Merge pull request #64 from ikrommyd/feat-cruijff

feat: add `Cruijff` PDF

CHANGELOG.rst

@@ -8,6 +8,7 @@ Develop
 Major Features and Improvements
 -------------------------------
 - added CMSShape PDF
+- added Cruijff PDF
 
 Breaking changes
 ------------------

tests/test_pdf_cruijff.py

@@ -0,0 +1,88 @@
+"""Tests for Cruijff PDF."""
+
+import numpy as np
+import pytest
+import tensorflow as tf
+import zfit
+from numba_stats import cruijff as cruijff_numba
+from scipy import integrate
+from zfit.core.testing import tester
+
+import zfit_physics as zphys
+
+mu_true = 90.0
+sigmal_true = 5.0
+alphal_true = 3.0
+sigmar_true = 10.0
+alphar_true = 2.0
+
+
+def create_cruijff(mu, sigmal, alphal, sigmar, alphar, limits):
+    obs = zfit.Space("obs1", limits=limits)
+    cruijff = zphys.pdf.Cruijff(mu=mu, sigmal=sigmal, alphal=alphal, sigmar=sigmar, alphar=alphar, obs=obs, norm=False)
+    return cruijff, obs
+
+
+def test_cruijff_pdf():
+    # Test PDF here
+    cruijff, _ = create_cruijff(
+        mu=mu_true, sigmal=sigmal_true, alphal=alphal_true, sigmar=sigmar_true, alphar=alphar_true, limits=(50, 130)
+    )
+    assert cruijff.pdf(90.0).numpy() == pytest.approx(
+        cruijff_numba.density(
+            90.0,
+            beta_left=alphal_true,
+            beta_right=alphar_true,
+            loc=mu_true,
+            scale_left=sigmal_true,
+            scale_right=sigmar_true,
+        ).item(),
+        rel=1e-8,
+    )
+    np.testing.assert_allclose(
+        cruijff.pdf(tf.range(50.0, 130, 10_000)),
+        cruijff_numba.density(
+            tf.range(50.0, 130, 10_000).numpy(),
+            beta_left=alphal_true,
+            beta_right=alphar_true,
+            loc=mu_true,
+            scale_left=sigmal_true,
+            scale_right=sigmar_true,
+        ),
+        rtol=1e-8,
+    )
+    assert cruijff.pdf(tf.range(50.0, 130, 10_000)) <= cruijff.pdf(90.0)
+
+
+def test_cruihff_integral():
+    # Test CDF and integral here
+    cruijff, obs = create_cruijff(
+        mu=mu_true, sigmal=sigmal_true, alphal=alphal_true, sigmar=sigmar_true, alphar=alphar_true, limits=(50, 130)
+    )
+    full_interval_numeric = cruijff.numeric_integrate(obs, norm=False).numpy()
+    true_integral = 67.71494
+    numba_stats_full_integral = integrate.quad(
+        cruijff_numba.density, 50, 130, args=(alphal_true, alphar_true, mu_true, sigmal_true, sigmar_true)
+    )[0]
+    assert full_interval_numeric == pytest.approx(true_integral, 1e-7)
+    assert full_interval_numeric == pytest.approx(numba_stats_full_integral, 1e-7)
+
+    numeric_integral = cruijff.numeric_integrate(limits=(80, 100), norm=False).numpy()
+    numba_stats_integral = integrate.quad(
+        cruijff_numba.density, 80, 100, args=(alphal_true, alphar_true, mu_true, sigmal_true, sigmar_true)
+    )[0]
+    assert numeric_integral == pytest.approx(numba_stats_integral, 1e-7)
+
+
+# register the pdf here and provide sets of working parameter configurations
+def cruijff_params_factory():
+    mu = zfit.Parameter("mu", mu_true)
+    sigmal = zfit.Parameter("sigmal", sigmal_true)
+    alphal = zfit.Parameter("alphal", alphal_true)
+    sigmar = zfit.Parameter("sigmar", sigmar_true)
+    alphar = zfit.Parameter("alphar", alphar_true)
+
+    return {"mu": mu, "sigmal": sigmal, "alphal": alphal, "sigmar": sigmar, "alphar": alphar}
+
+
+tester.register_pdf(pdf_class=zphys.pdf.Cruijff, params_factories=cruijff_params_factory)

zfit_physics/models/pdf_cruijff.py

@@ -0,0 +1,105 @@
+from typing import Optional
+
+import zfit
+from zfit import z
+from zfit.util import ztyping
+from zfit.z import numpy as znp
+
+
+@z.function(wraps="tensor")
+def cruijff_pdf_func(x, mu, sigmal, alphal, sigmar, alphar):
+    """Calculate the Cruijff PDF.
+
+    Args:
+        x: value(s) for which the PDF will be calculated.
+        mu: Mean value
+        sigmal: Left width parameter.
+        alphal: Left tail acceleration parameter.
+        sigmar: Right width parameter.
+        alphar: Right tail acceleration parameter.
+
+    Returns:
+        `tf.Tensor`: The calculated PDF values.
+
+    Notes:
+        Implementation from https://arxiv.org/abs/1005.4087, citation 22, and http://hdl.handle.net/1794/9022, Equation IV.3
+    """
+    cond = znp.less(x, mu)
+
+    # compute only once (in graph this _may_ be optimized anyways, but surely not in eager)
+    xminmu = x - mu
+    tleft = znp.square(xminmu / sigmal)
+    tright = znp.square(xminmu / sigmar)
+    exponent = znp.where(
+        cond,
+        tleft / (1 + alphal * tleft),
+        tright / (1 + alphar * tright),
+    )
+    value = znp.exp(-0.5 * exponent)
+    return value
+
+
+class Cruijff(zfit.pdf.BasePDF):
+    _N_OBS = 1
+
+    def __init__(
+        self,
+        mu: ztyping.ParamTypeInput,
+        sigmal: ztyping.ParamTypeInput,
+        alphal: ztyping.ParamTypeInput,
+        sigmar: ztyping.ParamTypeInput,
+        alphar: ztyping.ParamTypeInput,
+        obs: ztyping.ObsTypeInput,
+        *,
+        extended: Optional[ztyping.ExtendedInputType] = False,
+        norm: Optional[ztyping.NormInputType] = None,
+        name: str = "Cruijff",
+    ):
+        """Cruijff PDF, a Gaussian with two width, left and right, and non-Gaussian tails.
+
+        Implementation from https://arxiv.org/abs/1005.4087, citation 22, and http://hdl.handle.net/1794/9022, Equation IV.3
+
+        .. math:
+
+            f(x; \\mu, \\sigma_{L}, \\alpha_{L}, \\sigma_{R}, \\alpha_{R}) = \\begin{cases}
+            \\exp(- \\frac{(x-\\mu)^2}{2 \\sigma_{L}^2 + \\alpha_{L} (x-\\mu)^2)}, \\mbox{for} x \\leqslant mu \\newline
+            \\exp(- \\frac{(x-\\mu)^2}{2 \\sigma_{R}^2 + \\alpha_{R} (x-\\mu)^2)}, \\mbox{for} x > mu
+
+        Args:
+            mu: Mean value
+            sigmal: Left width parameter.
+            alphal: Left tail acceleration parameter.
+            sigmar: Right width parameter.
+            alphar: Right tail acceleration parameter.
+        obs: |@doc:pdf.init.obs| Observables of the
+           model. This will be used as the default space of the PDF and,
+           if not given explicitly, as the normalization range.
+
+           The default space is used for example in the sample method: if no
+           sampling limits are given, the default space is used.
+
+           The observables are not equal to the domain as it does not restrict or
+           truncate the model outside this range. |@docend:pdf.init.obs|
+        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.name| Human-readable name
+           or label of
+           the PDF for better identification.
+           Has no programmatical functional purpose as identification. |@docend:pdf.init.name|
+        """
+        params = {"mu": mu, "sigmal": sigmal, "alphal": alphal, "sigmar": sigmar, "alphar": alphar}
+        super().__init__(obs=obs, params=params, extended=extended, norm=norm)
+
+    def _unnormalized_pdf(self, x):
+        mu = self.params["mu"]
+        sigmal = self.params["sigmal"].value()
+        alphal = self.params["alphal"].value()
+        sigmar = self.params["sigmar"].value()
+        alphar = self.params["alphar"].value()
+        x = z.unstack_x(x)
+        return cruijff_pdf_func(x=x, mu=mu, sigmal=sigmal, alphal=alphal, sigmar=sigmar, alphar=alphar)

zfit_physics/pdf.py

@@ -1,5 +1,6 @@
 from .models.pdf_argus import Argus
 from .models.pdf_cmsshape import CMSShape
+from .models.pdf_cruijff import Cruijff
 from .models.pdf_relbw import RelativisticBreitWigner
 
-__all__ = ["Argus", "RelativisticBreitWigner", "CMSShape"]
+__all__ = ["Argus", "RelativisticBreitWigner", "CMSShape", "Cruijff"]