From 56bf1057d67aafa9dc7364faaa1301a114f81aab Mon Sep 17 00:00:00 2001 From: iasonkrom Date: Tue, 19 Mar 2024 00:39:24 +0100 Subject: [PATCH] add cruijff pdf --- tests/test_pdf_cruijff.py | 91 ++++++++++++++++++++++++++ zfit_physics/models/pdf_cruijff.py | 100 +++++++++++++++++++++++++++++ zfit_physics/pdf.py | 3 +- 3 files changed, 193 insertions(+), 1 deletion(-) create mode 100644 tests/test_pdf_cruijff.py create mode 100644 zfit_physics/models/pdf_cruijff.py diff --git a/tests/test_pdf_cruijff.py b/tests/test_pdf_cruijff.py new file mode 100644 index 0000000..835c708 --- /dev/null +++ b/tests/test_pdf_cruijff.py @@ -0,0 +1,91 @@ +"""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 + +# Important, do the imports below +from zfit.core.testing import tester + +import zfit_physics as zphys + +# specify globals here. Do NOT add any TensorFlow but just pure python +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 hereA + cruijff, _ = create_cruijff( + mu=mu_true, sigmal=sigmal_true, alphal=alphal_true, sigmar=sigmar_true, alphar=alphar_true, limits=(50, 130) + ) + assert zfit.run(cruijff.pdf(90.0)) == 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 = zfit.run(cruijff.numeric_integrate(obs, norm_range=False)) + 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 = zfit.run(cruijff.numeric_integrate(limits=(80, 100), norm_range=False)) + 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) diff --git a/zfit_physics/models/pdf_cruijff.py b/zfit_physics/models/pdf_cruijff.py new file mode 100644 index 0000000..373b5da --- /dev/null +++ b/zfit_physics/models/pdf_cruijff.py @@ -0,0 +1,100 @@ +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): + """Caltulate 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 + """ + cond = znp.less(x, mu) + + func = znp.where( + cond, + znp.exp(-0.5 * znp.square((x - mu) / sigmal) / (1 + alphal * znp.square((x - mu) / sigmal))), + znp.exp(-0.5 * znp.square((x - mu) / sigmar) / (1 + alphar * znp.square((x - mu) / sigmar))), + ) + return func + + +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. + + Implementation from https://arxiv.org/abs/1005.4087 + + .. 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) diff --git a/zfit_physics/pdf.py b/zfit_physics/pdf.py index 2e7112b..3791657 100644 --- a/zfit_physics/pdf.py +++ b/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"]