Source code for zfit.models.histmodifier
# Copyright (c) 2024 zfit
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
pass # ruff: noqa: F401
from collections.abc import Mapping
import tensorflow as tf
import zfit.z.numpy as znp
from ..core.interfaces import ZfitBinnedPDF
from ..core.space 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 = pdf.space
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)