Skip to content

Commit

Permalink
BREAK: require n_events argument in lorentz classes (#381)
Browse files Browse the repository at this point in the history 
* FEAT: make `ArraySize` publically available
  • Loading branch information
@redeboer
redeboer committed Dec 22, 2023
1 parent 611ebd6 commit 3d82d5a
Show file tree
Hide file tree
Showing 4 changed files with 42 additions and 53 deletions.
12 changes: 6 additions & 6 deletions docs/_extend_docstrings.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
from sympy.printing.numpy import NumPyPrinter

from ampform.io import aslatex
from ampform.kinematics import FourMomentumSymbol, _ArraySize
from ampform.kinematics import ArraySize, FourMomentumSymbol
from ampform.sympy._array_expressions import ArrayMultiplication

if sys.version_info < (3, 8):
Expand Down Expand Up @@ -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,
)
Expand All @@ -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(
Expand Down Expand Up @@ -455,7 +455,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,
)
Expand Down
3 changes: 2 additions & 1 deletion docs/usage/kinematics.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@
"source": [
"from ampform.kinematics import (\n",
" ArrayMultiplication,\n",
" ArraySize,\n",
" BoostZMatrix,\n",
" Energy,\n",
" FourMomentumSymbol,\n",
Expand All @@ -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"
]
Expand Down
28 changes: 9 additions & 19 deletions src/ampform/kinematics/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -467,11 +467,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.Expr, **kwargs) -> BoostZMatrix:
return create_expression(cls, beta, n_events, **kwargs)

def as_explicit(self) -> sp.MutableDenseMatrix:
Expand Down Expand Up @@ -651,11 +647,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.Expr, **hints) -> RotationYMatrix:
return create_expression(cls, angle, n_events, **hints)

def as_explicit(self) -> sp.MutableDenseMatrix:
Expand Down Expand Up @@ -722,11 +714,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.Expr, **hints) -> RotationZMatrix:
return create_expression(cls, angle, n_events, **hints)

def as_explicit(self) -> sp.MutableDenseMatrix:
Expand Down Expand Up @@ -803,8 +791,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:
Expand Down Expand Up @@ -906,9 +896,9 @@ def __recursive_helicity_angles(

def _get_number_of_events(
four_momenta: Mapping[int, sp.Expr],
) -> _ArraySize:
) -> ArraySize:
sorted_momentum_symbols = sorted(four_momenta.values(), key=str)
return _ArraySize(sorted_momentum_symbols[0])
return ArraySize(sorted_momentum_symbols[0])


def compute_invariant_masses(
Expand Down
52 changes: 25 additions & 27 deletions tests/test_kinematics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

from ampform.helicity.decay import get_parent_id
from ampform.kinematics import (
ArraySize,
BoostMatrix,
BoostZMatrix,
Energy,
Expand All @@ -29,7 +30,6 @@
RotationZMatrix,
Theta,
ThreeMomentum,
_ArraySize,
_OnesArray,
_ZerosArray,
compute_boost_chain,
Expand Down Expand Up @@ -84,7 +84,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
Expand Down Expand Up @@ -130,7 +130,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)
Expand All @@ -148,9 +148,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)
Expand Down Expand Up @@ -306,27 +306,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)))
"""
Expand All @@ -345,27 +344,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)))
"""
Expand All @@ -391,7 +389,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)
Expand Down

0 comments on commit 3d82d5a

Please sign in to comment.