feat: add ErfExp pdf

CHANGELOG.rst

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

tests/test_pdf_cruijff.py

@@ -54,7 +54,7 @@ def test_cruijff_pdf():
     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,89 @@
+"""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
+
+alpha_true = 90.0
+beta_true = 0.08
+gamma_true = -1.0
+n_true = 0.2
+
+alpha_range = (65.0, 90.0)
+gamma_range = (-10, 10)
+beta_range = (0.01, 10)
+n_range = (0.1, 0.5)
+
+
+def _erfexp_numpy(x, alpha, beta, gamma, n):
+    return special.erfc((x - alpha) * beta) * np.exp(-gamma * (np.power(x, n) - np.power(alpha, n)))
+
+
+erfexp_numpy = np.vectorize(_erfexp_numpy, excluded=["alpha", "beta", "gamma", "n"])
+
+
+def create_erfexp(alpha, beta, gamma, n, limits):
+    obs = zfit.Space("obs1", limits=limits)
+    erfexp = zphys.pdf.ErfExp(alpha=alpha, beta=beta, gamma=gamma, n=n, obs=obs, norm=False)
+    return erfexp, obs
+
+
+def test_erfexp_pdf():
+    # Test PDF here
+    erfexp, _ = create_erfexp(alpha=alpha_true, beta=beta_true, gamma=gamma_true, n=n_true, limits=(50, 130))
+    assert erfexp.pdf(90.0).numpy().item() == pytest.approx(
+        erfexp_numpy(90.0, alpha=alpha_true, beta=beta_true, gamma=gamma_true, n=n_true), rel=1e-8
+    )
+    np.testing.assert_allclose(
+        erfexp.pdf(tf.range(50.0, 130, 10_000)),
+        erfexp_numpy(tf.range(50.0, 130, 10_000), alpha=alpha_true, beta=beta_true, gamma=gamma_true, n=n_true),
+        rtol=1e-8,
+    )
+
+
+def test_erfexp_pdf_random_params():
+    # Test PDF here in a loop with random parameters
+    for _ in range(1000):
+        alpha_true = np.random.uniform(*alpha_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(alpha=alpha_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)),
+            erfexp_numpy(tf.range(50.0, 130, 10_000), alpha=alpha_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(alpha=alpha_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=(alpha_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=(alpha_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():
+    alpha = zfit.Parameter("alpha", alpha_true)
+    beta = zfit.Parameter("beta", beta_true)
+    gamma = zfit.Parameter("gamma", gamma_true)
+    n = zfit.Parameter("n", n_true)
+
+    return {"alpha": alpha, "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,91 @@
+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, alpha, beta, gamma, n):
+    """Calculate the ErfExp PDF.
+
+    Args:
+        x: value(s) for which the PDF will be calculated.
+        alpha: 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 - alpha) * beta) * znp.exp(-gamma * (znp.power(x, n) - znp.power(alpha, n)))
+
+
+# Note: There is no analytic integral for the ErfExp PDF
+# We tried with sympy, Mathematica, Wolfram Alpha and https://www.integral-calculator.com/
+class ErfExp(zfit.pdf.BasePDF):
+    _N_OBS = 1
+
+    def __init__(
+        self,
+        alpha: 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 from https://gitlab.cern.ch/cms-muonPOG/spark_tnp/-/blob/Spark3/RooErfExp.cc
+
+        .. math:
+
+            f(x; \\alpha, \\beta, \\gamma, n) = \\text{erfc}(\\beta (x - \\alpha)) \\exp{(-\\gamma (x^n - \\alpha^n))}
+
+        Args:
+            alpha: Location parameter.
+            beta: Scale parameter.
+            gamma: Shape parameter.
+            n: Shape 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 = {"alpha": alpha, "beta": beta, "gamma": gamma, "n": n}
+        super().__init__(obs=obs, params=params, extended=extended, norm=norm)
+
+    def _unnormalized_pdf(self, x):
+        alpha = self.params["alpha"]
+        beta = self.params["beta"]
+        gamma = self.params["gamma"]
+        n = self.params["n"]
+        x = z.unstack_x(x)
+        return erfexp_pdf_func(x=x, alpha=alpha, 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"]