diff --git a/.cspell.json b/.cspell.json index 3dd99ec6..5b971deb 100644 --- a/.cspell.json +++ b/.cspell.json @@ -242,6 +242,7 @@ "vectorize", "venv", "vmap", + "vmapped", "weisskopf", "wirtinger", "xcode", diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c208a5ee..57b2e10d 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/benchmarks/ampform.py b/benchmarks/ampform.py index d6b64149..5cf3f151 100644 --- a/benchmarks/ampform.py +++ b/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: diff --git a/benchmarks/expression.py b/benchmarks/expression.py index 796b9c9a..f8a9606d 100644 --- a/benchmarks/expression.py +++ b/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: diff --git a/benchmarks/unbinned_nll.py b/benchmarks/unbinned_nll.py index c5ed4906..c7bd1a82 100644 --- a/benchmarks/unbinned_nll.py +++ b/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 diff --git a/docs/conf.py b/docs/conf.py index 2df67045..88ec9300 100644 --- a/docs/conf.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 5b40d91e..1f6014d7 100644 --- a/pyproject.toml +++ b/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]] diff --git a/src/tensorwaves/data/__init__.py b/src/tensorwaves/data/__init__.py index 7ac35528..a62951ff 100644 --- a/src/tensorwaves/data/__init__.py +++ b/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( diff --git a/src/tensorwaves/data/_data_sample.py b/src/tensorwaves/data/_data_sample.py index bbb35dee..34f937d5 100644 --- a/src/tensorwaves/data/_data_sample.py +++ b/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" diff --git a/src/tensorwaves/data/rng.py b/src/tensorwaves/data/rng.py index 3054ef71..68b8d6df 100644 --- a/src/tensorwaves/data/rng.py +++ b/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, diff --git a/src/tensorwaves/data/transform.py b/src/tensorwaves/data/transform.py index af244716..05e1d97c 100644 --- a/src/tensorwaves/data/transform.py +++ b/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) diff --git a/src/tensorwaves/estimator.py b/src/tensorwaves/estimator.py index 07470c3d..0a10ba9a 100644 --- a/src/tensorwaves/estimator.py +++ b/src/tensorwaves/estimator.py @@ -5,16 +5,19 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, overload from tensorwaves.config import _initialize_jax from tensorwaves.data.transform import SympyDataTransformer from tensorwaves.function._backend import find_function, raise_missing_module_error from tensorwaves.function.sympy import create_parametrized_function, prepare_caching from tensorwaves.interface import ( + Array, DataSample, DataTransformer, Estimator, + FloatArray, + ParameterType, ParameterValue, ParametrizedFunction, ) @@ -22,7 +25,6 @@ if TYPE_CHECKING: from collections.abc import Callable, Iterable, Mapping - import numpy as np import sympy as sp @@ -33,7 +35,7 @@ def create_cached_function( free_parameters: Iterable[sp.Basic], *, use_cse: bool = True, -) -> tuple[ParametrizedFunction[DataSample, np.ndarray], DataTransformer]: +) -> tuple[ParametrizedFunction[DataSample, Array], DataTransformer]: """Create a function and data transformer for cached computations. Once it is known which parameters in an expression are to be optimized, this @@ -85,14 +87,25 @@ def _determine_backend(function: ParametrizedFunction, backend: str | None) -> s def _coerce_parameter_types( - parameters: Mapping[str, ParameterValue], -) -> dict[str, ParameterValue]: - # normalize to float/complex so that JIT compilers see stable input types - # (an int value would otherwise trigger a re-trace once it becomes a float) - return { - name: complex(value) if isinstance(value, complex) else float(value) - for name, value in parameters.items() - } + parameters: Mapping[str, ParameterType], +) -> dict[str, ParameterType]: + return {name: _coerce_parameter_value(value) for name, value in parameters.items()} + + +def _coerce_parameter_value(value: ParameterType) -> ParameterType: + """Normalize values for stable JIT inputs and event-axis broadcasting. + + Scalars are converted to ``float`` or ``complex`` so that a change from an integer + value does not trigger retracing. Parameter arrays receive a new trailing axis so + that they broadcast against the data samples' event axis. + """ + if isinstance(value, complex): + return complex(value) + if isinstance(value, (int, float)): + return float(value) + if getattr(value, "ndim", 0) >= 1: + return value[..., None] + return value def _import_jax(): # ruff: ignore[missing-return-type-private-function] @@ -159,7 +172,7 @@ def _create_core_gradient(core: Callable, backend: str) -> Callable: def gradient( parameters: Mapping[str, ParameterValue], - *data_args: DataSample | np.ndarray | None, + *data_args: DataSample | Array | None, ) -> dict[str, ParameterValue]: return _conjugate_complex_gradient(raw_gradient(parameters, *data_args)) @@ -167,7 +180,7 @@ def gradient( def raise_gradient_not_implemented( parameters: Mapping[str, ParameterValue], - *data_args: DataSample | np.ndarray | None, + *data_args: DataSample | Array | None, ) -> dict[str, ParameterValue]: msg = f"Gradient not implemented for back-end {backend}." raise NotImplementedError(msg) @@ -201,10 +214,10 @@ class ChiSquared(Estimator): def __init__( self, - function: ParametrizedFunction[DataSample, np.ndarray], + function: ParametrizedFunction[DataSample, FloatArray], domain: DataSample, - observed_values: np.ndarray, - weights: np.ndarray | None = None, + observed_values: Array, + weights: Array | None = None, backend: str | None = None, ) -> None: backend = _determine_backend(function, backend) @@ -220,19 +233,27 @@ def __init__( sum_function = find_function("sum", backend) def estimator( - parameters: Mapping[str, ParameterValue], + parameters: Mapping[str, ParameterType], domain: DataSample, - observed_values: np.ndarray, - weights: np.ndarray, + observed_values: Array, + weights: Array, ) -> float: computed_values = function(domain, parameters) chi_squared = weights * (computed_values - observed_values) ** 2 - return sum_function(chi_squared) + return sum_function(chi_squared, axis=-1) self.__estimator = _jit_estimator_core(estimator, backend) self.__gradient = _create_core_gradient(estimator, backend) - def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: + @overload + def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: ... + @overload + def __call__(self, parameters: Mapping[str, Array]) -> FloatArray: ... + @overload + def __call__( + self, parameters: Mapping[str, ParameterType] + ) -> float | FloatArray: ... + def __call__(self, parameters: Mapping[str, ParameterType]) -> float | FloatArray: return self.__estimator(*self.__estimator_args(parameters)) def gradient( @@ -240,7 +261,7 @@ def gradient( ) -> dict[str, ParameterValue]: return self.__gradient(*self.__estimator_args(parameters)) - def __estimator_args(self, parameters: Mapping[str, ParameterValue]) -> tuple: + def __estimator_args(self, parameters: Mapping[str, ParameterType]) -> tuple: return ( _coerce_parameter_types(parameters), self.__domain, @@ -290,7 +311,7 @@ class UnbinnedNLL(Estimator): def __init__( self, - function: ParametrizedFunction[DataSample, np.ndarray], + function: ParametrizedFunction[DataSample, FloatArray], data: DataSample, phsp: DataSample, phsp_volume: float = 1.0, @@ -306,25 +327,37 @@ def __init__( log_function = find_function("log", backend) def estimator( - parameters: Mapping[str, ParameterValue], + parameters: Mapping[str, ParameterType], data: DataSample, phsp: DataSample, - phsp_weights: np.ndarray | None, + phsp_weights: Array | None, ) -> float: bare_intensities = function(data, parameters) phsp_intensities = function(phsp, parameters) if phsp_weights is not None: phsp_intensities *= phsp_weights - normalization_integral = phsp_volume * mean_function(phsp_intensities) - log_normalization = len(bare_intensities) * log_function( + normalization_integral = phsp_volume * mean_function( + phsp_intensities, axis=-1 + ) + log_normalization = bare_intensities.shape[-1] * log_function( normalization_integral ) - return log_normalization - sum_function(log_function(bare_intensities)) + return log_normalization - sum_function( + log_function(bare_intensities), axis=-1 + ) self.__estimator = _jit_estimator_core(estimator, backend) self.__gradient = _create_core_gradient(estimator, backend) - def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: + @overload + def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: ... + @overload + def __call__(self, parameters: Mapping[str, Array]) -> FloatArray: ... + @overload + def __call__( + self, parameters: Mapping[str, ParameterType] + ) -> float | FloatArray: ... + def __call__(self, parameters: Mapping[str, ParameterType]) -> float | FloatArray: return self.__estimator(*self.__estimator_args(parameters)) def gradient( @@ -332,7 +365,7 @@ def gradient( ) -> dict[str, ParameterValue]: return self.__gradient(*self.__estimator_args(parameters)) - def __estimator_args(self, parameters: Mapping[str, ParameterValue]) -> tuple: + def __estimator_args(self, parameters: Mapping[str, ParameterType]) -> tuple: return ( _coerce_parameter_types(parameters), self.__data, diff --git a/src/tensorwaves/function/__init__.py b/src/tensorwaves/function/__init__.py index e8f1ad9b..07350cf7 100644 --- a/src/tensorwaves/function/__init__.py +++ b/src/tensorwaves/function/__init__.py @@ -6,12 +6,13 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable import attrs -import numpy as np from attrs import field, frozen from tensorwaves.interface import ( + Array, DataSample, Function, + ParameterType, ParameterValue, ParametrizedFunction, ) @@ -42,7 +43,7 @@ class BackendFunction(Protocol): """ @property - def function(self) -> Callable[..., np.ndarray]: + def function(self) -> Callable[..., Array]: """Backend-native function that takes positional arguments only.""" @property @@ -101,7 +102,7 @@ def _to_tuple(argument_order: Iterable[str]) -> tuple[str, ...]: @frozen -class PositionalArgumentFunction(Function[DataSample, np.ndarray]): +class PositionalArgumentFunction(Function[DataSample, Array]): """Wrapper around a function with positional arguments. This class provides a :meth:`~.Function.__call__` that can take a `.DataSample` for @@ -113,7 +114,7 @@ class PositionalArgumentFunction(Function[DataSample, np.ndarray]): .. seealso:: :func:`.create_function` """ - function: Callable[..., np.ndarray] = field(validator=_validate_arguments) + function: Callable[..., Array] = field(validator=_validate_arguments) """A function with positional arguments only.""" argument_order: tuple[str, ...] = field( converter=_to_tuple, validator=[_all_str, _all_unique] @@ -122,12 +123,12 @@ class PositionalArgumentFunction(Function[DataSample, np.ndarray]): backend: str | None = None """Name of the computational backend that :attr:`function` was compiled for.""" - def __call__(self, data: DataSample) -> np.ndarray: + def __call__(self, data: DataSample) -> Array: args = [data[var_name] for var_name in self.argument_order] return self.function(*args) -class ParametrizedBackendFunction(ParametrizedFunction[DataSample, np.ndarray]): +class ParametrizedBackendFunction(ParametrizedFunction[DataSample, Array]): """Implements `.ParametrizedFunction` for a specific computational back-end. .. seealso:: :func:`.create_parametrized_function` @@ -135,7 +136,7 @@ class ParametrizedBackendFunction(ParametrizedFunction[DataSample, np.ndarray]): def __init__( self, - function: Callable[..., np.ndarray], + function: Callable[..., Array], argument_order: Iterable[str], parameters: Mapping[str, ParameterValue], backend: str | None = None, @@ -146,13 +147,16 @@ def __init__( def __call__( self, data: DataSample, - parameters: Mapping[str, ParameterValue] | None = None, - ) -> np.ndarray: - extended_data = {**data, **self.__merge_parameters(parameters)} - return self.__function(extended_data) # ty:ignore[invalid-argument-type] + parameters: Mapping[str, ParameterType] | None = None, + ) -> Array: + extended_data: dict = {**data, **self.__parameters} + if parameters is not None: + self.__validate_parameters(parameters) + extended_data.update(parameters) + return self.__function(extended_data) @property - def function(self) -> Callable[..., np.ndarray]: + def function(self) -> Callable[..., Array]: return self.__function.function @property @@ -170,18 +174,15 @@ def parameters(self) -> dict[str, ParameterValue]: def with_parameters( self, parameters: Mapping[str, ParameterValue] ) -> ParametrizedBackendFunction: + self.__validate_parameters(parameters) return ParametrizedBackendFunction( function=self.function, argument_order=self.argument_order, - parameters=self.__merge_parameters(parameters), + parameters={**self.__parameters, **parameters}, backend=self.backend, ) - def __merge_parameters( - self, parameters: Mapping[str, ParameterValue] | None - ) -> dict[str, ParameterValue]: - if parameters is None: - return self.__parameters + def __validate_parameters(self, parameters: Mapping[str, ParameterType]) -> None: over_defined = set(parameters) - set(self.__parameters) if over_defined: sep = "\n " @@ -191,7 +192,6 @@ def __merge_parameters( f" Expecting one of:{sep}{parameter_listing}" ) raise ValueError(msg) - return {**self.__parameters, **parameters} def get_source_code(function: Function) -> str: diff --git a/src/tensorwaves/interface.py b/src/tensorwaves/interface.py index 85c5e740..69733b5c 100644 --- a/src/tensorwaves/interface.py +++ b/src/tensorwaves/interface.py @@ -2,15 +2,21 @@ from __future__ import annotations +import sys from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Generic, TypeAlias, overload import attrs import numpy as np from attrs import field, frozen from attrs.validators import instance_of, optional +if sys.version_info >= (3, 13): + from typing import TypeVar +else: + from typing_extensions import TypeVar # https://peps.python.org/pep-0696 + if TYPE_CHECKING: # pragma: no cover from IPython.lib.pretty import PrettyPrinter @@ -36,10 +42,29 @@ class Function(ABC, Generic[InputType, OutputType]): def __call__(self, data: InputType) -> OutputType: ... -DataSample = dict[str, np.ndarray] +ScalarT = TypeVar("ScalarT", bound=np.generic, default=Any) +"""The scalar type (dtype) of an `Array`.""" +Array: TypeAlias = np.ndarray[Any, np.dtype[ScalarT]] +"""Type representing numerical arrays. + +The alias is generic in its scalar type, so `Array` is dtype-agnostic, while +`FloatArray` narrows to an array of floats. +""" +FloatArray: TypeAlias = Array[np.floating] +"""An `Array` of real-valued numbers.""" +DataSample = dict[str, Array] """Mapping of variable names to a sequence of data points, used by `Function`.""" ParameterValue = complex | float -"""Allowed types for parameter values.""" +"""Allowed types for scalar parameter values.""" +ParameterType = ParameterValue | Array +"""Types for parameter values in an evaluation, including arrays of values. + +An array of parameter values represents several parameter points that are evaluated in +one call through `broadcasting +`_ against the event axis of +a `.DataSample`. This can be used to propagate fit uncertainties by evaluating over e.g. +bootstrapped parameter samples in a single, backend-parallelized call. +""" class ParametrizedFunction(Function[InputType, OutputType]): @@ -62,12 +87,13 @@ class ParametrizedFunction(Function[InputType, OutputType]): def __call__( self, data: InputType, - parameters: Mapping[str, ParameterValue] | None = None, + parameters: Mapping[str, ParameterType] | None = None, ) -> OutputType: """Evaluate the function over :code:`data` for these parameter values. Given parameter values are merged with the defaults in :attr:`parameters` for - this evaluation only. + this evaluation only. Parameter values may be arrays, as long as they broadcast + against the event arrays in :code:`data` (see `.ParameterType`). """ @property @@ -90,7 +116,7 @@ class DataTransformer(Function[DataSample, DataSample]): """ -class Estimator(Function[Mapping[str, ParameterValue], float]): +class Estimator(Function[Mapping[str, ParameterType], float | FloatArray]): """Estimator for discrepancy model and data. See the :mod:`.estimator` module for different implementations of this interface. @@ -98,9 +124,22 @@ class Estimator(Function[Mapping[str, ParameterValue], float]): .. automethod:: __call__ """ + @overload + def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: ... + @overload + def __call__(self, parameters: Mapping[str, Array]) -> FloatArray: ... + @overload + def __call__( + self, parameters: Mapping[str, ParameterType] + ) -> float | FloatArray: ... @abstractmethod - def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: # ty:ignore[invalid-method-override] - """Compute estimator value for this combination of parameter values.""" + def __call__(self, parameters: Mapping[str, ParameterType]) -> float | FloatArray: # ty:ignore[invalid-method-override] + """Compute estimator value for this combination of parameter values. + + Parameter values may be one-dimensional arrays of shape :code:`(p,)`, in which + case the estimator returns an array of :code:`p` estimator values, one for each + parameter point (see `.ParameterType`). + """ @abstractmethod def gradient( @@ -224,7 +263,7 @@ class RealNumberGenerator(ABC): @abstractmethod def __call__( self, size: int, min_value: float = 0.0, max_value: float = 1.0 - ) -> np.ndarray: + ) -> Array: """Generate random floats in the range [min_value, max_value).""" @property diff --git a/src/tensorwaves/optimizer/scipy.py b/src/tensorwaves/optimizer/scipy.py index 915ea2b2..7e033dfd 100644 --- a/src/tensorwaves/optimizer/scipy.py +++ b/src/tensorwaves/optimizer/scipy.py @@ -91,7 +91,7 @@ def wrapped_function(pars: list) -> float: logs=_create_log( optimizer=type(self), estimator_type=type(estimator), - estimator_value=estimator(parameters), + estimator_value=estimator_value, function_call=n_function_calls, parameters=parameters, ), @@ -113,7 +113,7 @@ def wrapped_callback(pars: Iterable[float]) -> None: logs=_create_log( optimizer=type(self), estimator_type=type(estimator), - estimator_value=float(estimator_value), + estimator_value=estimator_value, function_call=n_function_calls, parameters=create_parameter_dict(pars), ), diff --git a/tests/data/test_data.py b/tests/data/test_data.py index 77efecf5..3ba7cd81 100644 --- a/tests/data/test_data.py +++ b/tests/data/test_data.py @@ -17,6 +17,7 @@ ) from tensorwaves.function.sympy import create_function from tensorwaves.interface import ( + Array, DataGenerator, DataSample, Function, @@ -27,8 +28,8 @@ from _pytest.capture import CaptureFixture -class FlatDistribution(Function[DataSample, np.ndarray]): - def __call__(self, data: DataSample) -> np.ndarray: +class FlatDistribution(Function[DataSample, Array]): + def __call__(self, data: DataSample) -> Array: some_key = next(iter(data)) sample_size = len(data[some_key]) return np.ones(sample_size) diff --git a/tests/optimizer/test_fit_simple_model.py b/tests/optimizer/test_fit_simple_model.py index 97c14576..c6531046 100644 --- a/tests/optimizer/test_fit_simple_model.py +++ b/tests/optimizer/test_fit_simple_model.py @@ -23,7 +23,7 @@ import iminuit - from tensorwaves.interface import DataSample, Function + from tensorwaves.interface import Array, DataSample, Function def generate_domain( @@ -40,7 +40,7 @@ def generate_domain( def generate_data( size: int, boundaries: dict[str, tuple[float, float]], - function: Function[DataSample, np.ndarray], + function: Function[DataSample, Array], rng: np.random.Generator, bunch_size: int = 10_000, ) -> DataSample: diff --git a/tests/optimizer/test_minuit.py b/tests/optimizer/test_minuit.py index c6413e5f..f786e223 100644 --- a/tests/optimizer/test_minuit.py +++ b/tests/optimizer/test_minuit.py @@ -1,10 +1,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, overload import pytest -from tensorwaves.interface import Estimator, ParameterValue +from tensorwaves.interface import ( + Array, + Estimator, + FloatArray, + ParameterType, + ParameterValue, +) from tensorwaves.optimizer.minuit import Minuit2 from . import CallbackMock, assert_invocations @@ -19,7 +25,15 @@ class Polynomial1DMinimaEstimator(Estimator): def __init__(self, polynomial: Callable) -> None: self.__polynomial = polynomial - def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: + @overload + def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: ... + @overload + def __call__(self, parameters: Mapping[str, Array]) -> FloatArray: ... + @overload + def __call__( + self, parameters: Mapping[str, ParameterType] + ) -> float | FloatArray: ... + def __call__(self, parameters: Mapping[str, ParameterType]) -> float | FloatArray: x = parameters["x"] return self.__polynomial(x) @@ -33,7 +47,15 @@ class Polynomial2DMinimaEstimator(Estimator): def __init__(self, polynomial: Callable) -> None: self.__polynomial = polynomial - def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: + @overload + def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: ... + @overload + def __call__(self, parameters: Mapping[str, Array]) -> FloatArray: ... + @overload + def __call__( + self, parameters: Mapping[str, ParameterType] + ) -> float | FloatArray: ... + def __call__(self, parameters: Mapping[str, ParameterType]) -> float | FloatArray: x = parameters["x"] y = parameters["y"] return self.__polynomial(x, y) diff --git a/tests/optimizer/test_scipy.py b/tests/optimizer/test_scipy.py index 03f6ec88..5872d422 100644 --- a/tests/optimizer/test_scipy.py +++ b/tests/optimizer/test_scipy.py @@ -1,10 +1,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, overload import pytest -from tensorwaves.interface import Estimator, ParameterValue +from tensorwaves.interface import ( + Array, + Estimator, + FloatArray, + ParameterType, + ParameterValue, +) from tensorwaves.optimizer.scipy import ScipyMinimizer from . import CallbackMock, assert_invocations @@ -19,7 +25,15 @@ class Polynomial1DMinimaEstimator(Estimator): def __init__(self, polynomial: Callable) -> None: self.__polynomial = polynomial - def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: + @overload + def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: ... + @overload + def __call__(self, parameters: Mapping[str, Array]) -> FloatArray: ... + @overload + def __call__( + self, parameters: Mapping[str, ParameterType] + ) -> float | FloatArray: ... + def __call__(self, parameters: Mapping[str, ParameterType]) -> float | FloatArray: x = parameters["x"] return self.__polynomial(x) @@ -33,7 +47,15 @@ class Polynomial2DMinimaEstimator(Estimator): def __init__(self, polynomial: Callable) -> None: self.__polynomial = polynomial - def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: + @overload + def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: ... + @overload + def __call__(self, parameters: Mapping[str, Array]) -> FloatArray: ... + @overload + def __call__( + self, parameters: Mapping[str, ParameterType] + ) -> float | FloatArray: ... + def __call__(self, parameters: Mapping[str, ParameterType]) -> float | FloatArray: x = parameters["x"] y = parameters["y"] return self.__polynomial(x, y) diff --git a/tests/test_estimator.py b/tests/test_estimator.py index 1b44ddb9..a1183b29 100644 --- a/tests/test_estimator.py +++ b/tests/test_estimator.py @@ -42,6 +42,18 @@ def test_call(self, backend): ) assert estimator({"a": 0, "b": 2}) == 2.5 + def test_array_valued_parameters(self): + x_data = {"x": np.array([0.0, 1.0, 2.0])} + y_data = np.array([0.0, 1.0, 2.0]) + function = ParametrizedBackendFunction( + function=lambda a, b, x: a + b * x, + argument_order=("a", "b", "x"), + parameters={"a": 0, "b": 1}, + ) + estimator = ChiSquared(function, x_data, y_data) + b_values = np.array([1.0, 2.0]) + np.testing.assert_allclose(estimator({"b": b_values}), [0.0, 5.0]) + def test_jit_compiled_once(self): trace_count = 0 @@ -173,6 +185,43 @@ def test_create_cached_function(backend): np.testing.assert_allclose(intensities, cached_intensities) +@pytest.mark.parametrize("backend", ["jax", "numba", "numpy", "tf"]) +def test_unbinned_nll_with_array_valued_parameters(backend: str): + x, mu, sigma = sp.symbols("x mu sigma") + function = create_parametrized_function( + expression=sp.exp(-(((x - mu) / sigma) ** 2) / 2), + parameters={mu: 0.5, sigma: 0.1}, + backend=backend, + ) + rng = np.random.default_rng(seed=0) + data = {"x": rng.normal(0.5, 0.1, size=2_000)} + phsp = {"x": rng.uniform(-2.0, 5.0, size=5_000)} + estimator = UnbinnedNLL(function, data, phsp, phsp_volume=7.0) + mu_values = np.array([0.4, 0.5, 0.6]) + batched_output = estimator({"mu": mu_values}) + scalar_outputs = [estimator({"mu": value}) for value in mu_values] + assert batched_output.shape == mu_values.shape + np.testing.assert_allclose(batched_output, scalar_outputs, rtol=1e-8) + + +def test_unbinned_nll_batched_evaluation_equals_jax_vmap(): + jax = pytest.importorskip("jax") + x, mu, sigma = sp.symbols("x mu sigma") + function = create_parametrized_function( + expression=sp.exp(-(((x - mu) / sigma) ** 2) / 2), + parameters={mu: 0.5, sigma: 0.1}, + backend="jax", + ) + rng = np.random.default_rng(seed=0) + data = {"x": rng.normal(0.5, 0.1, size=2_000)} + phsp = {"x": rng.uniform(-2.0, 5.0, size=5_000)} + estimator = UnbinnedNLL(function, data, phsp, phsp_volume=7.0) + mu_values = np.array([0.4, 0.5, 0.6]) + batched_output = estimator({"mu": mu_values}) + vmapped_output = jax.vmap(lambda value: estimator({"mu": value}))(mu_values) + np.testing.assert_allclose(batched_output, vmapped_output, rtol=1e-8) + + NUMPY_RNG = np.random.default_rng(12345) diff --git a/uv.lock b/uv.lock index cf132881..ad7fbaf5 100644 --- a/uv.lock +++ b/uv.lock @@ -1624,7 +1624,7 @@ dependencies = [ { name = "comm" }, { name = "debugpy" }, { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "jupyter-client" }, { name = "jupyter-core" }, { name = "matplotlib-inline" }, @@ -1646,7 +1646,7 @@ version = "0.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "ipywidgets" }, { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -1688,7 +1688,7 @@ wheels = [ [[package]] name = "ipython" -version = "9.17.0" +version = "9.17.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", @@ -1720,9 +1720,9 @@ dependencies = [ { name = "traitlets" }, { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c4/bc/e05ae123712ce4e1fde4408eedca1791fc1ff832684565132ea1dc646092/ipython-9.17.0.tar.gz", hash = "sha256:1dc69e6966b270fb259f676c71a21450e63607729b14a672b942914a54e8b730", size = 4538547, upload-time = "2026-08-28T09:00:58.233Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/32/99451b1283ec5d92ad77073f12e1c667dc10775384d8f15c2914207149dd/ipython-9.17.1.tar.gz", hash = "sha256:8919be8c27f20a6f4423145028063f6637b42a03ce57665bb12015ee1f073529", size = 4539289, upload-time = "2026-09-01T08:29:32.6Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/18/5f/f992b57e8deb8fa6c2e614b422d80698adb5107902bbc89832e203665bd1/ipython-9.17.0-py3-none-any.whl", hash = "sha256:ce647713be8fef3fab2418c515a0def4d45d6705dd102be2c6d1f3015d7368b0", size = 638698, upload-time = "2026-08-28T09:00:56.174Z" }, + { url = "https://files.pythonhosted.org/packages/2d/1e/65b59cf518c106aa755e7f7da3099027738687a862ec785060702a481320/ipython-9.17.1-py3-none-any.whl", hash = "sha256:6d1645743cfd1a07eb695d85aa2b5fa66721f8cbae9431d4049f7084bbf06509", size = 639038, upload-time = "2026-09-01T08:29:30.673Z" }, ] [[package]] @@ -1744,7 +1744,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "comm" }, { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "jupyterlab-widgets" }, { name = "traitlets" }, { name = "widgetsnbextension" }, @@ -3019,7 +3019,7 @@ dependencies = [ { name = "importlib-metadata" }, { name = "ipykernel" }, { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "jupyter-cache" }, { name = "myst-parser", version = "4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "myst-parser", version = "5.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -5364,7 +5364,7 @@ wheels = [ [package.optional-dependencies] ipython = [ { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] [[package]] @@ -5736,6 +5736,7 @@ dependencies = [ { name = "rich" }, { name = "sympy" }, { name = "tqdm" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] [package.optional-dependencies] @@ -5776,7 +5777,7 @@ dev = [ { name = "black" }, { name = "ipympl" }, { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "jupyter-ruff" }, { name = "jupyterlab" }, { name = "jupyterlab-git" }, @@ -5876,7 +5877,7 @@ style = [ { name = "black" }, { name = "ipympl" }, { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "numba" }, @@ -5895,7 +5896,7 @@ style = [ ] test = [ { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pyarrow" }, { name = "pytest" }, { name = "pytest-benchmark" }, @@ -5904,7 +5905,7 @@ test = [ ] test-types = [ { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pytest" }, { name = "pytest-mock" }, ] @@ -5912,7 +5913,7 @@ types = [ { name = "black" }, { name = "ipympl" }, { name = "ipython", version = "8.39.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "ipython", version = "9.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ipython", version = "9.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "numba" }, @@ -5948,6 +5949,7 @@ requires-dist = [ { name = "tensorwaves", extras = ["phsp"], marker = "extra == 'pwa'" }, { name = "tensorwaves", extras = ["tf"], marker = "extra == 'phsp'" }, { name = "tqdm", specifier = ">=4.24.0" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'", specifier = ">=4.4.0" }, ] provides-extras = ["jax", "numba", "phsp", "pwa", "scipy", "tf"]