From 9bda9a7904f69faf4fb2ff77910d81af549381c7 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 30 Aug 2021 11:17:42 +0200 Subject: [PATCH 01/13] Add relativistic Breit-Wigner PDF without analytic CDF --- tests/test_pdf_relbw.py | 38 ++++++++++++++++ zfit_physics/models/pdf_relbw.py | 77 ++++++++++++++++++++++++++++++++ zfit_physics/pdf.py | 3 +- 3 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 tests/test_pdf_relbw.py create mode 100644 zfit_physics/models/pdf_relbw.py diff --git a/tests/test_pdf_relbw.py b/tests/test_pdf_relbw.py new file mode 100644 index 0000000..c177577 --- /dev/null +++ b/tests/test_pdf_relbw.py @@ -0,0 +1,38 @@ +"""Tests for relativistic Breit-Wigner PDF.""" +import numpy as np +import pytest +import tensorflow as tf +import zfit + +# 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 +param1_true = 0.3 +param2_true = 1.2 + + +def test_pdf(): + # Test PDF values here + obs = zfit.Space("obs1", (0, 200)) + + relbw = zphys.pdf.RelativisticBreitWigner(m=125., Gamma=0.05, obs=obs) + assert zfit.run(relbw.pdf(125.)) == pytest.approx(12.7, rel=1e-3) + + #analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False)) + #numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False)) + #assert pytest.approx(analytic_integral, 4e-2) == numeric_integral + + +# register the pdf here and provide sets of working parameter configurations + + +def relbw_params_factory(): + m = zfit.Parameter("m", 4.5) + Gamma = zfit.Parameter("Gamma", -2.3) + return {"m": m, "Gamma": Gamma} + + +tester.register_pdf(pdf_class=zphys.pdf.RelativisticBreitWigner, params_factories=relbw_params_factory) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py new file mode 100644 index 0000000..b8fc43e --- /dev/null +++ b/zfit_physics/models/pdf_relbw.py @@ -0,0 +1,77 @@ +import zfit +from zfit.core.space import ANY_LOWER, ANY_UPPER, Space +from zfit import z +import tensorflow as tf +import numpy as np + + +class RelativisticBreitWigner(zfit.pdf.ZPDF): + _N_OBS = 1 # dimension, can be omitted + _PARAMS = ['m', 'Gamma'] # the name of the parameters + + def _unnormalized_pdf(self, x): + """ + Calculate the PDF at value(s) x. + + Parameters + ---------- + x : tf.Tensor + Energies. Either one value or an array + """ + x = zfit.z.unstack_x(x) + alpha = self.params['Gamma'] / self.params['m'] + gamma = self.params['m']**2 * (1. + alpha**2)**0.5 + k = 2.**(3. / 2.) * self.params['m']**2 * alpha * gamma / (np.pi * (self.params['m']**2 + gamma)**0.5) + + return k / ((x**2 - self.params['m']**2)**2 + self.params['m']**4 * alpha**2) + + +def relbw_cdf_func(x, m, Gamma): + """ + Parameters + ---------- + x + m : float + mass of resonance + Gamma : float + width of resonance + Returns + ------- + cdf : float + CDF of Breit-Wigner distribution + The CDf was found by Mathematica: + pdf = k/((m^2 - m^2)^2 + m^4*alpha^2) + cdf = Integrate[pdf, m] + >>> BW = breit_wigner(m=125., Gamma=0.05) + >>> BW.cdf(125.) + 0.50052268648248666 + >>> BW.cdf(1E10) + 1.0000000000000004 + >>> BW.cdf(0.) + 0.0 + """ + alpha = Gamma / m + gamma = m ** 2 * (1. + alpha ** 2) ** 0.5 + k = 2. ** (3. / 2.) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) + + arg_1 = z.to_complex(-1) ** (1. / 4.) / (-1j + alpha) ** 0.5 * x / m + arg_2 = z.to_complex(-1) ** (3. / 4.) / (1j + alpha) ** 0.5 * x / m + + shape = -1j * tf.math.atan(arg_1) / (-1j + alpha) ** 0.5 - tf.math.atan(arg_2) / (1j + alpha) ** 0.5 + norm = z.to_complex(-1) ** (1. / 4.) * k / (2. * alpha * m ** 3) + + cdf_ = shape * norm + cdf_ = z.to_real(cdf_) + return cdf_ + + +def relbw_integral(limits, params, model): + lower, upper = limits.rect_limits + lower_cdf = relbw_cdf_func(x=limits[0], m=params['m'], Gamma=params['Gamma']) + upper_cdf = relbw_cdf_func(x=limits[1], m=params['m'], Gamma=params['Gamma']) + return upper - lower + + +relbw_integral_limits = Space(axes=(0,), limits=(((ANY_LOWER,),), ((ANY_UPPER,),))) +# TODO analytic integral does not work +#RelativisticBreitWigner.register_analytic_integral(func=relbw_integral, limits=relbw_integral_limits) diff --git a/zfit_physics/pdf.py b/zfit_physics/pdf.py index 96c5a62..875e0bd 100644 --- a/zfit_physics/pdf.py +++ b/zfit_physics/pdf.py @@ -1,3 +1,4 @@ from .models.pdf_argus import Argus +from .models.pdf_relbw import RelativisticBreitWigner -__all__ = ["Argus"] +__all__ = ["Argus", "RelativisticBreitWigner"] From 2c40a598b3792d5009bf8f3e067409b2e81600b4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 30 Aug 2021 16:22:49 +0000 Subject: [PATCH 02/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_pdf_relbw.py | 10 ++++---- zfit_physics/models/pdf_relbw.py | 41 ++++++++++++++++++-------------- 2 files changed, 28 insertions(+), 23 deletions(-) diff --git a/tests/test_pdf_relbw.py b/tests/test_pdf_relbw.py index c177577..b9570c1 100644 --- a/tests/test_pdf_relbw.py +++ b/tests/test_pdf_relbw.py @@ -18,12 +18,12 @@ def test_pdf(): # Test PDF values here obs = zfit.Space("obs1", (0, 200)) - relbw = zphys.pdf.RelativisticBreitWigner(m=125., Gamma=0.05, obs=obs) - assert zfit.run(relbw.pdf(125.)) == pytest.approx(12.7, rel=1e-3) + relbw = zphys.pdf.RelativisticBreitWigner(m=125.0, Gamma=0.05, obs=obs) + assert zfit.run(relbw.pdf(125.0)) == pytest.approx(12.7, rel=1e-3) - #analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False)) - #numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False)) - #assert pytest.approx(analytic_integral, 4e-2) == numeric_integral + # analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False)) + # numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False)) + # assert pytest.approx(analytic_integral, 4e-2) == numeric_integral # register the pdf here and provide sets of working parameter configurations diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index b8fc43e..c89ec05 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -1,17 +1,16 @@ +import numpy as np +import tensorflow as tf import zfit -from zfit.core.space import ANY_LOWER, ANY_UPPER, Space from zfit import z -import tensorflow as tf -import numpy as np +from zfit.core.space import ANY_LOWER, ANY_UPPER, Space class RelativisticBreitWigner(zfit.pdf.ZPDF): _N_OBS = 1 # dimension, can be omitted - _PARAMS = ['m', 'Gamma'] # the name of the parameters + _PARAMS = ["m", "Gamma"] # the name of the parameters def _unnormalized_pdf(self, x): - """ - Calculate the PDF at value(s) x. + """Calculate the PDF at value(s) x. Parameters ---------- @@ -19,11 +18,17 @@ def _unnormalized_pdf(self, x): Energies. Either one value or an array """ x = zfit.z.unstack_x(x) - alpha = self.params['Gamma'] / self.params['m'] - gamma = self.params['m']**2 * (1. + alpha**2)**0.5 - k = 2.**(3. / 2.) * self.params['m']**2 * alpha * gamma / (np.pi * (self.params['m']**2 + gamma)**0.5) + alpha = self.params["Gamma"] / self.params["m"] + gamma = self.params["m"] ** 2 * (1.0 + alpha ** 2) ** 0.5 + k = ( + 2.0 ** (3.0 / 2.0) + * self.params["m"] ** 2 + * alpha + * gamma + / (np.pi * (self.params["m"] ** 2 + gamma) ** 0.5) + ) - return k / ((x**2 - self.params['m']**2)**2 + self.params['m']**4 * alpha**2) + return k / ((x ** 2 - self.params["m"] ** 2) ** 2 + self.params["m"] ** 4 * alpha ** 2) def relbw_cdf_func(x, m, Gamma): @@ -51,14 +56,14 @@ def relbw_cdf_func(x, m, Gamma): 0.0 """ alpha = Gamma / m - gamma = m ** 2 * (1. + alpha ** 2) ** 0.5 - k = 2. ** (3. / 2.) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) + gamma = m ** 2 * (1.0 + alpha ** 2) ** 0.5 + k = 2.0 ** (3.0 / 2.0) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) - arg_1 = z.to_complex(-1) ** (1. / 4.) / (-1j + alpha) ** 0.5 * x / m - arg_2 = z.to_complex(-1) ** (3. / 4.) / (1j + alpha) ** 0.5 * x / m + arg_1 = z.to_complex(-1) ** (1.0 / 4.0) / (-1j + alpha) ** 0.5 * x / m + arg_2 = z.to_complex(-1) ** (3.0 / 4.0) / (1j + alpha) ** 0.5 * x / m shape = -1j * tf.math.atan(arg_1) / (-1j + alpha) ** 0.5 - tf.math.atan(arg_2) / (1j + alpha) ** 0.5 - norm = z.to_complex(-1) ** (1. / 4.) * k / (2. * alpha * m ** 3) + norm = z.to_complex(-1) ** (1.0 / 4.0) * k / (2.0 * alpha * m ** 3) cdf_ = shape * norm cdf_ = z.to_real(cdf_) @@ -67,11 +72,11 @@ def relbw_cdf_func(x, m, Gamma): def relbw_integral(limits, params, model): lower, upper = limits.rect_limits - lower_cdf = relbw_cdf_func(x=limits[0], m=params['m'], Gamma=params['Gamma']) - upper_cdf = relbw_cdf_func(x=limits[1], m=params['m'], Gamma=params['Gamma']) + lower_cdf = relbw_cdf_func(x=limits[0], m=params["m"], Gamma=params["Gamma"]) + upper_cdf = relbw_cdf_func(x=limits[1], m=params["m"], Gamma=params["Gamma"]) return upper - lower relbw_integral_limits = Space(axes=(0,), limits=(((ANY_LOWER,),), ((ANY_UPPER,),))) # TODO analytic integral does not work -#RelativisticBreitWigner.register_analytic_integral(func=relbw_integral, limits=relbw_integral_limits) +# RelativisticBreitWigner.register_analytic_integral(func=relbw_integral, limits=relbw_integral_limits) From 24ea041d850e451cb8632b7e0ade9a89106b7004 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 30 Aug 2021 19:37:45 +0200 Subject: [PATCH 03/13] Add analytic integral and corresponding tests --- tests/test_pdf_relbw.py | 41 ++++++++++++------- zfit_physics/models/pdf_relbw.py | 70 ++++++++++++++++++++------------ 2 files changed, 71 insertions(+), 40 deletions(-) diff --git a/tests/test_pdf_relbw.py b/tests/test_pdf_relbw.py index b9570c1..a98bf2e 100644 --- a/tests/test_pdf_relbw.py +++ b/tests/test_pdf_relbw.py @@ -1,5 +1,4 @@ """Tests for relativistic Breit-Wigner PDF.""" -import numpy as np import pytest import tensorflow as tf import zfit @@ -10,28 +9,42 @@ import zfit_physics as zphys # specify globals here. Do NOT add any TensorFlow but just pure python -param1_true = 0.3 -param2_true = 1.2 +m_true = 125. +Gamma_true = 0.05 -def test_pdf(): - # Test PDF values here - obs = zfit.Space("obs1", (0, 200)) +def create_relbw(m, Gamma, limits): + obs = zfit.Space("obs1", limits) + relbw = zphys.pdf.RelativisticBreitWigner(m=m, Gamma=Gamma, obs=obs) + return relbw, obs - relbw = zphys.pdf.RelativisticBreitWigner(m=125.0, Gamma=0.05, obs=obs) - assert zfit.run(relbw.pdf(125.0)) == pytest.approx(12.7, rel=1e-3) - # analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False)) - # numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False)) - # assert pytest.approx(analytic_integral, 4e-2) == numeric_integral +def test_relbw_pdf(): + # Test PDF here + relbw, _ = create_relbw(m_true, Gamma_true, limits=(0, 200)) + assert zfit.run(relbw.pdf(125.)) == pytest.approx(12.732396211295313, rel=1e-4) + assert relbw.pdf(tf.range(0., 200, 10_000)) <= relbw.pdf(125.) + sample = relbw.sample(1000) + tf.debugging.assert_all_finite(sample, 'Some samples from the relbw PDF are NaN or infinite') + assert sample.n_events == 1000 + assert all(tf.logical_and(0 <= sample, sample <= 200)) -# register the pdf here and provide sets of working parameter configurations +def test_relbw_integral(): + # Test CDF and integral here + relbw, obs = create_relbw(m_true, Gamma_true, limits=(0, 200)) + + analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False))[0] + numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False))[0] + assert analytic_integral == pytest.approx(1., 1e-4) + assert numeric_integral == pytest.approx(1., 2e-3) + +# register the pdf here and provide sets of working parameter configurations def relbw_params_factory(): - m = zfit.Parameter("m", 4.5) - Gamma = zfit.Parameter("Gamma", -2.3) + m = zfit.Parameter("m", m_true) + Gamma = zfit.Parameter("Gamma", Gamma_true) return {"m": m, "Gamma": Gamma} diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index c89ec05..4759595 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -1,16 +1,36 @@ -import numpy as np -import tensorflow as tf import zfit -from zfit import z from zfit.core.space import ANY_LOWER, ANY_UPPER, Space +from zfit import z +import tensorflow as tf +import numpy as np + + +def arctan_complex(x): + """ + Function that evaluates arctan(x) using tensorflow but also supports complex numbers. + + Args: x + + Returns: arctan(x) + + Notes + ----- + Formula used: https://www.wolframalpha.com/input/?i=arctan%28a%2Bb*i%29 + """ + return 1/2 * 1j * tf.math.log(1 - 1j*x) - 1/2 * 1j * tf.math.log(1 + 1j*x) class RelativisticBreitWigner(zfit.pdf.ZPDF): - _N_OBS = 1 # dimension, can be omitted - _PARAMS = ["m", "Gamma"] # the name of the parameters + """ + Relativistic Breit-Wigner distribution. + Formula for PDF and CDF are based on https://gist.github.com/andrewfowlie/cd0ed7e6c96f7c9e88f85eb3b9665b97 + """ + _N_OBS = 1 + _PARAMS = ['m', 'Gamma'] def _unnormalized_pdf(self, x): - """Calculate the PDF at value(s) x. + """ + Calculate the PDF at value(s) x. Parameters ---------- @@ -18,17 +38,11 @@ def _unnormalized_pdf(self, x): Energies. Either one value or an array """ x = zfit.z.unstack_x(x) - alpha = self.params["Gamma"] / self.params["m"] - gamma = self.params["m"] ** 2 * (1.0 + alpha ** 2) ** 0.5 - k = ( - 2.0 ** (3.0 / 2.0) - * self.params["m"] ** 2 - * alpha - * gamma - / (np.pi * (self.params["m"] ** 2 + gamma) ** 0.5) - ) + alpha = self.params['Gamma'] / self.params['m'] + gamma = self.params['m']**2 * (1. + alpha**2)**0.5 + k = 2.**(3. / 2.) * self.params['m']**2 * alpha * gamma / (np.pi * (self.params['m']**2 + gamma)**0.5) - return k / ((x ** 2 - self.params["m"] ** 2) ** 2 + self.params["m"] ** 4 * alpha ** 2) + return k / ((x**2 - self.params['m']**2)**2 + self.params['m']**4 * alpha**2) def relbw_cdf_func(x, m, Gamma): @@ -55,15 +69,19 @@ def relbw_cdf_func(x, m, Gamma): >>> BW.cdf(0.) 0.0 """ + Gamma = z.to_complex(Gamma) + m = z.to_complex(m) + x = z.to_complex(x) + alpha = Gamma / m - gamma = m ** 2 * (1.0 + alpha ** 2) ** 0.5 - k = 2.0 ** (3.0 / 2.0) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) + gamma = m ** 2 * (1. + alpha ** 2) ** 0.5 + k = 2. ** (3. / 2.) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) - arg_1 = z.to_complex(-1) ** (1.0 / 4.0) / (-1j + alpha) ** 0.5 * x / m - arg_2 = z.to_complex(-1) ** (3.0 / 4.0) / (1j + alpha) ** 0.5 * x / m + arg_1 = z.to_complex(-1) ** (1. / 4.) / (-1j + alpha) ** 0.5 * x / m + arg_2 = z.to_complex(-1) ** (3. / 4.) / (1j + alpha) ** 0.5 * x / m - shape = -1j * tf.math.atan(arg_1) / (-1j + alpha) ** 0.5 - tf.math.atan(arg_2) / (1j + alpha) ** 0.5 - norm = z.to_complex(-1) ** (1.0 / 4.0) * k / (2.0 * alpha * m ** 3) + shape = -1j * arctan_complex(arg_1) / (-1j + alpha) ** 0.5 - arctan_complex(arg_2) / (1j + alpha) ** 0.5 + norm = z.to_complex(-1) ** (1. / 4.) * k / (2. * alpha * m ** 3) cdf_ = shape * norm cdf_ = z.to_real(cdf_) @@ -72,11 +90,11 @@ def relbw_cdf_func(x, m, Gamma): def relbw_integral(limits, params, model): lower, upper = limits.rect_limits - lower_cdf = relbw_cdf_func(x=limits[0], m=params["m"], Gamma=params["Gamma"]) - upper_cdf = relbw_cdf_func(x=limits[1], m=params["m"], Gamma=params["Gamma"]) - return upper - lower + lower_cdf = relbw_cdf_func(x=lower, m=params['m'], Gamma=params['Gamma']) + upper_cdf = relbw_cdf_func(x=upper, m=params['m'], Gamma=params['Gamma']) + return upper_cdf - lower_cdf relbw_integral_limits = Space(axes=(0,), limits=(((ANY_LOWER,),), ((ANY_UPPER,),))) # TODO analytic integral does not work -# RelativisticBreitWigner.register_analytic_integral(func=relbw_integral, limits=relbw_integral_limits) +RelativisticBreitWigner.register_analytic_integral(func=relbw_integral, limits=relbw_integral_limits) From d144bb42aea92ecba302ffaaed6e62d2f0380219 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 30 Aug 2021 19:00:52 +0000 Subject: [PATCH 04/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_pdf_relbw.py | 12 ++++---- zfit_physics/models/pdf_relbw.py | 49 ++++++++++++++++++-------------- 2 files changed, 33 insertions(+), 28 deletions(-) diff --git a/tests/test_pdf_relbw.py b/tests/test_pdf_relbw.py index a98bf2e..f12938f 100644 --- a/tests/test_pdf_relbw.py +++ b/tests/test_pdf_relbw.py @@ -9,7 +9,7 @@ import zfit_physics as zphys # specify globals here. Do NOT add any TensorFlow but just pure python -m_true = 125. +m_true = 125.0 Gamma_true = 0.05 @@ -22,11 +22,11 @@ def create_relbw(m, Gamma, limits): def test_relbw_pdf(): # Test PDF here relbw, _ = create_relbw(m_true, Gamma_true, limits=(0, 200)) - assert zfit.run(relbw.pdf(125.)) == pytest.approx(12.732396211295313, rel=1e-4) - assert relbw.pdf(tf.range(0., 200, 10_000)) <= relbw.pdf(125.) + assert zfit.run(relbw.pdf(125.0)) == pytest.approx(12.732396211295313, rel=1e-4) + assert relbw.pdf(tf.range(0.0, 200, 10_000)) <= relbw.pdf(125.0) sample = relbw.sample(1000) - tf.debugging.assert_all_finite(sample, 'Some samples from the relbw PDF are NaN or infinite') + tf.debugging.assert_all_finite(sample, "Some samples from the relbw PDF are NaN or infinite") assert sample.n_events == 1000 assert all(tf.logical_and(0 <= sample, sample <= 200)) @@ -37,8 +37,8 @@ def test_relbw_integral(): analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False))[0] numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False))[0] - assert analytic_integral == pytest.approx(1., 1e-4) - assert numeric_integral == pytest.approx(1., 2e-3) + assert analytic_integral == pytest.approx(1.0, 1e-4) + assert numeric_integral == pytest.approx(1.0, 2e-3) # register the pdf here and provide sets of working parameter configurations diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index 4759595..829e3de 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -1,13 +1,12 @@ +import numpy as np +import tensorflow as tf import zfit -from zfit.core.space import ANY_LOWER, ANY_UPPER, Space from zfit import z -import tensorflow as tf -import numpy as np +from zfit.core.space import ANY_LOWER, ANY_UPPER, Space def arctan_complex(x): - """ - Function that evaluates arctan(x) using tensorflow but also supports complex numbers. + """Function that evaluates arctan(x) using tensorflow but also supports complex numbers. Args: x @@ -17,20 +16,20 @@ def arctan_complex(x): ----- Formula used: https://www.wolframalpha.com/input/?i=arctan%28a%2Bb*i%29 """ - return 1/2 * 1j * tf.math.log(1 - 1j*x) - 1/2 * 1j * tf.math.log(1 + 1j*x) + return 1 / 2 * 1j * tf.math.log(1 - 1j * x) - 1 / 2 * 1j * tf.math.log(1 + 1j * x) class RelativisticBreitWigner(zfit.pdf.ZPDF): - """ - Relativistic Breit-Wigner distribution. + """Relativistic Breit-Wigner distribution. + Formula for PDF and CDF are based on https://gist.github.com/andrewfowlie/cd0ed7e6c96f7c9e88f85eb3b9665b97 """ + _N_OBS = 1 - _PARAMS = ['m', 'Gamma'] + _PARAMS = ["m", "Gamma"] def _unnormalized_pdf(self, x): - """ - Calculate the PDF at value(s) x. + """Calculate the PDF at value(s) x. Parameters ---------- @@ -38,11 +37,17 @@ def _unnormalized_pdf(self, x): Energies. Either one value or an array """ x = zfit.z.unstack_x(x) - alpha = self.params['Gamma'] / self.params['m'] - gamma = self.params['m']**2 * (1. + alpha**2)**0.5 - k = 2.**(3. / 2.) * self.params['m']**2 * alpha * gamma / (np.pi * (self.params['m']**2 + gamma)**0.5) + alpha = self.params["Gamma"] / self.params["m"] + gamma = self.params["m"] ** 2 * (1.0 + alpha ** 2) ** 0.5 + k = ( + 2.0 ** (3.0 / 2.0) + * self.params["m"] ** 2 + * alpha + * gamma + / (np.pi * (self.params["m"] ** 2 + gamma) ** 0.5) + ) - return k / ((x**2 - self.params['m']**2)**2 + self.params['m']**4 * alpha**2) + return k / ((x ** 2 - self.params["m"] ** 2) ** 2 + self.params["m"] ** 4 * alpha ** 2) def relbw_cdf_func(x, m, Gamma): @@ -74,14 +79,14 @@ def relbw_cdf_func(x, m, Gamma): x = z.to_complex(x) alpha = Gamma / m - gamma = m ** 2 * (1. + alpha ** 2) ** 0.5 - k = 2. ** (3. / 2.) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) + gamma = m ** 2 * (1.0 + alpha ** 2) ** 0.5 + k = 2.0 ** (3.0 / 2.0) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) - arg_1 = z.to_complex(-1) ** (1. / 4.) / (-1j + alpha) ** 0.5 * x / m - arg_2 = z.to_complex(-1) ** (3. / 4.) / (1j + alpha) ** 0.5 * x / m + arg_1 = z.to_complex(-1) ** (1.0 / 4.0) / (-1j + alpha) ** 0.5 * x / m + arg_2 = z.to_complex(-1) ** (3.0 / 4.0) / (1j + alpha) ** 0.5 * x / m shape = -1j * arctan_complex(arg_1) / (-1j + alpha) ** 0.5 - arctan_complex(arg_2) / (1j + alpha) ** 0.5 - norm = z.to_complex(-1) ** (1. / 4.) * k / (2. * alpha * m ** 3) + norm = z.to_complex(-1) ** (1.0 / 4.0) * k / (2.0 * alpha * m ** 3) cdf_ = shape * norm cdf_ = z.to_real(cdf_) @@ -90,8 +95,8 @@ def relbw_cdf_func(x, m, Gamma): def relbw_integral(limits, params, model): lower, upper = limits.rect_limits - lower_cdf = relbw_cdf_func(x=lower, m=params['m'], Gamma=params['Gamma']) - upper_cdf = relbw_cdf_func(x=upper, m=params['m'], Gamma=params['Gamma']) + lower_cdf = relbw_cdf_func(x=lower, m=params["m"], Gamma=params["Gamma"]) + upper_cdf = relbw_cdf_func(x=upper, m=params["m"], Gamma=params["Gamma"]) return upper_cdf - lower_cdf From 0efd4cc0748ff1c04be713828c0305d95e2a6f7a Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 30 Aug 2021 21:37:17 +0200 Subject: [PATCH 05/13] Added docstrings and JIT --- zfit_physics/models/pdf_relbw.py | 86 +++++++++++++++++++------------- 1 file changed, 51 insertions(+), 35 deletions(-) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index 829e3de..f688494 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -1,40 +1,53 @@ import numpy as np import tensorflow as tf import zfit +from zfit.util import ztyping from zfit import z from zfit.core.space import ANY_LOWER, ANY_UPPER, Space +@z.function def arctan_complex(x): - """Function that evaluates arctan(x) using tensorflow but also supports complex numbers. + r"""Function that evaluates arctan(x) using tensorflow but also supports complex numbers. + It is defined as + .. math:: - Args: x + \mathrm{arctan}(x) = \frac{i}{2} \left(\ln(1-ix) - \ln(1+ix)\right) - Returns: arctan(x) + Args: + x: tf.Tensor - Notes - ----- - Formula used: https://www.wolframalpha.com/input/?i=arctan%28a%2Bb*i%29 + Returns: + .. math:: \mathrm{arctan}(x) + + Notes: + Formula is taken from https://www.wolframalpha.com/input/?i=arctan%28a%2Bb*i%29 + TODO: move somewhere? """ - return 1 / 2 * 1j * tf.math.log(1 - 1j * x) - 1 / 2 * 1j * tf.math.log(1 + 1j * x) + return 1 / 2 * 1j * (tf.math.log(1 - 1j * x) - tf.math.log(1 + 1j * x)) class RelativisticBreitWigner(zfit.pdf.ZPDF): """Relativistic Breit-Wigner distribution. Formula for PDF and CDF are based on https://gist.github.com/andrewfowlie/cd0ed7e6c96f7c9e88f85eb3b9665b97 + + Args: + m: the average value + Gamma: the width of the distribution """ _N_OBS = 1 _PARAMS = ["m", "Gamma"] - def _unnormalized_pdf(self, x): + def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: """Calculate the PDF at value(s) x. - Parameters - ---------- - x : tf.Tensor - Energies. Either one value or an array + Args: + x : Either one value or an array + + Returns: + `tf.Tensor`: The value(s) of the unnormalized PDF at x. """ x = zfit.z.unstack_x(x) alpha = self.params["Gamma"] / self.params["m"] @@ -50,29 +63,21 @@ def _unnormalized_pdf(self, x): return k / ((x ** 2 - self.params["m"] ** 2) ** 2 + self.params["m"] ** 4 * alpha ** 2) +@z.function def relbw_cdf_func(x, m, Gamma): """ - Parameters - ---------- - x - m : float - mass of resonance - Gamma : float - width of resonance - Returns - ------- - cdf : float - CDF of Breit-Wigner distribution - The CDf was found by Mathematica: - pdf = k/((m^2 - m^2)^2 + m^4*alpha^2) - cdf = Integrate[pdf, m] - >>> BW = breit_wigner(m=125., Gamma=0.05) - >>> BW.cdf(125.) - 0.50052268648248666 - >>> BW.cdf(1E10) - 1.0000000000000004 - >>> BW.cdf(0.) - 0.0 + Analytical function for the CDF of the relativistic Breit-Wigner distribution. + + Args: + x: value(s) for which the CDF will be calculated. + m: Mean value + Gamma: width + + Returns: + `tf.Tensor`: The calculated CDF values. + + Notes: + Based on code from this [github gist](https://gist.github.com/andrewfowlie/cd0ed7e6c96f7c9e88f85eb3b9665b97#file-bw-py-L112-L154) """ Gamma = z.to_complex(Gamma) m = z.to_complex(m) @@ -93,13 +98,24 @@ def relbw_cdf_func(x, m, Gamma): return cdf_ -def relbw_integral(limits, params, model): +def relbw_integral(limits: ztyping.SpaceType, params: dict, model) -> tf.Tensor: + """ + Calculates the analytic integral of the relativistic Breit-Wigner PDF. + + Args: + limits: An object with attribute rect_limits. + params: A hashmap from which the parameters that defines the PDF will be extracted. + model: Will be ignored. + + Returns: + The calculated integral. + """ lower, upper = limits.rect_limits lower_cdf = relbw_cdf_func(x=lower, m=params["m"], Gamma=params["Gamma"]) upper_cdf = relbw_cdf_func(x=upper, m=params["m"], Gamma=params["Gamma"]) return upper_cdf - lower_cdf +# These lines of code adds the analytic integral function to RelativisticBreitWigner PDF. relbw_integral_limits = Space(axes=(0,), limits=(((ANY_LOWER,),), ((ANY_UPPER,),))) -# TODO analytic integral does not work RelativisticBreitWigner.register_analytic_integral(func=relbw_integral, limits=relbw_integral_limits) From 09a7842477d937eb014661c8b7b05451662a09c8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 30 Aug 2021 19:39:58 +0000 Subject: [PATCH 06/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- zfit_physics/models/pdf_relbw.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index f688494..5a91948 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -1,9 +1,9 @@ import numpy as np import tensorflow as tf import zfit -from zfit.util import ztyping from zfit import z from zfit.core.space import ANY_LOWER, ANY_UPPER, Space +from zfit.util import ztyping @z.function @@ -65,8 +65,7 @@ def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: @z.function def relbw_cdf_func(x, m, Gamma): - """ - Analytical function for the CDF of the relativistic Breit-Wigner distribution. + """Analytical function for the CDF of the relativistic Breit-Wigner distribution. Args: x: value(s) for which the CDF will be calculated. @@ -99,8 +98,7 @@ def relbw_cdf_func(x, m, Gamma): def relbw_integral(limits: ztyping.SpaceType, params: dict, model) -> tf.Tensor: - """ - Calculates the analytic integral of the relativistic Breit-Wigner PDF. + """Calculates the analytic integral of the relativistic Breit-Wigner PDF. Args: limits: An object with attribute rect_limits. From 715946f1648de02902602851739ae287b4181755 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 1 Sep 2021 20:34:22 +0200 Subject: [PATCH 07/13] Rename Gamma to gamma --- tests/test_pdf_relbw.py | 14 +++++++------- zfit_physics/models/pdf_relbw.py | 30 ++++++++++++++++-------------- 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/tests/test_pdf_relbw.py b/tests/test_pdf_relbw.py index f12938f..2648283 100644 --- a/tests/test_pdf_relbw.py +++ b/tests/test_pdf_relbw.py @@ -10,18 +10,18 @@ # specify globals here. Do NOT add any TensorFlow but just pure python m_true = 125.0 -Gamma_true = 0.05 +gamma_true = 0.05 -def create_relbw(m, Gamma, limits): +def create_relbw(m, gamma, limits): obs = zfit.Space("obs1", limits) - relbw = zphys.pdf.RelativisticBreitWigner(m=m, Gamma=Gamma, obs=obs) + relbw = zphys.pdf.RelativisticBreitWigner(m=m, gamma=gamma, obs=obs) return relbw, obs def test_relbw_pdf(): # Test PDF here - relbw, _ = create_relbw(m_true, Gamma_true, limits=(0, 200)) + relbw, _ = create_relbw(m_true, gamma_true, limits=(0, 200)) assert zfit.run(relbw.pdf(125.0)) == pytest.approx(12.732396211295313, rel=1e-4) assert relbw.pdf(tf.range(0.0, 200, 10_000)) <= relbw.pdf(125.0) @@ -33,7 +33,7 @@ def test_relbw_pdf(): def test_relbw_integral(): # Test CDF and integral here - relbw, obs = create_relbw(m_true, Gamma_true, limits=(0, 200)) + relbw, obs = create_relbw(m_true, gamma_true, limits=(0, 200)) analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False))[0] numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False))[0] @@ -44,8 +44,8 @@ def test_relbw_integral(): # register the pdf here and provide sets of working parameter configurations def relbw_params_factory(): m = zfit.Parameter("m", m_true) - Gamma = zfit.Parameter("Gamma", Gamma_true) - return {"m": m, "Gamma": Gamma} + gamma = zfit.Parameter("gamma", gamma_true) + return {"m": m, "gamma": gamma} tester.register_pdf(pdf_class=zphys.pdf.RelativisticBreitWigner, params_factories=relbw_params_factory) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index 5a91948..d2befaf 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -1,9 +1,9 @@ import numpy as np import tensorflow as tf import zfit +from zfit.util import ztyping from zfit import z from zfit.core.space import ANY_LOWER, ANY_UPPER, Space -from zfit.util import ztyping @z.function @@ -34,11 +34,11 @@ class RelativisticBreitWigner(zfit.pdf.ZPDF): Args: m: the average value - Gamma: the width of the distribution + gamma: the width of the distribution """ _N_OBS = 1 - _PARAMS = ["m", "Gamma"] + _PARAMS = ["m", "gamma"] def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: """Calculate the PDF at value(s) x. @@ -50,7 +50,7 @@ def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: `tf.Tensor`: The value(s) of the unnormalized PDF at x. """ x = zfit.z.unstack_x(x) - alpha = self.params["Gamma"] / self.params["m"] + alpha = self.params["gamma"] / self.params["m"] gamma = self.params["m"] ** 2 * (1.0 + alpha ** 2) ** 0.5 k = ( 2.0 ** (3.0 / 2.0) @@ -64,13 +64,14 @@ def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: @z.function -def relbw_cdf_func(x, m, Gamma): - """Analytical function for the CDF of the relativistic Breit-Wigner distribution. +def relbw_cdf_func(x, m, gamma): + """ + Analytical function for the CDF of the relativistic Breit-Wigner distribution. Args: x: value(s) for which the CDF will be calculated. m: Mean value - Gamma: width + gamma: width Returns: `tf.Tensor`: The calculated CDF values. @@ -78,13 +79,13 @@ def relbw_cdf_func(x, m, Gamma): Notes: Based on code from this [github gist](https://gist.github.com/andrewfowlie/cd0ed7e6c96f7c9e88f85eb3b9665b97#file-bw-py-L112-L154) """ - Gamma = z.to_complex(Gamma) + gamma = z.to_complex(gamma) m = z.to_complex(m) x = z.to_complex(x) - alpha = Gamma / m - gamma = m ** 2 * (1.0 + alpha ** 2) ** 0.5 - k = 2.0 ** (3.0 / 2.0) * m ** 2 * alpha * gamma / (np.pi * (m ** 2 + gamma) ** 0.5) + alpha = gamma / m + gamma2 = m ** 2 * (1.0 + alpha ** 2) ** 0.5 + k = 2.0 ** (3.0 / 2.0) * m ** 2 * alpha * gamma2 / (np.pi * (m ** 2 + gamma2) ** 0.5) arg_1 = z.to_complex(-1) ** (1.0 / 4.0) / (-1j + alpha) ** 0.5 * x / m arg_2 = z.to_complex(-1) ** (3.0 / 4.0) / (1j + alpha) ** 0.5 * x / m @@ -98,7 +99,8 @@ def relbw_cdf_func(x, m, Gamma): def relbw_integral(limits: ztyping.SpaceType, params: dict, model) -> tf.Tensor: - """Calculates the analytic integral of the relativistic Breit-Wigner PDF. + """ + Calculates the analytic integral of the relativistic Breit-Wigner PDF. Args: limits: An object with attribute rect_limits. @@ -109,8 +111,8 @@ def relbw_integral(limits: ztyping.SpaceType, params: dict, model) -> tf.Tensor: The calculated integral. """ lower, upper = limits.rect_limits - lower_cdf = relbw_cdf_func(x=lower, m=params["m"], Gamma=params["Gamma"]) - upper_cdf = relbw_cdf_func(x=upper, m=params["m"], Gamma=params["Gamma"]) + lower_cdf = relbw_cdf_func(x=lower, m=params["m"], gamma=params["gamma"]) + upper_cdf = relbw_cdf_func(x=upper, m=params["m"], gamma=params["gamma"]) return upper_cdf - lower_cdf From 3bf8797f1756a77ec1c3c740f791e769877d5154 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 1 Sep 2021 20:41:19 +0200 Subject: [PATCH 08/13] Update CHANGELOG.rst --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c1a07e1..f6a8cdf 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -7,6 +7,7 @@ Develop Major Features and Improvements ------------------------------- +- Added relativistic Breit-Wigner PDF Breaking changes ------------------ From 1c18c37cd46f091bcce8cfa39d68c32a87928e82 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 1 Sep 2021 20:48:19 +0200 Subject: [PATCH 09/13] Unstack x in cdf function --- tests/test_pdf_relbw.py | 4 ++-- zfit_physics/models/pdf_relbw.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_pdf_relbw.py b/tests/test_pdf_relbw.py index 2648283..81bb066 100644 --- a/tests/test_pdf_relbw.py +++ b/tests/test_pdf_relbw.py @@ -35,8 +35,8 @@ def test_relbw_integral(): # Test CDF and integral here relbw, obs = create_relbw(m_true, gamma_true, limits=(0, 200)) - analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False))[0] - numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False))[0] + analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False)) + numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False)) assert analytic_integral == pytest.approx(1.0, 1e-4) assert numeric_integral == pytest.approx(1.0, 2e-3) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index d2befaf..c0304d6 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -49,7 +49,7 @@ def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: Returns: `tf.Tensor`: The value(s) of the unnormalized PDF at x. """ - x = zfit.z.unstack_x(x) + x = z.unstack_x(x) alpha = self.params["gamma"] / self.params["m"] gamma = self.params["m"] ** 2 * (1.0 + alpha ** 2) ** 0.5 k = ( @@ -81,7 +81,7 @@ def relbw_cdf_func(x, m, gamma): """ gamma = z.to_complex(gamma) m = z.to_complex(m) - x = z.to_complex(x) + x = z.to_complex(z.unstack_x(x)) alpha = gamma / m gamma2 = m ** 2 * (1.0 + alpha ** 2) ** 0.5 From aafdced04cba24e7c23f86fcdb67f38afe4678b1 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 1 Sep 2021 21:55:38 +0200 Subject: [PATCH 10/13] Add another test to test integral on non-full interval --- tests/test_pdf_relbw.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_pdf_relbw.py b/tests/test_pdf_relbw.py index 81bb066..d62bcd9 100644 --- a/tests/test_pdf_relbw.py +++ b/tests/test_pdf_relbw.py @@ -35,10 +35,14 @@ def test_relbw_integral(): # Test CDF and integral here relbw, obs = create_relbw(m_true, gamma_true, limits=(0, 200)) - analytic_integral = zfit.run(relbw.analytic_integrate(obs, norm_range=False)) - numeric_integral = zfit.run(relbw.numeric_integrate(obs, norm_range=False)) - assert analytic_integral == pytest.approx(1.0, 1e-4) - assert numeric_integral == pytest.approx(1.0, 2e-3) + full_interval_analytic = zfit.run(relbw.analytic_integrate(obs, norm_range=False)) + full_interval_numeric = zfit.run(relbw.numeric_integrate(obs, norm_range=False)) + assert full_interval_analytic == pytest.approx(1.0, 1e-4) + assert full_interval_numeric == pytest.approx(1.0, 2e-3) + + analytic_integral = zfit.run(relbw.analytic_integrate(limits=(50, 100), norm_range=False)) + numeric_integral = zfit.run(relbw.numeric_integrate(limits=(50, 100), norm_range=False)) + assert analytic_integral == pytest.approx(numeric_integral, 0.01) # register the pdf here and provide sets of working parameter configurations From 00330edd9ae29db4063eba9a6fa1eafc7f609ee0 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Wed, 1 Sep 2021 22:06:41 +0200 Subject: [PATCH 11/13] Add wraps argument to z.function --- zfit_physics/models/pdf_relbw.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index c0304d6..4ef8dd7 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -6,7 +6,7 @@ from zfit.core.space import ANY_LOWER, ANY_UPPER, Space -@z.function +@z.function(wraps='tensor) def arctan_complex(x): r"""Function that evaluates arctan(x) using tensorflow but also supports complex numbers. It is defined as @@ -63,7 +63,7 @@ def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: return k / ((x ** 2 - self.params["m"] ** 2) ** 2 + self.params["m"] ** 4 * alpha ** 2) -@z.function +@z.function(wraps='tensor') def relbw_cdf_func(x, m, gamma): """ Analytical function for the CDF of the relativistic Breit-Wigner distribution. From 09e0d29f3006dffc287ee7606c6f9e9c5d10ba13 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Wed, 1 Sep 2021 22:09:15 +0200 Subject: [PATCH 12/13] fix typo --- zfit_physics/models/pdf_relbw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index 4ef8dd7..3ed2e90 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -6,7 +6,7 @@ from zfit.core.space import ANY_LOWER, ANY_UPPER, Space -@z.function(wraps='tensor) +@z.function(wraps='tensor') def arctan_complex(x): r"""Function that evaluates arctan(x) using tensorflow but also supports complex numbers. It is defined as From 5c9ccec5f7b186e6131cdad85c86efd0d8e31956 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 1 Sep 2021 22:39:24 +0200 Subject: [PATCH 13/13] Refactor to pdf to relbw_pdf_func --- zfit_physics/models/pdf_relbw.py | 63 +++++++++++++++++++++----------- 1 file changed, 41 insertions(+), 22 deletions(-) diff --git a/zfit_physics/models/pdf_relbw.py b/zfit_physics/models/pdf_relbw.py index 3ed2e90..a0a9db3 100644 --- a/zfit_physics/models/pdf_relbw.py +++ b/zfit_physics/models/pdf_relbw.py @@ -7,24 +7,33 @@ @z.function(wraps='tensor') -def arctan_complex(x): - r"""Function that evaluates arctan(x) using tensorflow but also supports complex numbers. - It is defined as - .. math:: - - \mathrm{arctan}(x) = \frac{i}{2} \left(\ln(1-ix) - \ln(1+ix)\right) +def relbw_pdf_func(x, m, gamma): + """ + Calculate the relativistic Breit-Wigner PDF. Args: - x: tf.Tensor + x: value(s) for which the CDF will be calculated. + m: Mean value + gamma: width Returns: - .. math:: \mathrm{arctan}(x) + `tf.Tensor`: The calculated PDF values. Notes: - Formula is taken from https://www.wolframalpha.com/input/?i=arctan%28a%2Bb*i%29 - TODO: move somewhere? + Based on code from this [github gist](https://gist.github.com/andrewfowlie/cd0ed7e6c96f7c9e88f85eb3b9665b97#file-bw-py-L87-L110) """ - return 1 / 2 * 1j * (tf.math.log(1 - 1j * x) - tf.math.log(1 + 1j * x)) + x = z.unstack_x(x) + alpha = gamma / m + gamma2 = m ** 2 * (1.0 + alpha ** 2) ** 0.5 + k = ( + 2.0 ** (3.0 / 2.0) + * m ** 2 + * alpha + * gamma2 + / (np.pi * (m ** 2 + gamma2) ** 0.5) + ) + + return k / ((x ** 2 - m ** 2) ** 2 + m ** 4 * alpha ** 2) class RelativisticBreitWigner(zfit.pdf.ZPDF): @@ -49,18 +58,28 @@ def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor: Returns: `tf.Tensor`: The value(s) of the unnormalized PDF at x. """ - x = z.unstack_x(x) - alpha = self.params["gamma"] / self.params["m"] - gamma = self.params["m"] ** 2 * (1.0 + alpha ** 2) ** 0.5 - k = ( - 2.0 ** (3.0 / 2.0) - * self.params["m"] ** 2 - * alpha - * gamma - / (np.pi * (self.params["m"] ** 2 + gamma) ** 0.5) - ) + return relbw_pdf_func(x, m=self.params['m'], gamma=self.params['gamma']) + - return k / ((x ** 2 - self.params["m"] ** 2) ** 2 + self.params["m"] ** 4 * alpha ** 2) +@z.function(wraps='tensor') +def arctan_complex(x): + r"""Function that evaluates arctan(x) using tensorflow but also supports complex numbers. + It is defined as + .. math:: + + \mathrm{arctan}(x) = \frac{i}{2} \left(\ln(1-ix) - \ln(1+ix)\right) + + Args: + x: tf.Tensor + + Returns: + .. math:: \mathrm{arctan}(x) + + Notes: + Formula is taken from https://www.wolframalpha.com/input/?i=arctan%28a%2Bb*i%29 + TODO: move somewhere? + """ + return 1 / 2 * 1j * (tf.math.log(1 - 1j * x) - tf.math.log(1 + 1j * x)) @z.function(wraps='tensor')