Merge pull request #67 from zfit/ikrommyd/feat-erfexp

feat: add `ErfExp` pdf

CHANGELOG.rst

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

tests/test_pdf_cruijff.py

@@ -19,7 +19,7 @@
 
 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)
+    cruijff = zphys.pdf.Cruijff(mu=mu, sigmal=sigmal, alphal=alphal, sigmar=sigmar, alphar=alphar, obs=obs)
     return cruijff, obs
 
 
@@ -28,7 +28,7 @@ def test_cruijff_pdf():
     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(
+    assert cruijff.pdf(90.0, norm=False).numpy() == pytest.approx(
         cruijff_numba.density(
             90.0,
             beta_left=alphal_true,
@@ -39,8 +39,20 @@ def test_cruijff_pdf():
         ).item(),
         rel=1e-8,
     )
+    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()
+        / 67.71494,
+        rel=1e-7,
+    )
     np.testing.assert_allclose(
-        cruijff.pdf(tf.range(50.0, 130, 10_000)),
+        cruijff.pdf(tf.range(50.0, 130, 10_000), norm=False),
         cruijff_numba.density(
             tf.range(50.0, 130, 10_000).numpy(),
             beta_left=alphal_true,
@@ -51,10 +63,23 @@ def test_cruijff_pdf():
         ),
         rtol=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,
+        )
+        / 67.71494,
+        rtol=1e-8,
+    )
     assert cruijff.pdf(tf.range(50.0, 130, 10_000)) <= cruijff.pdf(90.0)
 
 
-def test_cruihff_integral():
+def test_cruijff_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)

tests/test_pdf_erfexp.py

@@ -0,0 +1,98 @@
+"""Tests for ErfExp PDF."""
+
+import numpy as np
+import pytest
+import tensorflow as tf
+import zfit
+from scipy import integrate, special
+from zfit.core.testing import tester
+
+import zfit_physics as zphys
+
+mu_true = 90.0
+beta_true = 0.08
+gamma_true = -1.0
+n_true = 0.2
+
+mu_range = (65.0, 90.0)
+gamma_range = (-10, 10)
+beta_range = (0.01, 10)
+n_range = (0.1, 0.5)
+
+
+def _erfexp_numpy(x, mu, beta, gamma, n):
+    return special.erfc((x - mu) * beta) * np.exp(-gamma * (np.power(x, n) - np.power(mu, n)))
+
+
+erfexp_numpy = np.vectorize(_erfexp_numpy, excluded=["mu", "beta", "gamma", "n"])
+
+
+def create_erfexp(mu, beta, gamma, n, limits):
+    obs = zfit.Space("obs1", limits=limits)
+    erfexp = zphys.pdf.ErfExp(mu=mu, beta=beta, gamma=gamma, n=n, obs=obs)
+    return erfexp, obs
+
+
+def test_erfexp_pdf():
+    # Test PDF here
+    erfexp, _ = create_erfexp(mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true, limits=(50, 130))
+    assert erfexp.pdf(90.0, norm=False).numpy().item() == pytest.approx(
+        erfexp_numpy(90.0, mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true), rel=1e-8
+    )
+    assert erfexp.pdf(90.0).numpy().item() == pytest.approx(
+        erfexp_numpy(90.0, mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true) / 71.18838, rel=1e-8
+    )
+    np.testing.assert_allclose(
+        erfexp.pdf(tf.range(50.0, 130, 10_000), norm=False),
+        erfexp_numpy(tf.range(50.0, 130, 10_000), mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true),
+        rtol=1e-8,
+    )
+    np.testing.assert_allclose(
+        erfexp.pdf(tf.range(50.0, 130, 10_000)),
+        erfexp_numpy(tf.range(50.0, 130, 10_000), mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true) / 71.18838,
+        rtol=1e-8,
+        atol=1e-8,
+    )
+
+
+def test_erfexp_pdf_random_params():
+    # Test PDF here in a loop with random parameters
+    for _ in range(1000):
+        mu_true = np.random.uniform(*mu_range)
+        beta_true = np.random.uniform(*beta_range)
+        gamma_true = np.random.uniform(*gamma_range)
+        n_true = np.random.uniform(*n_range)
+
+        erfexp, __ = create_erfexp(mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true, limits=(50, 130))
+        np.testing.assert_allclose(
+            erfexp.pdf(tf.range(50.0, 130, 10_000), norm=False),
+            erfexp_numpy(tf.range(50.0, 130, 10_000), mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true),
+            rtol=1e-5,
+        )
+
+
+def test_erfexp_integral():
+    # Test CDF and integral here
+    erfexp, obs = create_erfexp(mu=mu_true, beta=beta_true, gamma=gamma_true, n=n_true, limits=(50, 130))
+    full_interval_numeric = erfexp.numeric_integrate(obs, norm=False).numpy()
+    true_integral = 71.18838
+    numpy_full_integral = integrate.quad(erfexp_numpy, 50, 130, args=(mu_true, beta_true, gamma_true, n_true))[0]
+    assert full_interval_numeric == pytest.approx(true_integral, 1e-7)
+    assert full_interval_numeric == pytest.approx(numpy_full_integral, 1e-7)
+
+    numeric_integral = erfexp.numeric_integrate(limits=(80, 100), norm=False).numpy()
+    numpy_integral = integrate.quad(erfexp_numpy, 80, 100, args=(mu_true, beta_true, gamma_true, n_true))[0]
+    assert numeric_integral == pytest.approx(numpy_integral, 1e-7)
+
+
+# register the pdf here and provide sets of working parameter configurations
+def erfexp_params_factory():
+    mu = zfit.Parameter("mu", mu_true)
+    beta = zfit.Parameter("beta", beta_true)
+    gamma = zfit.Parameter("gamma", gamma_true)
+    n = zfit.Parameter("n", n_true)
+
+    return {"mu": mu, "beta": beta, "gamma": gamma, "n": n}
+
+
+tester.register_pdf(pdf_class=zphys.pdf.ErfExp, params_factories=erfexp_params_factory)

zfit_physics/models/pdf_cruijff.py

@@ -62,8 +62,9 @@ def __init__(
         .. 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 \\end{cases}
+            \\exp{\\left(- \\frac{(x-\\mu)^2}{2 \\sigma_{L}^2 + \\alpha_{L} (x-\\mu)^2}\\right)}, \\mbox{for} x \\leqslant mu \\newline
+            \\exp{\\left(- \\frac{(x-\\mu)^2}{2 \\sigma_{R}^2 + \\alpha_{R} (x-\\mu)^2}\\right)}, \\mbox{for} x > mu
+            \\end{cases}
 
         Args:
             mu: Mean value

zfit_physics/models/pdf_erfexp.py

@@ -0,0 +1,99 @@
+from typing import Optional
+
+import tensorflow as tf
+import zfit
+from zfit import z
+from zfit.util import ztyping
+from zfit.z import numpy as znp
+
+
+@z.function(wraps="tensor")
+def erfexp_pdf_func(x, mu, beta, gamma, n):
+    """Calculate the ErfExp PDF.
+
+    Args:
+        x: value(s) for which the PDF will be calculated.
+        mu: Location parameter.
+        beta: Scale parameter.
+        gamma: Shape parameter.
+        n: Shape parameter.
+
+    Returns:
+        `tf.Tensor`: The calculated PDF values.
+
+    Notes:
+        Implementation from https://gitlab.cern.ch/cms-muonPOG/spark_tnp/-/blob/Spark3/RooErfExp.cc
+        The parameters beta and gamma are given in reverse order in this c++ implementation.
+    """
+    return tf.math.erfc((x - mu) * beta) * znp.exp(-gamma * (znp.power(x, n) - znp.power(mu, n)))
+
+
+# Note: There is no analytic integral for the ErfExp PDF
+# We tried with sympy, Mathematica, Wolfram Alpha and https://www.integral-calculator.com/
+# import sympy as sp
+#
+# # Define symbols
+# x, mu, beta, gamma, n = sp.symbols('x mu beta gamma n', real=True)
+#
+# # Define the function
+# func = sp.erfc((x - mu) * beta) * sp.exp(-gamma * (x**n - mu**n))
+# sp.integrate(func, x)
+class ErfExp(zfit.pdf.BasePDF):
+    _N_OBS = 1
+
+    def __init__(
+        self,
+        mu: ztyping.ParamTypeInput,
+        beta: ztyping.ParamTypeInput,
+        gamma: ztyping.ParamTypeInput,
+        n: ztyping.ParamTypeInput,
+        obs: ztyping.ObsTypeInput,
+        *,
+        extended: Optional[ztyping.ExtendedInputType] = False,
+        norm: Optional[ztyping.NormInputType] = None,
+        name: str = "ErfExp",
+    ):
+        """ErfExp PDF, the product of a complementary error function and an exponential function.
+
+        Implementation following closely `C++ version of custom RooErfExp.cc <https://gitlab.cern.ch/cms-muonPOG/spark_tnp/-/blob/Spark3/RooErfExp.cc>`_
+
+        .. math:
+
+            f(x; \\mu, \\beta, \\gamma, n) = \\text{erfc}(\\beta (x - \\mu)) \\exp{(-\\gamma (x^n - \\mu^n))}
+
+        Args:
+            mu: Location parameter.
+            beta: Scale parameter.
+            gamma: Shape parameter, scale of exponential term.
+            n: Shape parameter, power in exponential term.
+            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, "beta": beta, "gamma": gamma, "n": n}
+        super().__init__(obs=obs, params=params, extended=extended, norm=norm)
+
+    def _unnormalized_pdf(self, x):
+        mu = self.params["mu"]
+        beta = self.params["beta"]
+        gamma = self.params["gamma"]
+        n = self.params["n"]
+        x = z.unstack_x(x)
+        return erfexp_pdf_func(x=x, mu=mu, beta=beta, gamma=gamma, n=n)

zfit_physics/pdf.py

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