enh: add compwa PDF wrapper

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

zfit_physics/compwa/__init__.py

@@ -0,0 +1,3 @@
+from . import data, pdf, variables
+
+__all__ = ["data", "pdf", "variables"]

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))

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)