Module flarefly.fitter

Module containing the class used to perform mass fits

Classes

class F2MassFitter (data_handler, name_signal_pdf, name_background_pdf, **kwargs)
Expand source code
class F2MassFitter:
    """
    Class used to perform mass fits with the zfit library
    https://github.com/zfit/zfit
    """

    # pylint: disable=too-many-statements, too-many-branches
    def __init__(self, data_handler, name_signal_pdf, name_background_pdf, **kwargs):
        """
        Initialize the F2MassFitter class
        Parameters
        -------------------------------------------------
        data_handler: flarefly.DataHandler
            The data handler containing the data to fit

        name_signal_pdf: list
            The list of names for the signal pdfs. The possible options are:

            - 'nosignal'

            - 'gaussian'

            - 'doublegaus'

            - 'gausexptail'

            - 'genergausexptail'

            - 'genergausexptailsymm'

            - 'bifurgaus'

            - 'crystalball'

            - 'doublecb'

            - 'doublecbsymm'

            - 'genercrystalball'

            - 'cauchy'

            - 'voigtian'

            - 'kde_exact' (requires to set the datasample and options)

            - 'kde_grid' (requires to set the datasample and options)

            - 'kde_fft' (requires to set the datasample and options)

            - 'kde_isj' (requires to set the datasample and options)

            - 'hist' (only for binned fits, requires to set the datasample)

        name_background_pdf: list
            The list of names of the background pdfs. The possible options are:

            - 'nobkg'

            - 'expo'

            - 'powlaw'

            - 'expopow'

            - 'expopowext'

            - 'chebpolN' (N is the order of the polynomial)

            - 'kde_exact' (requires to set the datasample and options)

            - 'kde_grid' (requires to set the datasample and options)

            - 'kde_fft' (requires to set the datasample and options)

            - 'kde_isj' (requires to set the datasample and options)

            - 'hist' (only for binned fits, requires to set the datasample)

        **kwargs: dict
            Additional optional arguments:

            - name_refl_pdf: list
                The list of names of the signal pdfs. It must have the same length as the signal list.
                The possible options are:

                - 'kde_exact' (requires to set the datasample and options)

                - 'kde_grid' (requires to set the datasample and options)

                - 'kde_fft' (requires to set the datasample and options)

                - 'kde_isj' (requires to set the datasample and options)

                - 'hist' (only for binned fits, requires to set the datasample)

            - name: str
                Optional name for the fitter,
                needed in case of multiple fitters defined in the same script

            - chi2_loss: bool
                chi2 minimization if True, nll minmization else,
                default value to False

            - minuit_mode:
                A number used by minuit to define the internal minimization strategy, either 0, 1 or 2.
                0 is the fastest, 2 is the slowest
                (see more details in
                https://zfit.readthedocs.io/en/latest/user_api/minimize/_generated/minimizers/zfit.minimize.Minuit.html#zfit.minimize.Minuit)
                Default value to 0

            - limits: list
                list of fit limits to include in the fit

            - tol: float
                Termination value for the convergence/stopping criterion of the algorithm in order to determine
                if the minimum has been found.
                Default value to 0.001

            - verbosity: int
                verbosity level (from 0 to 10)
                Default value to 0

            - signal_at_threshold: list
                list of booleans which indicate whether the signal PDFs are at threshold or not.
                Each element corresponds to a signal pdf

            - label_signal_pdf: list
                list of labels for signal pdfs

            - label_bkg_pdf: list
                list of labels for background pdfs
        """

        self._data_handler_ = data_handler
        self._name_signal_pdf_ = name_signal_pdf
        self._name_background_pdf_ = name_background_pdf
        self.label_signal_pdf = kwargs.get(
            'label_signal_pdf', [f'signal {idx}' for idx in range(len(name_signal_pdf))])
        self.label_bkg_pdf = kwargs.get(
            'label_bkg_pdf', [f'background {idx}' for idx in range(len(name_background_pdf))])
        self._name_refl_pdf_ = kwargs.get('name_refl_pdf', [None for _ in name_signal_pdf])
        if len(self._name_refl_pdf_) != len(self._name_signal_pdf_):
            Logger('List of pdfs for signals and reflections different! Exit', 'FATAL')
        if self._name_signal_pdf_[0] == 'nosignal':
            self._signal_pdf_ = []
            self._hist_signal_sample_ = []
            self._kde_signal_sample_ = []
            self._kde_signal_option_ = []
            self._name_refl_pdf_ = []
        else:
            self._signal_pdf_ = [None for _ in name_signal_pdf]
            self._hist_signal_sample_ = [None for _ in name_signal_pdf]
            self._kde_signal_sample_ = [None for _ in name_signal_pdf]
            self._kde_signal_option_ = [None for _ in name_signal_pdf]
        if self._name_background_pdf_[0] == 'nobkg':
            self._background_pdf_ = []
            self._hist_bkg_sample_ = []
            self._kde_bkg_sample_ = []
            self._kde_bkg_option_ = []
        else:
            self._background_pdf_ = [None for _ in name_background_pdf]
            self._hist_bkg_sample_ = [None for _ in name_background_pdf]
            self._kde_bkg_sample_ = [None for _ in name_background_pdf]
            self._kde_bkg_option_ = [None for _ in name_background_pdf]

        self._refl_pdf_ = [None for _ in self._name_refl_pdf_]
        self._hist_refl_sample_ = [None for _ in self._name_refl_pdf_]
        self._kde_refl_sample_ = [None for _ in self._name_refl_pdf_]
        self._kde_refl_option_ = [None for _ in self._name_refl_pdf_]
        self._refl_over_sgn_ = [0. for _ in self._name_refl_pdf_]
        self._total_pdf_ = None
        self._total_pdf_binned_ = None
        self._fit_result_ = None
        self._init_sgn_pars_ = [{} for _ in name_signal_pdf]
        self._init_bkg_pars_ = [{} for _ in name_background_pdf]
        self._limits_sgn_pars_ = [{} for _ in name_signal_pdf]
        self._limits_bkg_pars_ = [{} for _ in name_background_pdf]
        self._fix_sgn_pars_ = [{} for _ in name_signal_pdf]
        self._fix_fracs_to_pdfs_ = []
        self._fix_bkg_pars_ = [{} for _ in name_background_pdf]
        self._sgn_pars_ = [{} for _ in name_signal_pdf]
        self._bkg_pars_ = [{} for _ in name_background_pdf]
        if self._name_signal_pdf_[0] != 'nosignal' and self._name_background_pdf_[0] == 'nobkg':
            if self._name_refl_pdf_[0] is not None:
                Logger('Not possible to use reflections without background pdf', 'FATAL')
            self._fracs_ = [None for _ in range(len(name_signal_pdf) - 1)]
        elif self._name_signal_pdf_[0] == 'nosignal' and self._name_background_pdf_[0] != 'nobkg':
            self._fracs_ = [None for _ in range(len(name_background_pdf) - 1)]
        elif self._name_signal_pdf_[0] != 'nosignal' and self._name_background_pdf_[0] != 'nobkg':
            self._fracs_ = [None for _ in range(2 * len(name_signal_pdf) + len(name_background_pdf) - 1)]
        else:
            Logger('No signal nor background pdf defined', 'FATAL')
        self._total_pdf_norm_ = None
        self._ratio_truncated_ = None
        self._rawyield_ = [0. for _ in name_signal_pdf]
        self._rawyield_err_ = [0. for _ in name_signal_pdf]
        self._minimizer_ = zfit.minimize.Minuit(
            verbosity=kwargs.get('verbosity', 7),
            mode=kwargs.get('minuit_mode', 0),
            tol=kwargs.get('tol', 0.001)
        )
        if 'limits' in kwargs and self._data_handler_.get_is_binned():
            Logger('Restriction of fit limits is not yet implemented in binned fits!', 'FATAL')
        if 'limits' in kwargs and 'nosignal' in self._name_signal_pdf_:
            Logger('Using signal PDFs while restricting fit limits!', 'WARNING')
        self._limits_ = kwargs.get(
            'limits', self._data_handler_.get_limits())
        if not isinstance(self._limits_[0], list):
            self._limits_ = [self._limits_]
        self._is_truncated_ = self._limits_ != [self._data_handler_.get_limits()]
        self._name_ = kwargs.get('name', 'fitter')
        self._ndf_ = None
        self._chi2_loss_ = kwargs.get('chi2_loss', False)
        self._base_sgn_cmap_ = plt.colormaps.get_cmap('viridis')
        self._sgn_cmap_ = ListedColormap(self._base_sgn_cmap_(np.linspace(0.4, 0.65, len(self._signal_pdf_))))
        n_bkg_colors = len(self._background_pdf_)
        if len(self._background_pdf_) == 0:
            n_bkg_colors = 1  # to avoid crash
        self._base_bkg_cmap_ = plt.colormaps.get_cmap('Reds')
        self._bkg_cmap_ = ListedColormap(self._base_bkg_cmap_(np.linspace(0.8, 0.2, n_bkg_colors)))
        self._base_refl_cmap_ = plt.colormaps.get_cmap('summer')
        self._refl_cmap_ = ListedColormap(self._base_refl_cmap_(np.linspace(0., 0.6, len(self._refl_pdf_))))

        self._raw_residuals_ = []
        self._raw_residual_variances_ = []
        self._std_residuals_ = []
        self._std_residual_variances_ = []

        self._signal_at_threshold = kwargs.get('signal_at_threshold', [False for _ in name_signal_pdf])
        self._signalthr_pdf_ = [None for _ in name_signal_pdf]

        zfit.settings.advanced_warnings.all = False
        zfit.settings.changed_warnings.all = False
        self.pdg_api = pdg.connect()

    def cleanup(self):
        """
        Cleanup function to clear the graph cache of zfit
        """
        zfit.run.clear_graph_cache()

    def __del__(self):
        try:
            self.cleanup()
        except Exception:  # pylint: disable=broad-exception-caught
            pass

    # pylint: disable=too-many-branches, too-many-statements
    def __build_signal_pdfs(self, obs):
        """
        Helper function to compose the signal pdfs
        """

        for ipdf, (pdf_name, at_threshold) in enumerate(zip(self._name_signal_pdf_, self._signal_at_threshold)):
            if pdf_name == 'nosignal':
                Logger('Performing fit with no signal pdf', 'WARNING')
                break
            if 'kde' in pdf_name:
                self._signal_pdf_[ipdf] = PDFBuilder.build_signal_kde(
                    pdf_name,
                    self._kde_signal_sample_[ipdf],
                    self._name_,
                    ipdf,
                    self._kde_signal_option_[ipdf]
                )
            elif pdf_name == 'hist':
                self._signal_pdf_[ipdf] = PDFBuilder.build_signal_hist(
                    obs,
                    self._hist_signal_sample_[ipdf],
                    self._name_,
                    ipdf,
                )
            else:
                self._signal_pdf_[ipdf], self._sgn_pars_[ipdf] = PDFBuilder.build_signal_pdf(
                    pdf_name,
                    obs,
                    self._name_,
                    ipdf,
                    self._init_sgn_pars_[ipdf],
                    self._limits_sgn_pars_[ipdf],
                    self._fix_sgn_pars_[ipdf]
                )

            if at_threshold:
                # pion mass as default
                self._init_sgn_pars_[ipdf].setdefault('powerthr', 1.)
                self._limits_sgn_pars_[ipdf].setdefault('massthr', [None, None])
                self._fix_sgn_pars_[ipdf].setdefault('powerthr', False)
                self._fix_sgn_pars_[ipdf].setdefault('massthr', True)
                self._limits_sgn_pars_[ipdf].setdefault('powerthr', [None, None])
                self._init_sgn_pars_[ipdf].setdefault('massthr', self.pdg_api.get_particle_by_mcid(211).mass)
                self._sgn_pars_[ipdf][f'{self._name_}_massthr_signal{ipdf}'] = zfit.Parameter(
                    f'{self._name_}_massthr_signal{ipdf}', self._init_sgn_pars_[ipdf]['massthr'],
                    self._limits_sgn_pars_[ipdf]['massthr'][0], self._limits_sgn_pars_[ipdf]['massthr'][1],
                    floating=not self._fix_sgn_pars_[ipdf]['massthr'])
                self._sgn_pars_[ipdf][f'{self._name_}_powerthr_signal{ipdf}'] = zfit.Parameter(
                    f'{self._name_}_powerthr_signal{ipdf}', self._init_sgn_pars_[ipdf]['powerthr'],
                    self._limits_sgn_pars_[ipdf]['powerthr'][0], self._limits_sgn_pars_[ipdf]['powerthr'][1],
                    floating=not self._fix_sgn_pars_[ipdf]['powerthr'])
                self._signalthr_pdf_[ipdf] = cpdf.Pow(
                    obs=obs,
                    mass=self._sgn_pars_[ipdf][f'{self._name_}_massthr_signal{ipdf}'],
                    power=self._sgn_pars_[ipdf][f'{self._name_}_powerthr_signal{ipdf}']
                )
                self._signal_pdf_[ipdf] = zfit.pdf.ProductPDF(
                    [self._signal_pdf_[ipdf], self._signalthr_pdf_[ipdf]], obs=obs)

    def __build_background_pdfs(self, obs):
        """
        Helper function to compose the background pdfs
        """

        for ipdf, pdf_name in enumerate(self._name_background_pdf_):
            if pdf_name == 'nobkg':
                Logger('Performing fit with no bkg pdf', 'WARNING')
                break
            if 'kde' in pdf_name:
                self._background_pdf_[ipdf] = PDFBuilder.build_bkg_kde(
                    pdf_name,
                    self._kde_bkg_sample_[ipdf],
                    self._name_,
                    ipdf,
                    self._kde_bkg_option_[ipdf]
                )

            elif pdf_name == 'hist':
                self._background_pdf_[ipdf] = PDFBuilder.build_bkg_hist(
                    obs,
                    self._hist_bkg_sample_[ipdf],
                    self._name_,
                    ipdf,
                )
            else:
                self._background_pdf_[ipdf], self._bkg_pars_[ipdf] = PDFBuilder.build_bkg_pdf(
                    pdf_name,
                    obs,
                    self._name_,
                    ipdf,
                    self._init_bkg_pars_[ipdf],
                    self._limits_bkg_pars_[ipdf],
                    self._fix_bkg_pars_[ipdf]
                )
                if pdf_name in ['powlaw', 'expopow', 'expopowext'] and\
                        self._data_handler_.get_limits()[0] < self._init_bkg_pars_[ipdf]["mass"]:
                    Logger(
                        'The mass parameter in powlaw cannot be smaller than the lower fit limit, '
                        'please fix it.',
                        'FATAL'
                    )

    # pylint: disable=too-many-branches, too-many-statements
    def __build_reflection_pdfs(self, obs):
        """
        Helper function to compose the reflection pdfs
        """

        for ipdf, pdf_name in enumerate(self._name_refl_pdf_):
            if pdf_name is None:  # by default we put a dummy pdf
                low = zfit.Parameter(f'{self._name_}_low_refl{ipdf}',
                                     self._data_handler_.get_limits()[0],
                                     self._data_handler_.get_limits()[0] - 0.01,
                                     self._data_handler_.get_limits()[0] + 0.01,
                                     floating=False)
                high = zfit.Parameter(f'{self._name_}_high_refl{ipdf}',
                                      self._data_handler_.get_limits()[1],
                                      self._data_handler_.get_limits()[1] - 0.01,
                                      self._data_handler_.get_limits()[1] + 0.01,
                                      floating=False)
                self._refl_pdf_[ipdf] = zfit.pdf.Uniform(obs=obs, low=low, high=high)
            elif 'kde' in pdf_name:
                if self._kde_refl_sample_[ipdf]:
                    if pdf_name == 'kde_exact':
                        self._refl_pdf_[ipdf] = zfit.pdf.KDE1DimExact(self._kde_refl_sample_[ipdf].get_data(),
                                                                      obs=self._kde_refl_sample_[ipdf].get_obs(),
                                                                      name=f'{self._name_}_kde_refl{ipdf}',
                                                                      **self._kde_refl_option_[ipdf])
                    elif pdf_name == 'kde_grid':
                        self._refl_pdf_[ipdf] = zfit.pdf.KDE1DimGrid(self._kde_refl_sample_[ipdf].get_data(),
                                                                     obs=self._kde_refl_sample_[ipdf].get_obs(),
                                                                     name=f'{self._name_}_kde_refl{ipdf}',
                                                                     **self._kde_refl_option_[ipdf])
                    elif pdf_name == 'kde_fft':
                        self._refl_pdf_[ipdf] = zfit.pdf.KDE1DimFFT(self._kde_refl_sample_[ipdf].get_data(),
                                                                    obs=self._kde_refl_sample_[ipdf].get_obs(),
                                                                    name=f'{self._name_}_kde_refl{ipdf}',
                                                                    **self._kde_refl_option_[ipdf])
                    elif pdf_name == 'kde_isj':
                        self._refl_pdf_[ipdf] = zfit.pdf.KDE1DimISJ(self._kde_refl_sample_[ipdf].get_data(),
                                                                    obs=self._kde_refl_sample_[ipdf].get_obs(),
                                                                    name=f'{self._name_}_kde_refl{ipdf}',
                                                                    **self._kde_refl_option_[ipdf])
                else:
                    Logger(f'Missing datasample for Kernel Density Estimation of reflection {ipdf}!', 'FATAL')
            elif pdf_name == 'hist':
                if self._hist_refl_sample_[ipdf]:
                    self._refl_pdf_[ipdf] = zfit.pdf.SplinePDF(
                        zfit.pdf.HistogramPDF(self._hist_refl_sample_[ipdf].get_binned_data(),
                                              name=f'{self._name_}_hist_refl{ipdf}'),
                        order=3,
                        obs=obs
                    )
                else:
                    Logger(f'Missing datasample for histogram template of reflections {ipdf}!', 'FATAL')
            else:
                Logger(f'Reflection pdf {pdf_name} not supported', 'FATAL')

    def __get_constrained_frac_par(self, frac_par, factor_par, refl=False):
        """
        Helper function to create a parameter constrained to a value
        """

        def par_func(par, factor):
            return par * factor

        if refl:
            name_composed_param = f'{self._name_}_{factor_par.name.replace("factor", "frac_refl")}'
        else:
            name_composed_param = f'{self._name_}_{factor_par.name.replace("factor", "frac")}'

        frac_par_constrained = zfit.ComposedParameter(name_composed_param,
                                                      par_func,
                                                      params=[frac_par, factor_par],
                                                      unpack_params=True
                                                      )
        return frac_par_constrained

    def __set_frac_constraints(self):
        for frac_fix_info in self._fix_fracs_to_pdfs_:
            if frac_fix_info['fixed_pdf_type'] == 'signal':
                if frac_fix_info['target_pdf_type'] == 'signal':
                    self._fracs_[frac_fix_info['fixed_pdf_idx']] = self.__get_constrained_frac_par(
                        self._fracs_[frac_fix_info['target_pdf_idx']],
                        frac_fix_info['factor']
                    )
                    # also set for reflections
                    fix_factor = zfit.Parameter(
                        frac_fix_info['factor'].name.replace('factor', 'factor_refl'),
                        float(
                            frac_fix_info['factor'].value()
                        ) * self._refl_over_sgn_[frac_fix_info['fixed_pdf_idx']],
                        floating=False
                    )
                    self._fracs_[frac_fix_info['fixed_pdf_idx'] + len(self._signal_pdf_)] = \
                        self.__get_constrained_frac_par(
                            self._fracs_[frac_fix_info['target_pdf_idx']],
                            fix_factor,
                            refl=True
                        )
                else:  # fix to background PDF
                    self._fracs_[frac_fix_info['fixed_pdf_idx']] = self.__get_constrained_frac_par(
                        self._fracs_[frac_fix_info['target_pdf_idx'] + 2 * len(self._signal_pdf_)],
                        frac_fix_info['factor']
                    )
                    # also set for reflections
                    fix_factor = zfit.Parameter(
                        frac_fix_info['factor'].name.replace('factor', 'factor_refl'),
                        float(
                            frac_fix_info['factor'].value()
                        ) * self._refl_over_sgn_[frac_fix_info['fixed_pdf_idx']],
                        floating=False
                    )
                    self._fracs_[frac_fix_info['fixed_pdf_idx'] + len(self._signal_pdf_)] = \
                        self.__get_constrained_frac_par(
                            self._fracs_[frac_fix_info['target_pdf_idx'] + 2 * len(self._signal_pdf_)],
                            fix_factor,
                            refl=True
                        )
            else:  # fix background fraction
                if frac_fix_info['target_pdf_type'] == 'signal':
                    self._fracs_[frac_fix_info['fixed_pdf_idx'] + 2 * len(self._signal_pdf_)] = \
                        self.__get_constrained_frac_par(
                            self._fracs_[frac_fix_info['target_pdf_idx']],
                            frac_fix_info['factor']
                        )
                else:  # fix to background PDF
                    self._fracs_[frac_fix_info['fixed_pdf_idx'] + 2 * len(self._signal_pdf_)] = \
                        self.__get_constrained_frac_par(
                            self._fracs_[frac_fix_info['target_pdf_idx'] + 2 * len(self._signal_pdf_)],
                            frac_fix_info['factor']
                        )

    def __build_total_pdf(self):
        """
        Helper function to compose the total pdf
        """

        if self._data_handler_.get_is_binned():  # we need unbinned pdfs to sum them
            obs = self._data_handler_.get_unbinned_obs_from_binned_data()
        else:
            obs = self._data_handler_.get_obs()

        # order of the pdfs is signal, background

        self.__build_signal_pdfs(obs)
        self.__build_background_pdfs(obs)
        if len(self._background_pdf_) != 0:
            self.__build_reflection_pdfs(obs)
        else:
            self._refl_pdf_ = []

        self.__get_total_pdf_norm()

        if len(self._signal_pdf_) + len(self._background_pdf_) == 1:
            if len(self._signal_pdf_) == 0:
                self._total_pdf_ = self._background_pdf_[0]
                if self._is_truncated_:
                    self._total_pdf_ = self._total_pdf_.to_truncated(limits=self._limits_, obs=obs)
                return
            if len(self._background_pdf_) == 0:
                self._total_pdf_ = self._signal_pdf_[0]
                if self._is_truncated_:
                    self._total_pdf_ = self._total_pdf_.to_truncated(limits=self._limits_, obs=obs)
                return

        for ipdf, _ in enumerate(self._signal_pdf_):
            self._init_sgn_pars_[ipdf].setdefault('frac', 0.1)
            self._fix_sgn_pars_[ipdf].setdefault('frac', False)
            self._limits_sgn_pars_[ipdf].setdefault('frac', [0, 1.])
            if len(self._background_pdf_) == 0 and ipdf == len(self._signal_pdf_) - 1:
                continue
            self._fracs_[ipdf] = zfit.Parameter(f'{self._name_}_frac_signal{ipdf}',
                                                self._init_sgn_pars_[ipdf]['frac'],
                                                self._limits_sgn_pars_[ipdf]['frac'][0],
                                                self._limits_sgn_pars_[ipdf]['frac'][1],
                                                floating=not self._fix_sgn_pars_[ipdf]['frac'])

            # normalisation of reflection fixed to the one of the signal
            if len(self._background_pdf_) != 0:
                def func_mult(params):
                    return params['ros'] * params['s']
                self._fracs_[ipdf + len(self._signal_pdf_)] = zfit.ComposedParameter(
                    f'{self._name_}_frac_refl{ipdf}',
                    func_mult, params={'ros': self._refl_over_sgn_[ipdf],
                                       's': self._fracs_[ipdf]}
                )

        if len(self._background_pdf_) > 1:
            for ipdf, _ in enumerate(self._background_pdf_[:-1]):
                self._init_bkg_pars_[ipdf].setdefault('frac', 0.1)
                self._fix_bkg_pars_[ipdf].setdefault('frac', False)
                self._limits_bkg_pars_[ipdf].setdefault('frac', [0, 1.])
                self._fracs_[ipdf + 2 * len(self._signal_pdf_)] = zfit.Parameter(
                    f'{self._name_}_frac_bkg{ipdf}',
                    self._init_bkg_pars_[ipdf]['frac'],
                    self._limits_bkg_pars_[ipdf]['frac'][0],
                    self._limits_bkg_pars_[ipdf]['frac'][1],
                    floating=not self._fix_bkg_pars_[ipdf]['frac'])

        self.__set_frac_constraints()

        self._total_pdf_ = zfit.pdf.SumPDF(self._signal_pdf_+self._refl_pdf_+self._background_pdf_,
                                           self._fracs_)

        if not self._data_handler_.get_is_binned() and self._is_truncated_:
            self._total_pdf_ = self._total_pdf_.to_truncated(limits=self._limits_, obs=obs, norm=obs)

    def __build_total_pdf_binned(self):
        """
        Helper function to compose the total pdf binned from unbinned
        """

        # for binned data, obs already contains the wanted binning
        if self._data_handler_.get_is_binned():
            obs = self._data_handler_.get_obs()
        # for unbinned data, one needs to impose a binning
        else:
            obs = self._data_handler_.get_binned_obs_from_unbinned_data()

        if len(self._signal_pdf_) + len(self._background_pdf_) == 1:
            if len(self._signal_pdf_) == 0:
                self._total_pdf_binned_ = zfit.pdf.BinnedFromUnbinnedPDF(self._background_pdf_[0], obs)
                return
            if len(self._background_pdf_) == 0:
                self._total_pdf_binned_ = zfit.pdf.BinnedFromUnbinnedPDF(self._signal_pdf_[0], obs)
                return

        self._total_pdf_binned_ = zfit.pdf.BinnedFromUnbinnedPDF(zfit.pdf.SumPDF(
            self._signal_pdf_+self._refl_pdf_+self._background_pdf_, self._fracs_), obs)

    def __get_frac_and_err(self, par_name, frac_par, frac_type):
        if 'constrained' in par_name:
            split_par_name = par_name.split(sep='_constrained_to_')
            target_frac_name = split_par_name[0].replace(split_par_name[0].split('_')[-1], split_par_name[1])
            # find the parameter containing target_frac_name
            target_par_name, target_par = [
                (par.name, par) for par in self._fracs_ if target_frac_name in par.name
            ][0]

            if 'constrained' in target_par_name:
                if f'{self._name_}_frac_signal' in target_par_name:
                    frac_type = 'signal'
                else:
                    frac_type = 'bkg'
                frac_temp, err_temp, isgn = self.__get_frac_and_err(target_par_name, target_par, frac_type)
                return (
                    frac_temp * float(frac_par.params['param_1'].value()),
                    err_temp * float(frac_par.params['param_1'].value()),
                    isgn
                )
            frac = self._fit_result_.params[target_frac_name]['value']
            frac_err = self._fit_result_.params[target_frac_name]['hesse']['error']

            return (
                frac * float(frac_par.params['param_1'].value()),
                frac_err * float(frac_par.params['param_1'].value()),
                int(split_par_name[0][-1])
            )

        return (
            self._fit_result_.params[par_name]['value'],
            self._fit_result_.params[par_name]['hesse']['error'],
            int(par_name.split(sep=f'{self._name_}_frac_{frac_type}')[-1])
        )

    def __get_all_fracs(self):
        """
        Helper function to get all fractions

        Returns
        -------------------------------------------------
        signal_fracs: list
            fractions of the signal pdfs
        background_fracs: list
            fractions of the background pdfs
        refl_fracs: list
            fractions of the reflected signal pdfs
        signal_err_fracs: list
            errors of fractions of the signal pdfs
        bkg_err_fracs: list
            errors of fractions of the background pdfs
        refl_err_fracs: list
            errors of fractions of the reflected signal pdfs
        """
        signal_fracs, bkg_fracs, refl_fracs, signal_err_fracs, bkg_err_fracs, refl_err_fracs = ([] for _ in range(6))
        for frac_par in self._fracs_:
            if frac_par is None:
                continue
            par_name = frac_par.name
            if f'{self._name_}_frac_signal' in par_name:
                signal_frac, signal_err, isgn = self.__get_frac_and_err(par_name, frac_par, 'signal')
                signal_fracs.append(signal_frac)
                signal_err_fracs.append(signal_err)
                refl_fracs.append(signal_frac * self._refl_over_sgn_[isgn])
                refl_err_fracs.append(signal_err * self._refl_over_sgn_[isgn])
            elif f'{self._name_}_frac_bkg' in par_name:
                bkg_frac, bkg_err, isgn = self.__get_frac_and_err(par_name, frac_par, 'bkg')
                bkg_fracs.append(bkg_frac)
                bkg_err_fracs.append(bkg_err)

        if len(signal_fracs) == len(bkg_fracs) == len(refl_fracs) == 0:
            if len(self._background_pdf_) == 0:
                signal_fracs.append(1.)
                signal_err_fracs.append(0.)
                refl_fracs.append(0.)
                refl_err_fracs.append(0.)
            elif len(self._signal_pdf_) == 0:
                bkg_fracs.append(1.)
                bkg_err_fracs.append(0.)
        else:
            if len(self._background_pdf_) == 0:
                signal_fracs.append(1 - sum(signal_fracs) - sum(refl_fracs) - sum(bkg_fracs))
                signal_err_fracs.append(np.sqrt(sum(list(err**2 for err in signal_err_fracs + bkg_err_fracs))))
            else:
                bkg_fracs.append(1 - sum(signal_fracs) - sum(refl_fracs) - sum(bkg_fracs))
                bkg_err_fracs.append(np.sqrt(sum(list(err**2 for err in signal_err_fracs + bkg_err_fracs))))

        return signal_fracs, bkg_fracs, refl_fracs, signal_err_fracs, bkg_err_fracs, refl_err_fracs

    def __get_total_pdf_norm(self):
        """
        Get the normalization of the total pdf
        """
        if self._data_handler_.get_is_binned():
            self._total_pdf_norm_ = float(self._data_handler_.get_norm())
            return

        norm_total_pdf = 0.
        for lim in self._limits_:
            norm_obs = self._data_handler_.get_obs().with_limits(lim)
            norm_total_pdf += float(self._data_handler_.get_data().with_obs(norm_obs).n_events)
        self._total_pdf_norm_ = norm_total_pdf

    def __get_ratio_truncated(self):
        """
        Get the ratio of the sidebands integral over the total integral
        """
        if not self._is_truncated_:
            self._ratio_truncated_ = 1.
            return

        signal_fracs, bkg_fracs, refl_fracs, _, _, _ = self.__get_all_fracs()
        fracs = signal_fracs + refl_fracs + bkg_fracs
        limits = self._data_handler_.get_limits()
        sidebands_integral, total_integral = 0., 0.
        for i_pdf, pdf in enumerate(self._signal_pdf_ + self._refl_pdf_ + self._background_pdf_):
            sidebands_integral += float(sum(pdf.integrate(lim) * fracs[i_pdf] for lim in self._limits_))
            total_integral += float(pdf.integrate(limits) * fracs[i_pdf])

        self._ratio_truncated_ = sidebands_integral / total_integral

    def __get_raw_residuals(self):
        """
        Get the raw residuals (data_value - bkg_model_value) for all bins
        """

        bins = self._data_handler_.get_nbins()
        norm = self._total_pdf_norm_
        self._raw_residuals_ = [None]*bins
        background_pdf_binned_ = [None for _ in enumerate(self._name_background_pdf_)]
        model_bkg_values = [None for _ in enumerate(self._name_background_pdf_)]

        # access normalized data values and errors for all bins
        if self._data_handler_.get_is_binned():
            binned_data = self._data_handler_.get_binned_data()
            data_values = binned_data.values()
            self._raw_residual_variances_ = binned_data.variances()
            obs = self._data_handler_.get_obs()
        else:
            data_values = self._data_handler_.get_binned_data_from_unbinned_data()
            self._raw_residual_variances_ = data_values  # poissonian errors
            obs = self._data_handler_.get_binned_obs_from_unbinned_data()

        # get background fractions
        if len(self._background_pdf_) > 0:
            if len(self._background_pdf_) == 1:
                _, bkg_fracs, _, _, _, _ = self.__get_all_fracs()
            else:
                _, bkg_fracs, _, _, _, _ = self.__get_all_fracs()
            # access model predicted values for background
            for ipdf, bkg_name in enumerate(self._name_background_pdf_):
                background_pdf_binned_[ipdf] = zfit.pdf.BinnedFromUnbinnedPDF(self._background_pdf_[ipdf], obs)
                norm_bkg = norm
                if "hist" in bkg_name:
                    norm_bkg /= float(sum(background_pdf_binned_[ipdf].values()))
                model_bkg_values[ipdf] = background_pdf_binned_[ipdf].values() * bkg_fracs[ipdf]
                model_bkg_values[ipdf] *= norm_bkg / self._ratio_truncated_
            # compute residuals
            for ibin, data in enumerate(data_values):
                self._raw_residuals_[ibin] = float(data)
                for ipdf, _ in enumerate(self._name_background_pdf_):
                    self._raw_residuals_[ibin] -= model_bkg_values[ipdf][ibin]
        else:
            for ibin, data in enumerate(data_values):
                self._raw_residuals_[ibin] = float(data)

    def __get_std_residuals(self):
        """
        Get the standardized residuals
        (data_value - bkg_model_value)/ sigma_data for all bins
        """

        bins = self._data_handler_.get_nbins()
        norm = self._total_pdf_norm_
        self._std_residuals_ = [None]*bins
        self._std_residual_variances_ = [None]*bins

        # access normalized data values and errors for all bins
        if self._data_handler_.get_is_binned():
            binned_data = self._data_handler_.get_binned_data()
            data_values = binned_data.values()
            variances = binned_data.variances()
        else:
            if self._limits_ != [self._data_handler_.get_limits()]:  # restricted fit limits
                Logger('Standard residuals not yet implemented with restricted fit limits', 'FATAL')
            data_values = self._data_handler_.get_binned_data_from_unbinned_data()
            variances = data_values  # poissonian errors

        # access model predicted values for background
        self.__build_total_pdf_binned()
        model_values = self._total_pdf_binned_.values()*norm/self._ratio_truncated_
        for ibin, (data, model, variance) in enumerate(zip(data_values, model_values, variances)):
            if variance == 0:
                Logger('Null variance. Consider enlarging the bins.', 'FATAL')
            self._std_residuals_[ibin] = float((data - model)/np.sqrt(variance))
            self._std_residual_variances_[ibin] = float(variance/np.sqrt(variance))

    def __prefit(self, excluded_regions):
        """
        Perform a prefit to the background distributions with the zfit library.
        The expected signal regions must be provided.

        Parameters
        -------------------------------------------------
        excluded_regions: list
            List or list of lists with limits for the excluded regions

        """

        if self._name_background_pdf_[0] == "nobkg":
            Logger('Prefit cannot be performed in case of no background, skip', 'WARNING')
            return

        if self._data_handler_.get_is_binned():
            Logger('Prefit not yet implemented for binned fits, skip', 'WARNING')
            return

        if len(self._limits_) > 1:
            Logger('Prefit not supported for fits on disjoint intervals, skip', 'WARNING')
            return

        obs = self._data_handler_.get_obs()

        # we build the limits
        if not isinstance(excluded_regions[0], list):
            excluded_regions = [excluded_regions]

        # check if excluded regions overlap
        limits_sb = []
        if excluded_regions[0][0] > self._limits_[0][0]:
            limits_sb.append([self._limits_[0][0], excluded_regions[0][0]])
        for iregion, region in enumerate(excluded_regions[:-1]):
            if region[1] > self._limits_[0][0] and \
                excluded_regions[iregion+1][0] < self._limits_[-1][-1] and \
                    region[1] < excluded_regions[iregion+1][0]:
                limits_sb.append([region[1], excluded_regions[iregion+1][0]])
        if excluded_regions[-1][1] < self._limits_[-1][-1]:
            limits_sb.append([excluded_regions[-1][1], self._limits_[-1][-1]])

        if len(limits_sb) == 0:
            Logger('Cannot perform prefit with the excluded regions set since no sidebands left, skip', 'WARNING')
            return

        fracs_bkg_prefit = []
        for ipdf, _ in enumerate(self._background_pdf_):
            fracs_bkg_prefit.append(zfit.Parameter(f'{self._name_}_frac_bkg{ipdf}_prefit', 0.1, 0., 1.))

        prefit_background_pdf = None
        if len(self._background_pdf_) > 1:
            prefit_background_pdf = zfit.pdf.SumPDF(self._background_pdf_, fracs_bkg_prefit).to_truncated(
                limits=limits_sb, obs=obs, norm=obs)
        else:
            prefit_background_pdf = self._background_pdf_[0].to_truncated(
                limits=limits_sb, obs=obs, norm=obs)

        prefit_loss = zfit.loss.UnbinnedNLL(model=prefit_background_pdf,
                                            data=self._data_handler_.get_data())

        Logger("Performing background only PREFIT", "INFO")
        res_prefit = self._minimizer_.minimize(loss=prefit_loss)
        Logger(res_prefit, 'RESULT')

        # re-initialise the parameters from those obtained in the prefit
        for par in res_prefit.params:
            which_pdf = int(par.name[-1])
            self._bkg_pars_[which_pdf][par.name].set_value(res_prefit.params[par.name]['value'])

    # pylint: disable=too-many-nested-blocks
    def mass_zfit(self, do_prefit=False, **kwargs):
        """
        Perform a mass fit with the zfit library

        Parameters
        -------------------------------------------------
        do_prefit: bool
            Flag to enable prefit to background only (supported for unbinned fits only)
            The expected signal regions must be provided.

        **kwargs: dict
            Additional optional arguments:

            - prefit_excluded_regions: list
                List or list of lists with limits for the excluded regions

            - prefit_exclude_nsigma:
                Alternative argument to prefit_excluded_regions that excludes the regions around the signals,
                using the initialised mean and sigma parameters

        Returns
        -------------------------------------------------
        fit_result: zfit.minimizers.fitresult.FitResult
            The fit result
        """

        if self._data_handler_ is None:
            Logger('Data handler not specified', 'FATAL')

        self._raw_residuals_ = []
        self._raw_residual_variances_ = []
        self._std_residuals_ = []
        self._std_residual_variances_ = []

        self.__build_total_pdf()
        self.__build_total_pdf_binned()

        # do prefit
        if do_prefit:
            skip_prefit = False
            excluded_regions = []
            exclude_signals_nsigma = kwargs.get('prefit_exclude_nsigma', None)
            if exclude_signals_nsigma is not None:
                for ipdf, pdf_name in enumerate(self._name_signal_pdf_):
                    if pdf_name not in ['gaussian', 'cauchy', 'voigtian']:
                        Logger(f"Sigma parameter not defined for {pdf_name}, "
                               "cannot use nsigma for excluded regions in prefit", "Error")
                        skip_prefit = True
                    else:
                        mass_name = 'm' if pdf_name in ['cauchy', 'voigtian'] else 'mu'
                        mass = self._init_sgn_pars_[ipdf][mass_name]
                        sigma = self._init_sgn_pars_[ipdf]['sigma']
                        excluded_regions.append(
                            [mass - exclude_signals_nsigma * sigma, mass + exclude_signals_nsigma * sigma])
            else:
                excluded_regions = kwargs.get('prefit_excluded_regions', None)
                if excluded_regions is None:
                    Logger("To perform the prefit, you should set the excluded regions, skip", "Error")
                    skip_prefit = True

            if not skip_prefit:
                self.__prefit(excluded_regions)

        if self._data_handler_.get_is_binned():
            # chi2 loss
            if self._chi2_loss_:
                loss = zfit.loss.BinnedChi2(self._total_pdf_binned_, self._data_handler_.get_binned_data())
            # nll loss
            else:
                loss = zfit.loss.BinnedNLL(self._total_pdf_binned_, self._data_handler_.get_binned_data())
        else:
            loss = zfit.loss.UnbinnedNLL(model=self._total_pdf_, data=self._data_handler_.get_data())

        Logger("Performing FIT", "INFO")
        self._fit_result_ = self._minimizer_.minimize(loss=loss)
        Logger(self._fit_result_, 'RESULT')

        if self._fit_result_.hesse() == {}:
            if self._fit_result_.hesse(method='hesse_np') == {}:
                Logger('Impossible to compute hesse error', 'FATAL')

        self.__get_ratio_truncated()
        signal_fracs, _, _, signal_frac_errs, _, _ = self.__get_all_fracs()

        norm = self._total_pdf_norm_

        if len(self._fracs_) == 0:
            self._rawyield_[0] = self._data_handler_.get_norm()
            self._rawyield_err_[0] = np.sqrt(self._rawyield_[0])
        else:
            for i_pdf, _ in enumerate(self._signal_pdf_):
                self._rawyield_[i_pdf] = signal_fracs[i_pdf] * norm / self._ratio_truncated_
                self._rawyield_err_[i_pdf] = signal_frac_errs[i_pdf] * norm / self._ratio_truncated_

        return self._fit_result_

    # pylint: disable=too-many-statements, too-many-locals
    def plot_mass_fit(self, **kwargs):
        """
        Plot the mass fit

        Parameters
        -------------------------------------------------
        **kwargs: dict
            Additional optional arguments:

            - style: str
                style to be used (see https://github.com/scikit-hep/mplhep for more details)

            - logy: bool
                log scale in y axis

            - figsize: tuple
                size of the figure

            - axis_title: str
                x-axis title

            - show_extra_info: bool
                show mu, sigma, chi2/ndf, signal, bkg, signal/bkg, significance

            - extra_info_massnsigma: float
                number of sigmas for extra info

            - extra_info_massnhwhm: float
                number of hwhms for extra info (alternative to extra_info_massnsigma)

            - extra_info_massrange: list
                mass range limits for extra info (alternative to extra_info_massnsigma)

            - extra_info_loc: list
                location of extra info (one for chi2 and one for other info)

            - legend_loc: str
                location of the legend

            - num: int
                number of bins to plot pdfs converted into histograms

        Returns
        -------------------------------------------------
        fig: matplotlib.figure.Figure
            figure containing the mass fit plot
        """

        style = kwargs.get('style', 'LHCb2')
        logy = kwargs.get('logy', False)
        figsize = kwargs.get('figsize', (7, 7))
        bins = self._data_handler_.get_nbins()
        axis_title = kwargs.get('axis_title', self._data_handler_.get_var_name())
        show_extra_info = kwargs.get('show_extra_info', False)
        num = kwargs.get('num', 10000)
        mass_range = kwargs.get('extra_info_massrange', None)
        nhwhm = kwargs.get('extra_info_massnhwhm', None)
        nsigma = kwargs.get('extra_info_massnsigma', 3)
        info_loc = kwargs.get('extra_info_loc', ['lower right', 'upper left'])
        legend_loc = kwargs.get('legend_loc', 'best')

        mplhep.style.use(style)

        obs = self._data_handler_.get_obs()
        limits = self._data_handler_.get_limits()

        fig, axs = plt.subplots(figsize=figsize)

        hdata = self._data_handler_.to_hist(lower_edge=limits[0],
                                            upper_edge=limits[1],
                                            nbins=bins,
                                            varname=self._data_handler_.get_var_name())

        hdata.plot(yerr=True, color='black', histtype='errorbar', label='data')
        bin_sigma = (limits[1] - limits[0]) / bins
        norm_total_pdf = self._total_pdf_norm_ * bin_sigma

        x_plot = np.linspace(limits[0], limits[1], num=num)
        total_func = zfit.run(self._total_pdf_.pdf(x_plot, norm_range=obs))
        signal_funcs, refl_funcs, bkg_funcs = ([] for _ in range(3))

        for signal_pdf in self._signal_pdf_:
            signal_funcs.append(zfit.run(signal_pdf.pdf(x_plot, norm_range=obs)))
        for refl_pdf in self._refl_pdf_:
            refl_funcs.append(zfit.run(refl_pdf.pdf(x_plot, norm_range=obs)))
        for bkg_pdf in self._background_pdf_:
            bkg_funcs.append(zfit.run(bkg_pdf.pdf(x_plot, norm_range=obs)))

        signal_fracs, bkg_fracs, refl_fracs, _, _, _ = self.__get_all_fracs()

        # first draw backgrounds
        for ibkg, (bkg_func, bkg_frac) in enumerate(zip(bkg_funcs, bkg_fracs)):
            plt.plot(x_plot, bkg_func * norm_total_pdf * bkg_frac / self._ratio_truncated_,
                     color=self._bkg_cmap_(ibkg), ls='--', label=self.label_bkg_pdf[ibkg])
        # then draw signals
        for isgn, (signal_func, frac) in enumerate(zip(signal_funcs, signal_fracs)):
            plt.plot(x_plot, signal_func * norm_total_pdf * frac / self._ratio_truncated_, color=self._sgn_cmap_(isgn))
            plt.fill_between(x_plot, signal_func * norm_total_pdf * frac / self._ratio_truncated_,
                             color=self._sgn_cmap_(isgn),
                             alpha=0.5, label=self.label_signal_pdf[isgn])

        # finally draw reflected signals (if any)
        for irefl, (refl_func, frac) in enumerate(zip(refl_funcs, refl_fracs)):
            if self._name_refl_pdf_[irefl] is None:
                continue
            plt.plot(x_plot, refl_func * norm_total_pdf * frac / self._ratio_truncated_, color=self._refl_cmap_(irefl))
            plt.fill_between(x_plot, refl_func * norm_total_pdf * frac / self._ratio_truncated_,
                             color=self._refl_cmap_(irefl),
                             alpha=0.5, label=f'reflected signal {irefl}')

        plt.plot(x_plot, total_func * norm_total_pdf, color='xkcd:blue', label='total fit')
        plt.xlim(limits[0], limits[1])
        plt.xlabel(axis_title)
        plt.ylabel(rf'counts / {(limits[1]-limits[0])/bins*1000:0.1f} MeV/$c^2$')
        if logy:
            plt.yscale('log')
            plt.ylim(min(total_func) * norm_total_pdf / 5, max(total_func) * norm_total_pdf * 5)
        else:
            plt.ylim(0., max(total_func) * norm_total_pdf * 1.5)

        if show_extra_info:
            # info on chi2/ndf
            chi2 = self.get_chi2()
            ndf = self.get_ndf()
            anchored_text_chi2 = AnchoredText(fr'$\chi^2 / \mathrm{{ndf}} =${chi2:.2f} / {ndf}',
                                              loc=info_loc[0],
                                              frameon=False)
            # signal and background info for all signals
            text = []
            for idx, _ in enumerate(self._name_signal_pdf_):
                mass, mass_unc = self.get_mass(idx)
                sigma, sigma_unc = None, None
                gamma, gamma_unc = None, None
                rawyield, rawyield_err = self.get_raw_yield(idx=idx)
                if self._name_signal_pdf_[idx] in [
                    'gausexptail', 'genergausexptailsymm', 'gaussian', 'crystalball',
                    'doublecb', 'doublecbsymm', 'voigtian', 'hist'
                    ]:
                    sigma, sigma_unc = self.get_sigma(idx)
                if self._name_signal_pdf_[idx] in ['cauchy', 'voigtian']:
                    gamma, gamma_unc = self.get_signal_parameter(idx, 'gamma')
                extra_info = fr'{self.label_signal_pdf[idx]}''\n' \
                    + fr'  $\mu = {mass*1000:.1f}\pm{mass_unc*1000:.1f}$ MeV$/c^2$''\n'
                if sigma is not None:
                    extra_info += fr'  $\sigma = {sigma*1000:.1f}\pm{sigma_unc*1000:.1f}$ MeV$/c^2$''\n'
                if gamma is not None:
                    extra_info += fr'  $\Gamma/2 = {gamma*1000:.1f}\pm{gamma_unc*1000:.1f}$ MeV$/c^2$''\n'
                if mass_range is not None:
                    bkg, bkg_err = self.get_background(idx=idx, min=mass_range[0], max=mass_range[1])
                    s_over_b, s_over_b_err = self.get_signal_over_background(idx=idx, min=mass_range[0],
                                                                             max=mass_range[1])
                    signif, signif_err = self.get_significance(idx=idx, min=mass_range[0], max=mass_range[1])
                    interval = f'[{mass_range[0]:.3f}, {mass_range[1]:.3f}]'
                    extra_info += fr'  $S={rawyield:.0f} \pm {rawyield_err:.0f}$''\n'
                    extra_info += fr'  $B({interval})={bkg:.0f} \pm {bkg_err:.0f}$''\n'
                    extra_info += fr'  $S/B({interval})={s_over_b:.2f} \pm {s_over_b_err:.2f}$''\n'
                    extra_info += fr'  Signif.$({interval})={signif:.1f} \pm {signif_err:.1f}$'
                elif nhwhm is not None:
                    bkg, bkg_err = self.get_background(idx=idx, nhwhm=nhwhm)
                    s_over_b, s_over_b_err = self.get_signal_over_background(idx=idx, nhwhm=nhwhm)
                    signif, signif_err = self.get_significance(idx=idx, nhwhm=nhwhm)
                    extra_info += fr'  $S=${rawyield:.0f} $\pm$ {rawyield_err:.0f}''\n'
                    extra_info += fr'  $B({nhwhm}~\mathrm{{HWHM}})=${bkg:.0f} $\pm$ {bkg_err:.0f}''\n'
                    extra_info += fr'  $S/B({nhwhm}~\mathrm{{HWHM}})=${s_over_b:.2f} $\pm$ {s_over_b_err:.2f}''\n'
                    extra_info += fr'  Signif.$({nhwhm}~\mathrm{{HWHM}})=${signif:.1f} $\pm$ {signif_err:.1f}'
                else:
                    bkg, bkg_err = self.get_background(idx=idx, nsigma=nsigma)
                    s_over_b, s_over_b_err = self.get_signal_over_background(idx=idx, nsigma=nsigma)
                    signif, signif_err = self.get_significance(idx=idx, nsigma=nsigma)
                    extra_info += fr'  $S=${rawyield:.0f} $\pm$ {rawyield_err:.0f}''\n'
                    extra_info += fr'  $B({nsigma}\sigma)=${bkg:.0f} $\pm$ {bkg_err:.0f}''\n'
                    extra_info += fr'  $S/B({nsigma}\sigma)=${s_over_b:.2f} $\pm$ {s_over_b_err:.2f}''\n'
                    extra_info += fr'  Signif.$({nsigma}\sigma)=${signif:.1f} $\pm$ {signif_err:.1f}'
                text.append(extra_info)
            concatenated_text = '\n'.join(text)
            anchored_text_signal = AnchoredText(concatenated_text, loc=info_loc[1], frameon=False)

            axs.add_artist(anchored_text_chi2)
            axs.add_artist(anchored_text_signal)

        plt.legend(loc=legend_loc)

        Logger('plot_mass_fit now returns a tuple (fig, axs) !', 'WARNING')

        return fig, axs

    # pylint: disable=too-many-statements, too-many-locals
    def dump_to_root(self, filename, **kwargs):
        """
        Plot the mass fit

        Parameters
        -------------------------------------------------
        filename: str
            Name of output ROOT file

        **kwargs: dict
            Additional optional arguments:

            - axis_title: str
                x-axis title

            - num: int
                number of bins to plot pdfs converted into histograms

            - option: str
                option (recreate or update)


            - suffix: str
                suffix to append to objects
        """

        num = kwargs.get('num', 10000)
        suffix = kwargs.get('suffix', '')
        option = kwargs.get('option', 'recreate')
        bins = self._data_handler_.get_nbins()
        obs = self._data_handler_.get_obs()
        limits = self._data_handler_.get_limits()

        hdata = self._data_handler_.to_hist(lower_edge=limits[0],
                                            upper_edge=limits[1],
                                            nbins=bins,
                                            varname=self._data_handler_.get_var_name())
        # write data
        self.__write_data(hdata, f'hdata{suffix}', filename, option)

        bin_sigma = (limits[1] - limits[0]) / bins
        norm = self._total_pdf_norm_ * bin_sigma

        x_plot = np.linspace(limits[0], limits[1], num=num)

        total_func = zfit.run(self._total_pdf_.pdf(x_plot, norm_range=obs))
        # write total_func
        self.__write_pdf(histname=f'total_func{suffix}', weight=total_func * norm, num=num,
                         filename=filename, option='update')

        signal_funcs, bkg_funcs, refl_funcs = ([] for _ in range(3))
        for signal_pdf in self._signal_pdf_:
            signal_funcs.append(zfit.run(signal_pdf.pdf(x_plot, norm_range=obs)))
        for refl_pdf in self._refl_pdf_:
            refl_funcs.append(zfit.run(refl_pdf.pdf(x_plot, norm_range=obs)))
        for bkg_pdf in self._background_pdf_:
            bkg_funcs.append(zfit.run(bkg_pdf.pdf(x_plot, norm_range=obs)))

        signal_fracs, bkg_fracs, refl_fracs, _, _, _ = self.__get_all_fracs()

        # first write backgrounds
        for ibkg, (bkg_func, bkg_frac) in enumerate(zip(bkg_funcs, bkg_fracs)):
            self.__write_pdf(histname=f'bkg_{ibkg}{suffix}',
                             weight=bkg_func * norm * bkg_frac / self._ratio_truncated_,
                             num=num, filename=filename, option='update')
        # then write signals
        for isgn, (frac, signal_func) in enumerate(zip(signal_funcs, signal_fracs)):
            self.__write_pdf(histname=f'signal_{isgn}{suffix}',
                             weight=signal_func * norm * frac / self._ratio_truncated_,
                             num=num, filename=filename, option='update')

        # finally write reflected signals
        for irefl, (frac, refl_func) in enumerate(zip(refl_funcs, refl_fracs)):
            if self._name_refl_pdf_[irefl] is None:
                continue
            self.__write_pdf(histname=f'refl_{irefl}{suffix}',
                             weight=refl_func * norm * frac / self._ratio_truncated_,
                             num=num, filename=filename, option='update')

    @property
    def get_fit_result(self):
        """
        Get the fit result

        Returns
        -------------------------------------------------
        fit_result: zfit.minimizers.fitresult.FitResult
            The fit result
        """
        return self._fit_result_

    def get_ndf(self):
        """
        Get the number of degrees of freedom for chi2 fit
        ndf = nbins - nfreeparams - 1
        -1 because the data sample size is fixed

        Returns
        -------------------------------------------------
        ndf: int
            The number of degrees of freedom
        """
        nbins = self._data_handler_.get_nbins()
        nfreeparams = len(self._fit_result_.params)
        self._ndf_ = nbins - nfreeparams - 1
        return self._ndf_

    def get_chi2(self):
        """
        Get chi2 for binned data

        Returns
        -------------------------------------------------
        chi2: float
            chi2
        """

        chi2 = 0
        norm = self._data_handler_.get_norm()

        if self._data_handler_.get_is_binned():
            # for chi2 loss, just retrieve loss value in fit result
            if self._chi2_loss_:
                return float(self._fit_result_.loss.value())

            # for nll loss, compute chi2 "by hand"
            # access normalized data values and errors for all bins
            binned_data = self._data_handler_.get_binned_data()
            data_values = binned_data.values()
            data_variances = binned_data.variances()
            # access model predicted values
            model_values = self._total_pdf_binned_.values() * norm
            # compute chi2
            for (data, model, data_variance) in zip(data_values, model_values, data_variances):
                chi2 += (data - model)**2/data_variance
            return chi2

        if self._limits_ != [self._data_handler_.get_limits()]:  # restricted fit limits
            Logger('chi2 not yet implemented with restricted fit limits', 'FATAL')
        # for unbinned data
        data_values = self._data_handler_.get_binned_data_from_unbinned_data()
        # access model predicted values
        model_values = self._total_pdf_binned_.values() * norm
        # compute chi2
        for (data, model) in zip(data_values, model_values):
            if data == 0:
                continue
            chi2 += (data - model)**2/data

        return float(chi2)

    def get_chi2_ndf(self):
        """
        Get the reduced chi2 (chi2 divided by number of degrees of freedom)
        for binned data

        Returns
        -------------------------------------------------
        chi2_ndf: float
            The reduced chi2
        """
        return self.get_chi2()/self.get_ndf()

    def plot_raw_residuals(self, **kwargs):
        """
        Plot the raw residuals

        Parameters
        -------------------------------------------------
        **kwargs: dict
            Additional optional arguments:

            - style: str
                style to be used (see https://github.com/scikit-hep/mplhep for more details)

            - figsize: tuple
                size of the figure

            - axis_title: str
                x-axis title

        Returns
        -------------------------------------------------
        fig: matplotlib.figure.Figure
            figure containing the raw residuals plot
        """

        style = kwargs.get('style', 'LHCb2')
        figsize = kwargs.get('figsize', (7, 7))
        axis_title = kwargs.get('axis_title', self._data_handler_.get_var_name())

        mplhep.style.use(style)

        obs = self._data_handler_.get_obs()
        limits = self._data_handler_.get_limits()

        fig = plt.figure(figsize=figsize)

        if len(self._raw_residuals_) == 0:
            self.__get_raw_residuals()

        # draw residuals
        plt.errorbar(
            self._data_handler_.get_bin_center(),
            self._raw_residuals_,
            xerr=None,
            yerr=np.sqrt(self._raw_residual_variances_),
            linestyle="None",
            elinewidth=1,
            capsize=0,
            color="black",
            marker="o",
            markersize=5,
            label="residuals"
        )
        bins = self._data_handler_.get_nbins()
        bin_sigma = (limits[1] - limits[0]) / bins
        norm = self._total_pdf_norm_ * bin_sigma

        x_plot = np.linspace(limits[0], limits[1], num=1000)
        signal_funcs, refl_funcs = ([] for _ in range(2))
        for signal_pdf in self._signal_pdf_:
            signal_funcs.append(zfit.run(signal_pdf.pdf(x_plot, norm_range=obs)))
        for refl_pdf in self._refl_pdf_:
            refl_funcs.append(zfit.run(refl_pdf.pdf(x_plot, norm_range=obs)))

        signal_fracs, _, refl_fracs, _, _, _ = self.__get_all_fracs()

        # draw signals
        for isgn, (signal_func, frac) in enumerate(zip(signal_funcs, signal_fracs)):
            plt.plot(x_plot, signal_func * norm * frac / self._ratio_truncated_, color=self._sgn_cmap_(isgn))
            plt.fill_between(x_plot, signal_func * norm * frac / self._ratio_truncated_, color=self._sgn_cmap_(isgn),
                             alpha=0.5, label=self.label_signal_pdf[isgn])

        # finally draw reflected signals (if any)
        is_there_refl = False
        for irefl, (refl_func, frac) in enumerate(zip(refl_funcs, refl_fracs)):
            if self._name_refl_pdf_[irefl] is None:
                continue
            is_there_refl = True
            plt.plot(x_plot, refl_func * norm * frac / self._ratio_truncated_, color=self._refl_cmap_(irefl))
            plt.fill_between(x_plot, refl_func * norm * frac / self._ratio_truncated_, color=self._refl_cmap_(irefl),
                             alpha=0.5, label=f'reflected signal {irefl}')

        # draw signal + reflected signals (if any)
        if is_there_refl:
            for isgn, (signal_func, refl_func, frac_sgn, frac_refl) in enumerate(
                    zip(signal_funcs, refl_funcs, signal_fracs, refl_fracs)):
                plt.plot(x_plot, (signal_func * frac_sgn + frac_refl * refl_func) * norm / self._ratio_truncated_,
                         color='xkcd:blue', label='total - bkg')

        plt.xlim(limits[0], limits[1])
        plt.xlabel(axis_title)
        plt.ylabel(rf'(data - fitted bkg) / {(limits[1]-limits[0])/bins*1000:0.1f} MeV/$c^2$')
        plt.legend(loc='best')

        return fig

    def plot_std_residuals(self, **kwargs):
        """
        Plot the raw residuals

        Parameters
        -------------------------------------------------
        **kwargs: dict
            Additional optional arguments:

            - style: str
                style to be used (see https://github.com/scikit-hep/mplhep for more details)

            - figsize: tuple
                size of the figure

            - axis_title: str
                x-axis title

        Returns
        -------------------------------------------------
        fig: matplotlib.figure.Figure
            figure containing the raw residuals plot
        """

        style = kwargs.get('style', 'LHCb2')
        figsize = kwargs.get('figsize', (7, 7))
        axis_title = kwargs.get('axis_title', self._data_handler_.get_var_name())

        mplhep.style.use(style)

        limits = self._data_handler_.get_limits()
        bins = self._data_handler_.get_nbins()
        bin_center = self._data_handler_.get_bin_center()

        fig = plt.figure(figsize=figsize)

        if len(self._std_residuals_) == 0:
            self.__get_std_residuals()
        # draw residuals
        plt.errorbar(bin_center,
                     self._std_residuals_,
                     xerr=None,
                     yerr=np.sqrt(self._std_residual_variances_),
                     linestyle="None",
                     elinewidth=1,
                     capsize=0,
                     color="black",
                     marker="o",
                     markersize=5,
                     label=None)

        # line at 0
        plt.plot([bin_center[0], bin_center[-1]], [0., 0.], lw=2, color='xkcd:blue')
        # line at -3 sigma
        plt.plot([bin_center[0], bin_center[-1]], [-3., -3.], lw=2, color='xkcd:red')
        # line at 3 sigma
        plt.plot([bin_center[0], bin_center[-1]], [3., 3.], lw=2, color='xkcd:red')

        plt.xlim(limits[0], limits[1])
        plt.xlabel(axis_title)
        plt.ylabel(fr"$\dfrac{{ \mathrm{{data}} - \mathrm{{total \ fit}} }}{{ \sigma_{{ \mathrm{{data}} }} }}$"
                   fr"/ {(limits[1]-limits[0])/bins*1000:0.1f} MeV/$c^2$")

        return fig

    def get_raw_yield(self, idx=0):
        """
        Get the raw yield and its error

        Parameters
        -------------------------------------------------
        idx: int
            Index of the raw yield to be returned (default: 0)

        Returns
        -------------------------------------------------
        raw_yield: float
            The raw yield obtained from the fit
        raw_yield_err: float
            The raw yield error obtained from the fit
        """
        return self._rawyield_[idx], self._rawyield_err_[idx]

    def get_raw_yield_bincounting(self, idx=0, **kwargs):
        """
        Get the raw yield and its error via the bin-counting method

        Parameters
        -------------------------------------------------
        idx: int
            Index of the raw yield to be returned (default: 0)

        **kwargs: dict
            Additional optional arguments:

            - nsigma: float
                nsigma invariant-mass window around mean for signal counting

            - nhwhm: float
                number of hwhm invariant-mass window around mean for signal counting
                (alternative to nsigma)

            - min: float
                minimum value of invariant-mass for signal counting (alternative to nsigma)

            - max: float
                maximum value of invariant-mass for signal counting (alternative to nsigma)

        Returns
        -------------------------------------------------
        raw_yield: float
            The raw yield obtained from the bin counting
        raw_yield_err: float
            The raw yield error obtained from the bin counting
        """

        nsigma = kwargs.get('nsigma', 3.)
        nhwhm = kwargs.get('nhwhm', None)
        min_value = kwargs.get('min', None)
        max_value = kwargs.get('max', None)
        use_nsigma = True

        if nhwhm is not None and (min_value is not None or max_value is not None):
            Logger('I cannot compute the signal within a fixed mass interval and a number of HWFM', 'ERROR')
            return 0., 0.

        if min_value is not None and max_value is not None:
            use_nsigma = False

        if nhwhm is not None:
            use_nsigma = False
            if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
                Logger('HWHM not defined, I cannot compute the signal for this pdf', 'ERROR')
                return 0., 0.
            mass, _ = self.get_mass(idx)
            hwhm, _ = self.get_hwhm(idx)
            min_value = mass - nhwhm * hwhm
            max_value = mass + nhwhm * hwhm

        if use_nsigma:
            if self._name_signal_pdf_[idx] not in [
                'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
                'doublecb', 'doublecbsymm', 'voigtian', 'hist'
                ]:
                Logger('Sigma not defined, I cannot compute the signal for this pdf', 'ERROR')
                return 0., 0.
            mass, _ = self.get_mass(idx)
            sigma, _ = self.get_sigma(idx)
            min_value = mass - nsigma * sigma
            max_value = mass + nsigma * sigma

        if len(self._raw_residuals_) == 0:
            self.__get_raw_residuals()

        bin_centers = self._data_handler_.get_bin_center()
        bin_width = (bin_centers[1] - bin_centers[0]) / 2
        raw_yield, raw_yield_err = 0., 0.
        for residual, variance, bin_center in zip(self._raw_residuals_,
                                                  self._raw_residual_variances_, bin_centers):
            if bin_center - bin_width >= min_value and bin_center + bin_width <= max_value:
                raw_yield += residual
                raw_yield_err += variance
        raw_yield = float(raw_yield)
        raw_yield_err = np.sqrt(float(raw_yield_err))

        return raw_yield, raw_yield_err

    def get_mass(self, idx=0):
        """
        Get the mass and its error

        Parameters
        -------------------------------------------------
        idx: int
            Index of the mass to be returned (default: 0)

        Returns
        -------------------------------------------------
        mass: float
            The mass value obtained from the fit
        mass_err: float
            The mass error obtained from the fit
        """
        if 'hist' in self._name_signal_pdf_[idx]:
            hist = self._signal_pdf_[idx].to_hist()
            bin_limits = hist.to_numpy()[1]
            centres = [0.5 * (minn + maxx) for minn, maxx in zip(bin_limits[1:],  bin_limits[:-1])]
            counts = hist.values()
            mass = np.average(centres, weights=counts)
            mass_err = 0.
        else:
            mass_name = 'm' if self._name_signal_pdf_[idx] in ['cauchy', 'voigtian'] else 'mu'
            if self._fix_sgn_pars_[idx][mass_name]:
                mass = self._init_sgn_pars_[idx][mass_name]
                mass_err = 0.
            else:
                mass = self._fit_result_.params[f'{self._name_}_{mass_name}_signal{idx}']['value']
                mass_err = self._fit_result_.params[f'{self._name_}_{mass_name}_signal{idx}']['hesse']['error']

        return mass, mass_err

    def get_sigma(self, idx=0):
        """
        Get the sigma and its error

        Parameters
        -------------------------------------------------
        idx: int
            Index of the sigma to be returned (default: 0)

        Returns
        -------------------------------------------------
        sigma: float
            The sigma value obtained from the fit
        sigma_err: float
            The sigma error obtained from the fit
        """
        if self._name_signal_pdf_[idx] not in [
            'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
            'doublecb', 'doublecbsymm', 'voigtian', 'hist'
            ]:
            Logger(f'Sigma parameter not defined for {self._name_signal_pdf_[idx]} pdf!', 'ERROR')
            return 0., 0.

        # if histogram, the rms is used as proxy
        if 'hist' in self._name_signal_pdf_[idx]:
            Logger(f'RMS used as proxy for sigma parameter of {self._name_signal_pdf_[idx]} pdf!', 'WARNING')
            mean = self.get_mass(idx)[0]
            hist = self._signal_pdf_[idx].to_hist()
            bin_limits = hist.to_numpy()[1]
            centres = [0.5 * (minn + maxx) for minn, maxx in zip(bin_limits[1:],  bin_limits[:-1])]
            counts = hist.values()
            sigma = np.sqrt(np.average((centres - mean)**2, weights=counts))
            sigma_err = 0.
        else:
            if self._fix_sgn_pars_[idx]['sigma']:
                sigma = self._init_sgn_pars_[idx]['sigma']
                sigma_err = 0.
            else:
                sigma = self._fit_result_.params[f'{self._name_}_sigma_signal{idx}']['value']
                sigma_err = self._fit_result_.params[f'{self._name_}_sigma_signal{idx}']['hesse']['error']

        return sigma, sigma_err

    def get_hwhm(self, idx=0):
        """
        Get the half width half maximum and its error

        Parameters
        -------------------------------------------------
        idx: int
            Index of the sigma to be returned (default: 0)

        Returns
        -------------------------------------------------
        hwhm: float
            The sigma value obtained from the fit
        hwhm_err: float
            The sigma error obtained from the fit
        """
        if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
            Logger(f'HFWM parameter not defined for {self._name_signal_pdf_[idx]} pdf!', 'ERROR')
            return 0., 0.

        if self._name_signal_pdf_[idx] == 'gaussian':
            mult_fact = np.sqrt(2 * np.log(2))
            hwhm, hwhm_err = self.get_sigma(idx)
            hwhm *= mult_fact
            hwhm_err *= mult_fact
        elif self._name_signal_pdf_[idx] == 'cauchy':
            hwhm, hwhm_err = self.get_signal_parameter(idx, 'gamma')
        elif self._name_signal_pdf_[idx] == 'voigtian':
            mult_fact = np.sqrt(2 * np.log(2))
            sigma, sigma_err = self.get_sigma(idx)
            sigma *= mult_fact
            sigma_err *= mult_fact
            gamma, gamma_err = self.get_signal_parameter(idx, 'gamma')
            hwhm = 0.5346 * gamma + np.sqrt(0.2166 * gamma**2 + sigma**2)
            # we neglect the correlation between sigma and gamma
            der_sigma = sigma / np.sqrt(0.0721663 + sigma**2)
            der_gamma = 0.5346 + (0.2166 * gamma) / np.sqrt(0.2166 * gamma**2 + sigma**2)
            hwhm_err = np.sqrt(der_sigma**2 * sigma_err**2 + der_gamma**2 * gamma_err**2)

        return hwhm, hwhm_err

    def get_signal_parameter(self, idx, par_name):
        """
        Get a signal parameter and its error

        Parameters
        -------------------------------------------------
        idx: int
            Index of the parameter to be returned (default: 0)
        par_name: str
            parameter to return

        Returns
        -------------------------------------------------
        parameter: float
            The parameter value obtained from the fit
        parameter_err: float
            The parameter error obtained from the fit

        """

        if par_name == 'gamma':
            Logger('The gamma parameter that you are getting is half of the width of a resonance,'
                   ' for more info check the Cauchy pdf defined here '
                   'https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/basic/zfit.pdf.Cauchy.html',
                   'WARNING')

        if self._fix_sgn_pars_[idx][par_name]:
            parameter = self._init_sgn_pars_[idx][par_name]
            parameter_err = 0.
        else:
            parameter = self._fit_result_.params[f'{self._name_}_{par_name}_signal{idx}']['value']
            parameter_err = self._fit_result_.params[f'{self._name_}_{par_name}_signal{idx}']['hesse']['error']

        return parameter, parameter_err

    def get_background_parameter(self, idx, par_name):
        """
        Get a background parameter and its error

        Parameters
        -------------------------------------------------
        idx: int
            Index of the parameter to be returned (default: 0)
        par_name: str
            parameter to return

        Returns
        -------------------------------------------------
        parameter: float
            The parameter value obtained from the fit
        parameter_err: float
            The parameter error obtained from the fit

        """

        if self._fix_bkg_pars_[idx][par_name]:
            parameter = self._init_bkg_pars_[idx][par_name]
            parameter_err = 0.
        else:
            parameter = self._fit_result_.params[f'{self._name_}_{par_name}_bkg{idx}']['value']
            parameter_err = self._fit_result_.params[f'{self._name_}_{par_name}_bkg{idx}']['hesse']['error']

        return parameter, parameter_err

    def get_signal(self, idx=0, **kwargs):
        """
        Get the signal and its error in a given invariant-mass region

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal to be returned
        **kwargs: dict
            Additional optional arguments:

            - nsigma: float
                nsigma invariant-mass window around mean for signal computation

            - nhwhm: float
                number of hwhm invariant-mass window around mean for signal and background computation
                (alternative to nsigma)

            - min: float
                minimum value of invariant-mass for signal computation (alternative to nsigma)

            - max: float
                maximum value of invariant-mass for signal computation (alternative to nsigma)

        Returns
        -------------------------------------------------
        signal: float
            The signal value obtained from the fit
        signal_err: float
            The signal error obtained from the fit
        """

        nsigma = kwargs.get('nsigma', 3.)
        nhwhm = kwargs.get('nhwhm', None)
        min_value = kwargs.get('min', None)
        max_value = kwargs.get('max', None)
        use_nsigma = True

        if nhwhm is not None and (min_value is not None or max_value is not None):
            Logger('I cannot compute the signal within a fixed mass interval and a number of HWFM', 'ERROR')
            return 0., 0.

        if min_value is not None and max_value is not None:
            use_nsigma = False

        if nhwhm is not None:
            use_nsigma = False
            if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
                Logger('HWHM not defined, I cannot compute the signal for this pdf', 'ERROR')
                return 0., 0.
            mass, _ = self.get_mass(idx)
            hwhm, _ = self.get_hwhm(idx)
            min_value = mass - nhwhm * hwhm
            max_value = mass + nhwhm * hwhm

        if use_nsigma:
            if self._name_signal_pdf_[idx] not in [
                'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
                'doublecb', 'doublecbsymm', 'voigtian', 'hist'
                ]:
                Logger('Sigma not defined, I cannot compute the signal for this pdf', 'ERROR')
                return 0., 0.
            mass, _ = self.get_mass(idx)
            sigma, _ = self.get_sigma(idx)
            min_value = mass - nsigma * sigma
            max_value = mass + nsigma * sigma

        # pylint: disable=missing-kwoa
        signal = self._signal_pdf_[idx].integrate((min_value, max_value))

        signal_fracs, _, refl_fracs, signal_err_fracs, _, _ = self.__get_all_fracs()

        if len(self._background_pdf_) > 0:
            frac = signal_fracs[idx]
            frac_err = signal_err_fracs[idx]
        else:
            if len(self._signal_pdf_) == 1:
                frac = 1.
                frac_err = 0.
            if idx < len(signal_fracs):
                frac = signal_fracs[idx]
                frac_err = signal_err_fracs[idx]
            else:
                frac = 1. - sum(signal_fracs) - sum(refl_fracs)
                frac_err = np.sqrt(sum(list(err**2 for err in signal_err_fracs)))

        norm = self._data_handler_.get_norm()
        norm_err = norm * frac_err
        norm *= frac

        return float(signal * norm), float(signal * norm_err)

    def get_background(self, idx=0, **kwargs):
        """
        Get the background and its error in a given invariant-mass region

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal to be used to compute nsigma window
        **kwargs: dict
            Additional optional arguments:

            - nsigma: float
                nsigma invariant-mass window around mean for background computation

            - nhwhm: float
                number of hwhm invariant-mass window around mean for signal and background computation
                (alternative to nsigma)

            - min: float
                minimum value of invariant-mass for background computation (alternative to nsigma)

            - max: float
                maximum value of invariant-mass for background computation (alternative to nsigma)

        Returns
        -------------------------------------------------
        background: float
            The background value obtained from the fit
        background_err: float
            The background error obtained from the fit
        """

        if not self._background_pdf_:
            Logger('Background not fitted', 'ERROR')
            return 0., 0.

        nsigma = kwargs.get('nsigma', 3.)
        nhwhm = kwargs.get('nhwhm', None)
        min_value = kwargs.get('min', None)
        max_value = kwargs.get('max', None)
        use_nsigma = True

        if nhwhm is not None and (min_value is not None or max_value is not None):
            Logger('I cannot compute the signal within a fixed mass interval and a number of HWFM', 'ERROR')
            return 0., 0.

        if min_value is not None and max_value is not None:
            use_nsigma = False

        if nhwhm is not None:
            use_nsigma = False
            if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
                Logger('HWHM not defined, I cannot compute the signal for this pdf', 'ERROR')
                return 0., 0.
            mass, _ = self.get_mass(idx)
            hwhm, _ = self.get_hwhm(idx)
            min_value = mass - nhwhm * hwhm
            max_value = mass + nhwhm * hwhm

        if use_nsigma:
            if self._name_signal_pdf_[idx] not in [
                'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
                'doublecb', 'doublecbsymm', 'voigtian', 'hist'
                ]:
                Logger('Sigma not defined, I cannot compute the signal for this pdf', 'ERROR')
                return 0., 0.
            mass, _ = self.get_mass(idx)
            sigma, _ = self.get_sigma(idx)
            min_value = mass - nsigma * sigma
            max_value = mass + nsigma * sigma

        _, bkg_fracs, _, _, bkg_err_fracs, _ = self.__get_all_fracs()

        limits = self._data_handler_.get_limits()
        min_value = max(min_value, limits[0])
        max_value = min(max_value, limits[1])

        # pylint: disable=missing-kwoa
        background, background_err = 0., 0.
        for idx2, bkg in enumerate(self._background_pdf_):

            norm = self._total_pdf_norm_ * bkg_fracs[idx2]
            norm_err = norm * bkg_err_fracs[idx2]

            bkg_int = float(bkg.integrate((min_value, max_value)))
            background += bkg_int * norm
            background_err += (bkg_int * norm_err)**2

        background /= self._ratio_truncated_
        background_err = np.sqrt(background_err) / self._ratio_truncated_

        return float(background), float(background_err)

    def get_signal_over_background(self, idx=0, **kwargs):
        """
        Get the S/B ratio and its error in a given invariant-mass region

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal to be used to compute nsigma window
        **kwargs: dict
            Additional optional arguments:

            - nsigma: float
                nsigma invariant-mass window around mean for signal and background computation

            - nhwhm: float
                number of hwhm invariant-mass window around mean for signal and background computation
                (alternative to nsigma)

            - min: float
                minimum value of invariant-mass for signal and background computation (alternative to nsigma)

            - max: float
                maximum value of invariant-mass for signal and background computation (alternative to nsigma)

        Returns
        -------------------------------------------------
        signal_over_background: float
            The S/B value obtained from the fit
        signal_over_background_err: float
            The S/B error obtained from the fit
        """

        signal = self.get_signal(idx, **kwargs)
        bkg = self.get_background(idx, **kwargs)
        signal_over_background = signal[0]/bkg[0]
        signal_over_background_err = np.sqrt(signal[1]**2/signal[0]**2 + bkg[1]**2/bkg[0]**2)
        signal_over_background_err *= signal_over_background

        return signal_over_background, signal_over_background_err

    def get_significance(self, idx=0, **kwargs):
        """
        Get the significance and its error in a given invariant-mass region

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal to be used to compute nsigma window
        **kwargs: dict
            Additional optional arguments:

            - nsigma: float
                nsigma invariant-mass window around mean for signal and background computation

            - nhwhm: float
                number of hwhm invariant-mass window around mean for signal and background computation
                (alternative to nsigma)

            - min: float
                minimum value of invariant-mass for signal and background computation (alternative to nsigma)

            - max: float
                maximum value of invariant-mass for signal and background computation (alternative to nsigma)

        Returns
        -------------------------------------------------
        significance: float
            The significance value obtained from the fit
        significance_err: float
            The significance error obtained from the fit
        """

        signal = self.get_signal(idx, **kwargs)
        bkg = self.get_background(idx, **kwargs)
        significance = signal[0]/np.sqrt(signal[0]+bkg[0])
        sig_plus_bkg = signal[0] + bkg[0]

        significance_err = significance*np.sqrt(
            (signal[1]**2 + bkg[1]**2) / (4. * sig_plus_bkg**2) + (
                bkg[0]/sig_plus_bkg) * signal[1]**2 / signal[0]**2)

        return significance, significance_err

    def get_sweights(self):
        """
        Calculate sWeights for signal and background components.

        Returns:
            dict: A dictionary containing sWeights for 'signal' and 'bkg' components.
        """

        signal_pdf_extended = []
        refl_pdf_extended = []
        bkg_pdf_extended = []

        norm = self._data_handler_.get_norm()

        signal_fracs, bkg_fracs, refl_fracs, _, _, _ = self.__get_all_fracs()
        bkg_fracs.append(1-sum(bkg_fracs)-sum(signal_fracs)-sum(refl_fracs))

        total_signal_yields = 0.
        total_bkg_yields = 0.
        for pdf, frac in zip(self._signal_pdf_, signal_fracs):
            signal_pdf_extended.append(pdf.create_extended(frac * norm))
            total_signal_yields += frac * norm
        for pdf, frac in zip(self._refl_pdf_, refl_fracs):
            refl_pdf_extended.append(pdf.create_extended(frac * norm))
            total_signal_yields += frac * norm
        for pdf, frac in zip(self._background_pdf_, bkg_fracs):
            bkg_pdf_extended.append(pdf.create_extended(frac * norm))
            total_bkg_yields += frac * norm

        total_signal_yields_par = zfit.Parameter('signal_yield', total_signal_yields)
        total_bkg_yields_par = zfit.Parameter('bkg_yield', total_bkg_yields)

        if len(signal_pdf_extended + refl_pdf_extended) > 1:
            signal_pdf_for_sweights = zfit.pdf.SumPDF(
                signal_pdf_extended + refl_pdf_extended,
                extended=total_signal_yields_par
            )
        else:
            signal_pdf_for_sweights = signal_pdf_extended[0]

        if len(bkg_pdf_extended) > 1:
            bkg_pdf_for_sweights = zfit.pdf.SumPDF(
                bkg_pdf_extended,
                extended=total_bkg_yields_par
            )
        else:
            bkg_pdf_for_sweights = bkg_pdf_extended[0]

        total_pdf_for_sweights = zfit.pdf.SumPDF([signal_pdf_for_sweights, bkg_pdf_for_sweights])
        sweights = compute_sweights(total_pdf_for_sweights, self._data_handler_.get_data())
        sweights_labels = list(sweights.keys())

        return {'signal': sweights[sweights_labels[0]], 'bkg': sweights[sweights_labels[1]]}

    def set_particle_mass(self, idx, **kwargs):
        """
        Set the particle mass

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal
        **kwargs: dict
            Additional optional arguments:

            - mass: float
                The mass of the particle

            - pdg_id: int
                PDG ID of the particle (alternative to mass)

            - pdg_name: str
                Name of the particle (alternative to mass)

            - limits: list
                minimum and maximum limits for the mass parameter

            - fix: bool
                fix the mass parameter
        """
        mass_name = 'm' if self._name_signal_pdf_[idx] in ['cauchy', 'voigtian'] else 'mu'
        mass = 0.
        if 'mass' in kwargs:
            mass = kwargs['mass']
        elif 'pdg_id' in kwargs:
            mass = self.pdg_api.get_particle_by_mcid(kwargs['pdg_id']).mass
        elif 'pdg_name' in kwargs:
            mass = self.pdg_api.get_particle_by_name(kwargs['pdg_name']).mass
        else:
            Logger(f'"mass", "pdg_id", and "pdg_name" not provided, mass value for signal {idx} will not be set',
                   'ERROR')
        self._init_sgn_pars_[idx][mass_name] = mass
        if 'limits' in kwargs:
            self._limits_sgn_pars_[idx][mass_name] = kwargs['limits']
        if 'fix' in kwargs:
            self._fix_sgn_pars_[idx][mass_name] = kwargs['fix']

    def set_signal_initpar(self, idx, par_name, init_value, **kwargs):
        """
        Set a signal parameter

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal
        par_name: str
            The name of the parameter to be set
        init_value: float
            The value of parameter to be set
        **kwargs: dict
            Additional optional arguments:

            - limits: list
                minimum and maximum limits for the parameter

            - fix: bool
                fix the parameter to init_value
        """
        if par_name == 'gamma':
            Logger('The gamma parameter that you are setting is half of the width of a resonance,'
                   ' for more info check the Cauchy pdf defined here '
                   'https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/basic/zfit.pdf.Cauchy.html',
                   'WARNING')

        self._init_sgn_pars_[idx][par_name] = init_value
        if 'limits' in kwargs:
            self._limits_sgn_pars_[idx][par_name] = kwargs['limits']
        if 'fix' in kwargs:
            self._fix_sgn_pars_[idx][par_name] = kwargs['fix']

    def set_background_initpar(self, idx, par_name, init_value, **kwargs):
        """
        Set a background parameter

        Parameters
        -------------------------------------------------
        idx: int
            Index of the background
        par_name: str
            The name of the parameter to be set
        init_value: float
            The value of parameter to be set
        **kwargs: dict
            Additional optional arguments:

            - limits: list
                minimum and maximum limits for the parameter

            - fix: bool
                fix the mass parameter
        """
        self._init_bkg_pars_[idx][par_name] = init_value
        if 'limits' in kwargs:
            self._limits_bkg_pars_[idx][par_name] = kwargs['limits']
        if 'fix' in kwargs:
            self._fix_bkg_pars_[idx][par_name] = kwargs['fix']

    def __get_parameter_fix_frac(self, factor, name):
        """
        Get the parameter fix for the fraction

        Parameters
        -------------------------------------------------
        idx_pdf: int
            Index of the pdf
        target_pdf: int
            Index of the target pdf
        factor: float
            Factor to multiply the fraction parameter
        """
        return zfit.Parameter(name, factor, floating=False)

    def __check_consistency_fix_frac(self, idx_pdf, target_pdf, idx_pdf_type, target_pdf_type):
        """
        Checks the consistency of fixing the fraction between PDFs.

        Parameters
        -------------------------------------------------
        idx_pdf: int
            The index of the PDF to check.
        target_pdf: int
            The target PDF to compare against.
        idx_pdf_type: 'signal' or 'bkg'
            The type of the PDF to check.
        target_pdf_type: 'signal' or 'bkg'
            The type of the target PDF to compare against.
        """
        if idx_pdf_type == target_pdf_type and idx_pdf == target_pdf:
            Logger(
                f'Index {idx_pdf} is the same as {target_pdf},'
                'cannot constrain the fraction to itself',
                'FATAL'
            )
        if target_pdf_type == 'signal' and target_pdf > len(self._signal_pdf_):
            Logger(
                f'Target signal index {target_pdf} is out of range',
                'FATAL'
            )
        if target_pdf_type == 'bkg' and target_pdf > len(self._background_pdf_):
            Logger(
                f'Target background index {target_pdf} is out of range',
                'FATAL'
            )

    def fix_signal_frac_to_signal_pdf(self, idx_pdf, target_pdf, factor=1):
        """
        Fix the frac parameter of a signal to the frac parameter of another signal

        Parameters
        -------------------------------------------------
        idx_pdf: int
            Index of the signal fraction to be fixed
        target_pdf: int
            Index of the signal fraction to be used as reference
        factor: float
            Factor to multiply the frac parameter of the target signal
        """
        self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'signal', 'signal')
        factor_par = self.__get_parameter_fix_frac(
            factor, f'factor_signal{idx_pdf}_constrained_to_signal{target_pdf}'
        )
        self._fix_fracs_to_pdfs_.append({
            'fixed_pdf_idx': idx_pdf,
            'target_pdf_idx': target_pdf,
            'factor': factor_par,
            'fixed_pdf_type': 'signal',
            'target_pdf_type': 'signal'
        })

    def fix_signal_frac_to_bkg_pdf(self, idx_pdf, target_pdf, factor=1):
        """
        Fix the frac parameter of a signal to the frac parameter of a background

        Parameters
        -------------------------------------------------
        idx_pdf: int
            Index of the signal fraction to be fixed
        target_pdf: int
            Index of the background fraction to be used as reference
        factor: float
            Factor to multiply the frac parameter of the target background
        """
        self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'signal', 'bkg')
        factor_par = self.__get_parameter_fix_frac(
            factor, f'factor_signal{idx_pdf}_constrained_to_bkg{target_pdf}'
        )
        self._fix_fracs_to_pdfs_.append({
            'fixed_pdf_idx': idx_pdf,
            'target_pdf_idx': target_pdf,
            'factor': factor_par,
            'fixed_pdf_type': 'signal',
            'target_pdf_type': 'bkg'
        })

    def fix_bkg_frac_to_signal_pdf(self, idx_pdf, target_pdf, factor=1):
        """
        Fix the frac parameter of a background to the frac parameter of a signal

        Parameters
        -------------------------------------------------
        idx_pdf: int
            Index of the background fraction to be fixed
        target_pdf: int
            Index of the signal fraction to be used as reference
        factor: float
            Factor to multiply the frac parameter of the target signal
        """
        self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'bkg', 'signal')
        factor_par = self.__get_parameter_fix_frac(
            factor, f'factor_bkg{idx_pdf}_constrained_to_signal{target_pdf}'
        )
        self._fix_fracs_to_pdfs_.append({
            'fixed_pdf_idx': idx_pdf,
            'target_pdf_idx': target_pdf,
            'factor': factor_par,
            'fixed_pdf_type': 'bkg',
            'target_pdf_type': 'signal'
        })

    def fix_bkg_frac_to_bkg_pdf(self, idx_pdf, target_pdf, factor=1):
        """
        Fix the frac parameter of a background to the frac parameter of another background

        Parameters
        -------------------------------------------------
        idx_pdf: int
            Index of the background fraction to be fixed
        target_pdf: int
            Index of the background fraction to be used as reference
        factor: float
            Factor to multiply the frac parameter of the target background
        """
        self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'bkg', 'bkg')
        factor_par = self.__get_parameter_fix_frac(
            factor, f'factor_bkg{idx_pdf}_constrained_to_bkg{target_pdf}'
        )
        self._fix_fracs_to_pdfs_.append({
            'fixed_pdf_idx': idx_pdf,
            'target_pdf_idx': target_pdf,
            'factor': factor_par,
            'fixed_pdf_type': 'bkg',
            'target_pdf_type': 'bkg'
        })

    # pylint: disable=line-too-long
    def set_signal_template(self, idx, sample):
        """
        Set sample and options for signal template

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal
        sample: flarefly.DataHandler
            Data sample for histogram template
        """

        self._hist_signal_sample_[idx] = sample

    # pylint: disable=line-too-long
    def set_signal_kde(self, idx, sample, **kwargs):
        """
        Set sample and options for signal kde

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal
        sample: flarefly.DataHandler
            Data sample for Kernel Density Estimation
        **kwargs: dict
            Arguments for kde options. See
            https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid
            for more details
        """

        self._kde_signal_sample_[idx] = sample
        self._kde_signal_option_[idx] = kwargs

    # pylint: disable=line-too-long
    def set_reflection_template(self, idx, sample, r_over_s):
        """
        Set sample and options for reflected signal template

        Parameters
        -------------------------------------------------
        idx: int
            Index of the reflected signal
        sample: flarefly.DataHandler
            Data sample for histogram template
        r_over_s: float
            R/S ratio
        """

        self._hist_refl_sample_[idx] = sample
        self._refl_over_sgn_[idx] = r_over_s

    # pylint: disable=line-too-long
    def set_reflection_kde(self, idx, sample, r_over_s, **kwargs):
        """
        Set sample and options for reflected signal kde

        Parameters
        -------------------------------------------------
        idx: int
            Index of the signal
        sample: flarefly.DataHandler
            Data sample for Kernel Density Estimation
        r_over_s: float
            R/S ratio
        **kwargs: dict
            Arguments for kde options. See
            https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid
            for more details
        """

        self._kde_signal_sample_[idx] = sample
        self._kde_signal_option_[idx] = kwargs
        self._refl_over_sgn_[idx] = r_over_s

    # pylint: disable=line-too-long
    def set_background_template(self, idx, sample):
        """
        Set sample and options for background template histogram

        Parameters
        -------------------------------------------------
        idx: int
            Index of the background
        sample: flarefly.DataHandler
            Data sample for template histogram
        """

        limits_bkg = sample.get_limits()
        limits_data = self._data_handler_.get_limits()
        if not np.isclose(limits_bkg[0], limits_data[0], rtol=1e-05, atol=1e-08, equal_nan=False) or \
                not np.isclose(limits_bkg[1], limits_data[1], rtol=1e-05, atol=1e-08, equal_nan=False):
            Logger(f'The data and the background template {idx} have different limits:'
                   f' \n       -> background template: {limits_bkg}, data -> {limits_data}', 'FATAL')

        self._hist_bkg_sample_[idx] = sample

    # pylint: disable=line-too-long
    def set_background_kde(self, idx, sample, **kwargs):
        """
        Set sample and options for background kde

        Parameters
        -------------------------------------------------
        idx: int
            Index of the background
        sample: flarefly.DataHandler
            Data sample for Kernel Density Estimation
        **kwargs: dict
            Arguments for kde options. See
            https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid
            for more details
        """

        self._kde_bkg_sample_[idx] = sample
        self._kde_bkg_option_[idx] = kwargs

    def __write_data(self, hdata, histname='hdata', filename='output.root', option='recreate'):
        """
        Helper method to save a data histogram in a .root file (TH1D format)

        Parameters
        -------------------------------------------------
        hdata: hist
            Histogram containing the data

        histname: str
            Name of the histogram

        filename: str
            Name of the ROOT file

        option: str
            Option (recreate or update)
        """
        if option not in ['recreate', 'update']:
            Logger('Illegal option to save outputs in ROOT file!', 'FATAL')

        if option == 'recreate':
            with uproot.recreate(filename) as ofile:
                ofile[histname] = hdata
        else:
            with uproot.update(filename) as ofile:
                ofile[histname] = hdata

    def __write_pdf(self, histname, weight, num, filename='output.root', option='recreate'):
        """
        Helper method to save a pdf histogram in a .root file (TH1D format)

        Parameters
        -------------------------------------------------
        histname: str
            Name of the histogram

        weight: array[float]
            Array of weights for histogram bins

        num: int
            Number of bins to plot pdfs converted into histograms

        filename: str
            ROOT file name

        option: str
            Option (recreate or update)
        """

        if option not in ['recreate', 'update']:
            Logger('Illegal option to save outputs in ROOT file!', 'FATAL')

        limits = self._data_handler_.get_limits()
        x_plot = np.linspace(limits[0], limits[1], num=num)
        histo = Hist.new.Reg(num, limits[0], limits[1], name=self._data_handler_.get_var_name()).Double()
        histo.fill(x_plot, weight=weight)

        if option == 'recreate':
            with uproot.recreate(filename) as ofile:
                ofile[histname] = histo
        else:
            with uproot.update(filename) as ofile:
                ofile[histname] = histo

    def get_name_signal_pdf(self):
        """
        Return the signal pdf names
        """
        return self._name_signal_pdf_.copy()

    def get_name_background_pdf(self):
        """
        Return the background pdf names
        """
        return self._name_background_pdf_.copy()

    def get_name_refl_pdf(self):
        """
        Return the reflection pdf names
        """
        return self._name_refl_pdf_.copy()

    def get_signal_pars(self):
        """
        Return the signal parameters
        """
        fracs = self.__get_all_fracs()
        signal_pars = []
        for i_sgn, sgn_par in enumerate(self._sgn_pars_):
            signal_pars.append({})
            for key, value in sgn_par.items():
                par_name = key.split(f'{self._name_}_')[-1]
                par_name = par_name.split(f'_signal{i_sgn}')[0]
                signal_pars[-1][par_name] = value.numpy()
            signal_pars[-1]['frac'] = fracs[0][i_sgn]
        return signal_pars

    def get_bkg_pars(self):
        """
        Return the background parameters
        """
        fracs = self.__get_all_fracs()
        bkg_pars = []
        for i_bkg, bkg_par in enumerate(self._bkg_pars_):
            bkg_pars.append({})
            for key, value in bkg_par.items():
                par_name = key.split(f'{self._name_}_')[-1]
                par_name = par_name.split(f'_bkg{i_bkg}')[0]
                bkg_pars[-1][par_name] = value.numpy()
            bkg_pars[-1]['frac'] = fracs[1][i_bkg]
        return bkg_pars

Class used to perform mass fits with the zfit library https://github.com/zfit/zfit

Initialize the F2MassFitter class Parameters


data_handler : flarefly.DataHandler
The data handler containing the data to fit
name_signal_pdf : list

The list of names for the signal pdfs. The possible options are:

  • 'nosignal'

  • 'gaussian'

  • 'doublegaus'

  • 'gausexptail'

  • 'genergausexptail'

  • 'genergausexptailsymm'

  • 'bifurgaus'

  • 'crystalball'

  • 'doublecb'

  • 'doublecbsymm'

  • 'genercrystalball'

  • 'cauchy'

  • 'voigtian'

  • 'kde_exact' (requires to set the datasample and options)

  • 'kde_grid' (requires to set the datasample and options)

  • 'kde_fft' (requires to set the datasample and options)

  • 'kde_isj' (requires to set the datasample and options)

  • 'hist' (only for binned fits, requires to set the datasample)

name_background_pdf : list

The list of names of the background pdfs. The possible options are:

  • 'nobkg'

  • 'expo'

  • 'powlaw'

  • 'expopow'

  • 'expopowext'

  • 'chebpolN' (N is the order of the polynomial)

  • 'kde_exact' (requires to set the datasample and options)

  • 'kde_grid' (requires to set the datasample and options)

  • 'kde_fft' (requires to set the datasample and options)

  • 'kde_isj' (requires to set the datasample and options)

  • 'hist' (only for binned fits, requires to set the datasample)

**kwargs : dict

Additional optional arguments:

  • name_refl_pdf: list The list of names of the signal pdfs. It must have the same length as the signal list. The possible options are:

    • 'kde_exact' (requires to set the datasample and options)

    • 'kde_grid' (requires to set the datasample and options)

    • 'kde_fft' (requires to set the datasample and options)

    • 'kde_isj' (requires to set the datasample and options)

    • 'hist' (only for binned fits, requires to set the datasample)

  • name: str Optional name for the fitter, needed in case of multiple fitters defined in the same script

  • chi2_loss: bool chi2 minimization if True, nll minmization else, default value to False

  • minuit_mode: A number used by minuit to define the internal minimization strategy, either 0, 1 or 2. 0 is the fastest, 2 is the slowest (see more details in https://zfit.readthedocs.io/en/latest/user_api/minimize/_generated/minimizers/zfit.minimize.Minuit.html#zfit.minimize.Minuit) Default value to 0

  • limits: list list of fit limits to include in the fit

  • tol: float Termination value for the convergence/stopping criterion of the algorithm in order to determine if the minimum has been found. Default value to 0.001

  • verbosity: int verbosity level (from 0 to 10) Default value to 0

  • signal_at_threshold: list list of booleans which indicate whether the signal PDFs are at threshold or not. Each element corresponds to a signal pdf

  • label_signal_pdf: list list of labels for signal pdfs

  • label_bkg_pdf: list list of labels for background pdfs

Instance variables

prop get_fit_result
Expand source code
@property
def get_fit_result(self):
    """
    Get the fit result

    Returns
    -------------------------------------------------
    fit_result: zfit.minimizers.fitresult.FitResult
        The fit result
    """
    return self._fit_result_

Get the fit result

Returns

fit_result : zfit.minimizers.fitresult.FitResult
The fit result

Methods

def cleanup(self)
Expand source code
def cleanup(self):
    """
    Cleanup function to clear the graph cache of zfit
    """
    zfit.run.clear_graph_cache()

Cleanup function to clear the graph cache of zfit

def dump_to_root(self, filename, **kwargs)
Expand source code
def dump_to_root(self, filename, **kwargs):
    """
    Plot the mass fit

    Parameters
    -------------------------------------------------
    filename: str
        Name of output ROOT file

    **kwargs: dict
        Additional optional arguments:

        - axis_title: str
            x-axis title

        - num: int
            number of bins to plot pdfs converted into histograms

        - option: str
            option (recreate or update)


        - suffix: str
            suffix to append to objects
    """

    num = kwargs.get('num', 10000)
    suffix = kwargs.get('suffix', '')
    option = kwargs.get('option', 'recreate')
    bins = self._data_handler_.get_nbins()
    obs = self._data_handler_.get_obs()
    limits = self._data_handler_.get_limits()

    hdata = self._data_handler_.to_hist(lower_edge=limits[0],
                                        upper_edge=limits[1],
                                        nbins=bins,
                                        varname=self._data_handler_.get_var_name())
    # write data
    self.__write_data(hdata, f'hdata{suffix}', filename, option)

    bin_sigma = (limits[1] - limits[0]) / bins
    norm = self._total_pdf_norm_ * bin_sigma

    x_plot = np.linspace(limits[0], limits[1], num=num)

    total_func = zfit.run(self._total_pdf_.pdf(x_plot, norm_range=obs))
    # write total_func
    self.__write_pdf(histname=f'total_func{suffix}', weight=total_func * norm, num=num,
                     filename=filename, option='update')

    signal_funcs, bkg_funcs, refl_funcs = ([] for _ in range(3))
    for signal_pdf in self._signal_pdf_:
        signal_funcs.append(zfit.run(signal_pdf.pdf(x_plot, norm_range=obs)))
    for refl_pdf in self._refl_pdf_:
        refl_funcs.append(zfit.run(refl_pdf.pdf(x_plot, norm_range=obs)))
    for bkg_pdf in self._background_pdf_:
        bkg_funcs.append(zfit.run(bkg_pdf.pdf(x_plot, norm_range=obs)))

    signal_fracs, bkg_fracs, refl_fracs, _, _, _ = self.__get_all_fracs()

    # first write backgrounds
    for ibkg, (bkg_func, bkg_frac) in enumerate(zip(bkg_funcs, bkg_fracs)):
        self.__write_pdf(histname=f'bkg_{ibkg}{suffix}',
                         weight=bkg_func * norm * bkg_frac / self._ratio_truncated_,
                         num=num, filename=filename, option='update')
    # then write signals
    for isgn, (frac, signal_func) in enumerate(zip(signal_funcs, signal_fracs)):
        self.__write_pdf(histname=f'signal_{isgn}{suffix}',
                         weight=signal_func * norm * frac / self._ratio_truncated_,
                         num=num, filename=filename, option='update')

    # finally write reflected signals
    for irefl, (frac, refl_func) in enumerate(zip(refl_funcs, refl_fracs)):
        if self._name_refl_pdf_[irefl] is None:
            continue
        self.__write_pdf(histname=f'refl_{irefl}{suffix}',
                         weight=refl_func * norm * frac / self._ratio_truncated_,
                         num=num, filename=filename, option='update')

Plot the mass fit

Parameters

filename : str
Name of output ROOT file
**kwargs : dict

Additional optional arguments:

  • axis_title: str x-axis title

  • num: int number of bins to plot pdfs converted into histograms

  • option: str option (recreate or update)

  • suffix: str suffix to append to objects

def fix_bkg_frac_to_bkg_pdf(self, idx_pdf, target_pdf, factor=1)
Expand source code
def fix_bkg_frac_to_bkg_pdf(self, idx_pdf, target_pdf, factor=1):
    """
    Fix the frac parameter of a background to the frac parameter of another background

    Parameters
    -------------------------------------------------
    idx_pdf: int
        Index of the background fraction to be fixed
    target_pdf: int
        Index of the background fraction to be used as reference
    factor: float
        Factor to multiply the frac parameter of the target background
    """
    self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'bkg', 'bkg')
    factor_par = self.__get_parameter_fix_frac(
        factor, f'factor_bkg{idx_pdf}_constrained_to_bkg{target_pdf}'
    )
    self._fix_fracs_to_pdfs_.append({
        'fixed_pdf_idx': idx_pdf,
        'target_pdf_idx': target_pdf,
        'factor': factor_par,
        'fixed_pdf_type': 'bkg',
        'target_pdf_type': 'bkg'
    })

Fix the frac parameter of a background to the frac parameter of another background

Parameters

idx_pdf : int
Index of the background fraction to be fixed
target_pdf : int
Index of the background fraction to be used as reference
factor : float
Factor to multiply the frac parameter of the target background
def fix_bkg_frac_to_signal_pdf(self, idx_pdf, target_pdf, factor=1)
Expand source code
def fix_bkg_frac_to_signal_pdf(self, idx_pdf, target_pdf, factor=1):
    """
    Fix the frac parameter of a background to the frac parameter of a signal

    Parameters
    -------------------------------------------------
    idx_pdf: int
        Index of the background fraction to be fixed
    target_pdf: int
        Index of the signal fraction to be used as reference
    factor: float
        Factor to multiply the frac parameter of the target signal
    """
    self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'bkg', 'signal')
    factor_par = self.__get_parameter_fix_frac(
        factor, f'factor_bkg{idx_pdf}_constrained_to_signal{target_pdf}'
    )
    self._fix_fracs_to_pdfs_.append({
        'fixed_pdf_idx': idx_pdf,
        'target_pdf_idx': target_pdf,
        'factor': factor_par,
        'fixed_pdf_type': 'bkg',
        'target_pdf_type': 'signal'
    })

Fix the frac parameter of a background to the frac parameter of a signal

Parameters

idx_pdf : int
Index of the background fraction to be fixed
target_pdf : int
Index of the signal fraction to be used as reference
factor : float
Factor to multiply the frac parameter of the target signal
def fix_signal_frac_to_bkg_pdf(self, idx_pdf, target_pdf, factor=1)
Expand source code
def fix_signal_frac_to_bkg_pdf(self, idx_pdf, target_pdf, factor=1):
    """
    Fix the frac parameter of a signal to the frac parameter of a background

    Parameters
    -------------------------------------------------
    idx_pdf: int
        Index of the signal fraction to be fixed
    target_pdf: int
        Index of the background fraction to be used as reference
    factor: float
        Factor to multiply the frac parameter of the target background
    """
    self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'signal', 'bkg')
    factor_par = self.__get_parameter_fix_frac(
        factor, f'factor_signal{idx_pdf}_constrained_to_bkg{target_pdf}'
    )
    self._fix_fracs_to_pdfs_.append({
        'fixed_pdf_idx': idx_pdf,
        'target_pdf_idx': target_pdf,
        'factor': factor_par,
        'fixed_pdf_type': 'signal',
        'target_pdf_type': 'bkg'
    })

Fix the frac parameter of a signal to the frac parameter of a background

Parameters

idx_pdf : int
Index of the signal fraction to be fixed
target_pdf : int
Index of the background fraction to be used as reference
factor : float
Factor to multiply the frac parameter of the target background
def fix_signal_frac_to_signal_pdf(self, idx_pdf, target_pdf, factor=1)
Expand source code
def fix_signal_frac_to_signal_pdf(self, idx_pdf, target_pdf, factor=1):
    """
    Fix the frac parameter of a signal to the frac parameter of another signal

    Parameters
    -------------------------------------------------
    idx_pdf: int
        Index of the signal fraction to be fixed
    target_pdf: int
        Index of the signal fraction to be used as reference
    factor: float
        Factor to multiply the frac parameter of the target signal
    """
    self.__check_consistency_fix_frac(idx_pdf, target_pdf, 'signal', 'signal')
    factor_par = self.__get_parameter_fix_frac(
        factor, f'factor_signal{idx_pdf}_constrained_to_signal{target_pdf}'
    )
    self._fix_fracs_to_pdfs_.append({
        'fixed_pdf_idx': idx_pdf,
        'target_pdf_idx': target_pdf,
        'factor': factor_par,
        'fixed_pdf_type': 'signal',
        'target_pdf_type': 'signal'
    })

Fix the frac parameter of a signal to the frac parameter of another signal

Parameters

idx_pdf : int
Index of the signal fraction to be fixed
target_pdf : int
Index of the signal fraction to be used as reference
factor : float
Factor to multiply the frac parameter of the target signal
def get_background(self, idx=0, **kwargs)
Expand source code
def get_background(self, idx=0, **kwargs):
    """
    Get the background and its error in a given invariant-mass region

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal to be used to compute nsigma window
    **kwargs: dict
        Additional optional arguments:

        - nsigma: float
            nsigma invariant-mass window around mean for background computation

        - nhwhm: float
            number of hwhm invariant-mass window around mean for signal and background computation
            (alternative to nsigma)

        - min: float
            minimum value of invariant-mass for background computation (alternative to nsigma)

        - max: float
            maximum value of invariant-mass for background computation (alternative to nsigma)

    Returns
    -------------------------------------------------
    background: float
        The background value obtained from the fit
    background_err: float
        The background error obtained from the fit
    """

    if not self._background_pdf_:
        Logger('Background not fitted', 'ERROR')
        return 0., 0.

    nsigma = kwargs.get('nsigma', 3.)
    nhwhm = kwargs.get('nhwhm', None)
    min_value = kwargs.get('min', None)
    max_value = kwargs.get('max', None)
    use_nsigma = True

    if nhwhm is not None and (min_value is not None or max_value is not None):
        Logger('I cannot compute the signal within a fixed mass interval and a number of HWFM', 'ERROR')
        return 0., 0.

    if min_value is not None and max_value is not None:
        use_nsigma = False

    if nhwhm is not None:
        use_nsigma = False
        if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
            Logger('HWHM not defined, I cannot compute the signal for this pdf', 'ERROR')
            return 0., 0.
        mass, _ = self.get_mass(idx)
        hwhm, _ = self.get_hwhm(idx)
        min_value = mass - nhwhm * hwhm
        max_value = mass + nhwhm * hwhm

    if use_nsigma:
        if self._name_signal_pdf_[idx] not in [
            'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
            'doublecb', 'doublecbsymm', 'voigtian', 'hist'
            ]:
            Logger('Sigma not defined, I cannot compute the signal for this pdf', 'ERROR')
            return 0., 0.
        mass, _ = self.get_mass(idx)
        sigma, _ = self.get_sigma(idx)
        min_value = mass - nsigma * sigma
        max_value = mass + nsigma * sigma

    _, bkg_fracs, _, _, bkg_err_fracs, _ = self.__get_all_fracs()

    limits = self._data_handler_.get_limits()
    min_value = max(min_value, limits[0])
    max_value = min(max_value, limits[1])

    # pylint: disable=missing-kwoa
    background, background_err = 0., 0.
    for idx2, bkg in enumerate(self._background_pdf_):

        norm = self._total_pdf_norm_ * bkg_fracs[idx2]
        norm_err = norm * bkg_err_fracs[idx2]

        bkg_int = float(bkg.integrate((min_value, max_value)))
        background += bkg_int * norm
        background_err += (bkg_int * norm_err)**2

    background /= self._ratio_truncated_
    background_err = np.sqrt(background_err) / self._ratio_truncated_

    return float(background), float(background_err)

Get the background and its error in a given invariant-mass region

Parameters

idx : int
Index of the signal to be used to compute nsigma window
**kwargs : dict

Additional optional arguments:

  • nsigma: float nsigma invariant-mass window around mean for background computation

  • nhwhm: float number of hwhm invariant-mass window around mean for signal and background computation (alternative to nsigma)

  • min: float minimum value of invariant-mass for background computation (alternative to nsigma)

  • max: float maximum value of invariant-mass for background computation (alternative to nsigma)

Returns

background : float
The background value obtained from the fit
background_err : float
The background error obtained from the fit
def get_background_parameter(self, idx, par_name)
Expand source code
def get_background_parameter(self, idx, par_name):
    """
    Get a background parameter and its error

    Parameters
    -------------------------------------------------
    idx: int
        Index of the parameter to be returned (default: 0)
    par_name: str
        parameter to return

    Returns
    -------------------------------------------------
    parameter: float
        The parameter value obtained from the fit
    parameter_err: float
        The parameter error obtained from the fit

    """

    if self._fix_bkg_pars_[idx][par_name]:
        parameter = self._init_bkg_pars_[idx][par_name]
        parameter_err = 0.
    else:
        parameter = self._fit_result_.params[f'{self._name_}_{par_name}_bkg{idx}']['value']
        parameter_err = self._fit_result_.params[f'{self._name_}_{par_name}_bkg{idx}']['hesse']['error']

    return parameter, parameter_err

Get a background parameter and its error

Parameters

idx : int
Index of the parameter to be returned (default: 0)
par_name : str
parameter to return

Returns

parameter : float
The parameter value obtained from the fit
parameter_err : float
The parameter error obtained from the fit
def get_bkg_pars(self)
Expand source code
def get_bkg_pars(self):
    """
    Return the background parameters
    """
    fracs = self.__get_all_fracs()
    bkg_pars = []
    for i_bkg, bkg_par in enumerate(self._bkg_pars_):
        bkg_pars.append({})
        for key, value in bkg_par.items():
            par_name = key.split(f'{self._name_}_')[-1]
            par_name = par_name.split(f'_bkg{i_bkg}')[0]
            bkg_pars[-1][par_name] = value.numpy()
        bkg_pars[-1]['frac'] = fracs[1][i_bkg]
    return bkg_pars

Return the background parameters

def get_chi2(self)
Expand source code
def get_chi2(self):
    """
    Get chi2 for binned data

    Returns
    -------------------------------------------------
    chi2: float
        chi2
    """

    chi2 = 0
    norm = self._data_handler_.get_norm()

    if self._data_handler_.get_is_binned():
        # for chi2 loss, just retrieve loss value in fit result
        if self._chi2_loss_:
            return float(self._fit_result_.loss.value())

        # for nll loss, compute chi2 "by hand"
        # access normalized data values and errors for all bins
        binned_data = self._data_handler_.get_binned_data()
        data_values = binned_data.values()
        data_variances = binned_data.variances()
        # access model predicted values
        model_values = self._total_pdf_binned_.values() * norm
        # compute chi2
        for (data, model, data_variance) in zip(data_values, model_values, data_variances):
            chi2 += (data - model)**2/data_variance
        return chi2

    if self._limits_ != [self._data_handler_.get_limits()]:  # restricted fit limits
        Logger('chi2 not yet implemented with restricted fit limits', 'FATAL')
    # for unbinned data
    data_values = self._data_handler_.get_binned_data_from_unbinned_data()
    # access model predicted values
    model_values = self._total_pdf_binned_.values() * norm
    # compute chi2
    for (data, model) in zip(data_values, model_values):
        if data == 0:
            continue
        chi2 += (data - model)**2/data

    return float(chi2)

Get chi2 for binned data

Returns

chi2 : float
chi2
def get_chi2_ndf(self)
Expand source code
def get_chi2_ndf(self):
    """
    Get the reduced chi2 (chi2 divided by number of degrees of freedom)
    for binned data

    Returns
    -------------------------------------------------
    chi2_ndf: float
        The reduced chi2
    """
    return self.get_chi2()/self.get_ndf()

Get the reduced chi2 (chi2 divided by number of degrees of freedom) for binned data

Returns

chi2_ndf : float
The reduced chi2
def get_hwhm(self, idx=0)
Expand source code
def get_hwhm(self, idx=0):
    """
    Get the half width half maximum and its error

    Parameters
    -------------------------------------------------
    idx: int
        Index of the sigma to be returned (default: 0)

    Returns
    -------------------------------------------------
    hwhm: float
        The sigma value obtained from the fit
    hwhm_err: float
        The sigma error obtained from the fit
    """
    if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
        Logger(f'HFWM parameter not defined for {self._name_signal_pdf_[idx]} pdf!', 'ERROR')
        return 0., 0.

    if self._name_signal_pdf_[idx] == 'gaussian':
        mult_fact = np.sqrt(2 * np.log(2))
        hwhm, hwhm_err = self.get_sigma(idx)
        hwhm *= mult_fact
        hwhm_err *= mult_fact
    elif self._name_signal_pdf_[idx] == 'cauchy':
        hwhm, hwhm_err = self.get_signal_parameter(idx, 'gamma')
    elif self._name_signal_pdf_[idx] == 'voigtian':
        mult_fact = np.sqrt(2 * np.log(2))
        sigma, sigma_err = self.get_sigma(idx)
        sigma *= mult_fact
        sigma_err *= mult_fact
        gamma, gamma_err = self.get_signal_parameter(idx, 'gamma')
        hwhm = 0.5346 * gamma + np.sqrt(0.2166 * gamma**2 + sigma**2)
        # we neglect the correlation between sigma and gamma
        der_sigma = sigma / np.sqrt(0.0721663 + sigma**2)
        der_gamma = 0.5346 + (0.2166 * gamma) / np.sqrt(0.2166 * gamma**2 + sigma**2)
        hwhm_err = np.sqrt(der_sigma**2 * sigma_err**2 + der_gamma**2 * gamma_err**2)

    return hwhm, hwhm_err

Get the half width half maximum and its error

Parameters

idx : int
Index of the sigma to be returned (default: 0)

Returns

hwhm : float
The sigma value obtained from the fit
hwhm_err : float
The sigma error obtained from the fit
def get_mass(self, idx=0)
Expand source code
def get_mass(self, idx=0):
    """
    Get the mass and its error

    Parameters
    -------------------------------------------------
    idx: int
        Index of the mass to be returned (default: 0)

    Returns
    -------------------------------------------------
    mass: float
        The mass value obtained from the fit
    mass_err: float
        The mass error obtained from the fit
    """
    if 'hist' in self._name_signal_pdf_[idx]:
        hist = self._signal_pdf_[idx].to_hist()
        bin_limits = hist.to_numpy()[1]
        centres = [0.5 * (minn + maxx) for minn, maxx in zip(bin_limits[1:],  bin_limits[:-1])]
        counts = hist.values()
        mass = np.average(centres, weights=counts)
        mass_err = 0.
    else:
        mass_name = 'm' if self._name_signal_pdf_[idx] in ['cauchy', 'voigtian'] else 'mu'
        if self._fix_sgn_pars_[idx][mass_name]:
            mass = self._init_sgn_pars_[idx][mass_name]
            mass_err = 0.
        else:
            mass = self._fit_result_.params[f'{self._name_}_{mass_name}_signal{idx}']['value']
            mass_err = self._fit_result_.params[f'{self._name_}_{mass_name}_signal{idx}']['hesse']['error']

    return mass, mass_err

Get the mass and its error

Parameters

idx : int
Index of the mass to be returned (default: 0)

Returns

mass : float
The mass value obtained from the fit
mass_err : float
The mass error obtained from the fit
def get_name_background_pdf(self)
Expand source code
def get_name_background_pdf(self):
    """
    Return the background pdf names
    """
    return self._name_background_pdf_.copy()

Return the background pdf names

def get_name_refl_pdf(self)
Expand source code
def get_name_refl_pdf(self):
    """
    Return the reflection pdf names
    """
    return self._name_refl_pdf_.copy()

Return the reflection pdf names

def get_name_signal_pdf(self)
Expand source code
def get_name_signal_pdf(self):
    """
    Return the signal pdf names
    """
    return self._name_signal_pdf_.copy()

Return the signal pdf names

def get_ndf(self)
Expand source code
def get_ndf(self):
    """
    Get the number of degrees of freedom for chi2 fit
    ndf = nbins - nfreeparams - 1
    -1 because the data sample size is fixed

    Returns
    -------------------------------------------------
    ndf: int
        The number of degrees of freedom
    """
    nbins = self._data_handler_.get_nbins()
    nfreeparams = len(self._fit_result_.params)
    self._ndf_ = nbins - nfreeparams - 1
    return self._ndf_

Get the number of degrees of freedom for chi2 fit ndf = nbins - nfreeparams - 1 -1 because the data sample size is fixed

Returns

ndf : int
The number of degrees of freedom
def get_raw_yield(self, idx=0)
Expand source code
def get_raw_yield(self, idx=0):
    """
    Get the raw yield and its error

    Parameters
    -------------------------------------------------
    idx: int
        Index of the raw yield to be returned (default: 0)

    Returns
    -------------------------------------------------
    raw_yield: float
        The raw yield obtained from the fit
    raw_yield_err: float
        The raw yield error obtained from the fit
    """
    return self._rawyield_[idx], self._rawyield_err_[idx]

Get the raw yield and its error

Parameters

idx : int
Index of the raw yield to be returned (default: 0)

Returns

raw_yield : float
The raw yield obtained from the fit
raw_yield_err : float
The raw yield error obtained from the fit
def get_raw_yield_bincounting(self, idx=0, **kwargs)
Expand source code
def get_raw_yield_bincounting(self, idx=0, **kwargs):
    """
    Get the raw yield and its error via the bin-counting method

    Parameters
    -------------------------------------------------
    idx: int
        Index of the raw yield to be returned (default: 0)

    **kwargs: dict
        Additional optional arguments:

        - nsigma: float
            nsigma invariant-mass window around mean for signal counting

        - nhwhm: float
            number of hwhm invariant-mass window around mean for signal counting
            (alternative to nsigma)

        - min: float
            minimum value of invariant-mass for signal counting (alternative to nsigma)

        - max: float
            maximum value of invariant-mass for signal counting (alternative to nsigma)

    Returns
    -------------------------------------------------
    raw_yield: float
        The raw yield obtained from the bin counting
    raw_yield_err: float
        The raw yield error obtained from the bin counting
    """

    nsigma = kwargs.get('nsigma', 3.)
    nhwhm = kwargs.get('nhwhm', None)
    min_value = kwargs.get('min', None)
    max_value = kwargs.get('max', None)
    use_nsigma = True

    if nhwhm is not None and (min_value is not None or max_value is not None):
        Logger('I cannot compute the signal within a fixed mass interval and a number of HWFM', 'ERROR')
        return 0., 0.

    if min_value is not None and max_value is not None:
        use_nsigma = False

    if nhwhm is not None:
        use_nsigma = False
        if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
            Logger('HWHM not defined, I cannot compute the signal for this pdf', 'ERROR')
            return 0., 0.
        mass, _ = self.get_mass(idx)
        hwhm, _ = self.get_hwhm(idx)
        min_value = mass - nhwhm * hwhm
        max_value = mass + nhwhm * hwhm

    if use_nsigma:
        if self._name_signal_pdf_[idx] not in [
            'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
            'doublecb', 'doublecbsymm', 'voigtian', 'hist'
            ]:
            Logger('Sigma not defined, I cannot compute the signal for this pdf', 'ERROR')
            return 0., 0.
        mass, _ = self.get_mass(idx)
        sigma, _ = self.get_sigma(idx)
        min_value = mass - nsigma * sigma
        max_value = mass + nsigma * sigma

    if len(self._raw_residuals_) == 0:
        self.__get_raw_residuals()

    bin_centers = self._data_handler_.get_bin_center()
    bin_width = (bin_centers[1] - bin_centers[0]) / 2
    raw_yield, raw_yield_err = 0., 0.
    for residual, variance, bin_center in zip(self._raw_residuals_,
                                              self._raw_residual_variances_, bin_centers):
        if bin_center - bin_width >= min_value and bin_center + bin_width <= max_value:
            raw_yield += residual
            raw_yield_err += variance
    raw_yield = float(raw_yield)
    raw_yield_err = np.sqrt(float(raw_yield_err))

    return raw_yield, raw_yield_err

Get the raw yield and its error via the bin-counting method

Parameters

idx : int
Index of the raw yield to be returned (default: 0)
**kwargs : dict

Additional optional arguments:

  • nsigma: float nsigma invariant-mass window around mean for signal counting

  • nhwhm: float number of hwhm invariant-mass window around mean for signal counting (alternative to nsigma)

  • min: float minimum value of invariant-mass for signal counting (alternative to nsigma)

  • max: float maximum value of invariant-mass for signal counting (alternative to nsigma)

Returns

raw_yield : float
The raw yield obtained from the bin counting
raw_yield_err : float
The raw yield error obtained from the bin counting
def get_sigma(self, idx=0)
Expand source code
def get_sigma(self, idx=0):
    """
    Get the sigma and its error

    Parameters
    -------------------------------------------------
    idx: int
        Index of the sigma to be returned (default: 0)

    Returns
    -------------------------------------------------
    sigma: float
        The sigma value obtained from the fit
    sigma_err: float
        The sigma error obtained from the fit
    """
    if self._name_signal_pdf_[idx] not in [
        'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
        'doublecb', 'doublecbsymm', 'voigtian', 'hist'
        ]:
        Logger(f'Sigma parameter not defined for {self._name_signal_pdf_[idx]} pdf!', 'ERROR')
        return 0., 0.

    # if histogram, the rms is used as proxy
    if 'hist' in self._name_signal_pdf_[idx]:
        Logger(f'RMS used as proxy for sigma parameter of {self._name_signal_pdf_[idx]} pdf!', 'WARNING')
        mean = self.get_mass(idx)[0]
        hist = self._signal_pdf_[idx].to_hist()
        bin_limits = hist.to_numpy()[1]
        centres = [0.5 * (minn + maxx) for minn, maxx in zip(bin_limits[1:],  bin_limits[:-1])]
        counts = hist.values()
        sigma = np.sqrt(np.average((centres - mean)**2, weights=counts))
        sigma_err = 0.
    else:
        if self._fix_sgn_pars_[idx]['sigma']:
            sigma = self._init_sgn_pars_[idx]['sigma']
            sigma_err = 0.
        else:
            sigma = self._fit_result_.params[f'{self._name_}_sigma_signal{idx}']['value']
            sigma_err = self._fit_result_.params[f'{self._name_}_sigma_signal{idx}']['hesse']['error']

    return sigma, sigma_err

Get the sigma and its error

Parameters

idx : int
Index of the sigma to be returned (default: 0)

Returns

sigma : float
The sigma value obtained from the fit
sigma_err : float
The sigma error obtained from the fit
def get_signal(self, idx=0, **kwargs)
Expand source code
def get_signal(self, idx=0, **kwargs):
    """
    Get the signal and its error in a given invariant-mass region

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal to be returned
    **kwargs: dict
        Additional optional arguments:

        - nsigma: float
            nsigma invariant-mass window around mean for signal computation

        - nhwhm: float
            number of hwhm invariant-mass window around mean for signal and background computation
            (alternative to nsigma)

        - min: float
            minimum value of invariant-mass for signal computation (alternative to nsigma)

        - max: float
            maximum value of invariant-mass for signal computation (alternative to nsigma)

    Returns
    -------------------------------------------------
    signal: float
        The signal value obtained from the fit
    signal_err: float
        The signal error obtained from the fit
    """

    nsigma = kwargs.get('nsigma', 3.)
    nhwhm = kwargs.get('nhwhm', None)
    min_value = kwargs.get('min', None)
    max_value = kwargs.get('max', None)
    use_nsigma = True

    if nhwhm is not None and (min_value is not None or max_value is not None):
        Logger('I cannot compute the signal within a fixed mass interval and a number of HWFM', 'ERROR')
        return 0., 0.

    if min_value is not None and max_value is not None:
        use_nsigma = False

    if nhwhm is not None:
        use_nsigma = False
        if self._name_signal_pdf_[idx] not in ['gaussian', 'cauchy', 'voigtian']:
            Logger('HWHM not defined, I cannot compute the signal for this pdf', 'ERROR')
            return 0., 0.
        mass, _ = self.get_mass(idx)
        hwhm, _ = self.get_hwhm(idx)
        min_value = mass - nhwhm * hwhm
        max_value = mass + nhwhm * hwhm

    if use_nsigma:
        if self._name_signal_pdf_[idx] not in [
            'gaussian', 'gausexptail', 'genergausexptailsymm', 'crystalball',
            'doublecb', 'doublecbsymm', 'voigtian', 'hist'
            ]:
            Logger('Sigma not defined, I cannot compute the signal for this pdf', 'ERROR')
            return 0., 0.
        mass, _ = self.get_mass(idx)
        sigma, _ = self.get_sigma(idx)
        min_value = mass - nsigma * sigma
        max_value = mass + nsigma * sigma

    # pylint: disable=missing-kwoa
    signal = self._signal_pdf_[idx].integrate((min_value, max_value))

    signal_fracs, _, refl_fracs, signal_err_fracs, _, _ = self.__get_all_fracs()

    if len(self._background_pdf_) > 0:
        frac = signal_fracs[idx]
        frac_err = signal_err_fracs[idx]
    else:
        if len(self._signal_pdf_) == 1:
            frac = 1.
            frac_err = 0.
        if idx < len(signal_fracs):
            frac = signal_fracs[idx]
            frac_err = signal_err_fracs[idx]
        else:
            frac = 1. - sum(signal_fracs) - sum(refl_fracs)
            frac_err = np.sqrt(sum(list(err**2 for err in signal_err_fracs)))

    norm = self._data_handler_.get_norm()
    norm_err = norm * frac_err
    norm *= frac

    return float(signal * norm), float(signal * norm_err)

Get the signal and its error in a given invariant-mass region

Parameters

idx : int
Index of the signal to be returned
**kwargs : dict

Additional optional arguments:

  • nsigma: float nsigma invariant-mass window around mean for signal computation

  • nhwhm: float number of hwhm invariant-mass window around mean for signal and background computation (alternative to nsigma)

  • min: float minimum value of invariant-mass for signal computation (alternative to nsigma)

  • max: float maximum value of invariant-mass for signal computation (alternative to nsigma)

Returns

signal : float
The signal value obtained from the fit
signal_err : float
The signal error obtained from the fit
def get_signal_over_background(self, idx=0, **kwargs)
Expand source code
def get_signal_over_background(self, idx=0, **kwargs):
    """
    Get the S/B ratio and its error in a given invariant-mass region

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal to be used to compute nsigma window
    **kwargs: dict
        Additional optional arguments:

        - nsigma: float
            nsigma invariant-mass window around mean for signal and background computation

        - nhwhm: float
            number of hwhm invariant-mass window around mean for signal and background computation
            (alternative to nsigma)

        - min: float
            minimum value of invariant-mass for signal and background computation (alternative to nsigma)

        - max: float
            maximum value of invariant-mass for signal and background computation (alternative to nsigma)

    Returns
    -------------------------------------------------
    signal_over_background: float
        The S/B value obtained from the fit
    signal_over_background_err: float
        The S/B error obtained from the fit
    """

    signal = self.get_signal(idx, **kwargs)
    bkg = self.get_background(idx, **kwargs)
    signal_over_background = signal[0]/bkg[0]
    signal_over_background_err = np.sqrt(signal[1]**2/signal[0]**2 + bkg[1]**2/bkg[0]**2)
    signal_over_background_err *= signal_over_background

    return signal_over_background, signal_over_background_err

Get the S/B ratio and its error in a given invariant-mass region

Parameters

idx : int
Index of the signal to be used to compute nsigma window
**kwargs : dict

Additional optional arguments:

  • nsigma: float nsigma invariant-mass window around mean for signal and background computation

  • nhwhm: float number of hwhm invariant-mass window around mean for signal and background computation (alternative to nsigma)

  • min: float minimum value of invariant-mass for signal and background computation (alternative to nsigma)

  • max: float maximum value of invariant-mass for signal and background computation (alternative to nsigma)

Returns

signal_over_background : float
The S/B value obtained from the fit
signal_over_background_err : float
The S/B error obtained from the fit
def get_signal_parameter(self, idx, par_name)
Expand source code
def get_signal_parameter(self, idx, par_name):
    """
    Get a signal parameter and its error

    Parameters
    -------------------------------------------------
    idx: int
        Index of the parameter to be returned (default: 0)
    par_name: str
        parameter to return

    Returns
    -------------------------------------------------
    parameter: float
        The parameter value obtained from the fit
    parameter_err: float
        The parameter error obtained from the fit

    """

    if par_name == 'gamma':
        Logger('The gamma parameter that you are getting is half of the width of a resonance,'
               ' for more info check the Cauchy pdf defined here '
               'https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/basic/zfit.pdf.Cauchy.html',
               'WARNING')

    if self._fix_sgn_pars_[idx][par_name]:
        parameter = self._init_sgn_pars_[idx][par_name]
        parameter_err = 0.
    else:
        parameter = self._fit_result_.params[f'{self._name_}_{par_name}_signal{idx}']['value']
        parameter_err = self._fit_result_.params[f'{self._name_}_{par_name}_signal{idx}']['hesse']['error']

    return parameter, parameter_err

Get a signal parameter and its error

Parameters

idx : int
Index of the parameter to be returned (default: 0)
par_name : str
parameter to return

Returns

parameter : float
The parameter value obtained from the fit
parameter_err : float
The parameter error obtained from the fit
def get_signal_pars(self)
Expand source code
def get_signal_pars(self):
    """
    Return the signal parameters
    """
    fracs = self.__get_all_fracs()
    signal_pars = []
    for i_sgn, sgn_par in enumerate(self._sgn_pars_):
        signal_pars.append({})
        for key, value in sgn_par.items():
            par_name = key.split(f'{self._name_}_')[-1]
            par_name = par_name.split(f'_signal{i_sgn}')[0]
            signal_pars[-1][par_name] = value.numpy()
        signal_pars[-1]['frac'] = fracs[0][i_sgn]
    return signal_pars

Return the signal parameters

def get_significance(self, idx=0, **kwargs)
Expand source code
def get_significance(self, idx=0, **kwargs):
    """
    Get the significance and its error in a given invariant-mass region

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal to be used to compute nsigma window
    **kwargs: dict
        Additional optional arguments:

        - nsigma: float
            nsigma invariant-mass window around mean for signal and background computation

        - nhwhm: float
            number of hwhm invariant-mass window around mean for signal and background computation
            (alternative to nsigma)

        - min: float
            minimum value of invariant-mass for signal and background computation (alternative to nsigma)

        - max: float
            maximum value of invariant-mass for signal and background computation (alternative to nsigma)

    Returns
    -------------------------------------------------
    significance: float
        The significance value obtained from the fit
    significance_err: float
        The significance error obtained from the fit
    """

    signal = self.get_signal(idx, **kwargs)
    bkg = self.get_background(idx, **kwargs)
    significance = signal[0]/np.sqrt(signal[0]+bkg[0])
    sig_plus_bkg = signal[0] + bkg[0]

    significance_err = significance*np.sqrt(
        (signal[1]**2 + bkg[1]**2) / (4. * sig_plus_bkg**2) + (
            bkg[0]/sig_plus_bkg) * signal[1]**2 / signal[0]**2)

    return significance, significance_err

Get the significance and its error in a given invariant-mass region

Parameters

idx : int
Index of the signal to be used to compute nsigma window
**kwargs : dict

Additional optional arguments:

  • nsigma: float nsigma invariant-mass window around mean for signal and background computation

  • nhwhm: float number of hwhm invariant-mass window around mean for signal and background computation (alternative to nsigma)

  • min: float minimum value of invariant-mass for signal and background computation (alternative to nsigma)

  • max: float maximum value of invariant-mass for signal and background computation (alternative to nsigma)

Returns

significance : float
The significance value obtained from the fit
significance_err : float
The significance error obtained from the fit
def get_sweights(self)
Expand source code
def get_sweights(self):
    """
    Calculate sWeights for signal and background components.

    Returns:
        dict: A dictionary containing sWeights for 'signal' and 'bkg' components.
    """

    signal_pdf_extended = []
    refl_pdf_extended = []
    bkg_pdf_extended = []

    norm = self._data_handler_.get_norm()

    signal_fracs, bkg_fracs, refl_fracs, _, _, _ = self.__get_all_fracs()
    bkg_fracs.append(1-sum(bkg_fracs)-sum(signal_fracs)-sum(refl_fracs))

    total_signal_yields = 0.
    total_bkg_yields = 0.
    for pdf, frac in zip(self._signal_pdf_, signal_fracs):
        signal_pdf_extended.append(pdf.create_extended(frac * norm))
        total_signal_yields += frac * norm
    for pdf, frac in zip(self._refl_pdf_, refl_fracs):
        refl_pdf_extended.append(pdf.create_extended(frac * norm))
        total_signal_yields += frac * norm
    for pdf, frac in zip(self._background_pdf_, bkg_fracs):
        bkg_pdf_extended.append(pdf.create_extended(frac * norm))
        total_bkg_yields += frac * norm

    total_signal_yields_par = zfit.Parameter('signal_yield', total_signal_yields)
    total_bkg_yields_par = zfit.Parameter('bkg_yield', total_bkg_yields)

    if len(signal_pdf_extended + refl_pdf_extended) > 1:
        signal_pdf_for_sweights = zfit.pdf.SumPDF(
            signal_pdf_extended + refl_pdf_extended,
            extended=total_signal_yields_par
        )
    else:
        signal_pdf_for_sweights = signal_pdf_extended[0]

    if len(bkg_pdf_extended) > 1:
        bkg_pdf_for_sweights = zfit.pdf.SumPDF(
            bkg_pdf_extended,
            extended=total_bkg_yields_par
        )
    else:
        bkg_pdf_for_sweights = bkg_pdf_extended[0]

    total_pdf_for_sweights = zfit.pdf.SumPDF([signal_pdf_for_sweights, bkg_pdf_for_sweights])
    sweights = compute_sweights(total_pdf_for_sweights, self._data_handler_.get_data())
    sweights_labels = list(sweights.keys())

    return {'signal': sweights[sweights_labels[0]], 'bkg': sweights[sweights_labels[1]]}

Calculate sWeights for signal and background components.

Returns

dict
A dictionary containing sWeights for 'signal' and 'bkg' components.
def mass_zfit(self, do_prefit=False, **kwargs)
Expand source code
def mass_zfit(self, do_prefit=False, **kwargs):
    """
    Perform a mass fit with the zfit library

    Parameters
    -------------------------------------------------
    do_prefit: bool
        Flag to enable prefit to background only (supported for unbinned fits only)
        The expected signal regions must be provided.

    **kwargs: dict
        Additional optional arguments:

        - prefit_excluded_regions: list
            List or list of lists with limits for the excluded regions

        - prefit_exclude_nsigma:
            Alternative argument to prefit_excluded_regions that excludes the regions around the signals,
            using the initialised mean and sigma parameters

    Returns
    -------------------------------------------------
    fit_result: zfit.minimizers.fitresult.FitResult
        The fit result
    """

    if self._data_handler_ is None:
        Logger('Data handler not specified', 'FATAL')

    self._raw_residuals_ = []
    self._raw_residual_variances_ = []
    self._std_residuals_ = []
    self._std_residual_variances_ = []

    self.__build_total_pdf()
    self.__build_total_pdf_binned()

    # do prefit
    if do_prefit:
        skip_prefit = False
        excluded_regions = []
        exclude_signals_nsigma = kwargs.get('prefit_exclude_nsigma', None)
        if exclude_signals_nsigma is not None:
            for ipdf, pdf_name in enumerate(self._name_signal_pdf_):
                if pdf_name not in ['gaussian', 'cauchy', 'voigtian']:
                    Logger(f"Sigma parameter not defined for {pdf_name}, "
                           "cannot use nsigma for excluded regions in prefit", "Error")
                    skip_prefit = True
                else:
                    mass_name = 'm' if pdf_name in ['cauchy', 'voigtian'] else 'mu'
                    mass = self._init_sgn_pars_[ipdf][mass_name]
                    sigma = self._init_sgn_pars_[ipdf]['sigma']
                    excluded_regions.append(
                        [mass - exclude_signals_nsigma * sigma, mass + exclude_signals_nsigma * sigma])
        else:
            excluded_regions = kwargs.get('prefit_excluded_regions', None)
            if excluded_regions is None:
                Logger("To perform the prefit, you should set the excluded regions, skip", "Error")
                skip_prefit = True

        if not skip_prefit:
            self.__prefit(excluded_regions)

    if self._data_handler_.get_is_binned():
        # chi2 loss
        if self._chi2_loss_:
            loss = zfit.loss.BinnedChi2(self._total_pdf_binned_, self._data_handler_.get_binned_data())
        # nll loss
        else:
            loss = zfit.loss.BinnedNLL(self._total_pdf_binned_, self._data_handler_.get_binned_data())
    else:
        loss = zfit.loss.UnbinnedNLL(model=self._total_pdf_, data=self._data_handler_.get_data())

    Logger("Performing FIT", "INFO")
    self._fit_result_ = self._minimizer_.minimize(loss=loss)
    Logger(self._fit_result_, 'RESULT')

    if self._fit_result_.hesse() == {}:
        if self._fit_result_.hesse(method='hesse_np') == {}:
            Logger('Impossible to compute hesse error', 'FATAL')

    self.__get_ratio_truncated()
    signal_fracs, _, _, signal_frac_errs, _, _ = self.__get_all_fracs()

    norm = self._total_pdf_norm_

    if len(self._fracs_) == 0:
        self._rawyield_[0] = self._data_handler_.get_norm()
        self._rawyield_err_[0] = np.sqrt(self._rawyield_[0])
    else:
        for i_pdf, _ in enumerate(self._signal_pdf_):
            self._rawyield_[i_pdf] = signal_fracs[i_pdf] * norm / self._ratio_truncated_
            self._rawyield_err_[i_pdf] = signal_frac_errs[i_pdf] * norm / self._ratio_truncated_

    return self._fit_result_

Perform a mass fit with the zfit library

Parameters

do_prefit : bool
Flag to enable prefit to background only (supported for unbinned fits only) The expected signal regions must be provided.
**kwargs : dict

Additional optional arguments:

  • prefit_excluded_regions: list List or list of lists with limits for the excluded regions

  • prefit_exclude_nsigma: Alternative argument to prefit_excluded_regions that excludes the regions around the signals, using the initialised mean and sigma parameters

Returns

fit_result : zfit.minimizers.fitresult.FitResult
The fit result
def plot_mass_fit(self, **kwargs)
Expand source code
def plot_mass_fit(self, **kwargs):
    """
    Plot the mass fit

    Parameters
    -------------------------------------------------
    **kwargs: dict
        Additional optional arguments:

        - style: str
            style to be used (see https://github.com/scikit-hep/mplhep for more details)

        - logy: bool
            log scale in y axis

        - figsize: tuple
            size of the figure

        - axis_title: str
            x-axis title

        - show_extra_info: bool
            show mu, sigma, chi2/ndf, signal, bkg, signal/bkg, significance

        - extra_info_massnsigma: float
            number of sigmas for extra info

        - extra_info_massnhwhm: float
            number of hwhms for extra info (alternative to extra_info_massnsigma)

        - extra_info_massrange: list
            mass range limits for extra info (alternative to extra_info_massnsigma)

        - extra_info_loc: list
            location of extra info (one for chi2 and one for other info)

        - legend_loc: str
            location of the legend

        - num: int
            number of bins to plot pdfs converted into histograms

    Returns
    -------------------------------------------------
    fig: matplotlib.figure.Figure
        figure containing the mass fit plot
    """

    style = kwargs.get('style', 'LHCb2')
    logy = kwargs.get('logy', False)
    figsize = kwargs.get('figsize', (7, 7))
    bins = self._data_handler_.get_nbins()
    axis_title = kwargs.get('axis_title', self._data_handler_.get_var_name())
    show_extra_info = kwargs.get('show_extra_info', False)
    num = kwargs.get('num', 10000)
    mass_range = kwargs.get('extra_info_massrange', None)
    nhwhm = kwargs.get('extra_info_massnhwhm', None)
    nsigma = kwargs.get('extra_info_massnsigma', 3)
    info_loc = kwargs.get('extra_info_loc', ['lower right', 'upper left'])
    legend_loc = kwargs.get('legend_loc', 'best')

    mplhep.style.use(style)

    obs = self._data_handler_.get_obs()
    limits = self._data_handler_.get_limits()

    fig, axs = plt.subplots(figsize=figsize)

    hdata = self._data_handler_.to_hist(lower_edge=limits[0],
                                        upper_edge=limits[1],
                                        nbins=bins,
                                        varname=self._data_handler_.get_var_name())

    hdata.plot(yerr=True, color='black', histtype='errorbar', label='data')
    bin_sigma = (limits[1] - limits[0]) / bins
    norm_total_pdf = self._total_pdf_norm_ * bin_sigma

    x_plot = np.linspace(limits[0], limits[1], num=num)
    total_func = zfit.run(self._total_pdf_.pdf(x_plot, norm_range=obs))
    signal_funcs, refl_funcs, bkg_funcs = ([] for _ in range(3))

    for signal_pdf in self._signal_pdf_:
        signal_funcs.append(zfit.run(signal_pdf.pdf(x_plot, norm_range=obs)))
    for refl_pdf in self._refl_pdf_:
        refl_funcs.append(zfit.run(refl_pdf.pdf(x_plot, norm_range=obs)))
    for bkg_pdf in self._background_pdf_:
        bkg_funcs.append(zfit.run(bkg_pdf.pdf(x_plot, norm_range=obs)))

    signal_fracs, bkg_fracs, refl_fracs, _, _, _ = self.__get_all_fracs()

    # first draw backgrounds
    for ibkg, (bkg_func, bkg_frac) in enumerate(zip(bkg_funcs, bkg_fracs)):
        plt.plot(x_plot, bkg_func * norm_total_pdf * bkg_frac / self._ratio_truncated_,
                 color=self._bkg_cmap_(ibkg), ls='--', label=self.label_bkg_pdf[ibkg])
    # then draw signals
    for isgn, (signal_func, frac) in enumerate(zip(signal_funcs, signal_fracs)):
        plt.plot(x_plot, signal_func * norm_total_pdf * frac / self._ratio_truncated_, color=self._sgn_cmap_(isgn))
        plt.fill_between(x_plot, signal_func * norm_total_pdf * frac / self._ratio_truncated_,
                         color=self._sgn_cmap_(isgn),
                         alpha=0.5, label=self.label_signal_pdf[isgn])

    # finally draw reflected signals (if any)
    for irefl, (refl_func, frac) in enumerate(zip(refl_funcs, refl_fracs)):
        if self._name_refl_pdf_[irefl] is None:
            continue
        plt.plot(x_plot, refl_func * norm_total_pdf * frac / self._ratio_truncated_, color=self._refl_cmap_(irefl))
        plt.fill_between(x_plot, refl_func * norm_total_pdf * frac / self._ratio_truncated_,
                         color=self._refl_cmap_(irefl),
                         alpha=0.5, label=f'reflected signal {irefl}')

    plt.plot(x_plot, total_func * norm_total_pdf, color='xkcd:blue', label='total fit')
    plt.xlim(limits[0], limits[1])
    plt.xlabel(axis_title)
    plt.ylabel(rf'counts / {(limits[1]-limits[0])/bins*1000:0.1f} MeV/$c^2$')
    if logy:
        plt.yscale('log')
        plt.ylim(min(total_func) * norm_total_pdf / 5, max(total_func) * norm_total_pdf * 5)
    else:
        plt.ylim(0., max(total_func) * norm_total_pdf * 1.5)

    if show_extra_info:
        # info on chi2/ndf
        chi2 = self.get_chi2()
        ndf = self.get_ndf()
        anchored_text_chi2 = AnchoredText(fr'$\chi^2 / \mathrm{{ndf}} =${chi2:.2f} / {ndf}',
                                          loc=info_loc[0],
                                          frameon=False)
        # signal and background info for all signals
        text = []
        for idx, _ in enumerate(self._name_signal_pdf_):
            mass, mass_unc = self.get_mass(idx)
            sigma, sigma_unc = None, None
            gamma, gamma_unc = None, None
            rawyield, rawyield_err = self.get_raw_yield(idx=idx)
            if self._name_signal_pdf_[idx] in [
                'gausexptail', 'genergausexptailsymm', 'gaussian', 'crystalball',
                'doublecb', 'doublecbsymm', 'voigtian', 'hist'
                ]:
                sigma, sigma_unc = self.get_sigma(idx)
            if self._name_signal_pdf_[idx] in ['cauchy', 'voigtian']:
                gamma, gamma_unc = self.get_signal_parameter(idx, 'gamma')
            extra_info = fr'{self.label_signal_pdf[idx]}''\n' \
                + fr'  $\mu = {mass*1000:.1f}\pm{mass_unc*1000:.1f}$ MeV$/c^2$''\n'
            if sigma is not None:
                extra_info += fr'  $\sigma = {sigma*1000:.1f}\pm{sigma_unc*1000:.1f}$ MeV$/c^2$''\n'
            if gamma is not None:
                extra_info += fr'  $\Gamma/2 = {gamma*1000:.1f}\pm{gamma_unc*1000:.1f}$ MeV$/c^2$''\n'
            if mass_range is not None:
                bkg, bkg_err = self.get_background(idx=idx, min=mass_range[0], max=mass_range[1])
                s_over_b, s_over_b_err = self.get_signal_over_background(idx=idx, min=mass_range[0],
                                                                         max=mass_range[1])
                signif, signif_err = self.get_significance(idx=idx, min=mass_range[0], max=mass_range[1])
                interval = f'[{mass_range[0]:.3f}, {mass_range[1]:.3f}]'
                extra_info += fr'  $S={rawyield:.0f} \pm {rawyield_err:.0f}$''\n'
                extra_info += fr'  $B({interval})={bkg:.0f} \pm {bkg_err:.0f}$''\n'
                extra_info += fr'  $S/B({interval})={s_over_b:.2f} \pm {s_over_b_err:.2f}$''\n'
                extra_info += fr'  Signif.$({interval})={signif:.1f} \pm {signif_err:.1f}$'
            elif nhwhm is not None:
                bkg, bkg_err = self.get_background(idx=idx, nhwhm=nhwhm)
                s_over_b, s_over_b_err = self.get_signal_over_background(idx=idx, nhwhm=nhwhm)
                signif, signif_err = self.get_significance(idx=idx, nhwhm=nhwhm)
                extra_info += fr'  $S=${rawyield:.0f} $\pm$ {rawyield_err:.0f}''\n'
                extra_info += fr'  $B({nhwhm}~\mathrm{{HWHM}})=${bkg:.0f} $\pm$ {bkg_err:.0f}''\n'
                extra_info += fr'  $S/B({nhwhm}~\mathrm{{HWHM}})=${s_over_b:.2f} $\pm$ {s_over_b_err:.2f}''\n'
                extra_info += fr'  Signif.$({nhwhm}~\mathrm{{HWHM}})=${signif:.1f} $\pm$ {signif_err:.1f}'
            else:
                bkg, bkg_err = self.get_background(idx=idx, nsigma=nsigma)
                s_over_b, s_over_b_err = self.get_signal_over_background(idx=idx, nsigma=nsigma)
                signif, signif_err = self.get_significance(idx=idx, nsigma=nsigma)
                extra_info += fr'  $S=${rawyield:.0f} $\pm$ {rawyield_err:.0f}''\n'
                extra_info += fr'  $B({nsigma}\sigma)=${bkg:.0f} $\pm$ {bkg_err:.0f}''\n'
                extra_info += fr'  $S/B({nsigma}\sigma)=${s_over_b:.2f} $\pm$ {s_over_b_err:.2f}''\n'
                extra_info += fr'  Signif.$({nsigma}\sigma)=${signif:.1f} $\pm$ {signif_err:.1f}'
            text.append(extra_info)
        concatenated_text = '\n'.join(text)
        anchored_text_signal = AnchoredText(concatenated_text, loc=info_loc[1], frameon=False)

        axs.add_artist(anchored_text_chi2)
        axs.add_artist(anchored_text_signal)

    plt.legend(loc=legend_loc)

    Logger('plot_mass_fit now returns a tuple (fig, axs) !', 'WARNING')

    return fig, axs

Plot the mass fit

Parameters

**kwargs : dict

Additional optional arguments:

  • style: str style to be used (see https://github.com/scikit-hep/mplhep for more details)

  • logy: bool log scale in y axis

  • figsize: tuple size of the figure

  • axis_title: str x-axis title

  • show_extra_info: bool show mu, sigma, chi2/ndf, signal, bkg, signal/bkg, significance

  • extra_info_massnsigma: float number of sigmas for extra info

  • extra_info_massnhwhm: float number of hwhms for extra info (alternative to extra_info_massnsigma)

  • extra_info_massrange: list mass range limits for extra info (alternative to extra_info_massnsigma)

  • extra_info_loc: list location of extra info (one for chi2 and one for other info)

  • legend_loc: str location of the legend

  • num: int number of bins to plot pdfs converted into histograms

Returns

fig : matplotlib.figure.Figure
figure containing the mass fit plot
def plot_raw_residuals(self, **kwargs)
Expand source code
def plot_raw_residuals(self, **kwargs):
    """
    Plot the raw residuals

    Parameters
    -------------------------------------------------
    **kwargs: dict
        Additional optional arguments:

        - style: str
            style to be used (see https://github.com/scikit-hep/mplhep for more details)

        - figsize: tuple
            size of the figure

        - axis_title: str
            x-axis title

    Returns
    -------------------------------------------------
    fig: matplotlib.figure.Figure
        figure containing the raw residuals plot
    """

    style = kwargs.get('style', 'LHCb2')
    figsize = kwargs.get('figsize', (7, 7))
    axis_title = kwargs.get('axis_title', self._data_handler_.get_var_name())

    mplhep.style.use(style)

    obs = self._data_handler_.get_obs()
    limits = self._data_handler_.get_limits()

    fig = plt.figure(figsize=figsize)

    if len(self._raw_residuals_) == 0:
        self.__get_raw_residuals()

    # draw residuals
    plt.errorbar(
        self._data_handler_.get_bin_center(),
        self._raw_residuals_,
        xerr=None,
        yerr=np.sqrt(self._raw_residual_variances_),
        linestyle="None",
        elinewidth=1,
        capsize=0,
        color="black",
        marker="o",
        markersize=5,
        label="residuals"
    )
    bins = self._data_handler_.get_nbins()
    bin_sigma = (limits[1] - limits[0]) / bins
    norm = self._total_pdf_norm_ * bin_sigma

    x_plot = np.linspace(limits[0], limits[1], num=1000)
    signal_funcs, refl_funcs = ([] for _ in range(2))
    for signal_pdf in self._signal_pdf_:
        signal_funcs.append(zfit.run(signal_pdf.pdf(x_plot, norm_range=obs)))
    for refl_pdf in self._refl_pdf_:
        refl_funcs.append(zfit.run(refl_pdf.pdf(x_plot, norm_range=obs)))

    signal_fracs, _, refl_fracs, _, _, _ = self.__get_all_fracs()

    # draw signals
    for isgn, (signal_func, frac) in enumerate(zip(signal_funcs, signal_fracs)):
        plt.plot(x_plot, signal_func * norm * frac / self._ratio_truncated_, color=self._sgn_cmap_(isgn))
        plt.fill_between(x_plot, signal_func * norm * frac / self._ratio_truncated_, color=self._sgn_cmap_(isgn),
                         alpha=0.5, label=self.label_signal_pdf[isgn])

    # finally draw reflected signals (if any)
    is_there_refl = False
    for irefl, (refl_func, frac) in enumerate(zip(refl_funcs, refl_fracs)):
        if self._name_refl_pdf_[irefl] is None:
            continue
        is_there_refl = True
        plt.plot(x_plot, refl_func * norm * frac / self._ratio_truncated_, color=self._refl_cmap_(irefl))
        plt.fill_between(x_plot, refl_func * norm * frac / self._ratio_truncated_, color=self._refl_cmap_(irefl),
                         alpha=0.5, label=f'reflected signal {irefl}')

    # draw signal + reflected signals (if any)
    if is_there_refl:
        for isgn, (signal_func, refl_func, frac_sgn, frac_refl) in enumerate(
                zip(signal_funcs, refl_funcs, signal_fracs, refl_fracs)):
            plt.plot(x_plot, (signal_func * frac_sgn + frac_refl * refl_func) * norm / self._ratio_truncated_,
                     color='xkcd:blue', label='total - bkg')

    plt.xlim(limits[0], limits[1])
    plt.xlabel(axis_title)
    plt.ylabel(rf'(data - fitted bkg) / {(limits[1]-limits[0])/bins*1000:0.1f} MeV/$c^2$')
    plt.legend(loc='best')

    return fig

Plot the raw residuals

Parameters

**kwargs : dict

Additional optional arguments:

Returns

fig : matplotlib.figure.Figure
figure containing the raw residuals plot
def plot_std_residuals(self, **kwargs)
Expand source code
def plot_std_residuals(self, **kwargs):
    """
    Plot the raw residuals

    Parameters
    -------------------------------------------------
    **kwargs: dict
        Additional optional arguments:

        - style: str
            style to be used (see https://github.com/scikit-hep/mplhep for more details)

        - figsize: tuple
            size of the figure

        - axis_title: str
            x-axis title

    Returns
    -------------------------------------------------
    fig: matplotlib.figure.Figure
        figure containing the raw residuals plot
    """

    style = kwargs.get('style', 'LHCb2')
    figsize = kwargs.get('figsize', (7, 7))
    axis_title = kwargs.get('axis_title', self._data_handler_.get_var_name())

    mplhep.style.use(style)

    limits = self._data_handler_.get_limits()
    bins = self._data_handler_.get_nbins()
    bin_center = self._data_handler_.get_bin_center()

    fig = plt.figure(figsize=figsize)

    if len(self._std_residuals_) == 0:
        self.__get_std_residuals()
    # draw residuals
    plt.errorbar(bin_center,
                 self._std_residuals_,
                 xerr=None,
                 yerr=np.sqrt(self._std_residual_variances_),
                 linestyle="None",
                 elinewidth=1,
                 capsize=0,
                 color="black",
                 marker="o",
                 markersize=5,
                 label=None)

    # line at 0
    plt.plot([bin_center[0], bin_center[-1]], [0., 0.], lw=2, color='xkcd:blue')
    # line at -3 sigma
    plt.plot([bin_center[0], bin_center[-1]], [-3., -3.], lw=2, color='xkcd:red')
    # line at 3 sigma
    plt.plot([bin_center[0], bin_center[-1]], [3., 3.], lw=2, color='xkcd:red')

    plt.xlim(limits[0], limits[1])
    plt.xlabel(axis_title)
    plt.ylabel(fr"$\dfrac{{ \mathrm{{data}} - \mathrm{{total \ fit}} }}{{ \sigma_{{ \mathrm{{data}} }} }}$"
               fr"/ {(limits[1]-limits[0])/bins*1000:0.1f} MeV/$c^2$")

    return fig

Plot the raw residuals

Parameters

**kwargs : dict

Additional optional arguments:

Returns

fig : matplotlib.figure.Figure
figure containing the raw residuals plot
def set_background_initpar(self, idx, par_name, init_value, **kwargs)
Expand source code
def set_background_initpar(self, idx, par_name, init_value, **kwargs):
    """
    Set a background parameter

    Parameters
    -------------------------------------------------
    idx: int
        Index of the background
    par_name: str
        The name of the parameter to be set
    init_value: float
        The value of parameter to be set
    **kwargs: dict
        Additional optional arguments:

        - limits: list
            minimum and maximum limits for the parameter

        - fix: bool
            fix the mass parameter
    """
    self._init_bkg_pars_[idx][par_name] = init_value
    if 'limits' in kwargs:
        self._limits_bkg_pars_[idx][par_name] = kwargs['limits']
    if 'fix' in kwargs:
        self._fix_bkg_pars_[idx][par_name] = kwargs['fix']

Set a background parameter

Parameters

idx : int
Index of the background
par_name : str
The name of the parameter to be set
init_value : float
The value of parameter to be set
**kwargs : dict

Additional optional arguments:

  • limits: list minimum and maximum limits for the parameter

  • fix: bool fix the mass parameter

def set_background_kde(self, idx, sample, **kwargs)
Expand source code
def set_background_kde(self, idx, sample, **kwargs):
    """
    Set sample and options for background kde

    Parameters
    -------------------------------------------------
    idx: int
        Index of the background
    sample: flarefly.DataHandler
        Data sample for Kernel Density Estimation
    **kwargs: dict
        Arguments for kde options. See
        https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid
        for more details
    """

    self._kde_bkg_sample_[idx] = sample
    self._kde_bkg_option_[idx] = kwargs

Set sample and options for background kde

Parameters

idx : int
Index of the background
sample : flarefly.DataHandler
Data sample for Kernel Density Estimation
**kwargs : dict
Arguments for kde options. See https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid for more details
def set_background_template(self, idx, sample)
Expand source code
def set_background_template(self, idx, sample):
    """
    Set sample and options for background template histogram

    Parameters
    -------------------------------------------------
    idx: int
        Index of the background
    sample: flarefly.DataHandler
        Data sample for template histogram
    """

    limits_bkg = sample.get_limits()
    limits_data = self._data_handler_.get_limits()
    if not np.isclose(limits_bkg[0], limits_data[0], rtol=1e-05, atol=1e-08, equal_nan=False) or \
            not np.isclose(limits_bkg[1], limits_data[1], rtol=1e-05, atol=1e-08, equal_nan=False):
        Logger(f'The data and the background template {idx} have different limits:'
               f' \n       -> background template: {limits_bkg}, data -> {limits_data}', 'FATAL')

    self._hist_bkg_sample_[idx] = sample

Set sample and options for background template histogram

Parameters

idx : int
Index of the background
sample : flarefly.DataHandler
Data sample for template histogram
def set_particle_mass(self, idx, **kwargs)
Expand source code
def set_particle_mass(self, idx, **kwargs):
    """
    Set the particle mass

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal
    **kwargs: dict
        Additional optional arguments:

        - mass: float
            The mass of the particle

        - pdg_id: int
            PDG ID of the particle (alternative to mass)

        - pdg_name: str
            Name of the particle (alternative to mass)

        - limits: list
            minimum and maximum limits for the mass parameter

        - fix: bool
            fix the mass parameter
    """
    mass_name = 'm' if self._name_signal_pdf_[idx] in ['cauchy', 'voigtian'] else 'mu'
    mass = 0.
    if 'mass' in kwargs:
        mass = kwargs['mass']
    elif 'pdg_id' in kwargs:
        mass = self.pdg_api.get_particle_by_mcid(kwargs['pdg_id']).mass
    elif 'pdg_name' in kwargs:
        mass = self.pdg_api.get_particle_by_name(kwargs['pdg_name']).mass
    else:
        Logger(f'"mass", "pdg_id", and "pdg_name" not provided, mass value for signal {idx} will not be set',
               'ERROR')
    self._init_sgn_pars_[idx][mass_name] = mass
    if 'limits' in kwargs:
        self._limits_sgn_pars_[idx][mass_name] = kwargs['limits']
    if 'fix' in kwargs:
        self._fix_sgn_pars_[idx][mass_name] = kwargs['fix']

Set the particle mass

Parameters

idx : int
Index of the signal
**kwargs : dict

Additional optional arguments:

  • mass: float The mass of the particle

  • pdg_id: int PDG ID of the particle (alternative to mass)

  • pdg_name: str Name of the particle (alternative to mass)

  • limits: list minimum and maximum limits for the mass parameter

  • fix: bool fix the mass parameter

def set_reflection_kde(self, idx, sample, r_over_s, **kwargs)
Expand source code
def set_reflection_kde(self, idx, sample, r_over_s, **kwargs):
    """
    Set sample and options for reflected signal kde

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal
    sample: flarefly.DataHandler
        Data sample for Kernel Density Estimation
    r_over_s: float
        R/S ratio
    **kwargs: dict
        Arguments for kde options. See
        https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid
        for more details
    """

    self._kde_signal_sample_[idx] = sample
    self._kde_signal_option_[idx] = kwargs
    self._refl_over_sgn_[idx] = r_over_s

Set sample and options for reflected signal kde

Parameters

idx : int
Index of the signal
sample : flarefly.DataHandler
Data sample for Kernel Density Estimation
r_over_s : float
R/S ratio
**kwargs : dict
Arguments for kde options. See https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid for more details
def set_reflection_template(self, idx, sample, r_over_s)
Expand source code
def set_reflection_template(self, idx, sample, r_over_s):
    """
    Set sample and options for reflected signal template

    Parameters
    -------------------------------------------------
    idx: int
        Index of the reflected signal
    sample: flarefly.DataHandler
        Data sample for histogram template
    r_over_s: float
        R/S ratio
    """

    self._hist_refl_sample_[idx] = sample
    self._refl_over_sgn_[idx] = r_over_s

Set sample and options for reflected signal template

Parameters

idx : int
Index of the reflected signal
sample : flarefly.DataHandler
Data sample for histogram template
r_over_s : float
R/S ratio
def set_signal_initpar(self, idx, par_name, init_value, **kwargs)
Expand source code
def set_signal_initpar(self, idx, par_name, init_value, **kwargs):
    """
    Set a signal parameter

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal
    par_name: str
        The name of the parameter to be set
    init_value: float
        The value of parameter to be set
    **kwargs: dict
        Additional optional arguments:

        - limits: list
            minimum and maximum limits for the parameter

        - fix: bool
            fix the parameter to init_value
    """
    if par_name == 'gamma':
        Logger('The gamma parameter that you are setting is half of the width of a resonance,'
               ' for more info check the Cauchy pdf defined here '
               'https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/basic/zfit.pdf.Cauchy.html',
               'WARNING')

    self._init_sgn_pars_[idx][par_name] = init_value
    if 'limits' in kwargs:
        self._limits_sgn_pars_[idx][par_name] = kwargs['limits']
    if 'fix' in kwargs:
        self._fix_sgn_pars_[idx][par_name] = kwargs['fix']

Set a signal parameter

Parameters

idx : int
Index of the signal
par_name : str
The name of the parameter to be set
init_value : float
The value of parameter to be set
**kwargs : dict

Additional optional arguments:

  • limits: list minimum and maximum limits for the parameter

  • fix: bool fix the parameter to init_value

def set_signal_kde(self, idx, sample, **kwargs)
Expand source code
def set_signal_kde(self, idx, sample, **kwargs):
    """
    Set sample and options for signal kde

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal
    sample: flarefly.DataHandler
        Data sample for Kernel Density Estimation
    **kwargs: dict
        Arguments for kde options. See
        https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid
        for more details
    """

    self._kde_signal_sample_[idx] = sample
    self._kde_signal_option_[idx] = kwargs

Set sample and options for signal kde

Parameters

idx : int
Index of the signal
sample : flarefly.DataHandler
Data sample for Kernel Density Estimation
**kwargs : dict
Arguments for kde options. See https://zfit.readthedocs.io/en/latest/user_api/pdf/_generated/kde_api/zfit.pdf.KDE1DimGrid.html#zfit.pdf.KDE1DimGrid for more details
def set_signal_template(self, idx, sample)
Expand source code
def set_signal_template(self, idx, sample):
    """
    Set sample and options for signal template

    Parameters
    -------------------------------------------------
    idx: int
        Index of the signal
    sample: flarefly.DataHandler
        Data sample for histogram template
    """

    self._hist_signal_sample_[idx] = sample

Set sample and options for signal template

Parameters

idx : int
Index of the signal
sample : flarefly.DataHandler
Data sample for histogram template