Source code for zfit.models.histmodifier

#  Copyright (c) 2025 zfit

from __future__ import annotations

import typing
from collections.abc import Mapping

import tensorflow as tf

import zfit.z.numpy as znp
from zfit._interfaces import ZfitBinnedPDF

from ..core.space import supports
from ..util import ztyping
from ..util.exception import SpecificFunctionNotImplemented
from .binned_functor import BaseBinnedFunctorPDF

if typing.TYPE_CHECKING:
    import zfit  # noqa: F401




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  # noqa: PLC0415

            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  # noqa: PLC0415

                def sumfunc(params):
                    del params  # unused
                    values = self.counts()
                    return znp.sum(values)

                from zfit.core.parameter import get_auto_number  # noqa: PLC0415

                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)