Module flarefly.pdf_builder

PDFBuilder class to create signal and background PDFs using zfit.

Classes

class PDFBuilder
Expand source code 
class PDFBuilder:
    """Class to build signal and background PDFs using zfit."""

    def __init__(self):
        self.pdg_api = pdg.connect()

    @staticmethod
    def build_signal_pdf(
        pdf: F2PDFBase,
        obs: zfit.Space,
        name: str,
        ipdf: int,
    ) -> zfit.pdf.BasePDF:
        """Build a signal PDF with configurable parameters.

        Args:
            pdf: The PDF to build
            obs: The observable space for the PDF
            name: Base name for parameters
            ipdf: Index of the PDF
        """
        if pdf.is_kde():
            pdf = PDFBuilder.build_signal_kde(
                pdf,
                name,
                ipdf
            )
            return
        if pdf.is_hist():
            pdf = PDFBuilder.build_signal_hist(
                pdf,
                obs,
                name,
                ipdf
            )
            return

        config = get_signal_pdf_config(pdf.kind)

        # Update the input dictionaries with default values
        PDFBuilder._update_with_defaults(config, pdf)

        parameters = {}
        for par_name in config['parameters'].keys():
            param_name = f'{name}_{par_name}_signal{ipdf}'
            parameters[param_name] = zfit.Parameter(
                name=param_name,
                value=pdf.get_init_par(par_name),
                lower=pdf.get_limits_par(par_name)[0],
                upper=pdf.get_limits_par(par_name)[1],
                floating=not pdf.get_fix_par(par_name)
            )

        pdf_args = {'obs': obs}
        for arg_name in config['pdf_args']:
            if 'args_mapping' in config:  # flarefly -> zfit argument name mapping
                param_key = f"{name}_{config['args_mapping'][arg_name]}_signal{ipdf}"
            else:
                param_key = f'{name}_{arg_name}_signal{ipdf}'
            pdf_args[arg_name] = parameters[param_key]

        pdf.pdf = config['pdf_class'](**pdf_args)
        pdf.parameters = parameters

        if pdf.at_threshold:
            # pion mass as default
            pdf.set_default_par(
                'massthr',
                init=pdf.pdg_api.get_particle_by_mcid(211).mass,
                fix=True, limits=[None, None]
            )
            pdf.set_default_par('powerthr', init=1., fix=False, limits=[None, None])
            pdf.parameters[f'{name}_massthr_signal{ipdf}'] = zfit.Parameter(
                f'{name}_massthr_signal{ipdf}', pdf.get_init_par('massthr'),
                pdf.get_limits_par('massthr')[0], pdf.get_limits_par('massthr')[1],
                floating=not pdf.get_fix_par('massthr'))
            pdf.parameters[f'{name}_powerthr_signal{ipdf}'] = zfit.Parameter(
                f'{name}_powerthr_signal{ipdf}', pdf.get_init_par('powerthr'),
                pdf.get_limits_par('powerthr')[0], pdf.get_limits_par('powerthr')[1],
                floating=not pdf.get_fix_par('powerthr'))
            signalthr_pdf = cpdf.Pow(
                obs=obs,
                mass=pdf.parameters[f'{name}_massthr_signal{ipdf}'],
                power=pdf.parameters[f'{name}_powerthr_signal{ipdf}']
            )
            pdf.pdf = zfit.pdf.ProductPDF([pdf.pdf, signalthr_pdf], obs=obs)

    @staticmethod
    def build_signal_kde(
        pdf: F2PDFBase,
        name: str,
        ipdf: int,
    ) -> zfit.pdf.BasePDF:
        """Build a signal KDE PDF.

        Args:
            pdf: The PDF to build
            name: Base name for the PDF
            ipdf: Index of the PDF
        """
        if not pdf.kde_sample:
            Logger(f'Missing datasample for Kernel Density Estimation of signal {ipdf}!', 'FATAL')

        kde_options = pdf.kde_option or {}

        pdf.pdf = get_kde_pdf(pdf.kind)(
            data=pdf.kde_sample.get_data(),
            obs=pdf.kde_sample.get_obs(),
            name=f'{name}_kde_signal{ipdf}',
            **kde_options
        )

    @staticmethod
    def build_signal_hist(
        pdf: F2PDFBase,
        obs: zfit.Space,
        name: str,
        ipdf: int,
    ) -> zfit.pdf.BasePDF:
        """Build a signal PDF from a histogram template.

        Args:
            pdf: The PDF to build
            obs: The observable space for the PDF
            name: Base name for the PDF
            ipdf: Index of the PDF

        Returns:
            The constructed PDF
        """
        if not pdf.hist_sample:
            Logger(f'Missing datasample for histogram template of signal {ipdf}!', 'FATAL')

        pdf.pdf = zfit.pdf.SplinePDF(
            zfit.pdf.HistogramPDF(
                pdf.hist_sample.get_binned_data(),
                name=f'{name}_hist_signal{ipdf}'
            ),
            order=3,
            obs=obs
        )

    @staticmethod
    def build_bkg_pdf(
        pdf: F2PDFBase,
        obs: zfit.Space,
        name: str,
        ipdf: int,
    ) -> zfit.pdf.BasePDF:
        """Build a background PDF with configurable parameters.

        Args:
            pdf: The PDF to build
            obs: The observable space for the PDF
            name: Base name for parameters
            ipdf: Index of the PDF
        """
        if pdf.is_kde():
            pdf = PDFBuilder.build_bkg_kde(
                pdf,
                name,
                ipdf
            )
            return
        if pdf.is_hist():
            pdf = PDFBuilder.build_bkg_hist(
                pdf,
                obs,
                name,
                ipdf
            )
            return
        if pdf.kind == PDFType.CHEBPOL:
            # Handle Chebyshev polynomials specially
            PDFBuilder._build_chebyshev_pdf(
                pdf, obs, name, ipdf
            )
            return

        config = get_bkg_pdf_config(pdf.kind)

        # Update the input dictionaries with default values
        PDFBuilder._update_with_defaults(config, pdf)

        parameters = {}
        for par_name in config['parameters'].keys():
            param_name = f'{name}_{par_name}_bkg{ipdf}'
            parameters[param_name] = zfit.Parameter(
                name=param_name,
                value=pdf.get_init_par(par_name),
                lower=pdf.get_limits_par(par_name)[0],
                upper=pdf.get_limits_par(par_name)[1],
                floating=not pdf.get_fix_par(par_name)
            )

        pdf_args = {'obs': obs}
        for arg_name in config['pdf_args']:
            if 'args_mapping' in config:  # flarefly -> zfit argument name mapping
                param_key = f"{name}_{config['args_mapping'][arg_name]}_bkg{ipdf}"
            else:
                param_key = f'{name}_{arg_name}_bkg{ipdf}'
            pdf_args[arg_name] = parameters[param_key]

        pdf.pdf = config['pdf_class'](**pdf_args)
        pdf.parameters = parameters

    @staticmethod
    def _build_chebyshev_pdf(
        pdf: F2PDFBase,
        obs: zfit.Space,
        name: str,
        ipdf: int,
    ) -> Tuple[zfit.pdf.BasePDF, Dict[str, zfit.Parameter]]:
        """Build a Chebyshev polynomial background PDF."""

        # Create parameters for each coefficient
        parameters = {}
        for deg in range(pdf.kind.order + 1):
            par_name = f'c{deg}'
            pdf.set_default_init_par(par_name, 0.1)
            pdf.set_default_limits_par(par_name, [None, None])
            pdf.set_default_fix_par(par_name, False)

            param_name = f'{name}_{par_name}_bkg{ipdf}'
            parameters[param_name] = zfit.Parameter(
                name=param_name,
                value=pdf.get_init_par(par_name),
                lower=pdf.get_limits_par(par_name)[0],
                upper=pdf.get_limits_par(par_name)[1],
                floating=not pdf.get_fix_par(par_name)
            )

        # Prepare Chebyshev arguments
        coeff0 = parameters[f'{name}_c0_bkg{ipdf}']
        bkg_coeffs = [parameters[f'{name}_c{deg}_bkg{ipdf}'] for deg in range(1, pdf.kind.order + 1)]

        pdf.pdf = zfit.pdf.Chebyshev(obs=obs, coeff0=coeff0, coeffs=bkg_coeffs)
        pdf.parameters = parameters

    @staticmethod
    def build_bkg_kde(
        pdf: F2PDFBase,
        name: str,
        ipdf: int,
    ) -> zfit.pdf.BasePDF:
        """Build a background KDE PDF.

        Args:
            pdf_kind: Kind of KDE ('kde_exact', 'kde_grid', 'kde_fft', 'kde_isj')
            kde_sample: The sample data for KDE estimation
            name: Base name for the PDF
            ipdf: Index of the PDF
            kde_options: Additional options for the KDE

        Returns:
            The constructed KDE PDF
        """
        if not pdf.kde_sample:
            Logger(f'Missing datasample for Kernel Density Estimation of background {ipdf}!', 'FATAL')

        kde_options = pdf.kde_option or {}

        pdf.pdf = get_kde_pdf(pdf.kind)(
            data=pdf.kde_sample.get_data(),
            obs=pdf.kde_sample.get_obs(),
            name=f'{name}_kde_bkg{ipdf}',
            **kde_options
        )

    @staticmethod
    def build_bkg_hist(
        pdf: F2PDFBase,
        obs: zfit.Space,
        name: str,
        ipdf: int,
    ) -> zfit.pdf.BasePDF:
        """Build a bkg PDF from a histogram template.

        Args:
            pdf: The PDF to build
            obs: The observable space for the PDF
            name: Base name for the PDF
            ipdf: Index of the PDF
        """
        if not pdf.hist_sample:
            Logger(f'Missing datasample for histogram template of background {ipdf}!', 'FATAL')

        pdf.pdf = zfit.pdf.SplinePDF(
            zfit.pdf.HistogramPDF(
                pdf.hist_sample.get_binned_data(),
                name=f'{name}_hist_bkg{ipdf}'
            ),
            order=3,
            obs=obs
        )

    @staticmethod
    def _update_with_defaults(
        config: Dict[str, Any],
        pdf: F2PDFBase
    ) -> None:
        """Update the parameter dictionaries with default values from config."""
        for par_name, par_config in config['parameters'].items():
            pdf.set_default_init_par(par_name, par_config['init'])
            pdf.set_default_limits_par(par_name, par_config['limits'])
            pdf.set_default_fix_par(par_name, par_config['fix'])

Class to build signal and background PDFs using zfit.

Static methods

def build_bkg_hist(pdf: F2PDFBase,
obs: zfit.core.space.Space,
name: str,
ipdf: int) ‑> zfit.core.basepdf.BasePDF
Expand source code 
@staticmethod
def build_bkg_hist(
    pdf: F2PDFBase,
    obs: zfit.Space,
    name: str,
    ipdf: int,
) -> zfit.pdf.BasePDF:
    """Build a bkg PDF from a histogram template.

    Args:
        pdf: The PDF to build
        obs: The observable space for the PDF
        name: Base name for the PDF
        ipdf: Index of the PDF
    """
    if not pdf.hist_sample:
        Logger(f'Missing datasample for histogram template of background {ipdf}!', 'FATAL')

    pdf.pdf = zfit.pdf.SplinePDF(
        zfit.pdf.HistogramPDF(
            pdf.hist_sample.get_binned_data(),
            name=f'{name}_hist_bkg{ipdf}'
        ),
        order=3,
        obs=obs
    )

Build a bkg PDF from a histogram template.

Args

pdf
The PDF to build
obs
The observable space for the PDF
name
Base name for the PDF
ipdf
Index of the PDF
def build_bkg_kde(pdf: F2PDFBase,
name: str,
ipdf: int) ‑> zfit.core.basepdf.BasePDF
Expand source code 
@staticmethod
def build_bkg_kde(
    pdf: F2PDFBase,
    name: str,
    ipdf: int,
) -> zfit.pdf.BasePDF:
    """Build a background KDE PDF.

    Args:
        pdf_kind: Kind of KDE ('kde_exact', 'kde_grid', 'kde_fft', 'kde_isj')
        kde_sample: The sample data for KDE estimation
        name: Base name for the PDF
        ipdf: Index of the PDF
        kde_options: Additional options for the KDE

    Returns:
        The constructed KDE PDF
    """
    if not pdf.kde_sample:
        Logger(f'Missing datasample for Kernel Density Estimation of background {ipdf}!', 'FATAL')

    kde_options = pdf.kde_option or {}

    pdf.pdf = get_kde_pdf(pdf.kind)(
        data=pdf.kde_sample.get_data(),
        obs=pdf.kde_sample.get_obs(),
        name=f'{name}_kde_bkg{ipdf}',
        **kde_options
    )

Build a background KDE PDF.

Args

pdf_kind
Kind of KDE ('kde_exact', 'kde_grid', 'kde_fft', 'kde_isj')
kde_sample
The sample data for KDE estimation
name
Base name for the PDF
ipdf
Index of the PDF
kde_options
Additional options for the KDE

Returns

The constructed KDE PDF

def build_bkg_pdf(pdf: F2PDFBase,
obs: zfit.core.space.Space,
name: str,
ipdf: int) ‑> zfit.core.basepdf.BasePDF
Expand source code 
@staticmethod
def build_bkg_pdf(
    pdf: F2PDFBase,
    obs: zfit.Space,
    name: str,
    ipdf: int,
) -> zfit.pdf.BasePDF:
    """Build a background PDF with configurable parameters.

    Args:
        pdf: The PDF to build
        obs: The observable space for the PDF
        name: Base name for parameters
        ipdf: Index of the PDF
    """
    if pdf.is_kde():
        pdf = PDFBuilder.build_bkg_kde(
            pdf,
            name,
            ipdf
        )
        return
    if pdf.is_hist():
        pdf = PDFBuilder.build_bkg_hist(
            pdf,
            obs,
            name,
            ipdf
        )
        return
    if pdf.kind == PDFType.CHEBPOL:
        # Handle Chebyshev polynomials specially
        PDFBuilder._build_chebyshev_pdf(
            pdf, obs, name, ipdf
        )
        return

    config = get_bkg_pdf_config(pdf.kind)

    # Update the input dictionaries with default values
    PDFBuilder._update_with_defaults(config, pdf)

    parameters = {}
    for par_name in config['parameters'].keys():
        param_name = f'{name}_{par_name}_bkg{ipdf}'
        parameters[param_name] = zfit.Parameter(
            name=param_name,
            value=pdf.get_init_par(par_name),
            lower=pdf.get_limits_par(par_name)[0],
            upper=pdf.get_limits_par(par_name)[1],
            floating=not pdf.get_fix_par(par_name)
        )

    pdf_args = {'obs': obs}
    for arg_name in config['pdf_args']:
        if 'args_mapping' in config:  # flarefly -> zfit argument name mapping
            param_key = f"{name}_{config['args_mapping'][arg_name]}_bkg{ipdf}"
        else:
            param_key = f'{name}_{arg_name}_bkg{ipdf}'
        pdf_args[arg_name] = parameters[param_key]

    pdf.pdf = config['pdf_class'](**pdf_args)
    pdf.parameters = parameters

Build a background PDF with configurable parameters.

Args

pdf
The PDF to build
obs
The observable space for the PDF
name
Base name for parameters
ipdf
Index of the PDF
def build_signal_hist(pdf: F2PDFBase,
obs: zfit.core.space.Space,
name: str,
ipdf: int) ‑> zfit.core.basepdf.BasePDF
Expand source code 
@staticmethod
def build_signal_hist(
    pdf: F2PDFBase,
    obs: zfit.Space,
    name: str,
    ipdf: int,
) -> zfit.pdf.BasePDF:
    """Build a signal PDF from a histogram template.

    Args:
        pdf: The PDF to build
        obs: The observable space for the PDF
        name: Base name for the PDF
        ipdf: Index of the PDF

    Returns:
        The constructed PDF
    """
    if not pdf.hist_sample:
        Logger(f'Missing datasample for histogram template of signal {ipdf}!', 'FATAL')

    pdf.pdf = zfit.pdf.SplinePDF(
        zfit.pdf.HistogramPDF(
            pdf.hist_sample.get_binned_data(),
            name=f'{name}_hist_signal{ipdf}'
        ),
        order=3,
        obs=obs
    )

Build a signal PDF from a histogram template.

Args

pdf
The PDF to build
obs
The observable space for the PDF
name
Base name for the PDF
ipdf
Index of the PDF

Returns

The constructed PDF

def build_signal_kde(pdf: F2PDFBase,
name: str,
ipdf: int) ‑> zfit.core.basepdf.BasePDF
Expand source code 
@staticmethod
def build_signal_kde(
    pdf: F2PDFBase,
    name: str,
    ipdf: int,
) -> zfit.pdf.BasePDF:
    """Build a signal KDE PDF.

    Args:
        pdf: The PDF to build
        name: Base name for the PDF
        ipdf: Index of the PDF
    """
    if not pdf.kde_sample:
        Logger(f'Missing datasample for Kernel Density Estimation of signal {ipdf}!', 'FATAL')

    kde_options = pdf.kde_option or {}

    pdf.pdf = get_kde_pdf(pdf.kind)(
        data=pdf.kde_sample.get_data(),
        obs=pdf.kde_sample.get_obs(),
        name=f'{name}_kde_signal{ipdf}',
        **kde_options
    )

Build a signal KDE PDF.

Args

pdf
The PDF to build
name
Base name for the PDF
ipdf
Index of the PDF
def build_signal_pdf(pdf: F2PDFBase,
obs: zfit.core.space.Space,
name: str,
ipdf: int) ‑> zfit.core.basepdf.BasePDF
Expand source code 
@staticmethod
def build_signal_pdf(
    pdf: F2PDFBase,
    obs: zfit.Space,
    name: str,
    ipdf: int,
) -> zfit.pdf.BasePDF:
    """Build a signal PDF with configurable parameters.

    Args:
        pdf: The PDF to build
        obs: The observable space for the PDF
        name: Base name for parameters
        ipdf: Index of the PDF
    """
    if pdf.is_kde():
        pdf = PDFBuilder.build_signal_kde(
            pdf,
            name,
            ipdf
        )
        return
    if pdf.is_hist():
        pdf = PDFBuilder.build_signal_hist(
            pdf,
            obs,
            name,
            ipdf
        )
        return

    config = get_signal_pdf_config(pdf.kind)

    # Update the input dictionaries with default values
    PDFBuilder._update_with_defaults(config, pdf)

    parameters = {}
    for par_name in config['parameters'].keys():
        param_name = f'{name}_{par_name}_signal{ipdf}'
        parameters[param_name] = zfit.Parameter(
            name=param_name,
            value=pdf.get_init_par(par_name),
            lower=pdf.get_limits_par(par_name)[0],
            upper=pdf.get_limits_par(par_name)[1],
            floating=not pdf.get_fix_par(par_name)
        )

    pdf_args = {'obs': obs}
    for arg_name in config['pdf_args']:
        if 'args_mapping' in config:  # flarefly -> zfit argument name mapping
            param_key = f"{name}_{config['args_mapping'][arg_name]}_signal{ipdf}"
        else:
            param_key = f'{name}_{arg_name}_signal{ipdf}'
        pdf_args[arg_name] = parameters[param_key]

    pdf.pdf = config['pdf_class'](**pdf_args)
    pdf.parameters = parameters

    if pdf.at_threshold:
        # pion mass as default
        pdf.set_default_par(
            'massthr',
            init=pdf.pdg_api.get_particle_by_mcid(211).mass,
            fix=True, limits=[None, None]
        )
        pdf.set_default_par('powerthr', init=1., fix=False, limits=[None, None])
        pdf.parameters[f'{name}_massthr_signal{ipdf}'] = zfit.Parameter(
            f'{name}_massthr_signal{ipdf}', pdf.get_init_par('massthr'),
            pdf.get_limits_par('massthr')[0], pdf.get_limits_par('massthr')[1],
            floating=not pdf.get_fix_par('massthr'))
        pdf.parameters[f'{name}_powerthr_signal{ipdf}'] = zfit.Parameter(
            f'{name}_powerthr_signal{ipdf}', pdf.get_init_par('powerthr'),
            pdf.get_limits_par('powerthr')[0], pdf.get_limits_par('powerthr')[1],
            floating=not pdf.get_fix_par('powerthr'))
        signalthr_pdf = cpdf.Pow(
            obs=obs,
            mass=pdf.parameters[f'{name}_massthr_signal{ipdf}'],
            power=pdf.parameters[f'{name}_powerthr_signal{ipdf}']
        )
        pdf.pdf = zfit.pdf.ProductPDF([pdf.pdf, signalthr_pdf], obs=obs)

Build a signal PDF with configurable parameters.

Args

pdf
The PDF to build
obs
The observable space for the PDF
name
Base name for parameters
ipdf
Index of the PDF