Source code for zfit.models.interpolation

#  Copyright (c) 2023 zfit
from __future__ import annotations

import zfit.z.numpy as znp
from .functor import BaseFunctor
from ..core.interfaces import ZfitBinnedPDF
from ..core.space import supports
from ..util import ztyping
from ..util.exception import SpecificFunctionNotImplemented
from ..util.ztyping import ExtendedInputType, NormInputType
from ..z.interpolate_spline import interpolate_spline


[docs] class SplinePDF(BaseFunctor): def __init__( self, pdf: ZfitBinnedPDF, order: int | None = None, obs: ztyping.ObsTypeInput = None, extended: ExtendedInputType = None, norm: NormInputType = None, ) -> None: """Spline interpolate a binned PDF in order to get a smooth, unbinned PDF. Args: pdf: Binned PDF that will be interpolated. order: Spline interpolation order. Default is 3 obs: Unbinned observable. If not given, the observable of the pdf is used without the binning. 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| The normalization of the PDF. If this is parameter-like, it will be used as the """ if extended is None: extended = pdf.is_extended if extended is True: extended = pdf.get_yield() self._automatically_extended = True else: self._automatically_extended = False if obs is None: obs = pdf.space obs = obs.with_binning(None) super().__init__(pdfs=pdf, obs=obs, extended=extended, norm=norm) if order is None: order = 3 self._order = order @property def order(self): return self._order @supports(norm=True) def _ext_pdf(self, x, norm): if not self._automatically_extended: raise SpecificFunctionNotImplemented pdf = self.pdfs[0] density = pdf.ext_pdf(x.space, norm=norm) density_flat = znp.reshape(density, (-1,)) centers_list = znp.meshgrid(*pdf.space.binning.centers, indexing="ij") centers_list_flat = [znp.reshape(cent, (-1,)) for cent in centers_list] centers = znp.stack(centers_list_flat, axis=-1) # [None, :, None] # TODO: only 1 dim now probs = interpolate_spline( train_points=centers[None, ...], train_values=density_flat[None, :, None], query_points=x.value()[None, ...], order=self.order, ) return probs[0, ..., 0] @supports(norm=True) def _pdf(self, x, norm): pdf = self.pdfs[0] density = pdf.pdf( x.space, norm=norm ) # TODO: order? Give obs, pdf makes order and binning herself? centers = pdf.space.binning.centers[0][None, :, None] # TODO: only 1 dim now probs = interpolate_spline( train_points=centers, train_values=density[None, :, None], query_points=x.value()[None, ...], order=3, ) return probs[0, ..., 0]