WIP refactor!: implement HelicityKinematics.convert

Removes the need for SubSystem, ParticleReactionInfo, etc.

cspell.json

@@ -138,6 +138,7 @@
         "jupyterlab",
         "keras",
         "kernelspec",
+        "lambdification",
         "lambdify",
         "linestyle",
         "linkcheck",

src/tensorwaves/data/generate.py

@@ -4,6 +4,7 @@
 from typing import Callable, Optional
 
 import numpy as np
+from expertsystem.amplitude.kinematics import HelicityKinematics, ReactionInfo
 from tqdm import tqdm
 
 from tensorwaves.data.tf_phasespace import (
@@ -12,22 +13,17 @@
 )
 from tensorwaves.interfaces import (
     Function,
-    Kinematics,
     PhaseSpaceGenerator,
     UniformRealNumberGenerator,
 )
-from tensorwaves.physics.helicity_formalism.kinematics import (
-    HelicityKinematics,
-    ParticleReactionKinematicsInfo,
-)
 
 
 def _generate_data_bunch(
     bunch_size: int,
     phsp_generator: PhaseSpaceGenerator,
     random_generator: UniformRealNumberGenerator,
     intensity: Function,
-    kinematics: Kinematics,
+    kinematics: HelicityKinematics,
 ) -> np.ndarray:
     phsp_sample, weights = phsp_generator.generate(
         bunch_size, random_generator
@@ -38,17 +34,17 @@ def _generate_data_bunch(
 
     uniform_randoms = random_generator(bunch_size, max_value=maxvalue)
 
-    phsp_sample = phsp_sample.transpose(1, 0, 2)
-
-    return (phsp_sample[weights * intensities > uniform_randoms], maxvalue)
+    np_phsp_sample = np.array(phsp_sample.values())
+    np_phsp_sample = np_phsp_sample.transpose(1, 0, 2)
+    return (np_phsp_sample[weights * intensities > uniform_randoms], maxvalue)
 
 
 def generate_data(
     size: int,
     kinematics: HelicityKinematics,
     intensity: Function,
     phsp_generator: Callable[
-        [ParticleReactionKinematicsInfo], PhaseSpaceGenerator
+        [ReactionInfo], PhaseSpaceGenerator
     ] = TFPhaseSpaceGenerator,
     random_generator: Optional[UniformRealNumberGenerator] = None,
     bunch_size: int = 50000,
@@ -57,19 +53,17 @@ def generate_data(
 
     Args:
         size: Sample size to generate.
-        kinematics: A kinematics instance. Note that this instance must have a
-            property :attr:`~.HelicityKinematics.reaction_kinematics_info` of
-            the type `.ParticleReactionKinematicsInfo`, otherwise the phase
-            space generator instance cannot be constructed.
-        intensity: The intensity which will be sampled.
+        kinematics: A `~expertsystem.amplitude.kinematics.HelicityKinematics`
+            instance.
+        intensity: The intensity `.Function` that will be sampled.
         phsp_generator: Class of a phase space generator.
         random_generator: A uniform real random number generator. Defaults to
             `.TFUniformRealNumberGenerator` with **indeterministic** behavior.
         bunch_size: Adjusts size of a bunch. The requested sample size is
             generated from many smaller samples, aka bunches.
 
     """
-    phsp_gen_instance = phsp_generator(kinematics.reaction_kinematics_info)
+    phsp_gen_instance = phsp_generator(kinematics.reaction_info)
     if random_generator is None:
         random_generator = TFUniformRealNumberGenerator()
 
@@ -114,7 +108,7 @@ def generate_phsp(
     size: int,
     kinematics: HelicityKinematics,
     phsp_generator: Callable[
-        [ParticleReactionKinematicsInfo], PhaseSpaceGenerator
+        [ReactionInfo], PhaseSpaceGenerator
     ] = TFPhaseSpaceGenerator,
     random_generator: Optional[UniformRealNumberGenerator] = None,
     bunch_size: int = 50000,
@@ -124,17 +118,19 @@ def generate_phsp(
     Args:
         size: Sample size to generate.
         kinematics: A kinematics instance. Note that this instance must have a
-            property :attr:`~.HelicityKinematics.reaction_kinematics_info` of
-            the type `.ParticleReactionKinematicsInfo`, otherwise the phase
-            space generator instance cannot be constructed.
+            property
+            `~expertsystem.amplitude.kinematics.HelicityKinematics.reaction_info`
+            of the type
+            `expertsystem.amplitude.kinematics.ReactionInfo`,
+            otherwise the phase space generator instance cannot be constructed.
         phsp_generator: Class of a phase space generator.
         random_generator: A uniform real random number generator. Defaults to
             `.TFUniformRealNumberGenerator` with **indeterministic** behavior.
         bunch_size: Adjusts size of a bunch. The requested sample size is
             generated from many smaller samples, aka bunches.
 
     """
-    phsp_gen_instance = phsp_generator(kinematics.reaction_kinematics_info)
+    phsp_gen_instance = phsp_generator(kinematics.reaction_info)
     if random_generator is None:
         random_generator = TFUniformRealNumberGenerator()
 
@@ -145,14 +141,15 @@ def generate_phsp(
     )
     events = np.array([])
     while np.size(events, 0) < size:
-        four_momenta, weights = phsp_gen_instance.generate(
+        phsp_sample, weights = phsp_gen_instance.generate(
             bunch_size, random_generator
         )
-        four_momenta = four_momenta.transpose(1, 0, 2)
+        np_phsp_sample = np.array(phsp_sample.values())
+        np_phsp_sample = np_phsp_sample.transpose(1, 0, 2)
 
         hit_and_miss_randoms = random_generator(bunch_size)
 
-        bunch = four_momenta[weights > hit_and_miss_randoms]
+        bunch = np_phsp_sample[weights > hit_and_miss_randoms]
 
         if np.size(events, 0) > 0:
             events = np.vstack((events, bunch))

src/tensorwaves/data/tf_phasespace.py

@@ -1,7 +1,8 @@
 """Phase space generation using tensorflow."""
 
-from typing import Optional
+from typing import Dict, Optional, Tuple
 
+import expertsystem.amplitude.kinematics as es
 import numpy as np
 import phasespace
 import tensorflow as tf
@@ -11,25 +12,25 @@
     PhaseSpaceGenerator,
     UniformRealNumberGenerator,
 )
-from tensorwaves.physics.helicity_formalism.kinematics import (
-    ParticleReactionKinematicsInfo,
-)
 
 
 class TFPhaseSpaceGenerator(PhaseSpaceGenerator):
     """Implements a phase space generator using tensorflow."""
 
-    def __init__(
-        self, reaction_kinematics_info: ParticleReactionKinematicsInfo
-    ) -> None:
+    def __init__(self, reaction_info: es.ReactionInfo) -> None:
+        initial_states = reaction_info.initial_state.values()
+        if len(initial_states) != 1:
+            raise ValueError("Not a 1-to-n body decay")
+        initial_state = next(iter(initial_states))
         self.phsp_gen = phasespace.nbody_decay(
-            reaction_kinematics_info.total_invariant_mass,
-            reaction_kinematics_info.final_state_masses,
+            mass_top=initial_state.mass,
+            masses=[p.mass for p in reaction_info.final_state.values()],
+            names=list(map(str, reaction_info.final_state)),
         )
 
     def generate(
         self, size: int, rng: UniformRealNumberGenerator
-    ) -> np.ndarray:
+    ) -> Tuple[Dict[int, np.ndarray], np.ndarray]:
         if not isinstance(rng, TFUniformRealNumberGenerator):
             raise TypeError(
                 f"{TFPhaseSpaceGenerator.__name__} requires a "
@@ -39,10 +40,11 @@ def generate(
         weights, particles = self.phsp_gen.generate(
             n_events=size, seed=rng.generator
         )
-        particles = np.array(
-            tuple(particles[x].numpy() for x in particles.keys())
-        )
-        return particles, weights.numpy()
+        momentum_pool = {
+            int(label): momenta.numpy().T
+            for label, momenta in particles.items()
+        }
+        return momentum_pool, weights.numpy()
 
 
 class TFUniformRealNumberGenerator(UniformRealNumberGenerator):

src/tensorwaves/estimator.py

@@ -2,14 +2,12 @@
 
 All estimators have to implement the `~.interfaces.Estimator` interface.
 """
-from typing import Callable, Dict, List, Union
+from typing import Callable, Dict, List, Mapping, Union
+
+import sympy as sp
 
 from tensorwaves.interfaces import Estimator
-from tensorwaves.physics.amplitude import (
-    SympyModel,
-    get_backend_modules,
-    lambdify,
-)
+from tensorwaves.physics.amplitude import get_backend_modules, lambdify
 
 
 def gradient_creator(
@@ -53,7 +51,8 @@ class SympyUnbinnedNLL(  # pylint: disable=too-many-instance-attributes
 
     def __init__(
         self,
-        model: SympyModel,
+        expression: sp.Expr,
+        parameters: Mapping[sp.Symbol, Union[float, complex]],
         dataset: dict,
         phsp_dataset: dict,
         phsp_volume: float = 1.0,
@@ -62,11 +61,9 @@ def __init__(
         self.__gradient = gradient_creator(self.__call__, backend)
         backend_modules = get_backend_modules(backend)
 
-        self.__parameters: Dict[str, Union[float, complex]] = {
-            k.name: v for k, v in model.parameters.items()
-        }
+        self.__parameters = {s.name: v for s, v in parameters.items()}
 
-        model_expr = model.expression.doit()
+        model_expr = expression.doit()
 
         self.__bare_model = lambdify(
             tuple(model_expr.free_symbols),

src/tensorwaves/interfaces.py

@@ -3,6 +3,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Iterable, Optional, Tuple, Union
 
+import numpy as np
+
 
 class Function(ABC):
     """Interface of a callable function.
@@ -110,5 +112,11 @@ class PhaseSpaceGenerator(ABC):
     """Abstract class for generating phase space samples."""
 
     @abstractmethod
-    def generate(self, size: int, rng: UniformRealNumberGenerator) -> dict:
-        """Generate phase space sample."""
+    def generate(
+        self, size: int, rng: UniformRealNumberGenerator
+    ) -> Tuple[Dict[int, np.ndarray], np.ndarray]:
+        """Generate phase space sample.
+
+        Returns a `tuple` of a mapping of final state IDs to `numpy.array` s
+        with four-momentum tuples.
+        """

src/tensorwaves/physics/__init__.py

@@ -1,7 +1,7 @@
 """Evaluateable physics models for amplitude analysis."""
 
 __all__ = [
-    "helicity_formalism",
+    "amplitude",
 ]
 
-from . import helicity_formalism
+from . import amplitude

src/tensorwaves/physics/amplitude.py

@@ -1,39 +1,13 @@
-"""`~.Function` Adapter for `sympy` based models."""
+"""`.Function` Adapter for `sympy`-based models."""
 
 import logging
 from typing import Any, Callable, Dict, Tuple, Union
 
-import attr
-import sympy
+import sympy as sp
 
 from tensorwaves.interfaces import Function
 
 
-@attr.s(frozen=True)
-class SympyModel:
-    r"""Full definition of an arbitrary model based on sympy.
-
-    Note that input for particle physics amplitude models are based on four
-    momenta. However for reasons of convenience some models may define and use
-    a distinct set of kinematic variables (e.g. in the helicity formalism:
-    angles :math:`\\theta` and :math:`\\phi`). In this case a
-    `~.interfaces.Kinematics` instance (Adapter) is needed to convert four
-    momentum information into the custom set of kinematic variables.
-
-    Args:
-        expression : A sympy expression that contains the complete information
-          of the model based on some inputs. The inputs are defined via the
-          :code:`free_symbols` attribute of the sympy expression.
-        parameters: Defines which inputs of the model are parameters. The keys
-          represent the parameter set, while the values represent their default
-          values. Consequently the variables of the model are defined as the
-          intersection of the total input set with the parameter set.
-    """
-
-    expression: sympy.Expr = attr.ib()
-    parameters: Dict[sympy.Symbol, Union[float, complex]] = attr.ib()
-
-
 def get_backend_modules(
     backend: Union[str, tuple, dict],
 ) -> Union[str, tuple, dict]:
@@ -61,8 +35,8 @@ def get_backend_modules(
 
 
 def lambdify(
-    variables: Tuple[sympy.Symbol, ...],
-    expression: sympy.Expr,
+    variables: Tuple[sp.Symbol, ...],
+    expression: sp.Expr,
     backend: Union[str, tuple, dict],
 ) -> Callable:
     """Wrapper around `~sympy.utilities.lambdify.lambdify`.
@@ -76,7 +50,7 @@ def jax_lambdify() -> Callable:
         from jax import jit
 
         return jit(
-            sympy.lambdify(
+            sp.lambdify(
                 variables,
                 expression,
                 modules=backend_modules,
@@ -91,7 +65,7 @@ def jax_lambdify() -> Callable:
             from numba import jit
 
             return jit(
-                sympy.lambdify(
+                sp.lambdify(
                     variables,
                     expression,
                     modules="numpy",
@@ -102,7 +76,7 @@ def jax_lambdify() -> Callable:
         if any("jax" in x.__name__ for x in backend):
             return jax_lambdify()
 
-    return sympy.lambdify(
+    return sp.lambdify(
         variables,
         expression,
         modules=backend_modules,
@@ -124,20 +98,22 @@ class Intensity(Function):
     """
 
     def __init__(
-        self, model: SympyModel, backend: Union[str, tuple, dict] = "numpy"
+        self,
+        expression: sp.Expr,
+        parameters: Dict[sp.Symbol, Union[complex, float]],
+        backend: Union[str, tuple, dict] = "numpy",
     ):
-        full_sympy_model = model.expression.doit()
-        self.__input_variable_order = tuple(
+        full_sympy_model = expression.doit()
+        self.__input_variable_order: Tuple[str, ...] = tuple(
             x.name for x in full_sympy_model.free_symbols
         )
         self.__callable_model = lambdify(
             tuple(full_sympy_model.free_symbols),
             full_sympy_model,
             backend=backend,
         )
-
-        self.__parameters: Dict[str, Union[float, complex]] = {
-            k.name: v for k, v in model.parameters.items()
+        self.__parameters: Dict[str, Union[complex, float]] = {
+            s.name: v for s, v in parameters.items()
         }
 
     def __call__(self, dataset: Dict[str, Any]) -> Any: