Source code for zfit.models.histmodifier

#  Copyright (c) 2024 zfit

from __future__ import annotations

from typing import TYPE_CHECKING

    pass  # ruff: noqa: F401

from import Mapping

import tensorflow as tf

import zfit.z.numpy as znp

from ..core.interfaces import ZfitBinnedPDF
from import supports
from ..util import ztyping
from ..util.exception import SpecificFunctionNotImplemented
from .binned_functor import BaseBinnedFunctorPDF

[docs] class BinwiseScaleModifier(BaseBinnedFunctorPDF): def __init__( self, pdf: ZfitBinnedPDF, modifiers: bool | Mapping[str, ztyping.ParamTypeInput] | None = None, extended: ztyping.ExtendedInputType = None, norm: ztyping.NormInputType = None, name: str | None = "BinnedTemplatePDF", label: str | None = None, ) -> None: """Modifier that scales each bin separately of the *pdf*. Binwise modification can be used to account for uncorrelated or correlated uncertainties. Args: pdf: Binned pdf to be modified. modifiers: Modifiers for each bin. extended: |@doc:pdf.init.extended| The overall yield of the PDF. If this is parameter-like, it will be used as the yield, the expected number of events, and the PDF will be extended. An extended PDF has additional functionality, such as the ``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended| norm: |@doc:pdf.init.norm| Normalization of the PDF. By default, this is the same as the default space of the PDF. |@docend:pdf.init.norm| name: |@doc:pdf.init.label| Human-readable name or label of the PDF for a better description, to be used with plots etc. Has no programmatical functional purpose as identification. |@docend:pdf.init.label| label: |@doc:pdf.init.label| Human-readable name or label of the PDF for a better description, to be used with plots etc. Has no programmatical functional purpose as identification. |@docend:pdf.init.label| """ obs = if not isinstance(pdf, ZfitBinnedPDF): msg = "pdf must be a BinnedPDF" raise TypeError(msg) if extended is None: extended = pdf.is_extended if modifiers is None: modifiers = True if modifiers is True: import zfit modifiers = { f"sysshape_{i}": zfit.Parameter(f"auto_sysshape_{self}_{i}", 1.0) for i in range(pdf.counts(obs).shape.num_elements()) } if not isinstance(modifiers, dict): msg = "modifiers must be a dict-like object or True or None" raise TypeError(msg) params = modifiers.copy() self._binwise_modifiers = modifiers if extended is True: self._automatically_extended = True if modifiers: import zfit def sumfunc(params): del params # unused values = self.counts() return znp.sum(values) from zfit.core.parameter import get_auto_number params_sumfunc = modifiers.copy() dep_params = {} for p in pdf.get_params(): while True: i = get_auto_number() name_tmp = f"DUMMPY_PARAM_{i}" if name_tmp not in params_sumfunc: dep_params[name_tmp] = p break params_sumfunc.update(dep_params) extended = zfit.ComposedParameter( f"AUTO_binwise_modifier_{get_auto_number()}", sumfunc, params=params_sumfunc, ) else: extended = pdf.get_yield() elif extended is not False: self._automatically_extended = False super().__init__(obs=obs, name=name, params=params, models=pdf, extended=extended, norm=norm, label=label) @supports(norm=True) def _counts(self, x, norm=None): if not self._automatically_extended: raise SpecificFunctionNotImplemented return self._counts_with_modifiers(x, norm) def _counts_with_modifiers(self, x, norm): values = self.pdfs[0].counts(x, norm=norm) if modifiers := list(self._binwise_modifiers.values()): sysshape_flat = tf.stack(modifiers) modifiers = znp.reshape(sysshape_flat, values.shape) values = values * modifiers return values @supports(norm="space") def _rel_counts(self, x, norm=None): values = self._counts_with_modifiers(x, norm) return values / znp.sum(values)