Skip to content
Snippets Groups Projects
Commit 3d111d71 authored by Alexander Heidelbach's avatar Alexander Heidelbach
Browse files

Add constraints

parent 9bb8186d
Branches
No related tags found
No related merge requests found
from abc import ABC, abstractmethod
import numpy as np
from typing import Any, Dict, Mapping, Optional, Type, Union, List, TypeVar
from typing import Any, Dict, Mapping, Optional, Type, Union, List, TypeVar, Tuple
import zfit
import zfit.core
from hepstats.hypotests import Discovery
......@@ -25,7 +25,7 @@ class BaseFitter(ABC):
global_optimizer: str = "",
global_optimizer_args: Dict[Any, Any] = {},
truncated_limits: Optional[
tuple[float, float] | list[tuple[float, float]]
Union[Tuple[float, float], List[Tuple[float, float]]]
] = None,
) -> None:
self.model = model
......@@ -140,7 +140,11 @@ class BaseFitter(ABC):
def _single_fit(self) -> None:
assert isinstance(self.eval_model, zfit.core.interfaces.ZfitPDF)
nll = self.lossfunction(model=self.eval_model, data=self.data)
nll = self.lossfunction(
model=self.eval_model,
data=self.data,
constraints=self.model.constraints,
)
self.nll = nll
if self.global_optimizer:
minimizer_global = getattr(zfit.minimize, self.global_optimizer)(
......@@ -205,7 +209,7 @@ class BaseFitter(ABC):
self,
calculator: str = "AsymptoticCalculator",
calculator_kwargs: Dict[str, Any] = {"asimov_bins": 100},
) -> AsymptoticCalculator | FrequentistCalculator:
) -> Union[AsymptoticCalculator, FrequentistCalculator]:
assert self.nll is not None, "No NLL yet calculated, please check the fit"
if calculator == "AsymptoticCalculator":
return AsymptoticCalculator(self.nll, self.minimizer, **calculator_kwargs)
......@@ -220,7 +224,7 @@ class BaseFitter(ABC):
parameter_range: np.ndarray,
calculator: str = "AsymptoticCalculator",
calculator_kwargs: Dict[str, Any] = {"asimov_bins": 100},
) -> dict[str, float]:
) -> Dict[str, float]:
parameter = self.model.get_modelparameter(parameter_name)
calculator = self.get_calculator(
calculator=calculator, calculator_kwargs=calculator_kwargs
......@@ -241,7 +245,7 @@ class BaseFitter(ABC):
alpha: float = 0.05,
calculator: str = "AsymptoticCalculator",
calculator_kwargs: Dict[str, Any] = {"asimov_bins": 100},
) -> dict[str, float]:
) -> Dict[str, float]:
parameter = self.model.get_modelparameter(parameter_name)
calculator = self.get_calculator(
calculator=calculator, calculator_kwargs=calculator_kwargs
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment