WIP refactor: implement MomentumPool from expertsystem

src/tensorwaves/data/generate.py

@@ -4,6 +4,7 @@
 from typing import Callable, Optional
 
 import numpy as np
+from expertsystem.amplitude.data import MomentumPool
 from expertsystem.amplitude.kinematics import HelicityKinematics, ReactionInfo
 from tqdm import tqdm
 
@@ -36,7 +37,7 @@ def _generate_data_bunch(
 
     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)
+    return np_phsp_sample[weights * intensities > uniform_randoms], maxvalue
 
 
 def generate_data(
@@ -48,7 +49,7 @@ def generate_data(
     ] = TFPhaseSpaceGenerator,
     random_generator: Optional[UniformRealNumberGenerator] = None,
     bunch_size: int = 50000,
-) -> np.ndarray:
+) -> MomentumPool:
     """Facade function for creating data samples based on an intensities.
 
     Args:
@@ -101,7 +102,9 @@ def generate_data(
             events = bunch
         progress_bar.update()
     progress_bar.close()
-    return events[0:size].transpose(1, 0, 2)
+    events = events[0:size].transpose(1, 0, 2)
+    pos_to_state_id = dict(enumerate(kinematics.reaction_info.final_state))
+    return MomentumPool({pos_to_state_id[i]: events[i] for i in events})
 
 
 def generate_phsp(
@@ -112,7 +115,7 @@ def generate_phsp(
     ] = TFPhaseSpaceGenerator,
     random_generator: Optional[UniformRealNumberGenerator] = None,
     bunch_size: int = 50000,
-) -> np.ndarray:
+) -> MomentumPool:
     """Facade function for creating (unweighted) phase space samples.
 
     Args:
@@ -157,4 +160,6 @@ def generate_phsp(
             events = bunch
         progress_bar.update()
     progress_bar.close()
-    return events[0:size].transpose(1, 0, 2)
+    events = events[0:size].transpose(1, 0, 2)
+    pos_to_state_id = dict(enumerate(kinematics.reaction_info.final_state))
+    return MomentumPool({pos_to_state_id[i]: events[i] for i in events})

src/tensorwaves/data/tf_phasespace.py

@@ -1,11 +1,11 @@
 """Phase space generation using tensorflow."""
 
-from typing import Dict, Optional, Tuple
+from typing import Optional, Tuple
 
 import expertsystem.amplitude.kinematics as es
-import numpy as np
 import phasespace
 import tensorflow as tf
+from expertsystem.amplitude.data import MomentumPool, ScalarSequence
 from phasespace.random import get_rng
 
 from tensorwaves.interfaces import (
@@ -30,7 +30,7 @@ def __init__(self, reaction_info: es.ReactionInfo) -> None:
 
     def generate(
         self, size: int, rng: UniformRealNumberGenerator
-    ) -> Tuple[Dict[int, np.ndarray], np.ndarray]:
+    ) -> Tuple[MomentumPool, ScalarSequence]:
         if not isinstance(rng, TFUniformRealNumberGenerator):
             raise TypeError(
                 f"{TFPhaseSpaceGenerator.__name__} requires a "
@@ -40,11 +40,13 @@ def generate(
         weights, particles = self.phsp_gen.generate(
             n_events=size, seed=rng.generator
         )
-        momentum_pool = {
-            int(label): momenta.numpy().T
-            for label, momenta in particles.items()
-        }
-        return momentum_pool, weights.numpy()
+        momentum_pool = MomentumPool(
+            {
+                int(label): momenta.numpy().T
+                for label, momenta in particles.items()
+            }
+        )
+        return momentum_pool, ScalarSequence(weights.numpy())
 
 
 class TFUniformRealNumberGenerator(UniformRealNumberGenerator):
@@ -56,13 +58,15 @@ def __init__(self, seed: Optional[float] = None):
 
     def __call__(
         self, size: int, min_value: float = 0.0, max_value: float = 1.0
-    ) -> np.ndarray:
-        return self.generator.uniform(
-            shape=[size],
-            minval=min_value,
-            maxval=max_value,
-            dtype=self.dtype,
-        ).numpy()
+    ) -> ScalarSequence:
+        return ScalarSequence(
+            self.generator.uniform(
+                shape=[size],
+                minval=min_value,
+                maxval=max_value,
+                dtype=self.dtype,
+            ).numpy()
+        )
 
     @property
     def seed(self) -> Optional[float]:

src/tensorwaves/estimator.py

@@ -2,21 +2,23 @@
 
 All estimators have to implement the `~.interfaces.Estimator` interface.
 """
-from typing import Callable, Dict, List, Union
+from typing import Callable, Dict, List, Mapping, Union
+
+from expertsystem.amplitude.data import ScalarSequence
 
 from tensorwaves.interfaces import Estimator, Model
 from tensorwaves.physics.amplitude import get_backend_modules
 
 
 def gradient_creator(
-    function: Callable[[Dict[str, Union[float, complex]]], float],
+    function: Callable[[Mapping[str, Union[float, complex]]], float],
     backend: Union[str, tuple, dict],
 ) -> Callable[
-    [Dict[str, Union[float, complex]]], Dict[str, Union[float, complex]]
+    [Mapping[str, Union[float, complex]]], Dict[str, Union[float, complex]]
 ]:
     # pylint: disable=import-outside-toplevel
     def not_implemented(
-        parameters: Dict[str, Union[float, complex]]
+        parameters: Mapping[str, Union[float, complex]]
     ) -> Dict[str, Union[float, complex]]:
         raise NotImplementedError("Gradient not implemented.")
 
@@ -50,8 +52,8 @@ class SympyUnbinnedNLL(  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
         model: Model,
-        dataset: dict,
-        phsp_dataset: dict,
+        dataset: Mapping[str, Union[ScalarSequence, complex, float]],
+        phsp_dataset: Mapping[str, Union[ScalarSequence, complex, float]],
         phsp_volume: float = 1.0,
         backend: Union[str, tuple, dict] = "numpy",
     ) -> None:
@@ -77,14 +79,14 @@ def find_function_in_backend(name: str) -> Callable:
 
         self.__phsp_volume = phsp_volume
 
-        self.__data_args = []
-        self.__phsp_args = []
+        self.__data_args: Dict[str, Union[ScalarSequence, complex, float]] = {}
+        self.__phsp_args: Dict[str, Union[ScalarSequence, complex, float]] = {}
         self.__parameter_index_mapping: Dict[str, int] = {}
 
         for i, var_name in enumerate(model.variables):
             if var_name in dataset and var_name in phsp_dataset:
-                self.__data_args.append(dataset[var_name])
-                self.__phsp_args.append(phsp_dataset[var_name])
+                self.__data_args[var_name] = dataset[var_name]
+                self.__phsp_args[var_name] = phsp_dataset[var_name]
             elif var_name in dataset:
                 raise ValueError(
                     f"Datasets do not match! {var_name} exists in dataset but "
@@ -96,35 +98,36 @@ def find_function_in_backend(name: str) -> Callable:
                     "dataset but not in dataset."
                 )
             else:
-                self.__data_args.append(model.parameters[var_name])
-                self.__phsp_args.append(model.parameters[var_name])
+                self.__data_args[var_name] = model.parameters[var_name]
+                self.__phsp_args[var_name] = model.parameters[var_name]
                 self.__parameter_index_mapping[var_name] = i
 
-    def __call__(self, parameters: Dict[str, Union[float, complex]]) -> float:
+    def __call__(
+        self, parameters: Mapping[str, Union[float, complex]]
+    ) -> float:
         self.__update_parameters(parameters)
 
-        bare_intensities = self.__bare_model(*self.__data_args)
+        bare_intensities = self.__bare_model(self.__data_args)
         normalization_factor = 1.0 / (
             self.__phsp_volume
-            * self.__mean_function(self.__bare_model(*self.__phsp_args))
+            * self.__mean_function(self.__bare_model(self.__phsp_args))
         )
         likelihoods = normalization_factor * bare_intensities
         return -self.__sum_function(self.__log_function(likelihoods))
 
     def __update_parameters(
-        self, parameters: Dict[str, Union[float, complex]]
+        self, parameters: Mapping[str, Union[float, complex]]
     ) -> None:
         for par_name, value in parameters.items():
             if par_name in self.__parameter_index_mapping:
-                index = self.__parameter_index_mapping[par_name]
-                self.__data_args[index] = value
-                self.__phsp_args[index] = value
+                self.__data_args[par_name] = value
+                self.__phsp_args[par_name] = value
 
     @property
     def parameters(self) -> List[str]:
         return list(self.__parameter_index_mapping)
 
     def gradient(
-        self, parameters: Dict[str, Union[float, complex]]
+        self, parameters: Mapping[str, Union[float, complex]]
     ) -> Dict[str, Union[float, complex]]:
         return self.__gradient(parameters)

src/tensorwaves/interfaces.py

@@ -5,22 +5,18 @@
     Any,
     Dict,
     FrozenSet,
-    Generic,
     Iterable,
+    Mapping,
     Optional,
     Protocol,
     Tuple,
-    TypeVar,
     Union,
 )
 
-import numpy as np
+from expertsystem.amplitude.data import DataSet, MomentumPool, ScalarSequence
 
-DataType = TypeVar("DataType")
-"""Type of the data that is returned by `.Function.__call__`."""
 
-
-class Function(Protocol, Generic[DataType]):
+class Function(Protocol):
     """Interface of a callable function.
 
     The parameters of the model are separated from the domain variables. This
@@ -32,7 +28,9 @@ class Function(Protocol, Generic[DataType]):
     is to facilitate the events when parameters have changed.
     """
 
-    def __call__(self, dataset: Dict[str, DataType]) -> DataType:
+    def __call__(
+        self, dataset: Mapping[str, Union[ScalarSequence, complex, float]]
+    ) -> ScalarSequence:
         """Evaluate the function.
 
         Args:
@@ -60,13 +58,13 @@ def lambdify(self, backend: Union[str, tuple, dict]) -> Function:
         """
 
     @abstractmethod
-    def performance_optimize(self, fix_inputs: Dict[str, Any]) -> "Model":
+    def performance_optimize(self, fix_inputs: DataSet) -> "Model":
         """Create a performance optimized model, based on fixed inputs."""
 
     @property
     @abstractmethod
     def parameters(self) -> Dict[str, Union[float, complex]]:
-        """Get `dict` of parameters."""
+        """Get mapping of parameters to suggested initial values."""
 
     @property
     @abstractmethod
@@ -78,7 +76,9 @@ class Estimator(ABC):
     """Estimator for discrepancy model and data."""
 
     @abstractmethod
-    def __call__(self, parameters: Dict[str, Union[float, complex]]) -> float:
+    def __call__(
+        self, parameters: Mapping[str, Union[float, complex]]
+    ) -> float:
         """Evaluate discrepancy."""
 
     @property
@@ -88,7 +88,7 @@ def parameters(self) -> Iterable[str]:
 
     @abstractmethod
     def gradient(
-        self, parameters: Dict[str, Union[float, complex]]
+        self, parameters: Mapping[str, Union[float, complex]]
     ) -> Dict[str, Union[float, complex]]:
         """Calculate gradient for given parameter mapping."""
 
@@ -97,11 +97,11 @@ class Kinematics(ABC):
     """Abstract interface for computation of kinematic variables."""
 
     @abstractmethod
-    def convert(self, events: dict) -> dict:
+    def convert(self, events: MomentumPool) -> DataSet:
         """Convert a set of momentum tuples (events) to kinematic variables."""
 
     @abstractmethod
-    def is_within_phase_space(self, events: dict) -> Tuple[bool]:
+    def is_within_phase_space(self, events: MomentumPool) -> Tuple[bool]:
         """Check which events lie within phase space."""
 
     @property
@@ -114,7 +114,11 @@ class Optimizer(ABC):
     """Optimize a fit model to a data set."""
 
     @abstractmethod
-    def optimize(self, estimator: Estimator, initial_parameters: dict) -> dict:
+    def optimize(
+        self,
+        estimator: Estimator,
+        initial_parameters: Mapping[str, Union[float, complex]],
+    ) -> Dict[str, Any]:
         """Execute optimization."""
 
 
@@ -124,7 +128,7 @@ class UniformRealNumberGenerator(ABC):
     @abstractmethod
     def __call__(
         self, size: int, min_value: float = 0.0, max_value: float = 1.0
-    ) -> Union[float, list]:
+    ) -> ScalarSequence:
         """Generate random floats in the range from [min_value,max_value)."""
 
     @property  # type: ignore
@@ -144,7 +148,7 @@ class PhaseSpaceGenerator(ABC):
     @abstractmethod
     def generate(
         self, size: int, rng: UniformRealNumberGenerator
-    ) -> Tuple[Dict[int, np.ndarray], np.ndarray]:
+    ) -> Tuple[MomentumPool, ScalarSequence]:
         """Generate phase space sample.
 
         Returns a `tuple` of a mapping of final state IDs to `numpy.array` s

src/tensorwaves/optimizer/minuit.py

@@ -5,7 +5,7 @@
 import logging
 import time
 from datetime import datetime
-from typing import Dict, Iterable, Optional, Union
+from typing import Any, Dict, Iterable, Mapping, Optional, Union
 
 from iminuit import Minuit
 from tqdm import tqdm
@@ -16,7 +16,9 @@
 
 
 class ParameterFlattener:
-    def __init__(self, parameters: Dict[str, Union[float, complex]]) -> None:
+    def __init__(
+        self, parameters: Mapping[str, Union[float, complex]]
+    ) -> None:
         self.__real_imag_to_complex_name = {}
         self.__complex_to_real_imag_name = {}
         for name, val in parameters.items():
@@ -46,7 +48,7 @@ def unflatten(
         return parameters
 
     def flatten(
-        self, parameters: Dict[str, Union[float, complex]]
+        self, parameters: Mapping[str, Union[float, complex]]
     ) -> Dict[str, float]:
         flattened_parameters = {}
         for par_name, value in parameters.items():
@@ -80,8 +82,8 @@ def __init__(
     def optimize(  # pylint: disable=too-many-locals
         self,
         estimator: Estimator,
-        initial_parameters: Dict[str, Union[complex, float]],
-    ) -> dict:
+        initial_parameters: Mapping[str, Union[complex, float]],
+    ) -> Dict[str, Any]:
         parameter_handler = ParameterFlattener(initial_parameters)
         flattened_parameters = parameter_handler.flatten(initial_parameters)