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)),
p,
)
func = sp.lambdify([], expr.doit(), cse=True)
src = inspect.getsource(func)
expected_src = """
def _lambdifygenerated():
x0 = 1/sqrt(1 - beta**2)
x1 = len(p)
x2 = ones(x1)
x3 = zeros(x1)
return (einsum("...ij,...jk,...kl,...l->...i", array(
[
[x0, x3, x3, -beta*x0],
[x3, x2, x3, x3],
[x3, x3, x2, x3],
[-beta*x0, x3, x3, x0],
]
).transpose((2, 0, 1)), array(
[
[x2, x3, x3, x3],
[x3, cos(theta), x3, sin(theta)],
[x3, x3, x2, x3],
[x3, -sin(theta), x3, cos(theta)],
]
).transpose((2, 0, 1)), array(
[
[x2, x3, x3, x3],
[x3, cos(phi), -sin(phi), x3],
[x3, sin(phi), cos(phi), x3],
[x3, x3, x3, x2],
]
).transpose((2, 0, 1)), p))
"""
expected_src = textwrap.dedent(expected_src)
assert src.strip() == expected_src.strip()
def test_boost_into_own_rest_frame_gives_mass(self):
p = FourMomentumSymbol("p")
n_events = _ArraySize(p)
beta = ThreeMomentumNorm(p) / Energy(p)
expr = BoostZMatrix(beta, n_events)
func = sp.lambdify(p, expr.doit())
p_array = np.array([[5, 0, 0, 1]])
boost_z = func(p_array)[0]
boosted_array = np.einsum("...ij,...j->...i", boost_z, p_array)
mass = 4.89897949
assert pytest.approx(boosted_array[0]) == [mass, 0, 0, 0]
expr = InvariantMass(p)
func = sp.lambdify(p, expr.doit())
mass_array = func(p_array)
assert pytest.approx(mass_array[0]) == mass
def extend_BoostZMatrix() -> None:
from ampform.kinematics import BoostZMatrix
beta, n_events = sp.symbols("beta n")
expr = BoostZMatrix(beta, n_events)
_append_to_docstring(
BoostZMatrix,
f"""\n
This boost operates on a `FourMomentumSymbol` and looks like:
.. math:: {sp.latex(expr)} = {sp.latex(expr.as_explicit())}
:label: BoostZMatrix
""",
)
_append_to_docstring(
BoostZMatrix,
"""
In `TensorWaves <https://tensorwaves.rtfd.io>`_, this class is expressed in
a computational backend and it should operate on four-momentum arrays of
rank-2. As such, this boost matrix becomes a **rank-3** matrix. When using
`NumPy <https://numpy.org>`_ as backend, the computation looks as follows:
""",
)
b = sp.Symbol("b")
_append_code_rendering(
BoostZMatrix(b).doit(),
use_cse=True,
docstring_class=BoostZMatrix,
)
from ampform.kinematics import RotationYMatrix, RotationZMatrix
_append_to_docstring(
BoostZMatrix,
"""
Note that this code was generated with :func:`sympy.lambdify
<sympy.utilities.lambdify.lambdify>` with :code:`cse=True`. The repetition
of :func:`numpy.ones` is still bothersome, but these sub-nodes is also
extracted by :func:`sympy.cse <sympy.simplify.cse_main.cse>` if the
expression is nested further down in an :doc:`expression tree
<sympy:tutorial/manipulation>`, for instance when boosting a
`.FourMomentumSymbol` :math:`p` in the :math:`z`-direction:
""",
)
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)),
p,
)
_append_to_docstring(
BoostZMatrix,
f"""\n
.. math:: {sp.latex(expr)}
:label: boost-in-z-direction
which in :mod:`numpy` code becomes:
""",
)
_append_code_rendering(
expr.doit(), use_cse=True, docstring_class=BoostZMatrix
)
def rotation_func(self, rotation_expr):
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)