FEAT: accept array-valued parameters

.cspell.json

@@ -242,6 +242,7 @@
     "vectorize",
     "venv",
     "vmap",
+    "vmapped",
     "weisskopf",
     "wirtinger",
     "xcode",

.pre-commit-config.yaml

@@ -16,7 +16,7 @@ repos:
       - id: check-useless-excludes
 
   - repo: https://github.com/ComPWA/policy
-    rev: 0.9.6
+    rev: 0.9.7
     hooks:
       - id: check-dev-files
 

benchmarks/ampform.py

@@ -26,6 +26,7 @@
     from tensorwaves.interface import (
         DataSample,
         FitResult,
+        FloatArray,
         Function,
         ParameterValue,
         ParametrizedFunction,
@@ -72,7 +73,7 @@ def create_function(
 
 def generate_data(
     model: HelicityModel,
-    function: Function[DataSample, np.ndarray],
+    function: Function[DataSample, FloatArray],
     data_sample_size: int,
     phsp_sample_size: int,
     backend: str,
@@ -109,7 +110,7 @@ def generate_data(
 def fit(
     data: DataSample,
     phsp: DataSample,
-    function: ParametrizedFunction[DataSample, np.ndarray],
+    function: ParametrizedFunction[DataSample, FloatArray],
     initial_parameters: Mapping[str, ParameterValue],
     backend: str,
 ) -> FitResult:

benchmarks/expression.py

@@ -12,7 +12,7 @@
 from tensorwaves.optimizer.scipy import ScipyMinimizer
 
 if TYPE_CHECKING:
-    from tensorwaves.interface import DataSample, Function
+    from tensorwaves.interface import Array, DataSample, Function
 
 
 def gaussian(x: sp.Symbol, mu: sp.Symbol, sigma: sp.Symbol) -> sp.Expr:
@@ -64,7 +64,7 @@ def _generate_domain(
 
 def _generate_data(
     size: int,
-    function: Function[DataSample, np.ndarray],
+    function: Function[DataSample, Array],
     rng: np.random.Generator,
     bunch_size: int = 10_000,
 ) -> DataSample:

benchmarks/unbinned_nll.py

@@ -18,6 +18,8 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+    from tensorwaves.interface import Array, DataSample
+
     def prange(stop: int) -> range: ...
 
 else:
@@ -193,7 +195,7 @@ def intensities() -> tuple[np.ndarray, np.ndarray]:
 
 
 @pytest.fixture(scope="module")
-def estimator_samples() -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+def estimator_samples() -> tuple[DataSample, DataSample]:
     rng = np.random.default_rng(seed=0)
     data = {"x": rng.uniform(low=-2.0, high=2.0, size=1_000_000)}
     phsp = {"x": rng.uniform(low=-2.0, high=2.0, size=1_000_000)}
@@ -221,7 +223,7 @@ def tensorflow_intensities(
 
 @pytest.fixture(scope="module")
 def jax_estimator_samples(
-    estimator_samples: tuple[dict[str, np.ndarray], dict[str, np.ndarray]],
+    estimator_samples: tuple[DataSample, DataSample],
 ) -> tuple[dict[str, jax.Array], dict[str, jax.Array]]:
     configure(jax_precision="float64")
     data, phsp = estimator_samples
@@ -233,7 +235,7 @@ def jax_estimator_samples(
 
 @pytest.fixture(scope="module")
 def tensorflow_estimator_samples(
-    estimator_samples: tuple[dict[str, np.ndarray], dict[str, np.ndarray]],
+    estimator_samples: tuple[DataSample, DataSample],
 ) -> tuple[dict[str, tf.Tensor], dict[str, tf.Tensor]]:
     data, phsp = estimator_samples
     return {"x": tnp.asarray(data["x"])}, {"x": tnp.asarray(phsp["x"])}
@@ -267,8 +269,8 @@ def _create_estimator(
 
 
 def _compute_estimator_reference(
-    data: dict[str, np.ndarray],
-    phsp: dict[str, np.ndarray],
+    data: DataSample,
+    phsp: DataSample,
     center: float,
 ) -> float:
     data_intensities = _numpy_intensity(data["x"], center)
@@ -277,16 +279,16 @@ def _compute_estimator_reference(
 
 
 def _benchmark_estimator_numpy(
-    benchmark: Callable[[Callable[[], float]], float],
+    benchmark: Callable[[Callable[[], float | Array]], float | Array],
     backend: str,
-    data: dict[str, np.ndarray],
-    phsp: dict[str, np.ndarray],
+    data: DataSample,
+    phsp: DataSample,
     parameters: dict[str, float],
-) -> float:
+) -> float | Array:
     estimator = _create_estimator(backend, data, phsp)
     estimator(parameters)
 
-    def run() -> float:
+    def run() -> float | Array:
         return estimator(parameters)
 
     return benchmark(run)
@@ -413,7 +415,7 @@ def test_unbinned_nll_normalization_formula(
 def test_unbinned_nll_estimator(
     benchmark,
     backend: str,
-    estimator_samples: tuple[dict[str, np.ndarray], dict[str, np.ndarray]],
+    estimator_samples: tuple[DataSample, DataSample],
     request: pytest.FixtureRequest,
 ) -> None:
     data, phsp = estimator_samples

docs/conf.py

@@ -59,21 +59,28 @@ def get_tensorflow_url() -> str:
 add_module_names = False
 api_github_repo = f"{ORGANIZATION}/{REPO_NAME}"
 api_target_substitutions: dict[str, str | tuple[str, str]] = {
+    "Array": "tensorwaves.interface.Array",
     "DataSample": "tensorwaves.interface.DataSample",
+    "FloatArray": "tensorwaves.interface.FloatArray",
+    "np.floating": "numpy.floating",
     "np.ndarray": "numpy.ndarray",
     "ParameterValue": "tensorwaves.interface.ParameterValue",
     "Path": "pathlib.Path",
     "ProgressColumn": "rich.progress.ProgressColumn",
+    "ScalarT": "tensorwaves.interface.ScalarT",
     "sp.Basic": "sympy.core.basic.Basic",
     "sp.Expr": "sympy.core.expr.Expr",
     "sp.Symbol": "sympy.core.symbol.Symbol",
     "TypeAliasForwardRef": "typing.TypeAlias",
 }
 api_target_types: dict[str, str | tuple[str, str]] = {
+    "tensorwaves.interface.Array": "obj",
     "tensorwaves.interface.DataSample": "obj",
+    "tensorwaves.interface.FloatArray": "obj",
     "tensorwaves.interface.InputType": "obj",
     "tensorwaves.interface.OutputType": "obj",
     "tensorwaves.interface.ParameterValue": "obj",
+    "tensorwaves.interface.ScalarT": "obj",
 }
 author = "Common Partial Wave Analysis"
 autodoc_default_options = {
@@ -83,10 +90,13 @@ def get_tensorflow_url() -> str:
 }
 autodoc_member_order = "bysource"
 autodoc_type_aliases = {
+    "Array": "tensorwaves.interface.Array",
     "DataSample": "tensorwaves.interface.DataSample",
+    "FloatArray": "tensorwaves.interface.FloatArray",
     "InputType": "tensorwaves.interface.InputType",
     "OutputType": "tensorwaves.interface.OutputType",
     "ParameterValue": "tensorwaves.interface.ParameterValue",
+    "ScalarT": "tensorwaves.interface.ScalarT",
 }
 autodoc_typehints_format = "short"
 autosectionlabel_prefix_document = True

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "rich",
     "sympy >=1.9",  # lambdify cse
     "tqdm >=4.24.0",  # autonotebook
+    "typing-extensions >=4.4.0; python_version <'3.13.0'",  # TypeVar defaults
 ]
 dynamic = ["version"]
 
@@ -614,7 +615,7 @@ key-empty = "off"
 
 [[tool.tombi.schemas]]
 root = "tool.compwa.policy"
-path = "https://raw.githubusercontent.com/ComPWA/policy/0.9.6/compwa-policy.schema.json"
+path = "https://raw.githubusercontent.com/ComPWA/policy/0.9.7/compwa-policy.schema.json"
 include = ["pyproject.toml"]
 
 [[tool.ty.overrides]]

src/tensorwaves/data/__init__.py

@@ -11,6 +11,7 @@
     DataGenerator,
     DataSample,
     DataTransformer,
+    FloatArray,
     Function,
     RealNumberGenerator,
 )
@@ -75,7 +76,7 @@ class IntensityDistributionGenerator(DataGenerator):
     def __init__(
         self,
         domain_generator: DataGenerator,
-        function: Function[DataSample, np.ndarray],
+        function: Function[DataSample, FloatArray],
         domain_transformer: DataTransformer | None = None,
         bunch_size: int = 50_000,
     ) -> None:
@@ -123,7 +124,7 @@ def _generate_bunch(self, rng: RealNumberGenerator) -> tuple[DataSample, float]:
         )
         transformed_domain = self.__domain_transformer(domain)
         computed_intensities = self.__function(transformed_domain)
-        max_intensity: float = np.max(computed_intensities)
+        max_intensity = float(np.max(computed_intensities))
         random_intensities = rng(size=self.__bunch_size, max_value=max_intensity)
         weights = domain.get("weights", 1)
         hit_and_miss_sample = select_events(

src/tensorwaves/data/_data_sample.py

@@ -12,7 +12,7 @@
 
     from tqdm.auto import tqdm
 
-    from tensorwaves.interface import DataSample
+    from tensorwaves.interface import Array, DataSample
 
 
 def get_number_of_events(four_momenta: DataSample) -> int:
@@ -28,7 +28,7 @@ def merge_events(sample1: DataSample, sample2: DataSample) -> DataSample:
 
 def _determine_merge_method(
     sample: DataSample,
-) -> Callable[[tuple[np.ndarray, np.ndarray]], np.ndarray]:
+) -> Callable[[tuple[Array, Array]], Array]:
     if len(sample) == 0:
         return operator.itemgetter(1)
     some_array = next(iter(sample.values()))
@@ -44,7 +44,7 @@ def _determine_merge_method(
 def _merge_events(
     sample1: DataSample,
     sample2: DataSample,
-    merge_method: Callable[[tuple[np.ndarray, np.ndarray]], np.ndarray],
+    merge_method: Callable[[tuple[Array, Array]], Array],
 ) -> DataSample:
     if len(sample1) and len(sample2) and set(sample1) != set(sample2):
         msg = "Keys of data sets are not matching"

src/tensorwaves/data/rng.py

@@ -8,7 +8,7 @@
 
 from tensorwaves.config import _tensorflow_precision
 from tensorwaves.function._backend import raise_missing_module_error
-from tensorwaves.interface import RealNumberGenerator
+from tensorwaves.interface import Array, RealNumberGenerator
 
 if TYPE_CHECKING:  # pragma: no cover
     import tensorflow as tf
@@ -24,7 +24,7 @@ def __init__(self, seed: int | None = None) -> None:
 
     def __call__(
         self, size: int, min_value: float = 0.0, max_value: float = 1.0
-    ) -> np.ndarray:
+    ) -> Array:
         return self.generator.uniform(size=size, low=min_value, high=max_value)
 
     @property
@@ -50,7 +50,7 @@ def __init__(self, seed: int | None = None) -> None:
 
     def __call__(
         self, size: int, min_value: float = 0.0, max_value: float = 1.0
-    ) -> np.ndarray:
+    ) -> Array:
         return self.generator.uniform(
             shape=[size],
             minval=min_value,

src/tensorwaves/data/transform.py

@@ -8,14 +8,13 @@
 
 from tensorwaves.function import PositionalArgumentFunction
 from tensorwaves.function.sympy import _get_free_symbols, _lambdify_normal_or_fast
-from tensorwaves.interface import DataSample, DataTransformer, Function
+from tensorwaves.interface import Array, DataSample, DataTransformer, Function
 
 from ._attrs import to_tuple
 
 if TYPE_CHECKING:  # pragma: no cover
     from collections.abc import Mapping
 
-    import numpy as np
     import sympy as sp
 
 
@@ -55,9 +54,7 @@ def __call__(self, data: DataSample) -> DataSample:
 class SympyDataTransformer(DataTransformer):
     """Implementation of a `.DataTransformer`."""
 
-    def __init__(
-        self, functions: Mapping[str, Function[DataSample, np.ndarray]]
-    ) -> None:
+    def __init__(self, functions: Mapping[str, Function[DataSample, Array]]) -> None:
         if any(not isinstance(f, Function) for f in functions.values()):
             msg = (
                 f"Not all values in the mapping are an instance of {Function.__name__}"
@@ -66,7 +63,7 @@ def __init__(
         self.__functions = dict(functions)
 
     @property
-    def functions(self) -> dict[str, Function[DataSample, np.ndarray]]:
+    def functions(self) -> dict[str, Function[DataSample, Array]]:
         """Read-only access to the internal mapping of functions."""
         return dict(self.__functions)