data-manipulation-utilities 0.1.6__py3-none-any.whl → 0.1.9__py3-none-any.whl

dmu/stats/minimizers.py

@@ -0,0 +1,183 @@
+'''
+Module containing derived classes from ZFit minimizer
+'''
+import numpy
+
+import zfit
+from zfit.result                   import FitResult
+from zfit.core.basepdf             import BasePDF           as zpdf
+from zfit.minimizers.baseminimizer import FailMinimizeNaN
+from dmu.stats.gof_calculator      import GofCalculator
+from dmu.logging.log_store         import LogStore
+
+log = LogStore.add_logger('dmu:ml:minimizers')
+# ------------------------
+class AnealingMinimizer(zfit.minimize.Minuit):
+    '''
+    Class meant to minimizer zfit likelihoods by using multiple retries,
+    each retry is preceeded by the randomization of the fitting parameters
+    '''
+    # ------------------------
+    def __init__(self, ntries : int, pvalue : float = -1, chi2ndof : float = -1):
+        '''
+        ntries  : Try this number of times
+        pvalue  : Stop tries when this threshold is reached
+        chi2ndof: Use this value as a threshold to stop fits
+        '''
+        self._ntries   = ntries
+        self._pvalue   = pvalue
+        self._chi2ndof = chi2ndof
+
+        self._check_thresholds()
+
+        super().__init__()
+    # ------------------------
+    def _check_thresholds(self) -> None:
+        good_pvalue  = 0 <= self._pvalue < 1
+        good_chi2dof = self._chi2ndof > 0
+
+        if good_pvalue and good_chi2dof:
+            raise ValueError('Threshold for both chi2 and pvalue were specified')
+
+        if good_pvalue:
+            log.debug(f'Will use threshold on pvalue with value: {self._pvalue}')
+            return
+
+        if good_chi2dof:
+            log.debug(f'Will use threshold on chi2ndof with value: {self._chi2ndof}')
+            return
+
+        raise ValueError('Neither pvalue nor chi2 thresholds are valid')
+    # ------------------------
+    def _is_good_gof(self, ch2 : float, pvl : float) -> bool:
+        is_good_pval = pvl > self._pvalue   and self._pvalue   > 0
+        is_good_chi2 = ch2 < self._chi2ndof and self._chi2ndof > 0
+        is_good      = is_good_pval or is_good_chi2
+
+        if is_good_pval:
+            log.info(f'Stopping fit, found p-value: {pvl:.3f} > {self._pvalue:.3f}')
+
+        if is_good_chi2:
+            log.info(f'Stopping fit, found chi2/ndof: {ch2:.3f} > {self._chi2ndof:.3f}')
+
+        if not is_good:
+            log.debug(f'Could not read threshold, pvalue/chi2: {pvl:.3f}/{ch2:.3f}')
+
+        return is_good
+    # ------------------------
+    def _is_good_fit(self, res : FitResult) -> bool:
+        if not res.valid:
+            log.warning('Skipping invalid fit')
+            return False
+
+        if res.status != 0:
+            log.warning('Skipping fit with bad status')
+            return False
+
+        if not res.converged:
+            log.warning('Skipping non-converging fit')
+            return False
+
+        return True
+    # ------------------------
+    def _get_gof(self, nll) -> tuple[float, float]:
+        log.debug('Checking GOF')
+
+        gcl = GofCalculator(nll)
+        pvl = gcl.get_gof(kind='pvalue')
+        ch2 = gcl.get_gof(kind='chi2/ndof')
+
+        return ch2, pvl
+    # ------------------------
+    def _randomize_parameters(self, nll):
+        '''
+        Will move floating parameters of PDF according
+        to uniform PDF
+        '''
+
+        log.debug('Randomizing parameters')
+        l_model = nll.model
+        if len(l_model) != 1:
+            raise ValueError('Not found and and only one model')
+
+        model = l_model[0]
+        s_par = model.get_params(floating=True)
+        for par in s_par:
+            ival = par.value()
+            fval = numpy.random.uniform(par.lower, par.upper)
+            par.set_value(fval)
+            log.debug(f'{par.name:<20}{ival:<15.3f}{"->":<10}{fval:<15.3f}{"in":<5}{par.lower:<15.3e}{par.upper:<15.3e}')
+    # ------------------------
+    def _pick_best_fit(self, d_chi2_res : dict) -> FitResult:
+        nres = len(d_chi2_res)
+        if nres == 0:
+            raise ValueError('No fits found')
+
+        l_chi2_res= list(d_chi2_res.items())
+        l_chi2_res.sort()
+        chi2, res = l_chi2_res[0]
+
+        log.warning(f'Picking out best fit from {nres} fits with chi2: {chi2:.3f}')
+
+        return res
+    #------------------------------
+    def _set_pdf_pars(self, res : FitResult, pdf : zpdf) -> None:
+        '''
+        Will set the PDF floating parameter values as the result instance
+        '''
+        l_par_flt = list(pdf.get_params(floating= True))
+        l_par_fix = list(pdf.get_params(floating=False))
+        l_par     = l_par_flt + l_par_fix
+
+        d_val = { par.name : dc['value'] for par, dc in res.params.items()}
+
+        log.debug('Setting PDF parameters to best result')
+        for par in l_par:
+            if par.name not in d_val:
+                par_val = par.value().numpy()
+                log.debug(f'Skipping {par.name} = {par_val:.3e}')
+                continue
+
+            val = d_val[par.name]
+            log.debug(f'{"":<4}{par.name:<20}{"->":<10}{val:<20.3e}')
+            par.set_value(val)
+    # ------------------------
+    def _pdf_from_nll(self, nll) -> zpdf:
+        l_model = nll.model
+        if len(l_model) != 1:
+            raise ValueError('Cannot extract one and only one PDF from NLL')
+
+        return l_model[0]
+    # ------------------------
+    def minimize(self, nll, **kwargs) -> FitResult:
+        '''
+        Will run minimization and return FitResult object
+        '''
+
+        d_chi2_res : dict[float,FitResult] = {}
+        for i_try in range(self._ntries):
+            log.info(f'try {i_try:02}/{self._ntries:02}')
+            try:
+                res = super().minimize(nll, **kwargs)
+            except (FailMinimizeNaN, ValueError, RuntimeError) as exc:
+                log.warning(exc)
+                self._randomize_parameters(nll)
+                continue
+
+            if not self._is_good_fit(res):
+                continue
+
+            chi2, pvl = self._get_gof(nll)
+            d_chi2_res[chi2] = res
+
+            if self._is_good_gof(chi2, pvl):
+                return res
+
+            self._randomize_parameters(nll)
+
+        res = self._pick_best_fit(d_chi2_res)
+        pdf = self._pdf_from_nll(nll)
+        self._set_pdf_pars(res, pdf)
+
+        return res
+# ------------------------

dmu/stats/model_factory.py

@@ -0,0 +1,207 @@
+'''
+Module storing ZModel class
+'''
+# pylint: disable=too-many-lines, import-error
+
+from typing import Callable, Union
+
+import zfit
+from zfit.core.interfaces   import ZfitSpace as zobs
+from zfit.core.basepdf      import BasePDF   as zpdf
+from zfit.core.parameter    import Parameter as zpar
+from dmu.logging.log_store  import LogStore
+
+log=LogStore.add_logger('dmu:stats:model_factory')
+#-----------------------------------------
+class MethodRegistry:
+    '''
+    Class intended to store protected methods belonging to ModelFactory class
+    which is defined in this same module
+    '''
+    # Registry dictionary to hold methods
+    _d_method = {}
+
+    @classmethod
+    def register(cls, nickname : str):
+        '''
+        Decorator in charge of registering method for given nickname
+        '''
+        def decorator(method):
+            cls._d_method[nickname] = method
+            return method
+
+        return decorator
+
+    @classmethod
+    def get_method(cls, nickname : str) -> Union[Callable,None]:
+        '''
+        Will return method in charge of building PDF, for an input nickname
+        '''
+        return cls._d_method.get(nickname, None)
+#-----------------------------------------
+class ModelFactory:
+    '''
+    Class used to create Zfit PDFs by passing only the nicknames, e.g.:
+
+    ```python
+    from dmu.stats.model_factory import ModelFactory
+
+    l_pdf = ['dscb', 'gauss']
+    l_shr = ['mu']
+    mod   = ModelFactory(obs = obs, l_pdf = l_pdf, l_shared=l_shr)
+    pdf   = mod.get_pdf()
+    ```
+
+    where one can specify which parameters can be shared among the PDFs
+    '''
+    #-----------------------------------------
+    def __init__(self, obs : zobs, l_pdf : list[str], l_shared : list[str]):
+        '''
+        obs:      zfit obserbable
+        l_pdf:    List of PDF nicknames which are registered below
+        l_shared: List of parameter names that are shared
+        '''
+
+        self._l_pdf           = l_pdf
+        self._l_shr           = l_shared
+        self._l_can_be_shared = ['mu', 'sg']
+        self._obs             = obs
+
+        self._d_par : dict[str,zpar] = {}
+    #-----------------------------------------
+    def _get_name(self, name : str, suffix : str) -> str:
+        for can_be_shared in self._l_can_be_shared:
+            if name.startswith(f'{can_be_shared}_') and can_be_shared in self._l_shr:
+                return can_be_shared
+
+        return f'{name}{suffix}'
+    #-----------------------------------------
+    def _get_parameter(self,
+                       name   : str,
+                       suffix : str,
+                       val    : float,
+                       low    : float,
+                       high   : float) -> zpar:
+        name = self._get_name(name, suffix)
+        if name in self._d_par:
+            return self._d_par[name]
+
+        par  = zfit.param.Parameter(name, val, low, high)
+
+        self._d_par[name] = par
+
+        return par
+    #-----------------------------------------
+    @MethodRegistry.register('exp')
+    def _get_exponential(self, suffix : str = '') -> zpdf:
+        c   = self._get_parameter('c_exp', suffix, -0.005, -0.05, 0.00)
+        pdf = zfit.pdf.Exponential(c, self._obs)
+
+        return pdf
+    #-----------------------------------------
+    @MethodRegistry.register('pol1')
+    def _get_pol1(self, suffix : str = '') -> zpdf:
+        a   = self._get_parameter('a_pol1', suffix, -0.005, -0.95, 0.00)
+        pdf = zfit.pdf.Chebyshev(obs=self._obs, coeffs=[a])
+
+        return pdf
+    #-----------------------------------------
+    @MethodRegistry.register('pol2')
+    def _get_pol2(self, suffix : str = '') -> zpdf:
+        a   = self._get_parameter('a_pol2', suffix, -0.005, -0.95, 0.00)
+        b   = self._get_parameter('b_pol2', suffix,  0.000, -0.95, 0.95)
+        pdf = zfit.pdf.Chebyshev(obs=self._obs, coeffs=[a, b])
+
+        return pdf
+    #-----------------------------------------
+    @MethodRegistry.register('cbr')
+    def _get_cbr(self, suffix : str = '') -> zpdf:
+        mu  = self._get_parameter('mu_cbr', suffix, 5300, 5250, 5350)
+        sg  = self._get_parameter('sg_cbr', suffix,   10,    2,  300)
+        ar  = self._get_parameter('ac_cbr', suffix,   -2,  -4.,  -1.)
+        nr  = self._get_parameter('nc_cbr', suffix,    1,  0.5,  5.0)
+
+        pdf = zfit.pdf.CrystalBall(mu, sg, ar, nr, self._obs)
+
+        return pdf
+    #-----------------------------------------
+    @MethodRegistry.register('cbl')
+    def _get_cbl(self, suffix : str = '') -> zpdf:
+        mu  = self._get_parameter('mu_cbl', suffix, 5300, 5250, 5350)
+        sg  = self._get_parameter('sg_cbl', suffix,   10,    2,  300)
+        al  = self._get_parameter('ac_cbl', suffix,    2,   1.,   4.)
+        nl  = self._get_parameter('nc_cbl', suffix,    1,  0.5,  5.0)
+
+        pdf = zfit.pdf.CrystalBall(mu, sg, al, nl, self._obs)
+
+        return pdf
+    #-----------------------------------------
+    @MethodRegistry.register('gauss')
+    def _get_gauss(self, suffix : str = '') -> zpdf:
+        mu  = self._get_parameter('mu_gauss', suffix, 5300, 5250, 5350)
+        sg  = self._get_parameter('sg_gauss', suffix,   10,    2,  300)
+
+        pdf = zfit.pdf.Gauss(mu, sg, self._obs)
+
+        return pdf
+    #-----------------------------------------
+    @MethodRegistry.register('dscb')
+    def _get_dscb(self, suffix : str = '') -> zpdf:
+        mu  = self._get_parameter('mu_dscb', suffix, 5300, 5250, 5400)
+        sg  = self._get_parameter('sg_dscb', suffix,   10,    2,   30)
+        ar  = self._get_parameter('ar_dscb', suffix,    1,    0,    5)
+        al  = self._get_parameter('al_dscb', suffix,    1,    0,    5)
+        nr  = self._get_parameter('nr_dscb', suffix,    2,    1,    5)
+        nl  = self._get_parameter('nl_dscb', suffix,    2,    0,    5)
+
+        pdf = zfit.pdf.DoubleCB(mu, sg, al, nl, ar, nr, self._obs)
+
+        return pdf
+    #-----------------------------------------
+    def _get_pdf_types(self) -> list[tuple[str,str]]:
+        d_name_freq = {}
+
+        l_type = []
+        for name in self._l_pdf:
+            if name not in d_name_freq:
+                d_name_freq[name] = 1
+            else:
+                d_name_freq[name]+= 1
+
+            frq = d_name_freq[name]
+            frq = f'_{frq}'
+
+            l_type.append((name, frq))
+
+        return l_type
+    #-----------------------------------------
+    def _get_pdf(self, kind : str, preffix : str) -> zpdf:
+        fun = MethodRegistry.get_method(kind)
+        if fun is None:
+            raise NotImplementedError(f'PDF of type {kind} is not implemented')
+
+        return fun(self, preffix)
+    #-----------------------------------------
+    def _add_pdf(self, l_pdf : list[zpdf]) -> zpdf:
+        nfrc = len(l_pdf)
+        if nfrc == 1:
+            log.debug('Requested only one PDF, skipping sum')
+            return l_pdf[0]
+
+        l_frc= [ zfit.param.Parameter(f'frc_{ifrc + 1}', 0.5, 0, 1) for ifrc in range(nfrc - 1) ]
+
+        pdf = zfit.pdf.SumPDF(l_pdf, fracs=l_frc)
+
+        return pdf
+    #-----------------------------------------
+    def get_pdf(self) -> zpdf:
+        '''
+        Given a list of strings representing PDFs returns the a zfit PDF which is
+        the sum of them
+        '''
+        l_type=   self._get_pdf_types()
+        l_pdf = [ self._get_pdf(kind, preffix) for kind, preffix in l_type ]
+        pdf   =   self._add_pdf(l_pdf)
+
+        return pdf
+#-----------------------------------------