diff --git a/tests/compwa/test_wrapper.py b/tests/compwa/test_wrapper.py new file mode 100644 index 0000000..9577d4a --- /dev/null +++ b/tests/compwa/test_wrapper.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import numpy as np +import pandas as pd + +import zfit_physics.compwa as zcompwa + + +# @pytest.fixture() +def create_amplitude(): + import qrules + + reaction = qrules.generate_transitions( + initial_state=("J/psi(1S)", [-1, +1]), + final_state=["gamma", "pi0", "pi0"], + allowed_intermediate_particles=["f(0)"], + allowed_interaction_types=["strong", "EM"], + formalism="helicity", + ) + + import ampform + from ampform.dynamics.builder import create_non_dynamic_with_ff, create_relativistic_breit_wigner_with_ff + + model_builder = ampform.get_builder(reaction) + model_builder.scalar_initial_state_mass = True + model_builder.stable_final_state_ids = [0, 1, 2] + model_builder.set_dynamics("J/psi(1S)", create_non_dynamic_with_ff) + for name in reaction.get_intermediate_particles().names: + model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff) + model = model_builder.formulate() + + from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator + + rng = TFUniformRealNumberGenerator(seed=0) + phsp_generator = TFPhaseSpaceGenerator( + initial_state_mass=reaction.initial_state[-1].mass, + final_state_masses={i: p.mass for i, p in reaction.final_state.items()}, + ) + phsp_momenta = phsp_generator.generate(100_000, rng) + + unfolded_expression = model.expression.doit() + + return model, reaction + + +def test_wrapper_simple(): + import zfit + + model, reaction = create_amplitude() + + from tensorwaves.function.sympy import create_parametrized_function + + unfolded_expression = model.expression.doit() + intensity_func = create_parametrized_function( + expression=unfolded_expression, + parameters=model.parameter_defaults, + backend="tensorflow", + ) + from tensorwaves.data import SympyDataTransformer + + helicity_transformer = SympyDataTransformer.from_sympy( + model.kinematic_variables, backend="numpy" + ) + from tensorwaves.data import ( + IntensityDistributionGenerator, + TFPhaseSpaceGenerator, + TFUniformRealNumberGenerator, + TFWeightedPhaseSpaceGenerator, + ) + + rng = TFUniformRealNumberGenerator(seed=0) + phsp_generator = TFPhaseSpaceGenerator( + initial_state_mass=reaction.initial_state[-1].mass, + final_state_masses={i: p.mass for i, p in reaction.final_state.items()}, + ) + phsp_momenta = phsp_generator.generate(100_000, rng) + + weighted_phsp_generator = TFWeightedPhaseSpaceGenerator( + initial_state_mass=reaction.initial_state[-1].mass, + final_state_masses={i: p.mass for i, p in reaction.final_state.items()}, + ) + data_generator = IntensityDistributionGenerator( + domain_generator=weighted_phsp_generator, + function=intensity_func, + domain_transformer=helicity_transformer, + ) + data_momenta = data_generator.generate(10_000, rng) + + phsp = helicity_transformer(phsp_momenta) + data = helicity_transformer(data_momenta) + data_frame = pd.DataFrame(data) + phsp_frame = pd.DataFrame(phsp) + initial_parameters = { + R"C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}": ( + 1.0 + ), + "m_{f_{0}(500)}": 0.4, + "m_{f_{0}(980)}": 0.88, + "m_{f_{0}(1370)}": 1.22, + "m_{f_{0}(1500)}": 1.45, + "m_{f_{0}(1710)}": 1.83, + R"\Gamma_{f_{0}(500)}": 0.3, + R"\Gamma_{f_{0}(980)}": 0.1, + R"\Gamma_{f_{0}(1710)}": 0.3, + } + + # data conversion + # phsp_zfit = zfit.Data.from_pandas(phsp_frame) + # data_zfit = zfit.Data.from_pandas(data_frame) + data_frame = data_frame.astype(np.float64) + phsp_frame = phsp_frame.astype(np.float64) + intensity = intensity_func + + pdf = zcompwa.pdf.ComPWAPDF( + intensity=intensity, + norm=phsp_frame, + ) + # for p in pdf.get_params(): # returns the free params automatically + # p.set_value(p + np.random.normal(0, 0.01)) + # zfit.run.set_graph_mode(False) + zfit.run.set_autograd_mode(False) + loss = zfit.loss.UnbinnedNLL(pdf, data_frame) + + # ok, here I was caught up playing around :) Minuit seems to perform the best though + # minimizer = zfit.minimize.Minuit(verbosity=7, gradient=True) + # minimizer = zfit.minimize.Minuit(verbosity=7, gradient='zfit') + minimizer = zfit.minimize.ScipyLBFGSBV1(verbosity=8) + # minimizer = zfit.minimize.ScipyTrustKrylovV1(verbosity=8) + # minimizer = zfit.minimize.NLoptMMAV1(verbosity=9) + # minimizer = zfit.minimize.IpyoptV1(verbosity=8) + result = minimizer.minimize(loss) + print(result) + result.hesse() + print(result) + assert result.valid diff --git a/zfit_physics/compwa/__init__.py b/zfit_physics/compwa/__init__.py new file mode 100644 index 0000000..7ad62a3 --- /dev/null +++ b/zfit_physics/compwa/__init__.py @@ -0,0 +1,3 @@ +from . import data, pdf, variables + +__all__ = ["data", "pdf", "variables"] diff --git a/zfit_physics/compwa/data.py b/zfit_physics/compwa/data.py new file mode 100644 index 0000000..e69de29 diff --git a/zfit_physics/compwa/pdf.py b/zfit_physics/compwa/pdf.py new file mode 100644 index 0000000..4fe0dcb --- /dev/null +++ b/zfit_physics/compwa/pdf.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import numpy as np +import tensorflow as tf +import zfit # suppress tf warnings +import zfit.z.numpy as znp +from zfit import supports, z + +from .variables import obs_from_frame + + +def patched_call(self, data, *, params) -> np.ndarray: + # extended_data = {**self.__parameters, **data} # type: ignore[arg-type] + if params is not None: + self.update_parameters(params) + return self.__function(data) # type: ignore[arg-type] + + +class ComPWAPDF(zfit.pdf.BasePDF): + def __init__(self, intensity, norm, obs=None, params=None, extended=None, name="ComPWA"): + """ComPWA intensity normalized over the *norm* dataset.""" + if params is None: + params = { + name: zfit.param.convert_to_parameter(val, name=name, prefer_constant=False) + for name, val in intensity.parameters.items() + } + if obs is None: + obs = obs_from_frame(norm) + # intensity.__call__ = types.MethodType(patched_call, intensity) + super().__init__(obs, params=params, name=name, extended=extended) + self.intensity = intensity + norm = {ob: znp.array(ar) for ob, ar in zip(self.obs, z.unstack_x(norm))} + self.norm_sample = norm + + @supports(norm=True) + def _pdf(self, x, norm, params): + paramvalsfloat = [] + paramvalscomplex = [] + iscomplex = [] + # we need to split complex and floats to pass them to the numpy function, as it creates a tensor + for val in params.values(): + if val.dtype == znp.complex128: + iscomplex.append(True) + paramvalscomplex.append(val) + paramvalsfloat.append(znp.zeros_like(val, dtype=znp.float64)) + else: + iscomplex.append(False) + paramvalsfloat.append(val) + paramvalscomplex.append(znp.zeros_like(val, dtype=znp.complex128)) + + def unnormalized_pdf(x, paramvalsfloat, paramvalscomplex): + data = {ob: znp.array(ar) for ob, ar in zip(self.obs, x)} + paramsinternal = { + n: c if isc else f for n, f, c, isc in zip(params.keys(), paramvalsfloat, paramvalscomplex, iscomplex) + } + self.intensity.update_parameters(paramsinternal) + return self.intensity(data) + + xunstacked = z.unstack_x(x) + + probs = tf.numpy_function(unnormalized_pdf, [xunstacked, paramvalsfloat, paramvalscomplex], Tout=tf.float64) + if norm is not False: + normvalues = [znp.asarray(self.norm_sample[ob]) for ob in self.obs] + norm = znp.mean( + tf.numpy_function(unnormalized_pdf, [normvalues, paramvalsfloat, paramvalscomplex], Tout=tf.float64) + ) + norm.set_shape([]) + probs /= norm + probs.set_shape([None]) + return probs + + @z.function(wraps="tensorwaves") + def _jitted_normalization(self, norm, params): + return znp.mean(self._jitted_unnormalized_pdf(norm, params=params)) diff --git a/zfit_physics/compwa/variables.py b/zfit_physics/compwa/variables.py new file mode 100644 index 0000000..f468156 --- /dev/null +++ b/zfit_physics/compwa/variables.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from typing import Mapping + +import numpy as np +import pandas as pd +import zfit +from zfit.core.interfaces import ZfitUnbinnedData + + +def obs_from_frame(frame1, frame2=None, bufferfactor=0.01): + obs = [] + if frame2 is None: + frame2 = frame1 + + if isinstance(frame1, ZfitUnbinnedData) or isinstance(frame2, ZfitUnbinnedData): + return frame1.space + + if not isinstance(frame1, (Mapping, pd.DataFrame)) or not isinstance(frame2, (Mapping, pd.DataFrame)): + msg = "frame1 and frame2 have to be either a mapping or a pandas DataFrame, or a zfit Data object. They are currently of type: " + raise ValueError( + msg, + type(frame1), + type(frame2), + ) + for ob in frame2: + minimum = np.min([np.min(frame1[ob]), np.min(frame2[ob])]) + maximum = np.max([np.max(frame1[ob]), np.max(frame2[ob])]) + dist = maximum - minimum + buffer = bufferfactor * dist + obs.append( + zfit.Space( + ob, + limits=( + minimum - buffer, + maximum + buffer, + ), + ) + ) + return zfit.dimension.combine_spaces(*obs)