data-manipulation-utilities 0.0.1__py3-none-any.whl

dmu/stats/fitter.py

@@ -0,0 +1,515 @@
+'''
+Module holding zfitter class
+'''
+
+import pprint
+from typing                   import Union
+
+import numpy
+import zfit
+import pandas as pd
+
+from scipy                    import stats
+from zfit.minimizers.strategy import FailMinimizeNaN
+from zfit.result              import FitResult
+from zfit.core.data           import Data
+from dmu.logging.log_store    import LogStore
+
+log = LogStore.add_logger('dmu:statistics:fitter')
+#------------------------------
+class FitterGofError(Exception):
+    '''
+    Exception used when GoF cannot be calculated
+    '''
+#------------------------------
+class FitterFailedFit(Exception):
+    '''
+    Exception used when fitter fails
+    '''
+#------------------------------
+class Fitter:
+    '''
+    Class meant to be an interface to underlying fitters
+    '''
+    # pylint: disable=too-many-instance-attributes
+    #------------------------------
+    def __init__(self, pdf, data):
+        self._data_in = data
+        self._pdf     = pdf
+
+        self._data_zf : zfit.data.Data
+        self._data_np : numpy.ndarray
+        self._obs     : zfit.Space
+        self._d_par   : dict
+
+        self._ndof           = 10
+        self._pval_threshold = 0.01
+        self._initialized    = False
+    #------------------------------
+    def _initialize(self):
+        if self._initialized:
+            return
+
+        self._check_data()
+
+        self._intialized = True
+    #------------------------------
+    def _check_data(self):
+        if   isinstance(self._data_in, numpy.ndarray):
+            data_np = self._data_in
+        elif isinstance(self._data_in, zfit.Data):
+            data_np = zfit.run(zfit.z.unstack_x(self._data_in)) # convert original data to numpy array, needed by _calc_gof
+        elif isinstance(self._data_in, pd.DataFrame):
+            data_np = self._data_in.to_numpy()
+        elif isinstance(self._data_in, pd.Series):
+            self._data_in = pd.DataFrame(self._data_in)
+            data_np = self._data_in.to_numpy()
+        else:
+            data_type = str(type(self._data_in))
+            raise ValueError(f'Data is not a numpy array, zfit.Data or pandas.DataFrame, but {data_type}')
+
+        data_np       = self._check_numpy_data(data_np)
+        self._data_np = data_np
+        if not isinstance(self._data_in, zfit.Data):
+            self._data_zf = zfit.Data.from_numpy(obs=self._pdf.space, array=data_np)
+        else:
+            self._data_zf = self._data_in
+    #------------------------------
+    def _check_numpy_data(self, data):
+        shp = data.shape
+        if   len(shp) == 1:
+            pass
+        elif len(shp) == 2:
+            _, jval = shp
+            if jval != 1:
+                log.error(f'Invalid data shape: {shp}')
+                raise
+        else:
+            log.error(f'Invalid data shape: {shp}')
+            raise
+
+        ival = data.shape[0]
+
+        data = data[~numpy.isnan(data)]
+        data = data[~numpy.isinf(data)]
+
+        fval = data.shape[0]
+
+        if ival != fval:
+            log.warning(f'Data was trimmed for inf and nan: {ival} -> {fval}')
+
+        return data
+    #------------------------------
+    def _bin_pdf(self, nbins):
+        [[min_x]], [[max_x]] = self._pdf.space.limits
+        _, arr_edg = numpy.histogram(self._data_np, bins = nbins, range=(min_x, max_x))
+
+        size = arr_edg.size
+
+        l_bc = []
+        for i_edg in range(size - 1):
+            low = arr_edg[i_edg + 0]
+            hig = arr_edg[i_edg + 1]
+
+            var = self._pdf.integrate(limits = [low, hig])
+            val = var.numpy()[0]
+            l_bc.append(val * self._data_np.size)
+
+        return numpy.array(l_bc)
+    #------------------------------
+    def _get_binning(self):
+        min_x = numpy.min(self._data_np)
+        max_x = numpy.max(self._data_np)
+        nbins = self._ndof + self._get_float_pars()
+
+        return nbins, min_x, max_x
+    #------------------------------
+    def _calc_gof(self):
+        log.debug('Calculating GOF')
+        nbins, min_x, max_x = self._get_binning()
+
+        log.debug(f'Nbins: {nbins}')
+        log.debug(f'Range: [{min_x:.3f}, {max_x:.3f}]')
+
+        arr_data, _ = numpy.histogram(self._data_np, bins = nbins, range=(min_x, max_x))
+        arr_data    = arr_data.astype(float)
+        arr_modl    = self._bin_pdf(nbins)
+        norm        = numpy.sum(arr_data) / numpy.sum(arr_modl)
+        arr_modl    = norm * arr_modl
+        arr_res     = arr_modl - arr_data
+
+        arr_chi2    = numpy.divide(arr_res ** 2, arr_data, out=numpy.zeros_like(arr_data), where=arr_data!=0)
+        sum_chi2    = numpy.sum(arr_chi2)
+        pvalue      = 1 - stats.chi2.cdf(sum_chi2, self._ndof)
+
+        log.debug(f'{"Data":<20}{"Model":<20}{"chi2":<20}')
+        if pvalue < self._pval_threshold:
+            for data, modl, chi2 in zip(arr_data, arr_modl, arr_chi2):
+                log.debug(f'{data:<20.0f}{modl:<20.3f}{chi2:<20.3f}')
+
+        log.debug(f'Chi2: {sum_chi2:.3f}')
+        log.debug(f'Ndof: {self._ndof}')
+        log.debug(f'pval: {pvalue:<.3e}')
+
+        return (sum_chi2, self._ndof, pvalue)
+    #------------------------------
+    def _get_float_pars(self):
+        npar     = 0
+        s_par    = self._pdf.get_params()
+        for par in s_par:
+            if par.floating:
+                npar+=1
+
+        self._d_par = {par.name : par for par in s_par}
+
+        return npar
+    #------------------------------
+    def _reshuffle_pdf_pars(self):
+        '''
+        Will move floating parameters of PDF according
+        to uniform PDF
+        '''
+
+        s_par = self._pdf.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 _set_pdf_pars(self, res):
+        '''
+        Will set the PDF floating parameter values as the result instance
+        '''
+        l_par_flt = list(self._pdf.get_params(floating= True))
+        l_par_fix = list(self._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:
+                log.debug(f'Skipping {par.name} = {par.value().numpy():.3e}')
+                continue
+
+            val = d_val[par.name]
+            log.debug(f'{"":<4}{par.name:<20}{"->":<10}{val:<20.3e}')
+            par.set_value(val)
+    #------------------------------
+    def _get_constraints(self, cfg : dict) -> list[zfit.constraint.GaussianConstraint]:
+        '''
+        Takes dictionary of constraints as floats, returns list of GaussianConstraint objects
+        '''
+        if 'constraints' not in cfg:
+            log.debug('Not using any constraint')
+            return []
+
+        d_const = cfg['constraints']
+        s_par   = self._pdf.get_params(floating=True)
+        d_par   = { par.name : par for par in s_par}
+
+        log.info('Adding constraints:')
+        l_const = []
+        for par_name, (par_mu, par_sg) in d_const.items():
+            if par_name not in d_par:
+                log.error(s_par)
+                raise ValueError(f'Parameter {par_name} not found among floating parameters of model, above')
+
+            par = d_par[par_name]
+
+            if par_sg == 0:
+                par.floating = False
+                log.info(f'{"":<4}{par_name:<15}{par_mu:<15.3e}{par_sg:<15.3e}')
+                continue
+
+            const = zfit.constraint.GaussianConstraint(params=par, observation=float(par_mu), uncertainty=float(par_sg))
+            log.info(f'{"":<4}{par_name:<25}{par_mu:<15.3e}{par_sg:<15.3e}')
+            l_const.append(const)
+
+        return l_const
+    #------------------------------
+    def _get_ranges(self, cfg : dict) -> list:
+        if 'ranges' not in cfg:
+            return [None]
+
+        ranges = cfg['ranges']
+        log.info('-' * 30)
+        log.info(f'{"Low edge":>15}{"High edge":>15}')
+        log.info('-' * 30)
+        for rng in ranges:
+            log.info(f'{rng[0]:>15.3e}{rng[1]:>15.3e}')
+
+        return ranges
+    #------------------------------
+    def _get_subdataset(self, cfg : dict) -> Data:
+        if 'nentries' not in cfg:
+            return self._data_zf
+
+        nentries_out = cfg['nentries']
+        arr_inp      = self._data_zf.to_numpy().flatten()
+        nentries_inp = len(arr_inp)
+        if nentries_inp <= nentries_out:
+            log.warning(f'Input dataset in smaller than output dataset, {nentries_inp} < {nentries_out}')
+            return self._data_zf
+
+        has_weights = self._data_zf.weights is not None
+
+        if has_weights:
+            arr_wgt = self._data_zf.weights.numpy()
+            arr_inp = numpy.array([arr_inp, arr_wgt]).T
+
+        arr_out = numpy.random.choice(arr_inp, size=nentries_out, replace=False)
+        if has_weights:
+            arr_out = arr_out.T[0]
+            arr_wgt = arr_out.T[1]
+        else:
+            arr_wgt = None
+
+        data = zfit.data.from_numpy(array=arr_out, weights=arr_wgt, obs=self._data_zf.obs)
+
+        return data
+    #------------------------------
+    def _get_binned_observable(self, nbins : int):
+        obs = self._pdf.space
+        [[minx]], [[maxx]] = obs.limits
+
+        binning = zfit.binned.RegularBinning(nbins, minx, maxx, name=obs.label)
+        obs_bin = zfit.Space(obs.label, binning=binning)
+
+        return obs_bin
+    #------------------------------
+    def _get_nbins(self, cfg : dict) -> Union[None, int]:
+        if 'likelihood' not in cfg:
+            return None
+
+        if 'nbins' not in cfg['likelihood']:
+            return None
+
+        return cfg['likelihood']['nbins']
+    #------------------------------
+    def _get_nll(self, data_zf, constraints, frange, cfg):
+        nbins     = self._get_nbins(cfg)
+        if nbins is None:
+            log.info('No binning was specified, will do unbinned fit')
+            pdf = self._pdf
+        else:
+            log.info(f'Using {nbins} bins for fit')
+            obs     = self._get_binned_observable(nbins)
+            pdf     = zfit.pdf.BinnedFromUnbinnedPDF(self._pdf, obs)
+            data_zf = data_zf.to_binned(obs)
+
+        if not self._pdf.is_extended and nbins is None:
+            nll = zfit.loss.UnbinnedNLL(        model=pdf, data=data_zf, constraints=constraints, fit_range=frange)
+            return nll
+
+        if     self._pdf.is_extended and nbins is None:
+            nll = zfit.loss.ExtendedUnbinnedNLL(model=pdf, data=data_zf, constraints=constraints, fit_range=frange)
+            return nll
+
+        if frange is not None:
+            raise ValueError('Fit range cannot be defined for binned likelihoods')
+
+        if not self._pdf.is_extended:
+            nll = zfit.loss.BinnedNLL(          model=pdf, data=data_zf, constraints=constraints)
+            return nll
+
+        if     self._pdf.is_extended:
+            nll = zfit.loss.ExtendedBinnedNLL(  model=pdf, data=data_zf, constraints=constraints)
+            return nll
+
+        raise ValueError('Likelihood was neither Binned nor Unbinned nor Extended nor non-extended')
+    #------------------------------
+    def _get_full_nll(self, cfg : dict):
+        constraints = self._get_constraints(cfg)
+        ranges      = self._get_ranges(cfg)
+        data_zf     = self._get_subdataset(cfg)
+        l_nll       = [ self._get_nll(data_zf, constraints, frange, cfg) for frange in ranges ]
+        nll         = sum(l_nll[1:], l_nll[0])
+
+        return nll
+    #------------------------------
+    def _print_pars(self, cfg : dict):
+        '''
+        Will print current values parameters in cfg['print_pars'] list, if present
+        '''
+
+        if 'print_pars' not in cfg:
+            return
+
+        l_par_name = cfg['print_pars']
+        d_par_val  = { name : par.value().numpy() for name, par in self._d_par.items() if name in l_par_name}
+
+        l_name = list(d_par_val.keys())
+        l_value= list(d_par_val.values())
+
+        l_form = [   f'{var:<10}' for var in l_name]
+        header = ''.join(l_form)
+
+        l_form = [f'{val:<10.3f}' for val in l_value]
+        parval = ''.join(l_form)
+
+        log.info(header)
+        log.info(parval)
+    #------------------------------
+    def _minimize(self, nll, cfg : dict) -> tuple[FitResult, tuple]:
+        mnm = zfit.minimize.Minuit()
+        res = mnm.minimize(nll)
+
+        try:
+            gof = self._calc_gof()
+        except FitterGofError as exc:
+            raise FitterGofError('Cannot calculate GOF') from exc
+
+        chi2, _, pval  = gof
+        stat           = res.status
+
+        log.info(f'{chi2:<10.3f}{pval:<10.3e}{stat:<10}')
+        self._print_pars(cfg)
+
+        return res, gof
+    #------------------------------
+    def _fit_retries(self, cfg : dict) -> tuple[dict, FitResult]:
+        ntries       = cfg['strategy']['retry']['ntries']
+        pvalue_thresh= cfg['strategy']['retry']['pvalue_thresh']
+        ignore_status= cfg['strategy']['retry']['ignore_status']
+
+        nll        = self._get_full_nll(cfg = cfg)
+        d_pval_res = {}
+        last_res   = None
+        for i_try in range(ntries):
+            try:
+                res, gof = self._minimize(nll, cfg)
+            except (FailMinimizeNaN, FitterGofError, RuntimeError):
+                self._reshuffle_pdf_pars()
+                log.warning(f'Fit {i_try:03}/{ntries:03} failed due to exception')
+                continue
+
+            last_res = res
+            bad_fit  = res.status != 0 or not res.valid
+
+            if not ignore_status and bad_fit:
+                self._reshuffle_pdf_pars()
+                log.info(f'Fit {i_try:03}/{ntries:03} failed, status/validity: {res.status}/{res.valid}')
+                continue
+
+            chi2, _, pval   = gof
+            d_pval_res[chi2]=res
+
+            if pval > pvalue_thresh:
+                log.info(f'Reached {pval:.3f} (> {pvalue_thresh:.3f}) threshold after {i_try + 1} attempts')
+                return {chi2 : res}, res
+
+            self._reshuffle_pdf_pars()
+
+        if last_res is None:
+            raise FitterFailedFit('Cannot find any valid fit')
+
+        return d_pval_res, last_res
+    #------------------------------
+    def _pick_best_fit(self, d_pval_res : dict, last_res : FitResult) -> FitResult:
+        nsucc = len(d_pval_res)
+        if nsucc == 0:
+            log.warning('None of the fits succeeded, returning last result')
+            self._set_pdf_pars(last_res)
+
+            return last_res
+
+        l_pval_res= list(d_pval_res.items())
+        l_pval_res.sort()
+        _, res = l_pval_res[0]
+
+        log.debug('Picking out best fit from {nsucc} fits')
+        for chi2, _ in l_pval_res:
+            log.debug(f'{chi2:.3f}')
+
+        self._set_pdf_pars(res)
+
+        return res
+    #------------------------------
+    def _fit_in_steps(self, cfg : dict) -> FitResult:
+        l_nsample = cfg['strategy']['steps']['nsteps']
+        l_nsigma  = cfg['strategy']['steps']['nsigma']
+        l_yield   = cfg['strategy']['steps']['yields']
+
+        res = None
+        for nsample, nsigma in zip(l_nsample, l_nsigma):
+            log.info(f'Fitting with {nsample} samples')
+            cfg_step             = dict(cfg)
+            cfg_step['nentries'] = nsample
+
+            nll    = self._get_full_nll(cfg = cfg_step)
+            res, _ = self._minimize(nll, cfg_step)
+            res.hesse(method='minuit_hesse')
+            self._update_par_bounds(res, nsigma=nsigma, yields=l_yield)
+
+        log.info('Fitting full sample')
+        nll    = self._get_full_nll(cfg = cfg)
+        res, _ = self._minimize(nll, cfg)
+        res.hesse(method='minuit_hesse')
+
+        if res is None:
+            nsteps = len(l_nsample)
+            raise ValueError(f'No fit out of {nsteps} was done')
+
+        return res
+    #------------------------------
+    def _result_to_value_error(self, res : FitResult) -> dict[str, list[float]]:
+        d_par = {}
+        for par, d_val in res.params.items():
+            try:
+                val = d_val['value']
+                err = d_val['hesse']['error']
+            except KeyError as exc:
+                pprint.pprint(d_val)
+                raise KeyError(f'Cannot extract value, hesse or error from dictionary above') from exc
+
+            d_par[par.name] = [val, err]
+
+        return d_par
+    #------------------------------
+    def _update_par_bounds(self, res : FitResult, nsigma : float, yields : list[str]) -> None:
+        s_shape_par = self._pdf.get_params(is_yield=False, floating=True)
+        d_shp_par   = { par.name : par for par in s_shape_par if par.name not in yields}
+        d_fit_par   = self._result_to_value_error(res)
+
+        log.info(60 * '-')
+        log.info(f'{"Parameter":<20}{"Low bound":<20}{"High bound":<20}')
+        log.info(60 * '-')
+        for name, [val, err] in d_fit_par.items():
+            if name not in d_shp_par:
+                log.debug(f'Skipping {name} parameter')
+                continue
+
+            shape       = d_shp_par[name]
+            shape.lower = val - nsigma * err
+            shape.upper = val + nsigma * err
+
+            log.info(f'{name:<20}{val - err:<20.3e}{val + err:<20.3e}')
+    #------------------------------
+    def fit(self, cfg : Union[dict, None] = None):
+        '''
+        Runs the fit using the configuration specified by the cfg dictionary
+
+        Returns fit result
+        '''
+        self._initialize()
+
+        cfg = {} if cfg is None else cfg
+
+        log.info(f'{"chi2":<10}{"pval":<10}{"stat":<10}')
+        if 'strategy' not in cfg:
+            nll    = self._get_full_nll(cfg = cfg)
+            res, _ = self._minimize(nll, cfg)
+            res.hesse(method='minuit_hesse')
+        elif 'retry' in cfg['strategy']:
+            d_pval_res, last_res = self._fit_retries(cfg)
+            res = self._pick_best_fit(d_pval_res, last_res)
+        elif 'steps' in cfg['strategy']:
+            res = self._fit_in_steps(cfg)
+        else:
+            raise ValueError('Unsupported fitting strategy')
+
+
+        return res
+#------------------------------