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)