-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
6 changed files
with
160 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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()) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1 +1 @@ | ||
from . import pdf | ||
from . import pdf, func |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
from ..models.model_breit_wigner import RelativisticBreitWigner |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |