diff --git a/docs/_extend_docstrings.py b/docs/_extend_docstrings.py
index 830f3b8b0..7028ddb20 100644
--- a/docs/_extend_docstrings.py
+++ b/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,
)
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/lorentz.py b/src/ampform/kinematics/lorentz.py
index 3b240e8e3..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:
diff --git a/tests/kinematics/test_lorentz.py b/tests/kinematics/test_lorentz.py
index 44d6b0b4b..096ad2ae5 100644
--- a/tests/kinematics/test_lorentz.py
+++ b/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)