diff --git a/.cspell.json b/.cspell.json index 534e32c5..c9134483 100644 --- a/.cspell.json +++ b/.cspell.json @@ -240,6 +240,7 @@ "unbinned", "vectorize", "venv", + "vmap", "weisskopf", "wirtinger", "xcode", diff --git a/docs/amplitude-analysis.ipynb b/docs/amplitude-analysis.ipynb index 5447b28f..0c59c994 100644 --- a/docs/amplitude-analysis.ipynb +++ b/docs/amplitude-analysis.ipynb @@ -1907,7 +1907,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.13" + "version": "3.13.14" } }, "nbformat": 4, diff --git a/docs/usage.ipynb b/docs/usage.ipynb index 60d6a757..6a462c0a 100644 --- a/docs/usage.ipynb +++ b/docs/usage.ipynb @@ -805,7 +805,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.14" } }, "nbformat": 4, diff --git a/docs/usage/basics.ipynb b/docs/usage/basics.ipynb index 51635e3a..defae220 100644 --- a/docs/usage/basics.ipynb +++ b/docs/usage/basics.ipynb @@ -1263,7 +1263,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.13" + "version": "3.13.14" } }, "nbformat": 4, diff --git a/docs/usage/binned-fit.ipynb b/docs/usage/binned-fit.ipynb index 6dbf719e..a003e130 100644 --- a/docs/usage/binned-fit.ipynb +++ b/docs/usage/binned-fit.ipynb @@ -278,7 +278,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.14" } }, "nbformat": 4, diff --git a/docs/usage/chi-squared.ipynb b/docs/usage/chi-squared.ipynb index 65f8af9c..9b66b1ab 100644 --- a/docs/usage/chi-squared.ipynb +++ b/docs/usage/chi-squared.ipynb @@ -231,7 +231,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.14" } }, "nbformat": 4, diff --git a/docs/usage/unbinned-fit.ipynb b/docs/usage/unbinned-fit.ipynb index 8273d405..ae0de6c4 100644 --- a/docs/usage/unbinned-fit.ipynb +++ b/docs/usage/unbinned-fit.ipynb @@ -290,7 +290,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.14" } }, "nbformat": 4, diff --git a/src/tensorwaves/data/transform.py b/src/tensorwaves/data/transform.py index aae2b7ab..af244716 100644 --- a/src/tensorwaves/data/transform.py +++ b/src/tensorwaves/data/transform.py @@ -103,6 +103,6 @@ def from_sympy( max_complexity=max_complexity, ) functions[variable_name] = PositionalArgumentFunction( - function, argument_order + function, argument_order, backend ) return cls(functions) diff --git a/src/tensorwaves/estimator.py b/src/tensorwaves/estimator.py index dbc3f28c..052d0c8a 100644 --- a/src/tensorwaves/estimator.py +++ b/src/tensorwaves/estimator.py @@ -75,6 +75,15 @@ def create_cached_function( return cached_function, cache_transformer +def _determine_backend(function: ParametrizedFunction, backend: str | None) -> str: + if backend is not None: + return backend + function_backend = getattr(function, "backend", None) + if function_backend is None: + return "numpy" + return function_backend + + def gradient_creator( function: Callable[[Mapping[str, ParameterValue]], ParameterValue], backend: str, @@ -120,7 +129,8 @@ class ChiSquared(Estimator): (unweighted). A common choice is :math:`w_i = 1/\sigma_i^2`, with :math:`\sigma_i` the uncertainty in each measured value of :math:`y_i`. backend: Computational backend with which to compute the sum - :math:`\sum_{i=1}^n`. + :math:`\sum_{i=1}^n`. By default, this is the backend of the + :code:`function`, if it exposes one (see `.BackendFunction`). .. seealso:: :doc:`/usage/chi-squared` """ @@ -131,8 +141,9 @@ def __init__( domain: DataSample, observed_values: np.ndarray, weights: np.ndarray | None = None, - backend: str = "numpy", + backend: str | None = None, ) -> None: + backend = _determine_backend(function, backend) self.__function = function self.__domain = domain self.__observed_values = observed_values @@ -186,7 +197,8 @@ class UnbinnedNLL(Estimator): phsp_volume: Optional phase space volume :math:`V`, used in the normalization factor. Default: :math:`V=1`. backend: The computational back-end with which the sums and averages - should be computed. + should be computed. By default, this is the backend of the + :code:`function`, if it exposes one (see `.BackendFunction`). .. seealso:: :doc:`/usage/unbinned-fit` """ @@ -197,8 +209,9 @@ def __init__( data: DataSample, phsp: DataSample, phsp_volume: float = 1.0, - backend: str = "numpy", + backend: str | None = None, ) -> None: + backend = _determine_backend(function, backend) self.__data = dict(data) # shallow copy self.__phsp = {k: v for k, v in phsp.items() if k != "weights"} self.__phsp_weights = phsp.get("weights") diff --git a/src/tensorwaves/function/__init__.py b/src/tensorwaves/function/__init__.py index 325e2e4b..e8f1ad9b 100644 --- a/src/tensorwaves/function/__init__.py +++ b/src/tensorwaves/function/__init__.py @@ -3,7 +3,7 @@ from __future__ import annotations import inspect -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Protocol, runtime_checkable import attrs import numpy as np @@ -20,6 +20,40 @@ from collections.abc import Callable, Iterable, Mapping +@runtime_checkable +class BackendFunction(Protocol): + """A function object that exposes its backend-native kernel. + + Classes like `PositionalArgumentFunction` and `ParametrizedBackendFunction` wrap a + pure function that takes positional argument arrays. This protocol gives access to + that kernel and the backend it was compiled for, so that backend-native + transformations (such as :code:`jax.jit`, :code:`jax.grad`, or :code:`jax.vmap`) + can be applied to it and estimators can determine which computational backend to + use. + + >>> import sympy as sp + >>> from tensorwaves.function.sympy import create_function + >>> x, y = sp.symbols("x y") + >>> func = create_function(x**2 + y**2, backend="jax") + >>> func.backend + 'jax' + >>> func.argument_order + ('x', 'y') + """ + + @property + def function(self) -> Callable[..., np.ndarray]: + """Backend-native function that takes positional arguments only.""" + + @property + def argument_order(self) -> tuple[str, ...]: + """Name of each positional argument, with data variables before parameters.""" + + @property + def backend(self) -> str | None: + """Name of the computational backend, if known.""" + + def _all_str( _: PositionalArgumentFunction, __: attrs.Attribute, value: Iterable[str] ) -> None: @@ -85,6 +119,8 @@ class PositionalArgumentFunction(Function[DataSample, np.ndarray]): converter=_to_tuple, validator=[_all_str, _all_unique] ) """Ordered labels for each positional argument.""" + backend: str | None = None + """Name of the computational backend that :attr:`function` was compiled for.""" def __call__(self, data: DataSample) -> np.ndarray: args = [data[var_name] for var_name in self.argument_order] @@ -102,8 +138,9 @@ def __init__( function: Callable[..., np.ndarray], argument_order: Iterable[str], parameters: Mapping[str, ParameterValue], + backend: str | None = None, ) -> None: - self.__function = PositionalArgumentFunction(function, argument_order) + self.__function = PositionalArgumentFunction(function, argument_order, backend) self.__parameters = dict(parameters) def __call__( @@ -122,6 +159,10 @@ def function(self) -> Callable[..., np.ndarray]: def argument_order(self) -> tuple[str, ...]: return self.__function.argument_order + @property + def backend(self) -> str | None: + return self.__function.backend + @property def parameters(self) -> dict[str, ParameterValue]: return dict(self.__parameters) @@ -133,6 +174,7 @@ def with_parameters( function=self.function, argument_order=self.argument_order, parameters=self.__merge_parameters(parameters), + backend=self.backend, ) def __merge_parameters( diff --git a/src/tensorwaves/function/sympy/__init__.py b/src/tensorwaves/function/sympy/__init__.py index d11243e7..c470c373 100644 --- a/src/tensorwaves/function/sympy/__init__.py +++ b/src/tensorwaves/function/sympy/__init__.py @@ -77,6 +77,7 @@ def create_function( return PositionalArgumentFunction( function=lambdified_function, argument_order=tuple(map(str, sorted_symbols)), + backend=backend, ) @@ -143,6 +144,7 @@ def create_parametrized_function( # ruff:ignore[too-many-arguments] function=lambdified_function, argument_order=tuple(map(str, sorted_symbols)), parameters={str(symbol): value for symbol, value in parameters.items()}, + backend=backend, ) diff --git a/tests/function/test_function.py b/tests/function/test_function.py index 25015636..7c5a34de 100644 --- a/tests/function/test_function.py +++ b/tests/function/test_function.py @@ -5,6 +5,7 @@ import sympy as sp from tensorwaves.function import ( + BackendFunction, ParametrizedBackendFunction, PositionalArgumentFunction, get_source_code, @@ -58,6 +59,11 @@ def test_call( def test_function(self, function: ParametrizedBackendFunction): assert callable(function.function) + def test_backend(self, function: ParametrizedBackendFunction): + assert isinstance(function, BackendFunction) + assert function.backend == "numpy" + assert function.with_parameters({}).backend == "numpy" + def test_call_with_parameters(self): initial_parameter_values = {"a": 1.0, "b": 2.0} func = ParametrizedBackendFunction( @@ -87,6 +93,7 @@ def test_with_parameters(self): assert new_func is not func assert new_func.parameters == {"a": 2.0, "b": 2.0} assert new_func.function is func.function + assert func.backend is None assert func.parameters == initial_parameter_values data: DataSample = {"x": np.array([0.0, 1.0, 2.0])} np.testing.assert_array_equal(new_func(data), [2.0, 4.0, 6.0]) diff --git a/tests/test_estimator.py b/tests/test_estimator.py index ef21a52e..bf4ec966 100644 --- a/tests/test_estimator.py +++ b/tests/test_estimator.py @@ -42,6 +42,25 @@ def test_call(self, backend): ) assert estimator({"a": 0, "b": 2}) == 2.5 + @pytest.mark.parametrize("backend", ["jax", "numpy"]) + def test_backend_inferred_from_function(self, backend): + x_data = {"x": np.array([0.0, 1.0, 2.0])} + y_data = np.array([0.0, 1.0, 2.0]) + a, b, x = sp.symbols("a b x") + function = create_parametrized_function( + a + b * x, + parameters={a: 0.0, b: 1.0}, + backend=backend, + ) + estimator = ChiSquared(function, x_data, y_data) + if backend == "jax": + gradient = estimator.gradient({"a": 0.0, "b": 1.0}) + assert pytest.approx(gradient["a"]) == 0.0 + assert pytest.approx(gradient["b"]) == 0.0 + else: + with pytest.raises(NotImplementedError): + estimator.gradient({"a": 0.0, "b": 1.0}) + def gaussian(mu_: float, sigma_: float) -> ParametrizedBackendFunction: x, mu, sigma = sp.symbols("x, mu, sigma")