WIP: Add Breit wigner (+squared) Function

tests/test_model_breit_wigner.py

@@ -0,0 +1,39 @@
+"""Example test for a pdf or function"""
+import pytest
+import zfit
+import numpy as np
+# Important, do the imports below
+from zfit.core.testing import setup_function, teardown_function, 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
+obs1 = zfit.Space("obs1", limits=(-4, 5))
+
+
+def test_bw_pdf():
+    # test special properties  here
+    bw = zphys.unstable.pdf.RelativisticBreitWignerSquared(obs=obs1, mres=1., wres=0.3)
+
+    integral = bw.integrate(limits=obs1)
+    assert pytest.approx(1., rel=1e-3) == zfit.run(integral)
+
+
+def test_bw_func():
+    # test special properties  here
+    bw = zphys.unstable.func.RelativisticBreitWigner(obs=obs1, mres=1., wres=0.3)
+
+    vals = bw.func(x=np.random.uniform(size=(100, 1)))
+    assert 100 == len(zfit.run(vals))
+
+
+# register the pdf here and provide sets of working parameter configurations
+
+def _bw_params_factory():
+    mres = zfit.Parameter('mres', param1_true)
+    wres = zfit.Parameter('wres', param2_true)
+    return {"mres": mres, "wres": wres}
+
+# tester.register_func(func_class=zfit.pdf.Gauss, params_factories=_bw_params_factory())

zfit_physics/kinematics.py

@@ -0,0 +1,31 @@
+import tensorflow as tf
+import zfit
+from zfit import ztf
+
+
+def metric_tensor():
+    """
+    Metric tensor for Lorentz space (constant)
+    """
+    return ztf.constant([-1., -1., -1., 1.], dtype=zfit.settings.ztypes.float)
+
+
+def mass_squared(vector):
+    """Calculate the squared mass for a Lorentz 4-momentum."""
+    return ztf.reduce_sum(vector * vector * metric_tensor(), axis=1)
+
+
+def lorentz_vector(space, time):
+    """
+    Make a Lorentz vector from spatial and time components
+        space : 3-vector of spatial components
+        time  : time component
+    """
+    return tf.concat([space, tf.stack([time], axis=1)], axis=1)
+
+def vector(x, y, z):
+    """
+    Make a 3-vector from components
+    x, y, z : vector components
+    """
+    return tf.stack([x, y, z], axis=1)

zfit_physics/models/model_breit_wigner.py

@@ -0,0 +1,86 @@
+import tensorflow as tf
+import zfit
+from zfit import ztf
+from zfit.util import ztyping
+
+from .. import kinematics
+
+
+def relativistic_breit_wigner(m2, mres, wres):
+    """Relativistic Breit-Wigner function.
+    """
+    second_part = tf.complex(ztf.constant(0.), mres) * ztf.to_complex(wres)
+    below_div = ztf.to_complex(mres ** 2 - m2) - second_part
+    return 1. / below_div
+
+
+def _bw_func(self, x):
+    """The Breit Wigner function wrapped to be used within a PDF or a Func directly.
+
+    Args:
+        self:
+        x:
+
+    Returns:
+
+    """
+    var = x.unstack_x()
+    if isinstance(var, list):
+        m_sq = kinematics.mass_squared(tf.reduce_sum(
+            [kinematics.lorentz_vector(kinematics.vector(px, py, pz), pe)
+             for px, py, pz, pe in zip(*[iter(var)] * 4)],
+            axis=0))
+    elif self.using_m_squared:
+        m_sq = var
+    else:
+        m_sq = var * tf.math.conj(
+            var)  # TODO(Albert): this was squared, but should be mult with conj, right?
+    mres = self.params['mres']
+    wres = self.params['wres']
+    return relativistic_breit_wigner(m_sq, mres, wres)
+
+
+class RelativisticBreitWigner(zfit.func.BaseFunc):
+
+    def __init__(self, obs: ztyping.ObsTypeInput, mres: ztyping.ParamTypeInput, wres: ztyping.ParamTypeInput,
+                 using_m_squared: bool = False,
+                 name: str = "RelativisticBreitWigner"):
+        """Relativistic Breit Wigner Function
+
+        Args:
+            obs: Space the function is defined in
+            mres: Mass of the resonance
+            wres: width of the resonance mass
+            using_m_squared: Whether the input mass is already the squared mass. If not, will be squared on the
+                fly. If the input is a list, the kinematics of daughter particles is assumed and the mass squared is
+                calculated
+            name: Name of the Function
+        """
+        self.using_m_squared = using_m_squared
+        # HACK to make it usable in while loop
+        # zfit.run._enable_parameter_autoconversion = False
+
+        super().__init__(obs=obs, name=name, dtype=zfit.settings.ztypes.complex,
+                         params={'mres': mres, 'wres': wres})
+
+        # zfit.run._enable_parameter_autoconversion = True
+        # HACK end
+
+    def _func(self, x):
+        return _bw_func(self, x)
+
+
+class RelativisticBreitWignerSquared(zfit.pdf.BasePDF):
+
+    def __init__(self, obs: ztyping.ObsTypeInput, mres: ztyping.ParamTypeInput, wres: ztyping.ParamTypeInput,
+                 using_m_squared: bool = False, name="RelativisticBreitWignerPDF"):
+        self.using_m_squared = using_m_squared
+
+        super().__init__(obs=obs, name=name, dtype=zfit.settings.ztypes.float,
+                         params={'mres': mres, 'wres': wres})
+
+    def _unnormalized_pdf(self, x):
+        propagator = _bw_func(self, x)
+        val = propagator * tf.math.conj(propagator)
+        val = ztf.to_real(val)
+        return val

zfit_physics/unstable/__init__.py

@@ -1 +1 @@
-from . import pdf
+from . import pdf, func

zfit_physics/unstable/func.py

@@ -0,0 +1 @@
+from ..models.model_breit_wigner import RelativisticBreitWigner

zfit_physics/unstable/pdf.py

@@ -1,2 +1,3 @@
 from ..models.pdf_conv import NumConvPDFUnbinnedV1
-from ..models.pdf_kde import GaussianKDE
+from ..models.pdf_kde import GaussianKDE
+from ..models.model_breit_wigner import RelativisticBreitWignerSquared