From 27221187b84b7b83b3dc6996b728da8effbbb902 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Thu, 21 Dec 2023 21:25:04 +0100 Subject: [PATCH] BREAK: require `n_events` argument in lorentz classes (#381) * FEAT: make `ArraySize` publically available --- docs/_extend_docstrings.py | 12 +++---- docs/usage/kinematics.ipynb | 3 +- src/ampform/kinematics/angles.py | 6 ++-- src/ampform/kinematics/lorentz.py | 24 +++++--------- tests/kinematics/test_lorentz.py | 52 +++++++++++++++---------------- 5 files changed, 43 insertions(+), 54 deletions(-) diff --git a/docs/_extend_docstrings.py b/docs/_extend_docstrings.py index a00e5962e..7028ddb20 100644 --- a/docs/_extend_docstrings.py +++ b/docs/_extend_docstrings.py @@ -24,7 +24,7 @@ from sympy.printing.numpy import NumPyPrinter from ampform.io import aslatex -from ampform.kinematics.lorentz import FourMomentumSymbol, _ArraySize +from ampform.kinematics.lorentz import ArraySize, FourMomentumSymbol from ampform.sympy._array_expressions import ArrayMultiplication if sys.version_info < (3, 8): @@ -127,7 +127,7 @@ def extend_BoostZMatrix() -> None: ) b = sp.Symbol("b") _append_code_rendering( - BoostZMatrix(b).doit(), + BoostZMatrix(b, n_events=ArraySize(b)).doit(), use_cse=True, docstring_class=BoostZMatrix, ) @@ -147,9 +147,9 @@ def extend_BoostZMatrix() -> None: ) p, beta, phi, theta = sp.symbols("p beta phi theta") multiplication = ArrayMultiplication( - BoostZMatrix(beta, n_events=_ArraySize(p)), - RotationYMatrix(theta, n_events=_ArraySize(p)), - RotationZMatrix(phi, n_events=_ArraySize(p)), + BoostZMatrix(beta, n_events=ArraySize(p)), + RotationYMatrix(theta, n_events=ArraySize(p)), + RotationZMatrix(phi, n_events=ArraySize(p)), p, ) _append_to_docstring( @@ -460,7 +460,7 @@ def extend_RotationZMatrix() -> None: ) a = sp.Symbol("a") _append_code_rendering( - RotationZMatrix(a).doit(), + RotationZMatrix(a, n_events=ArraySize(a)).doit(), use_cse=True, docstring_class=RotationZMatrix, ) diff --git a/docs/usage/kinematics.ipynb b/docs/usage/kinematics.ipynb index 05b740200..a8b67e191 100644 --- a/docs/usage/kinematics.ipynb +++ b/docs/usage/kinematics.ipynb @@ -108,6 +108,7 @@ "source": [ "from ampform.kinematics.lorentz import (\n", " ArrayMultiplication,\n", + " ArraySize,\n", " BoostZMatrix,\n", " Energy,\n", " FourMomentumSymbol,\n", @@ -117,7 +118,7 @@ "p = FourMomentumSymbol(\"p\", shape=[])\n", "q = FourMomentumSymbol(\"q\", shape=[])\n", "beta = three_momentum_norm(p) / Energy(p)\n", - "Bz = BoostZMatrix(beta)\n", + "Bz = BoostZMatrix(beta, n_events=ArraySize(beta))\n", "Bz_expr = ArrayMultiplication(Bz, q)\n", "Bz_expr" ] diff --git a/src/ampform/kinematics/angles.py b/src/ampform/kinematics/angles.py index 6cfa1eb92..0813763cd 100644 --- a/src/ampform/kinematics/angles.py +++ b/src/ampform/kinematics/angles.py @@ -13,6 +13,7 @@ ) from ampform.helicity.naming import get_helicity_angle_symbols, get_helicity_suffix from ampform.kinematics.lorentz import ( + ArraySize, BoostMatrix, BoostZMatrix, Energy, @@ -23,7 +24,6 @@ NegativeMomentum, RotationYMatrix, RotationZMatrix, - _ArraySize, compute_boost_chain, three_momentum_norm, ) @@ -183,9 +183,9 @@ def __recursive_helicity_angles( return __recursive_helicity_angles(four_momenta, initial_state_edge.ending_node_id) -def _get_number_of_events(four_momenta: Mapping[int, sp.Expr]) -> _ArraySize: +def _get_number_of_events(four_momenta: Mapping[int, sp.Expr]) -> ArraySize: sorted_momentum_symbols = sorted(four_momenta.values(), key=str) - return _ArraySize(sorted_momentum_symbols[0]) + return ArraySize(sorted_momentum_symbols[0]) def compute_wigner_angles( diff --git a/src/ampform/kinematics/lorentz.py b/src/ampform/kinematics/lorentz.py index cd5b9cd72..56151d7c4 100644 --- a/src/ampform/kinematics/lorentz.py +++ b/src/ampform/kinematics/lorentz.py @@ -300,11 +300,7 @@ class BoostZMatrix(UnevaluatedExpression): :math:`n\times4\times4`. Defaults to the `len` of :code:`beta`. """ - def __new__( - cls, beta: sp.Basic, n_events: sp.Expr | None = None, **kwargs - ) -> BoostZMatrix: - if n_events is None: - n_events = _ArraySize(beta) + def __new__(cls, beta: sp.Basic, n_events: sp.Basic, **kwargs) -> BoostZMatrix: return create_expression(cls, beta, n_events, **kwargs) def as_explicit(self) -> sp.MutableDenseMatrix: @@ -484,11 +480,7 @@ class RotationYMatrix(UnevaluatedExpression): :math:`n\times4\times4`. Defaults to the `len` of :code:`angle`. """ - def __new__( - cls, angle: sp.Basic, n_events: sp.Expr | None = None, **hints - ) -> RotationYMatrix: - if n_events is None: - n_events = _ArraySize(angle) + def __new__(cls, angle: sp.Basic, n_events: sp.Basic, **hints) -> RotationYMatrix: return create_expression(cls, angle, n_events, **hints) def as_explicit(self) -> sp.MutableDenseMatrix: @@ -555,11 +547,7 @@ class RotationZMatrix(UnevaluatedExpression): :math:`n\times4\times4`. Defaults to the `len` of :code:`angle`. """ - def __new__( - cls, angle: sp.Basic, n_events: sp.Expr | None = None, **hints - ) -> RotationZMatrix: - if n_events is None: - n_events = _ArraySize(angle) + def __new__(cls, angle: sp.Basic, n_events: sp.Basic, **hints) -> RotationZMatrix: return create_expression(cls, angle, n_events, **hints) def as_explicit(self) -> sp.MutableDenseMatrix: @@ -636,8 +624,10 @@ def _numpycode(self, printer: NumPyPrinter, *args) -> str: return f"zeros({shape})" -class _ArraySize(NumPyPrintable): - def __new__(cls, array: sp.Basic, **kwargs) -> _ArraySize: +class ArraySize(NumPyPrintable): + """Symbolic expression for getting the size of a numerical array.""" + + def __new__(cls, array: sp.Basic, **kwargs) -> ArraySize: return create_expression(cls, array, **kwargs) def _numpycode(self, printer: NumPyPrinter, *args) -> str: diff --git a/tests/kinematics/test_lorentz.py b/tests/kinematics/test_lorentz.py index 178ee610f..096ad2ae5 100644 --- a/tests/kinematics/test_lorentz.py +++ b/tests/kinematics/test_lorentz.py @@ -11,6 +11,7 @@ from sympy.printing.numpy import NumPyPrinter from ampform.kinematics.lorentz import ( + ArraySize, BoostMatrix, BoostZMatrix, Energy, @@ -24,7 +25,6 @@ RotationYMatrix, RotationZMatrix, ThreeMomentum, - _ArraySize, _OnesArray, _ZerosArray, compute_boost_chain, @@ -58,7 +58,7 @@ def test_boost_in_z_direction_reduces_to_z_boost(self): ]) beta = three_momentum_norm(p) / Energy(p) - z_expr = BoostZMatrix(beta) + z_expr = BoostZMatrix(beta, n_events=ArraySize(p)) z_func = sp.lambdify(p, z_expr.doit(), cse=True) z_matrix = z_func(p_array)[0] assert pytest.approx(matrix) == z_matrix @@ -104,7 +104,7 @@ def test_boosting_back_gives_original_momentum( class TestBoostZMatrix: def test_boost_into_own_rest_frame_gives_mass(self): p = FourMomentumSymbol("p", shape=[]) - n_events = _ArraySize(p) + n_events = ArraySize(p) beta = three_momentum_norm(p) / Energy(p) expr = BoostZMatrix(beta, n_events) func = sp.lambdify(p, expr.doit(), cse=True) @@ -122,9 +122,9 @@ def test_boost_into_own_rest_frame_gives_mass(self): def test_numpycode_cse_in_expression_tree(self): p, beta, phi, theta = sp.symbols("p beta phi theta") expr = ArrayMultiplication( - BoostZMatrix(beta, n_events=_ArraySize(p)), - RotationYMatrix(theta, n_events=_ArraySize(p)), - RotationZMatrix(phi, n_events=_ArraySize(p)), + BoostZMatrix(beta, n_events=ArraySize(p)), + RotationYMatrix(theta, n_events=ArraySize(p)), + RotationZMatrix(phi, n_events=ArraySize(p)), p, ) func = sp.lambdify([], expr.doit(), cse=True) @@ -246,27 +246,26 @@ def test_same_as_inverse(self, data_sample: dict[int, np.ndarray]): class TestRotationYMatrix: @pytest.fixture(scope="session") def rotation_expr(self): - angle, n_events = sp.symbols("a n") - return RotationYMatrix(angle, n_events) + angle = sp.Symbol("a") + return RotationYMatrix(angle, n_events=ArraySize(angle)) @pytest.fixture(scope="session") - def rotation_func(self, rotation_expr): + def rotation_func(self, rotation_expr: RotationYMatrix): angle = sp.Symbol("a") - rotation_expr = rotation_expr.doit() - rotation_expr = rotation_expr.subs(sp.Symbol("n"), _ArraySize(angle)) - return sp.lambdify(angle, rotation_expr, cse=True) + return sp.lambdify(angle, rotation_expr.doit(), cse=True) def test_numpycode_cse(self, rotation_expr: RotationYMatrix): func = sp.lambdify([], rotation_expr.doit(), cse=True) src = inspect.getsource(func) expected_src = """ def _lambdifygenerated(): + x0 = len(a) return (array( [ - [ones(n), zeros(n), zeros(n), zeros(n)], - [zeros(n), cos(a), zeros(n), sin(a)], - [zeros(n), zeros(n), ones(n), zeros(n)], - [zeros(n), -sin(a), zeros(n), cos(a)], + [ones(x0), zeros(x0), zeros(x0), zeros(x0)], + [zeros(x0), cos(a), zeros(x0), sin(a)], + [zeros(x0), zeros(x0), ones(x0), zeros(x0)], + [zeros(x0), -sin(a), zeros(x0), cos(a)], ] ).transpose((2, 0, 1))) """ @@ -285,27 +284,26 @@ def test_rotation_over_pi_flips_xz(self, rotation_func): class TestRotationZMatrix: @pytest.fixture(scope="session") def rotation_expr(self): - angle, n_events = sp.symbols("a n") - return RotationZMatrix(angle, n_events) + angle = sp.Symbol("a") + return RotationZMatrix(angle, n_events=ArraySize(angle)) @pytest.fixture(scope="session") - def rotation_func(self, rotation_expr): + def rotation_func(self, rotation_expr: RotationZMatrix): angle = sp.Symbol("a") - rotation_expr = rotation_expr.doit() - rotation_expr = rotation_expr.subs(sp.Symbol("n"), _ArraySize(angle)) - return sp.lambdify(angle, rotation_expr, cse=True) + return sp.lambdify(angle, rotation_expr.doit(), cse=True) def test_numpycode_cse(self, rotation_expr: RotationZMatrix): func = sp.lambdify([], rotation_expr.doit(), cse=True) src = inspect.getsource(func) expected_src = """ def _lambdifygenerated(): + x0 = len(a) return (array( [ - [ones(n), zeros(n), zeros(n), zeros(n)], - [zeros(n), cos(a), -sin(a), zeros(n)], - [zeros(n), sin(a), cos(a), zeros(n)], - [zeros(n), zeros(n), zeros(n), ones(n)], + [ones(x0), zeros(x0), zeros(x0), zeros(x0)], + [zeros(x0), cos(a), -sin(a), zeros(x0)], + [zeros(x0), sin(a), cos(a), zeros(x0)], + [zeros(x0), zeros(x0), zeros(x0), ones(x0)], ] ).transpose((2, 0, 1))) """ @@ -331,7 +329,7 @@ def test_rotation_latex_repr_is_identical_with_doit(rotation): @pytest.mark.parametrize("rotation", [RotationYMatrix, RotationZMatrix]) def test_rotation_over_multiple_two_pi_is_identity(rotation): angle = sp.Symbol("a") - expr = rotation(angle) + expr = rotation(angle, n_events=ArraySize(angle)) func = sp.lambdify(angle, expr.doit(), cse=True) angle_array = np.arange(-2, 4, 1) * 2 * np.pi rotation_matrices = func(angle_array)