diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0597d43..273778c 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,6 +9,7 @@ Major Features and Improvements ------------------------------- - added CMSShape PDF - added Cruijff PDF +- added ErfExp PDF Breaking changes ------------------ diff --git a/tests/test_pdf_cruijff.py b/tests/test_pdf_cruijff.py index e81eca2..fdcce7a 100644 --- a/tests/test_pdf_cruijff.py +++ b/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) diff --git a/tests/test_pdf_erfexp.py b/tests/test_pdf_erfexp.py new file mode 100644 index 0000000..c3b8ef3 --- /dev/null +++ b/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) diff --git a/zfit_physics/models/pdf_cruijff.py b/zfit_physics/models/pdf_cruijff.py index b0ce2de..4619a47 100644 --- a/zfit_physics/models/pdf_cruijff.py +++ b/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 diff --git a/zfit_physics/models/pdf_erfexp.py b/zfit_physics/models/pdf_erfexp.py new file mode 100644 index 0000000..c948dbd --- /dev/null +++ b/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 `_ + + .. 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) diff --git a/zfit_physics/pdf.py b/zfit_physics/pdf.py index 3791657..c88e44a 100644 --- a/zfit_physics/pdf.py +++ b/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"]