BREAK: require n_events argument in boost classes

docs/_extend_docstrings.py

@@ -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,
     )
@@ -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,
     )

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"
    ]

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:

tests/kinematics/test_lorentz.py

@@ -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
@@ -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)