From 8b455390e41cae008f214a41d69283f45f5579f7 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Thu, 28 Dec 2023 00:18:20 +0100 Subject: [PATCH] BREAK: split kinematics module (#384) * FEAT: make `determine_indices()` publicly available --- .cspell.json | 2 + docs/_extend_docstrings.py | 39 +- docs/conf.py | 5 +- docs/usage/kinematics.ipynb | 4 +- src/ampform/dynamics/__init__.py | 5 +- src/ampform/dynamics/builder.py | 2 +- src/ampform/dynamics/phasespace.py | 46 +- src/ampform/helicity/__init__.py | 22 +- src/ampform/kinematics/__init__.py | 779 +----------------- src/ampform/kinematics/angles.py | 203 +++++ src/ampform/kinematics/lorentz.py | 612 ++++++++++++++ src/ampform/sympy/__init__.py | 34 + tests/external/test_sympy.py | 17 - tests/helicity/align/__init__.py | 1 + tests/kinematics/conftest.py | 23 + tests/kinematics/test_angles.py | 263 ++++++ .../test_lorentz.py} | 262 +----- 17 files changed, 1200 insertions(+), 1119 deletions(-) create mode 100644 src/ampform/kinematics/angles.py create mode 100644 src/ampform/kinematics/lorentz.py create mode 100644 tests/helicity/align/__init__.py create mode 100644 tests/kinematics/conftest.py create mode 100644 tests/kinematics/test_angles.py rename tests/{test_kinematics.py => kinematics/test_lorentz.py} (68%) diff --git a/.cspell.json b/.cspell.json index 6154ec396..1d9576d9c 100644 --- a/.cspell.json +++ b/.cspell.json @@ -41,6 +41,7 @@ ], "ignoreWords": [ "Autoupdate", + "Dalitzplot", "MAINT", "Minkowski", "adrs", @@ -254,6 +255,7 @@ "lineshapes", "mathbb", "matplotlib", + "Mikhasenko", "mypy", "nishijima", "numpy", diff --git a/docs/_extend_docstrings.py b/docs/_extend_docstrings.py index 3e04d1d88..35f71ef38 100644 --- a/docs/_extend_docstrings.py +++ b/docs/_extend_docstrings.py @@ -24,7 +24,7 @@ from sympy.printing.numpy import NumPyPrinter from ampform.io import aslatex -from ampform.kinematics import ArraySize, FourMomentumSymbol +from ampform.kinematics.lorentz import ArraySize, FourMomentumSymbol from ampform.sympy._array_expressions import ArrayMultiplication if sys.version_info < (3, 8): @@ -72,14 +72,14 @@ def extend_BlattWeisskopfSquared() -> None: def extend_BoostMatrix() -> None: - from ampform.kinematics import BoostMatrix + from ampform.kinematics.lorentz import BoostMatrix p = FourMomentumSymbol("p", shape=[]) expr = BoostMatrix(p) _append_to_docstring( BoostMatrix, f"""\n - This boost operates on a `FourMomentumSymbol` and looks like: + This boost operates on a `.FourMomentumSymbol` and looks like: .. math:: {sp.latex(expr)} = {sp.latex(expr.as_explicit())} :class: full-width @@ -103,14 +103,14 @@ def extend_BoostMatrix() -> None: def extend_BoostZMatrix() -> None: - from ampform.kinematics import BoostZMatrix + from ampform.kinematics.lorentz import BoostZMatrix beta, n_events = sp.symbols("beta n") matrix = BoostZMatrix(beta, n_events) _append_to_docstring( BoostZMatrix, f"""\n - This boost operates on a `FourMomentumSymbol` and looks like: + This boost operates on a `.FourMomentumSymbol` and looks like: .. math:: {sp.latex(matrix)} = {sp.latex(matrix.as_explicit())} :label: BoostZMatrix @@ -132,7 +132,7 @@ def extend_BoostZMatrix() -> None: docstring_class=BoostZMatrix, ) - from ampform.kinematics import RotationYMatrix, RotationZMatrix + from ampform.kinematics.lorentz import RotationYMatrix, RotationZMatrix _append_to_docstring( BoostZMatrix, @@ -259,7 +259,12 @@ def extend_EnergyDependentWidth() -> None: def extend_Energy_and_FourMomentumXYZ() -> None: - from ampform.kinematics import Energy, FourMomentumX, FourMomentumY, FourMomentumZ + from ampform.kinematics.lorentz import ( + Energy, + FourMomentumX, + FourMomentumY, + FourMomentumZ, + ) def _extend(component_class: type[sp.Expr]) -> None: _append_to_docstring(component_class, "\n\n") @@ -274,7 +279,7 @@ def _extend(component_class: type[sp.Expr]) -> None: def extend_EuclideanNorm() -> None: - from ampform.kinematics import EuclideanNorm + from ampform.kinematics.lorentz import EuclideanNorm vector = FourMomentumSymbol("v", shape=[]) expr = EuclideanNorm(vector) @@ -327,7 +332,7 @@ def extend_Kibble() -> None: def extend_InvariantMass() -> None: - from ampform.kinematics import InvariantMass + from ampform.kinematics.lorentz import InvariantMass p = FourMomentumSymbol("p", shape=[]) expr = InvariantMass(p) @@ -403,7 +408,7 @@ def extend_PhaseSpaceFactorSWave() -> None: def extend_Phi() -> None: - from ampform.kinematics import Phi + from ampform.kinematics.angles import Phi p = FourMomentumSymbol("p", shape=[]) expr = Phi(p) @@ -411,7 +416,7 @@ def extend_Phi() -> None: def extend_RotationYMatrix() -> None: - from ampform.kinematics import RotationYMatrix + from ampform.kinematics.lorentz import RotationYMatrix angle, n_events = sp.symbols("alpha n") expr = RotationYMatrix(angle, n_events) @@ -419,7 +424,7 @@ def extend_RotationYMatrix() -> None: RotationYMatrix, f"""\n The matrix for a rotation over angle :math:`\\alpha` around the :math:`y`-axis - operating on `FourMomentumSymbol` looks like: + operating on `.FourMomentumSymbol` looks like: .. math:: {sp.latex(expr)} = {sp.latex(expr.as_explicit())} :label: RotationYMatrix @@ -430,15 +435,15 @@ def extend_RotationYMatrix() -> None: def extend_RotationZMatrix() -> None: - from ampform.kinematics import RotationZMatrix + from ampform.kinematics.lorentz import RotationZMatrix angle, n_events = sp.symbols("alpha n") expr = RotationZMatrix(angle, n_events) _append_to_docstring( RotationZMatrix, f"""\n - The matrix for a rotation over angle :math:`\\alpha` around the :math:`z`-axis - operating on `FourMomentumSymbol` looks like: + The matrix for a rotation over angle :math:`\\alpha` around the + :math:`z`-axis operating on `.FourMomentumSymbol` looks like: .. math:: {sp.latex(expr)} = {sp.latex(expr.as_explicit())} :label: RotationZMatrix @@ -468,7 +473,7 @@ def extend_RotationZMatrix() -> None: def extend_Theta() -> None: - from ampform.kinematics import Theta + from ampform.kinematics.angles import Theta p = FourMomentumSymbol("p", shape=[]) expr = Theta(p) @@ -476,7 +481,7 @@ def extend_Theta() -> None: def extend_ThreeMomentum() -> None: - from ampform.kinematics import ThreeMomentum + from ampform.kinematics.lorentz import ThreeMomentum p = FourMomentumSymbol("p", shape=[]) expr = ThreeMomentum(p) diff --git a/docs/conf.py b/docs/conf.py index 86acc5f90..a7e4bce1f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -52,8 +52,8 @@ "BuilderReturnType": ("obj", "ampform.dynamics.builder.BuilderReturnType"), "DecoratedClass": ("obj", "ampform.sympy.deprecated.DecoratedClass"), "DecoratedExpr": ("obj", "ampform.sympy.deprecated.DecoratedExpr"), - "FourMomenta": ("obj", "ampform.kinematics.FourMomenta"), - "FourMomentumSymbol": ("obj", "ampform.kinematics.FourMomentumSymbol"), + "FourMomenta": ("obj", "ampform.kinematics.lorentz.FourMomenta"), + "FourMomentumSymbol": ("obj", "ampform.kinematics.lorentz.FourMomentumSymbol"), "InteractionProperties": "qrules.quantum_numbers.InteractionProperties", "LatexPrinter": "sympy.printing.printer.Printer", "Literal[(-1, 1)]": "typing.Literal", @@ -230,6 +230,7 @@ "show_navbar_depth": 2, "show_toc_level": 2, } +html_title = REPO_TITLE intersphinx_mapping = { "IPython": (f"https://ipython.readthedocs.io/en/{pin('IPython')}", None), "attrs": (f"https://www.attrs.org/en/{pin('attrs')}", None), diff --git a/docs/usage/kinematics.ipynb b/docs/usage/kinematics.ipynb index b3e71ee38..a8b67e191 100644 --- a/docs/usage/kinematics.ipynb +++ b/docs/usage/kinematics.ipynb @@ -106,7 +106,7 @@ }, "outputs": [], "source": [ - "from ampform.kinematics import (\n", + "from ampform.kinematics.lorentz import (\n", " ArrayMultiplication,\n", " ArraySize,\n", " BoostZMatrix,\n", @@ -225,7 +225,7 @@ "metadata": {}, "outputs": [], "source": [ - "from ampform.kinematics import BoostMatrix\n", + "from ampform.kinematics.lorentz import BoostMatrix\n", "\n", "B = BoostMatrix(p)\n", "B_expr = ArrayMultiplication(B, q)\n", diff --git a/src/ampform/dynamics/__init__.py b/src/ampform/dynamics/__init__.py index d2fe90f27..3238da74a 100644 --- a/src/ampform/dynamics/__init__.py +++ b/src/ampform/dynamics/__init__.py @@ -19,10 +19,9 @@ PhaseSpaceFactorComplex, # noqa: F401 PhaseSpaceFactorProtocol, PhaseSpaceFactorSWave, # noqa: F401 - _determine_indices, _indices_to_subscript, ) -from ampform.sympy import argument, unevaluated +from ampform.sympy import argument, determine_indices, unevaluated if TYPE_CHECKING: from sympy.printing.latex import LatexPrinter @@ -180,7 +179,7 @@ def evaluate(self) -> sp.Expr: def _latex_repr_(self, printer: LatexPrinter, *args) -> str: s = printer._print(self.args[0]) gamma0 = self.args[2] - subscript = _indices_to_subscript(_determine_indices(gamma0)) + subscript = _indices_to_subscript(determine_indices(gamma0)) name = Rf"\Gamma{subscript}" if self.name is None else self.name return Rf"{name}\left({s}\right)" diff --git a/src/ampform/dynamics/builder.py b/src/ampform/dynamics/builder.py index 646be50cd..4815f55ce 100644 --- a/src/ampform/dynamics/builder.py +++ b/src/ampform/dynamics/builder.py @@ -34,7 +34,7 @@ class TwoBodyKinematicVariableSet: """Data container for the essential variables of a two-body decay. This data container is inserted into a `.ResonanceDynamicsBuilder`, so that it can - build some lineshape expression from the `.dynamics` module. It also allows to + build some lineshape expression from the :mod:`.dynamics` module. It also allows to insert :doc:`custom dynamics ` into the amplitude model. """ diff --git a/src/ampform/dynamics/phasespace.py b/src/ampform/dynamics/phasespace.py index e7eb67da3..1e5902b31 100644 --- a/src/ampform/dynamics/phasespace.py +++ b/src/ampform/dynamics/phasespace.py @@ -12,14 +12,12 @@ from __future__ import annotations -import re import sys from typing import TYPE_CHECKING, Any, Sequence import sympy as sp -from sympy.printing.conventions import split_super_sub -from ampform.sympy import argument, unevaluated +from ampform.sympy import argument, determine_indices, unevaluated from ampform.sympy.math import ComplexSqrt if TYPE_CHECKING: @@ -87,7 +85,7 @@ def evaluate(self) -> sp.Expr: def _latex_repr_(self, printer: LatexPrinter, *args) -> str: s = self.args[0] s_latex = printer._print(self.args[0]) - subscript = _indices_to_subscript(_determine_indices(s)) + subscript = _indices_to_subscript(determine_indices(s)) name = "q^2" + subscript if self.name is None else self.name return Rf"{name}\left({s_latex}\right)" @@ -113,7 +111,7 @@ def evaluate(self) -> sp.Expr: def _latex_repr_(self, printer: LatexPrinter, *args) -> str: s_symbol = self.args[0] s_latex = printer._print(s_symbol) - subscript = _indices_to_subscript(_determine_indices(s_symbol)) + subscript = _indices_to_subscript(determine_indices(s_symbol)) name = R"\rho" + subscript if self.name is None else self.name return Rf"{name}\left({s_latex}\right)" @@ -144,7 +142,7 @@ def evaluate(self) -> sp.Expr: def _latex_repr_(self, printer: LatexPrinter, *args) -> str: s_symbol = self.args[0] s_latex = printer._print(s_symbol) - subscript = _indices_to_subscript(_determine_indices(s_symbol)) + subscript = _indices_to_subscript(determine_indices(s_symbol)) name = R"\hat{\rho}" + subscript if self.name is None else self.name return Rf"{name}\left({s_latex}\right)" @@ -171,7 +169,7 @@ def evaluate(self) -> sp.Expr: def _latex_repr_(self, printer: LatexPrinter, *args) -> str: s_symbol = self.args[0] s_latex = printer._print(s_symbol) - subscript = _indices_to_subscript(_determine_indices(s_symbol)) + subscript = _indices_to_subscript(determine_indices(s_symbol)) name = R"\rho^\mathrm{c}" + subscript if self.name is None else self.name return Rf"{name}\left({s_latex}\right)" @@ -197,7 +195,7 @@ def evaluate(self) -> sp.Expr: def _latex_repr_(self, printer: LatexPrinter, *args) -> str: s_symbol = self.args[0] s_latex = printer._print(s_symbol) - subscript = _indices_to_subscript(_determine_indices(s_symbol)) + subscript = _indices_to_subscript(determine_indices(s_symbol)) name = R"\rho^\mathrm{CM}" + subscript if self.name is None else self.name return Rf"{name}\left({s_latex}\right)" @@ -245,7 +243,7 @@ def evaluate(self) -> sp.Expr: def _latex_repr_(self, printer: LatexPrinter, *args) -> str: s_symbol = self.args[0] s_latex = printer._print(s_symbol) - subscript = _indices_to_subscript(_determine_indices(s_symbol)) + subscript = _indices_to_subscript(determine_indices(s_symbol)) name = R"\rho^\mathrm{eq}" + subscript if self.name is None else self.name return Rf"{name}\left({s_latex}\right)" @@ -285,33 +283,3 @@ def _indices_to_subscript(indices: Sequence[int]) -> str: return "" subscript = ",".join(map(str, indices)) return "_{" + subscript + "}" - - -def _determine_indices(symbol) -> list[int]: - r"""Extract any indices if available from a `~sympy.core.symbol.Symbol`. - - >>> _determine_indices(sp.Symbol("m1")) - [1] - >>> _determine_indices(sp.Symbol("m_a2")) - [2] - >>> _determine_indices(sp.Symbol(R"\alpha_{i2, 5}")) - [2, 5] - >>> _determine_indices(sp.Symbol("m")) - [] - - `~sympy.tensor.indexed.Indexed` instances can also be handled: - >>> m_a = sp.IndexedBase("m_a") - >>> _determine_indices(m_a[0]) - [0] - """ - _, _, subscripts = split_super_sub(sp.latex(symbol)) - if not subscripts: - return [] - subscript: str = subscripts[-1] - subscript = re.sub(r"[^0-9^\,]", "", subscript) - subscript = f"[{subscript}]" - try: - indices = eval(subscript) # noqa: PGH001, S307 - except SyntaxError: - return [] - return list(indices) diff --git a/src/ampform/helicity/__init__.py b/src/ampform/helicity/__init__.py index 04059b7d0..af20321f0 100644 --- a/src/ampform/helicity/__init__.py +++ b/src/ampform/helicity/__init__.py @@ -43,6 +43,15 @@ TwoBodyKinematicVariableSet, create_non_dynamic, ) +from ampform.helicity.decay import ( + TwoBodyDecay, + get_parent_id, + get_prefactor, + get_sibling_state_id, + group_by_spin_projection, + group_by_topology, + is_opposite_helicity_state, +) from ampform.helicity.naming import ( CanonicalAmplitudeNameGenerator, HelicityAmplitudeNameGenerator, @@ -53,19 +62,10 @@ get_topology_identifier, natural_sorting, ) -from ampform.kinematics import HelicityAdapter, get_invariant_mass_symbol +from ampform.kinematics import HelicityAdapter +from ampform.kinematics.lorentz import get_invariant_mass_symbol from ampform.sympy import PoolSum -from .decay import ( - TwoBodyDecay, - get_parent_id, - get_prefactor, - get_sibling_state_id, - group_by_spin_projection, - group_by_topology, - is_opposite_helicity_state, -) - if sys.version_info >= (3, 8): from functools import singledispatchmethod from typing import Literal diff --git a/src/ampform/kinematics/__init__.py b/src/ampform/kinematics/__init__.py index 069715b09..c11ef1755 100644 --- a/src/ampform/kinematics/__init__.py +++ b/src/ampform/kinematics/__init__.py @@ -1,4 +1,3 @@ -# cspell:ignore einsum """Classes and functions for relativistic four-momentum kinematics. .. autolink-preface:: @@ -12,36 +11,21 @@ import itertools from collections import abc from functools import singledispatch -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Mapping +from typing import TYPE_CHECKING, Iterable import attrs -import sympy as sp from qrules.topology import Topology from qrules.transition import ReactionInfo, StateTransition -from ampform.helicity.decay import ( - assert_isobar_topology, - determine_attached_final_state, - get_parent_id, - get_sibling_state_id, - is_opposite_helicity_state, - list_decay_chain_ids, +from ampform.helicity.decay import assert_isobar_topology, get_parent_id +from ampform.kinematics.angles import compute_helicity_angles, compute_wigner_angles +from ampform.kinematics.lorentz import ( + compute_invariant_masses, + create_four_momentum_symbols, ) -from ampform.helicity.naming import get_helicity_angle_symbols, get_helicity_suffix -from ampform.sympy import ExprClass, NumPyPrintable, unevaluated -from ampform.sympy._array_expressions import ( - ArrayAxisSum, - ArrayMultiplication, - ArraySlice, - ArraySum, - ArraySymbol, - MatrixMultiplication, -) -from ampform.sympy.math import ComplexSqrt if TYPE_CHECKING: - from sympy.printing.latex import LatexPrinter - from sympy.printing.numpy import NumPyPrinter + import sympy as sp class HelicityAdapter: @@ -151,752 +135,3 @@ def _(obj: Topology) -> Topology: @_get_topology.register(StateTransition) def _(obj: StateTransition) -> Topology: return obj.topology - - -def create_four_momentum_symbols(topology: Topology) -> FourMomenta: - """Create a set of array-symbols for a `~qrules.topology.Topology`. - - >>> from qrules.topology import create_isobar_topologies - >>> topologies = create_isobar_topologies(3) - >>> create_four_momentum_symbols(topologies[0]) - {0: p0, 1: p1, 2: p2} - """ - n_final_states = len(topology.outgoing_edge_ids) - return {i: FourMomentumSymbol(f"p{i}", shape=[]) for i in range(n_final_states)} - - -FourMomenta = Dict[int, "FourMomentumSymbol"] -"""A mapping of state IDs to their corresponding `FourMomentumSymbol`. - -It's best to create a `dict` of `FourMomenta` with :func:`create_four_momentum_symbols`. -""" -FourMomentumSymbol = ArraySymbol -r"""Array-`~sympy.core.symbol.Symbol` that represents an array of four-momenta. - -The array is assumed to be of shape :math:`n\times 4` with :math:`n` the number of -events. The four-momenta are assumed to be in the order :math:`\left(E,\vec{p}\right)`. -See also `Energy`, `FourMomentumX`, `FourMomentumY`, and `FourMomentumZ`. -""" - - -@unevaluated -class Energy(sp.Expr): - """Represents the energy-component of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - _latex_repr_ = R"E\left({momentum}\right)" - - def evaluate(self) -> ArraySlice: - return ArraySlice(self.momentum, (slice(None), 0)) - - -def _implement_latex_subscript( # pyright: ignore[reportUnusedFunction] - subscript: str, -) -> Callable[[type[ExprClass]], type[ExprClass]]: - def decorator(decorated_class: type[ExprClass]) -> type[ExprClass]: - def _latex_repr_(self: sp.Expr, printer: LatexPrinter, *args) -> str: - momentum = printer._print(self.momentum) # type: ignore[attr-defined] - if printer._needs_mul_brackets(self.momentum): # type: ignore[attr-defined] - momentum = Rf"\left({momentum}\right)" - else: - momentum = Rf"{{{momentum}}}" - return f"{momentum}_{subscript}" - - decorated_class._latex_repr_ = _latex_repr_ # type: ignore[assignment,attr-defined] - return decorated_class - - return decorator - - -@unevaluated -@_implement_latex_subscript(subscript="x") -class FourMomentumX(sp.Expr): - """Component :math:`x` of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - - def evaluate(self) -> ArraySlice: - return ArraySlice(self.momentum, (slice(None), 1)) - - -@unevaluated -@_implement_latex_subscript(subscript="y") -class FourMomentumY(sp.Expr): - """Component :math:`y` of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - - def evaluate(self) -> ArraySlice: - return ArraySlice(self.momentum, (slice(None), 2)) - - -@unevaluated -@_implement_latex_subscript(subscript="z") -class FourMomentumZ(sp.Expr): - """Component :math:`z` of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - - def evaluate(self) -> ArraySlice: - return ArraySlice(self.momentum, (slice(None), 3)) - - -@unevaluated -class ThreeMomentum(NumPyPrintable): - """Spatial components of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - _latex_repr_ = R"\vec{{{momentum}}}" - - def evaluate(self) -> ArraySlice: - return ArraySlice(self.momentum, (slice(None), slice(1, None))) - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - return printer._print(self.evaluate()) - - -@unevaluated -class EuclideanNorm(NumPyPrintable): - """Take the euclidean norm of an array over axis 1.""" - - vector: sp.Basic - _latex_repr = R"\left|{vector}\right|" - - def evaluate(self) -> ArraySlice: - return sp.sqrt(EuclideanNormSquared(self.vector)) - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - return printer._print(self.evaluate(), *args) - - -@unevaluated -class EuclideanNormSquared(sp.Expr): - """Take the squared euclidean norm of an array over axis 1.""" - - vector: sp.Basic - _latex_repr_ = R"\left|{vector}\right|^{{2}}" - - def evaluate(self) -> ArrayAxisSum: - return ArrayAxisSum(self.vector**2, axis=1) # type: ignore[operator] - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - return printer._print(self.evaluate(), *args) - - -def three_momentum_norm(momentum: sp.Basic) -> EuclideanNorm: - return EuclideanNorm(ThreeMomentum(momentum)) - - -@unevaluated -class InvariantMass(sp.Expr): - """Invariant mass of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - _latex_repr = "m_{{{momentum}}}" - - def evaluate(self) -> ComplexSqrt: - p = self.momentum - p_xyz = ThreeMomentum(p) - return ComplexSqrt(Energy(p) ** 2 - EuclideanNorm(p_xyz) ** 2) - - -@unevaluated -class Phi(sp.Expr): - r"""Azimuthal angle :math:`\phi` of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - _latex_repr_ = R"\phi\left({momentum}\right)" - - def evaluate(self) -> sp.Expr: - p = self.momentum - return sp.atan2(FourMomentumY(p), FourMomentumX(p)) - - -@unevaluated -class Theta(sp.Expr): - r"""Polar (elevation) angle :math:`\theta` of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - _latex_repr_ = R"\theta\left({momentum}\right)" - - def evaluate(self) -> sp.Expr: - p = self.momentum - return sp.acos(FourMomentumZ(p) / three_momentum_norm(p)) - - -@unevaluated -class NegativeMomentum(sp.Expr): - r"""Invert the spatial components of a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - _latex_repr = R"-\left({momentum}\right)" - - def evaluate(self) -> sp.Expr: - p = self.momentum - eta = MinkowskiMetric(p) - return ArrayMultiplication(eta, p) - - -@unevaluated(implement_doit=False) -class MinkowskiMetric(NumPyPrintable): - r"""Minkowski metric :math:`\eta = (1, -1, -1, -1)`.""" - - momentum: sp.Basic - _latex_repr_ = R"\boldsymbol{\eta}" - - def as_explicit(self) -> sp.MutableDenseMatrix: - return sp.Matrix([ - [1, 0, 0, 0], - [0, -1, 0, 0], - [0, 0, -1, 0], - [0, 0, 0, -1], - ]) - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - printer.module_imports[printer._module].update({"array", "ones", "zeros"}) - momentum = printer._print(self.momentum) - n_events = f"len({momentum})" - zeros = f"zeros({n_events})" - ones = f"ones({n_events})" - return f"""array( - [ - [{ones}, {zeros}, {zeros}, {zeros}], - [{zeros}, -{ones}, {zeros}, {zeros}], - [{zeros}, {zeros}, -{ones}, {zeros}], - [{zeros}, {zeros}, {zeros}, -{ones}], - ] - ).transpose((2, 0, 1))""" - - -@unevaluated(commutative=False) -class BoostZMatrix(sp.Expr): - r"""Represents a Lorentz boost matrix in the :math:`z`-direction. - - Args: - beta: Velocity in the :math:`z`-direction, :math:`\beta=p_z/E`. - n_events: Number of events :math:`n` for this matrix array of shape - :math:`n\times4\times4`. Defaults to the `len` of :code:`beta`. - """ - - beta: sp.Basic - n_events: sp.Basic - - def as_explicit(self) -> sp.MutableDenseMatrix: - beta = self.beta - gamma = 1 / ComplexSqrt(1 - beta**2) # type: ignore[operator] - return sp.Matrix([ - [gamma, 0, 0, -gamma * beta], - [0, 1, 0, 0], - [0, 0, 1, 0], - [-gamma * beta, 0, 0, gamma], - ]) - - def evaluate(self) -> _BoostZMatrixImplementation: - beta = self.beta - gamma = 1 / sp.sqrt(1 - beta**2) # type: ignore[operator] - n_events = self.n_events - return _BoostZMatrixImplementation( - beta=beta, - gamma=gamma, - gamma_beta=gamma * beta, - ones=_OnesArray(n_events), - zeros=_ZerosArray(n_events), - ) - - def _latex_repr_(self, printer: LatexPrinter, *args) -> str: - return printer._print(self.evaluate(), *args) - - -@unevaluated(implement_doit=False) -class _BoostZMatrixImplementation(NumPyPrintable): - beta: sp.Basic - gamma: sp.Basic - gamma_beta: sp.Basic - ones: _OnesArray - zeros: _ZerosArray - _latex_repr = R"\boldsymbol{{B_z}}\left({beta}\right)" - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - printer.module_imports[printer._module].add("array") - _, gamma, gamma_beta, ones, zeros = map(printer._print, self.args) - return f"""array( - [ - [{gamma}, {zeros}, {zeros}, -{gamma_beta}], - [{zeros}, {ones}, {zeros}, {zeros}], - [{zeros}, {zeros}, {ones}, {zeros}], - [-{gamma_beta}, {zeros}, {zeros}, {gamma}], - ] - ).transpose((2, 0, 1))""" - - -@unevaluated(commutative=False) -class BoostMatrix(sp.Expr): - r"""Compute a rank-3 Lorentz boost matrix from a `.FourMomentumSymbol`.""" - - momentum: sp.Basic - _latex_repr = R"\boldsymbol{{B}}\left({momentum}\right)" - - def as_explicit(self) -> sp.MutableDenseMatrix: - momentum = self.momentum - energy = Energy(momentum) - beta_sq = EuclideanNormSquared(ThreeMomentum(momentum)) / energy**2 - beta_x = FourMomentumX(momentum) / energy - beta_y = FourMomentumY(momentum) / energy - beta_z = FourMomentumZ(momentum) / energy - g = 1 / sp.sqrt(1 - beta_sq) - return sp.Matrix([ - [g, -g * beta_x, -g * beta_y, -g * beta_z], - [ - -g * beta_x, - 1 + (g - 1) * beta_x**2 / beta_sq, - (g - 1) * beta_y * beta_x / beta_sq, - (g - 1) * beta_z * beta_x / beta_sq, - ], - [ - -g * beta_y, - (g - 1) * beta_x * beta_y / beta_sq, - 1 + (g - 1) * beta_y**2 / beta_sq, - (g - 1) * beta_z * beta_y / beta_sq, - ], - [ - -g * beta_z, - (g - 1) * beta_x * beta_z / beta_sq, - (g - 1) * beta_y * beta_z / beta_sq, - 1 + (g - 1) * beta_z**2 / beta_sq, - ], - ]) - - def evaluate(self) -> _BoostMatrixImplementation: - p = self.momentum - energy = Energy(p) - beta_sq = EuclideanNormSquared(ThreeMomentum(p)) / energy**2 - beta_x = FourMomentumX(p) / energy - beta_y = FourMomentumY(p) / energy - beta_z = FourMomentumZ(p) / energy - gamma = 1 / sp.sqrt(1 - beta_sq) - return _BoostMatrixImplementation( - momentum=p, - b00=gamma, - b01=-gamma * beta_x, - b02=-gamma * beta_y, - b03=-gamma * beta_z, - b11=1 + (gamma - 1) * beta_x**2 / beta_sq, - b12=(gamma - 1) * beta_x * beta_y / beta_sq, - b13=(gamma - 1) * beta_x * beta_z / beta_sq, - b22=1 + (gamma - 1) * beta_y**2 / beta_sq, - b23=(gamma - 1) * beta_y * beta_z / beta_sq, - b33=1 + (gamma - 1) * beta_z**2 / beta_sq, - ) - - -@unevaluated(commutative=False, implement_doit=False) -class _BoostMatrixImplementation(NumPyPrintable): - momentum: sp.Basic - b00: sp.Basic - b01: sp.Basic - b02: sp.Basic - b03: sp.Basic - b11: sp.Basic - b12: sp.Basic - b13: sp.Basic - b22: sp.Basic - b23: sp.Basic - b33: sp.Basic - _latex_repr = R"\boldsymbol{{B}}\left({momentum}\right)" - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - _, b00, b01, b02, b03, b11, b12, b13, b22, b23, b33 = self.args - return f"""array( - [ - [{b00}, {b01}, {b02}, {b03}], - [{b01}, {b11}, {b12}, {b13}], - [{b02}, {b12}, {b22}, {b23}], - [{b03}, {b13}, {b23}, {b33}], - ] - ).transpose((2, 0, 1))""" - - -@unevaluated(commutative=False) -class RotationYMatrix(sp.Expr): - r"""Rotation matrix around the :math:`y`-axis for a `.FourMomentumSymbol`. - - Args: - angle: Angle with which to rotate, see e.g. `.Phi` and `.Theta`. - n_events: Number of events :math:`n` for this matrix array of shape - :math:`n\times4\times4`. Defaults to the `len` of :code:`angle`. - """ - - angle: sp.Basic - n_events: sp.Basic - - def as_explicit(self) -> sp.MutableDenseMatrix: - angle = self.angle - return sp.Matrix([ - [1, 0, 0, 0], - [0, sp.cos(angle), 0, sp.sin(angle)], - [0, 0, 1, 0], - [0, -sp.sin(angle), 0, sp.cos(angle)], - ]) - - def evaluate(self) -> _RotationYMatrixImplementation: - return _RotationYMatrixImplementation( - angle=self.angle, - cos_angle=sp.cos(self.angle), - sin_angle=sp.sin(self.angle), - ones=_OnesArray(self.n_events), - zeros=_ZerosArray(self.n_events), - ) - - def _latex_repr_(self, printer: LatexPrinter, *args) -> str: - return printer._print(self.evaluate(), *args) - - -@unevaluated(commutative=False, implement_doit=False) -class _RotationYMatrixImplementation(NumPyPrintable): - angle: sp.Basic - cos_angle: sp.Basic - sin_angle: sp.Basic - ones: _OnesArray - zeros: _ZerosArray - _latex_repr = R"\boldsymbol{{R_y}}\left({angle}\right)" - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - printer.module_imports[printer._module].add("array") - _, cos_angle, sin_angle, ones, zeros = map(printer._print, self.args) - return f"""array( - [ - [{ones}, {zeros}, {zeros}, {zeros}], - [{zeros}, {cos_angle}, {zeros}, {sin_angle}], - [{zeros}, {zeros}, {ones}, {zeros}], - [{zeros}, -{sin_angle}, {zeros}, {cos_angle}], - ] - ).transpose((2, 0, 1))""" - - -@unevaluated(commutative=False) -class RotationZMatrix(sp.Expr): - r"""Rotation matrix around the :math:`z`-axis for a `.FourMomentumSymbol`. - - Args: - angle: Angle with which to rotate, see e.g. `.Phi` and `.Theta`. - n_events: Number of events :math:`n` for this matrix array of shape - :math:`n\times4\times4`. Defaults to the `len` of :code:`angle`. - """ - - angle: sp.Basic - n_events: sp.Basic - - def as_explicit(self) -> sp.MutableDenseMatrix: - angle = self.angle - return sp.Matrix([ - [1, 0, 0, 0], - [0, sp.cos(angle), -sp.sin(angle), 0], - [0, sp.sin(angle), sp.cos(angle), 0], - [0, 0, 0, 1], - ]) - - def evaluate(self) -> _RotationZMatrixImplementation: - return _RotationZMatrixImplementation( - angle=self.angle, - cos_angle=sp.cos(self.angle), - sin_angle=sp.sin(self.angle), - ones=_OnesArray(self.n_events), - zeros=_ZerosArray(self.n_events), - ) - - def _latex_repr_(self, printer: LatexPrinter, *args) -> str: - return printer._print(self.evaluate(), *args) - - -@unevaluated(commutative=False, implement_doit=False) -class _RotationZMatrixImplementation(NumPyPrintable): - angle: sp.Basic - cos_angle: sp.Basic - sin_angle: sp.Basic - ones: _OnesArray - zeros: _ZerosArray - _latex_repr = R"\boldsymbol{{R_z}}\left({angle}\right)" - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - printer.module_imports[printer._module].add("array") - _, cos_angle, sin_angle, ones, zeros = map(printer._print, self.args) - return f"""array( - [ - [{ones}, {zeros}, {zeros}, {zeros}], - [{zeros}, {cos_angle}, -{sin_angle}, {zeros}], - [{zeros}, {sin_angle}, {cos_angle}, {zeros}], - [{zeros}, {zeros}, {zeros}, {ones}], - ] - ).transpose((2, 0, 1))""" - - -@unevaluated(implement_doit=False) -class _OnesArray(NumPyPrintable): - shape: Any - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - printer.module_imports[printer._module].add("ones") - shape = printer._print(self.shape) - return f"ones({shape})" - - -@unevaluated(implement_doit=False) -class _ZerosArray(NumPyPrintable): - shape: Any - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - printer.module_imports[printer._module].add("zeros") - shape = printer._print(self.shape) - return f"zeros({shape})" - - -@unevaluated(implement_doit=False) -class ArraySize(NumPyPrintable): - """Symbolic expression for getting the size of a numerical array.""" - - array: Any - _latex_repr_ = "N_{{{array}}}" - - def _numpycode(self, printer: NumPyPrinter, *args) -> str: - array = printer._print(self.array) - return f"len({array})" - - -def compute_helicity_angles( - four_momenta: Mapping[int, sp.Expr], topology: Topology -) -> dict[sp.Symbol, sp.Expr]: - """Formulate expressions for all helicity angles in a topology. - - Formulate expressions (`~sympy.core.expr.Expr`) for all helicity angles appearing in - a given `~qrules.topology.Topology`. The expressions are given in terms of - `FourMomenta` The expressions returned as values in a `dict`, where the keys are - defined by :func:`.get_helicity_angle_symbols`. - - Example - ------- - >>> from qrules.topology import create_isobar_topologies - >>> topologies = create_isobar_topologies(3) - >>> topology = topologies[0] - >>> four_momenta = create_four_momentum_symbols(topology) - >>> angles = compute_helicity_angles(four_momenta, topology) - >>> theta_symbol = sp.Symbol("theta_0", real=True) - >>> angles[theta_symbol] - Theta(p1 + p2) - """ - if topology.outgoing_edge_ids != set(four_momenta): - msg = ( - f"Momentum IDs {set(four_momenta)} do not match final state edge IDs" - f" {set(topology.outgoing_edge_ids)}" - ) - raise ValueError(msg) - - n_events = _get_number_of_events(four_momenta) - - def __recursive_helicity_angles( - four_momenta: Mapping[int, sp.Expr], node_id: int - ) -> dict[sp.Symbol, sp.Expr]: - helicity_angles: dict[sp.Symbol, sp.Expr] = {} - child_state_ids = sorted(topology.get_edge_ids_outgoing_from_node(node_id)) - if all(topology.edges[i].ending_node_id is None for i in child_state_ids): - state_id = child_state_ids[0] - if is_opposite_helicity_state(topology, state_id): - state_id = child_state_ids[1] - four_momentum: sp.Expr = four_momenta[state_id] - phi, theta = get_helicity_angle_symbols(topology, state_id) - helicity_angles[phi] = Phi(four_momentum) - helicity_angles[theta] = Theta(four_momentum) - for state_id in child_state_ids: - edge = topology.edges[state_id] - if edge.ending_node_id is not None: - # recursively determine all momenta ids in the list - sub_momenta_ids = determine_attached_final_state(topology, state_id) - if len(sub_momenta_ids) > 1: - # add all of these momenta together -> defines new subsystem - four_momentum = ArraySum(*[ - four_momenta[i] for i in sub_momenta_ids - ]) - - # boost all of those momenta into this new subsystem - phi_expr = Phi(four_momentum) - theta_expr = Theta(four_momentum) - p3_norm = three_momentum_norm(four_momentum) - beta = p3_norm / Energy(four_momentum) - new_momentum_pool: dict[int, sp.Expr] = { - k: ArrayMultiplication( - BoostZMatrix(beta, n_events), - RotationYMatrix(-theta_expr, n_events), - RotationZMatrix(-phi_expr, n_events), - p, - ) - for k, p in four_momenta.items() - if k in sub_momenta_ids - } - - # register current angle variables - if is_opposite_helicity_state(topology, state_id): - state_id = get_sibling_state_id(topology, state_id) - phi, theta = get_helicity_angle_symbols(topology, state_id) - helicity_angles[phi] = Phi(four_momentum) - helicity_angles[theta] = Theta(four_momentum) - - # call next recursion - angles = __recursive_helicity_angles( - new_momentum_pool, - edge.ending_node_id, - ) - helicity_angles.update(angles) - - return helicity_angles - - initial_state_id = next(iter(topology.incoming_edge_ids)) - initial_state_edge = topology.edges[initial_state_id] - assert initial_state_edge.ending_node_id is not None # noqa: S101 - return __recursive_helicity_angles(four_momenta, initial_state_edge.ending_node_id) - - -def _get_number_of_events( - four_momenta: Mapping[int, sp.Expr], -) -> ArraySize: - sorted_momentum_symbols = sorted(four_momenta.values(), key=str) - return ArraySize(sorted_momentum_symbols[0]) - - -def compute_invariant_masses( - four_momenta: FourMomenta, topology: Topology -) -> dict[sp.Symbol, sp.Expr]: - """Compute the invariant masses for all final state combinations.""" - if topology.outgoing_edge_ids != set(four_momenta): - msg = ( - f"Momentum IDs {set(four_momenta)} do not match final state edge IDs" - f" {set(topology.outgoing_edge_ids)}" - ) - raise ValueError(msg) - invariant_masses: dict[sp.Symbol, sp.Expr] = {} - for state_id in topology.edges: - attached_state_ids = determine_attached_final_state(topology, state_id) - total_momentum = ArraySum(*[four_momenta[i] for i in attached_state_ids]) - expr = InvariantMass(total_momentum) - symbol = get_invariant_mass_symbol(topology, state_id) - invariant_masses[symbol] = expr - return invariant_masses - - -def compute_wigner_angles( - topology: Topology, momenta: FourMomenta, state_id: int -) -> dict[sp.Symbol, sp.Expr]: - """Create an `~sympy.core.expr.Expr` for each angle in a Wigner rotation. - - Implementation of (B.2-4) in :cite:`marangottoHelicityAmplitudesGeneric2020`, with - :math:`x'_z` etc. taken from the result of :func:`compute_wigner_rotation_matrix`. - """ - wigner_rotation_matrix = compute_wigner_rotation_matrix(topology, momenta, state_id) - x_z = ArraySlice(wigner_rotation_matrix, (slice(None), 1, 3)) - y_z = ArraySlice(wigner_rotation_matrix, (slice(None), 2, 3)) - z_x = ArraySlice(wigner_rotation_matrix, (slice(None), 3, 1)) - z_y = ArraySlice(wigner_rotation_matrix, (slice(None), 3, 2)) - z_z = ArraySlice(wigner_rotation_matrix, (slice(None), 3, 3)) - suffix = get_helicity_suffix(topology, state_id) - alpha, beta, gamma = sp.symbols( - f"alpha{suffix} beta{suffix} gamma{suffix}", real=True - ) - return { - alpha: sp.atan2(z_y, z_x), - beta: sp.acos(z_z), - gamma: sp.atan2(y_z, -x_z), - } - - -def compute_wigner_rotation_matrix( - topology: Topology, momenta: FourMomenta, state_id: int -) -> MatrixMultiplication: - """Compute a Wigner rotation matrix. - - Implementation of Eq. (36) in :cite:`marangottoHelicityAmplitudesGeneric2020`. - """ - momentum = momenta[state_id] - inverted_direct_boost = BoostMatrix(NegativeMomentum(momentum)) - boost_chain = compute_boost_chain(topology, momenta, state_id) - return MatrixMultiplication(inverted_direct_boost, *boost_chain) - - -def compute_boost_chain( - topology: Topology, momenta: FourMomenta, state_id: int -) -> list[BoostMatrix]: - boost_matrices = [] - decay_chain_state_ids = __get_boost_chain_ids(topology, state_id) - boosted_momenta: dict[int, sp.Expr] = { - i: get_four_momentum_sum(topology, momenta, i) for i in decay_chain_state_ids - } - for current_state_id in decay_chain_state_ids: - current_momentum = boosted_momenta[current_state_id] - boost = BoostMatrix(current_momentum) - boosted_momenta = { - i: ArrayMultiplication(boost, p) for i, p in boosted_momenta.items() - } - boost_matrices.append(boost) - return boost_matrices - - -def __get_boost_chain_ids(topology: Topology, state_id: int) -> list[int]: - """Get the state IDs from first resonance to this final state. - - >>> from qrules.topology import create_isobar_topologies - >>> topology = create_isobar_topologies(3)[0] - >>> __get_boost_chain_ids(topology, state_id=0) - [0] - >>> __get_boost_chain_ids(topology, state_id=1) - [3, 1] - >>> __get_boost_chain_ids(topology, state_id=2) - [3, 2] - """ - decay_chain_state_ids = list(reversed(list_decay_chain_ids(topology, state_id))) - initial_state_id = next(iter(topology.incoming_edge_ids)) - decay_chain_state_ids.remove(initial_state_id) - return decay_chain_state_ids - - -def get_four_momentum_sum( - topology: Topology, momenta: FourMomenta, state_id: int -) -> ArraySum | FourMomentumSymbol: - """Get the `FourMomentumSymbol` or sum of momenta for **any** edge ID. - - If the edge ID is a final state ID, return its `FourMomentumSymbol`. If it's an - intermediate edge ID, return the sum of the momenta of the final states to which it - decays. - - >>> from qrules.topology import create_isobar_topologies - >>> topology = create_isobar_topologies(3)[0] - >>> momenta = create_four_momentum_symbols(topology) - >>> get_four_momentum_sum(topology, momenta, state_id=0) - p0 - >>> get_four_momentum_sum(topology, momenta, state_id=3) - p1 + p2 - """ - if state_id in topology.outgoing_edge_ids: - return momenta[state_id] - sub_momenta_ids = determine_attached_final_state(topology, state_id) - return ArraySum(*[momenta[i] for i in sub_momenta_ids]) - - -def get_invariant_mass_symbol(topology: Topology, state_id: int) -> sp.Symbol: - """Generate an invariant mass label for a state (edge on a topology). - - Example - ------- - In the case shown in Figure :ref:`one-to-five-topology-0`, the invariant mass of - state :math:`5` is :math:`m_{034}`, because :math:`p_5=p_0+p_3+p_4`: - - >>> from qrules.topology import create_isobar_topologies - >>> topologies = create_isobar_topologies(5) - >>> get_invariant_mass_symbol(topologies[0], state_id=5) - m_034 - - Naturally, the 'invariant' mass label for a final state is just the mass of the - state itself: - - >>> get_invariant_mass_symbol(topologies[0], state_id=1) - m_1 - """ - final_state_ids = determine_attached_final_state(topology, state_id) - mass_name = f"m_{''.join(map(str, sorted(final_state_ids)))}" - return sp.Symbol(mass_name, nonnegative=True) diff --git a/src/ampform/kinematics/angles.py b/src/ampform/kinematics/angles.py new file mode 100644 index 000000000..94af69dae --- /dev/null +++ b/src/ampform/kinematics/angles.py @@ -0,0 +1,203 @@ +"""Angle computations for (boosted) :mod:`.lorentz` vectors.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping + +import sympy as sp + +from ampform.helicity.decay import ( + determine_attached_final_state, + get_sibling_state_id, + is_opposite_helicity_state, +) +from ampform.helicity.naming import get_helicity_angle_symbols, get_helicity_suffix +from ampform.kinematics.lorentz import ( + ArraySize, + BoostMatrix, + BoostZMatrix, + Energy, + FourMomenta, + FourMomentumX, + FourMomentumY, + FourMomentumZ, + NegativeMomentum, + RotationYMatrix, + RotationZMatrix, + compute_boost_chain, + three_momentum_norm, +) +from ampform.sympy import unevaluated +from ampform.sympy._array_expressions import ( + ArrayMultiplication, + ArraySlice, + ArraySum, + MatrixMultiplication, +) + +if TYPE_CHECKING: + from qrules.topology import Topology + + +@unevaluated +class Phi(sp.Expr): + r"""Azimuthal angle :math:`\phi` of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + _latex_repr_ = R"\phi\left({momentum}\right)" + + def evaluate(self) -> sp.Expr: + p = self.momentum + return sp.atan2(FourMomentumY(p), FourMomentumX(p)) + + +@unevaluated +class Theta(sp.Expr): + r"""Polar (elevation) angle :math:`\theta` of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + _latex_repr_ = R"\theta\left({momentum}\right)" + + def evaluate(self) -> sp.Expr: + p = self.momentum + return sp.acos(FourMomentumZ(p) / three_momentum_norm(p)) + + +def compute_helicity_angles( + four_momenta: Mapping[int, sp.Expr], topology: Topology +) -> dict[sp.Symbol, sp.Expr]: + """Formulate expressions for all helicity angles in a topology. + + Formulate expressions (`~sympy.core.expr.Expr`) for all helicity angles appearing in + a given `~qrules.topology.Topology`. The expressions are given in terms of + `.FourMomenta` The expressions returned as values in a `dict`, where the keys are + defined by :func:`.get_helicity_angle_symbols`. + + Example + ------- + >>> from qrules.topology import create_isobar_topologies + >>> topologies = create_isobar_topologies(3) + >>> topology = topologies[0] + >>> from ampform.kinematics.lorentz import create_four_momentum_symbols + >>> four_momenta = create_four_momentum_symbols(topology) + >>> angles = compute_helicity_angles(four_momenta, topology) + >>> theta_symbol = sp.Symbol("theta_0", real=True) + >>> angles[theta_symbol] + Theta(p1 + p2) + """ + if topology.outgoing_edge_ids != set(four_momenta): + msg = ( + f"Momentum IDs {set(four_momenta)} do not match final state edge IDs" + f" {set(topology.outgoing_edge_ids)}" + ) + raise ValueError(msg) + + n_events = _get_number_of_events(four_momenta) + + def __recursive_helicity_angles( + four_momenta: Mapping[int, sp.Expr], node_id: int + ) -> dict[sp.Symbol, sp.Expr]: + helicity_angles: dict[sp.Symbol, sp.Expr] = {} + child_state_ids = sorted(topology.get_edge_ids_outgoing_from_node(node_id)) + if all(topology.edges[i].ending_node_id is None for i in child_state_ids): + state_id = child_state_ids[0] + if is_opposite_helicity_state(topology, state_id): + state_id = child_state_ids[1] + four_momentum: sp.Expr = four_momenta[state_id] + phi, theta = get_helicity_angle_symbols(topology, state_id) + helicity_angles[phi] = Phi(four_momentum) + helicity_angles[theta] = Theta(four_momentum) + for state_id in child_state_ids: + edge = topology.edges[state_id] + if edge.ending_node_id is not None: + # recursively determine all momenta ids in the list + sub_momenta_ids = determine_attached_final_state(topology, state_id) + if len(sub_momenta_ids) > 1: + # add all of these momenta together -> defines new subsystem + four_momentum = ArraySum(*[ + four_momenta[i] for i in sub_momenta_ids + ]) + + # boost all of those momenta into this new subsystem + phi_expr = Phi(four_momentum) + theta_expr = Theta(four_momentum) + p3_norm = three_momentum_norm(four_momentum) + beta = p3_norm / Energy(four_momentum) + new_momentum_pool: dict[int, sp.Expr] = { + k: ArrayMultiplication( + BoostZMatrix(beta, n_events), + RotationYMatrix(-theta_expr, n_events), + RotationZMatrix(-phi_expr, n_events), + p, + ) + for k, p in four_momenta.items() + if k in sub_momenta_ids + } + + # register current angle variables + if is_opposite_helicity_state(topology, state_id): + state_id = get_sibling_state_id(topology, state_id) + phi, theta = get_helicity_angle_symbols(topology, state_id) + helicity_angles[phi] = Phi(four_momentum) + helicity_angles[theta] = Theta(four_momentum) + + # call next recursion + angles = __recursive_helicity_angles( + new_momentum_pool, + edge.ending_node_id, + ) + helicity_angles.update(angles) + + return helicity_angles + + initial_state_id = next(iter(topology.incoming_edge_ids)) + initial_state_edge = topology.edges[initial_state_id] + if initial_state_edge.ending_node_id is None: + msg = "Edge does not end in a node" + raise ValueError(msg) + return __recursive_helicity_angles(four_momenta, initial_state_edge.ending_node_id) + + +def _get_number_of_events(four_momenta: Mapping[int, sp.Expr]) -> ArraySize: + sorted_momentum_symbols = sorted(four_momenta.values(), key=str) + return ArraySize(sorted_momentum_symbols[0]) + + +def compute_wigner_angles( + topology: Topology, momenta: FourMomenta, state_id: int +) -> dict[sp.Symbol, sp.Expr]: + """Create an `~sympy.core.expr.Expr` for each angle in a Wigner rotation. + + Implementation of (B.2-4) in :cite:`marangottoHelicityAmplitudesGeneric2020`, with + :math:`x'_z` etc. taken from the result of :func:`compute_wigner_rotation_matrix`. + See also `Wigner rotations `_. + """ + wigner_rotation_matrix = compute_wigner_rotation_matrix(topology, momenta, state_id) + x_z = ArraySlice(wigner_rotation_matrix, (slice(None), 1, 3)) + y_z = ArraySlice(wigner_rotation_matrix, (slice(None), 2, 3)) + z_x = ArraySlice(wigner_rotation_matrix, (slice(None), 3, 1)) + z_y = ArraySlice(wigner_rotation_matrix, (slice(None), 3, 2)) + z_z = ArraySlice(wigner_rotation_matrix, (slice(None), 3, 3)) + suffix = get_helicity_suffix(topology, state_id) + alpha, beta, gamma = sp.symbols( + f"alpha{suffix} beta{suffix} gamma{suffix}", real=True + ) + return { + alpha: sp.atan2(z_y, z_x), + beta: sp.acos(z_z), + gamma: sp.atan2(y_z, -x_z), + } + + +def compute_wigner_rotation_matrix( + topology: Topology, momenta: FourMomenta, state_id: int +) -> MatrixMultiplication: + """Compute a Wigner rotation matrix. + + Implementation of Eq. (36) in :cite:`marangottoHelicityAmplitudesGeneric2020`. See + also `Wigner rotations `_. + """ + momentum = momenta[state_id] + inverted_direct_boost = BoostMatrix(NegativeMomentum(momentum)) + boost_chain = compute_boost_chain(topology, momenta, state_id) + return MatrixMultiplication(inverted_direct_boost, *boost_chain) diff --git a/src/ampform/kinematics/lorentz.py b/src/ampform/kinematics/lorentz.py new file mode 100644 index 000000000..c075fdb43 --- /dev/null +++ b/src/ampform/kinematics/lorentz.py @@ -0,0 +1,612 @@ +"""Symbolic implementations for Lorentz vectors and boosts.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Dict + +import sympy as sp + +from ampform.helicity.decay import determine_attached_final_state, list_decay_chain_ids +from ampform.sympy import ExprClass, NumPyPrintable, unevaluated +from ampform.sympy._array_expressions import ( + ArrayAxisSum, + ArrayMultiplication, + ArraySlice, + ArraySum, + ArraySymbol, +) +from ampform.sympy.math import ComplexSqrt + +if TYPE_CHECKING: + from qrules.topology import Topology + from sympy.printing.latex import LatexPrinter + from sympy.printing.numpy import NumPyPrinter + + +def create_four_momentum_symbols(topology: Topology) -> FourMomenta: + """Create a set of array-symbols for a `~qrules.topology.Topology`. + + >>> from qrules.topology import create_isobar_topologies + >>> topologies = create_isobar_topologies(3) + >>> create_four_momentum_symbols(topologies[0]) + {0: p0, 1: p1, 2: p2} + """ + n_final_states = len(topology.outgoing_edge_ids) + return {i: FourMomentumSymbol(f"p{i}", shape=[]) for i in range(n_final_states)} + + +FourMomenta = Dict[int, "FourMomentumSymbol"] +"""A mapping of state IDs to their corresponding `FourMomentumSymbol`. + +It's best to create a `dict` of `.FourMomenta` with +:func:`create_four_momentum_symbols`. +""" +FourMomentumSymbol = ArraySymbol +r"""Array-`~sympy.core.symbol.Symbol` that represents an array of four-momenta. + +The array is assumed to be of shape :math:`n\times 4` with :math:`n` the number of +events. The four-momenta are assumed to be in the order :math:`\left(E,\vec{p}\right)`. +See also `Energy`, `FourMomentumX`, `FourMomentumY`, and `FourMomentumZ`. +""" + + +@unevaluated +class Energy(sp.Expr): + """Represents the energy-component of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + _latex_repr_ = R"E\left({momentum}\right)" + + def evaluate(self) -> ArraySlice: + return ArraySlice(self.momentum, (slice(None), 0)) + + +def _implement_latex_subscript( # pyright: ignore[reportUnusedFunction] + subscript: str, +) -> Callable[[type[ExprClass]], type[ExprClass]]: + def decorator(decorated_class: type[ExprClass]) -> type[ExprClass]: + def _latex_repr_(self: sp.Expr, printer: LatexPrinter, *args) -> str: + momentum = printer._print(self.momentum) # type: ignore[attr-defined] + if printer._needs_mul_brackets(self.momentum): # type: ignore[attr-defined] + momentum = Rf"\left({momentum}\right)" + else: + momentum = Rf"{{{momentum}}}" + return f"{momentum}_{subscript}" + + decorated_class._latex_repr_ = _latex_repr_ # type: ignore[assignment,attr-defined] + return decorated_class + + return decorator + + +@unevaluated +@_implement_latex_subscript(subscript="x") +class FourMomentumX(sp.Expr): + """Component :math:`x` of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + + def evaluate(self) -> ArraySlice: + return ArraySlice(self.momentum, (slice(None), 1)) + + +@unevaluated +@_implement_latex_subscript(subscript="y") +class FourMomentumY(sp.Expr): + """Component :math:`y` of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + + def evaluate(self) -> ArraySlice: + return ArraySlice(self.momentum, (slice(None), 2)) + + +@unevaluated +@_implement_latex_subscript(subscript="z") +class FourMomentumZ(sp.Expr): + """Component :math:`z` of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + + def evaluate(self) -> ArraySlice: + return ArraySlice(self.momentum, (slice(None), 3)) + + +@unevaluated +class ThreeMomentum(NumPyPrintable): + """Spatial components of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + _latex_repr_ = R"\vec{{{momentum}}}" + + def evaluate(self) -> ArraySlice: + return ArraySlice(self.momentum, (slice(None), slice(1, None))) + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + return printer._print(self.evaluate()) + + +@unevaluated +class EuclideanNorm(NumPyPrintable): + """Take the euclidean norm of an array over axis 1.""" + + vector: sp.Basic + _latex_repr = R"\left|{vector}\right|" + + def evaluate(self) -> ArraySlice: + return sp.sqrt(EuclideanNormSquared(self.vector)) + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + return printer._print(self.evaluate(), *args) + + +@unevaluated +class EuclideanNormSquared(sp.Expr): + """Take the squared euclidean norm of an array over axis 1.""" + + vector: sp.Basic + _latex_repr_ = R"\left|{vector}\right|^{{2}}" + + def evaluate(self) -> ArrayAxisSum: + return ArrayAxisSum(self.vector**2, axis=1) # type: ignore[operator] + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + return printer._print(self.evaluate(), *args) + + +def three_momentum_norm(momentum: sp.Basic) -> EuclideanNorm: + return EuclideanNorm(ThreeMomentum(momentum)) + + +@unevaluated +class InvariantMass(sp.Expr): + """Invariant mass of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + _latex_repr = "m_{{{momentum}}}" + + def evaluate(self) -> ComplexSqrt: + p = self.momentum + p_xyz = ThreeMomentum(p) + return ComplexSqrt(Energy(p) ** 2 - EuclideanNorm(p_xyz) ** 2) + + +@unevaluated +class NegativeMomentum(sp.Expr): + r"""Invert the spatial components of a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + _latex_repr = R"-\left({momentum}\right)" + + def evaluate(self) -> sp.Expr: + p = self.momentum + eta = MinkowskiMetric(p) + return ArrayMultiplication(eta, p) + + +@unevaluated(implement_doit=False) +class MinkowskiMetric(NumPyPrintable): + r"""Minkowski metric :math:`\eta = (1, -1, -1, -1)`.""" + + momentum: sp.Basic + _latex_repr_ = R"\boldsymbol{\eta}" + + def as_explicit(self) -> sp.MutableDenseMatrix: + return sp.Matrix([ + [1, 0, 0, 0], + [0, -1, 0, 0], + [0, 0, -1, 0], + [0, 0, 0, -1], + ]) + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + printer.module_imports[printer._module].update({"array", "ones", "zeros"}) + momentum = printer._print(self.momentum) + n_events = f"len({momentum})" + zeros = f"zeros({n_events})" + ones = f"ones({n_events})" + return f"""array( + [ + [{ones}, {zeros}, {zeros}, {zeros}], + [{zeros}, -{ones}, {zeros}, {zeros}], + [{zeros}, {zeros}, -{ones}, {zeros}], + [{zeros}, {zeros}, {zeros}, -{ones}], + ] + ).transpose((2, 0, 1))""" + + +@unevaluated(commutative=False) +class BoostZMatrix(sp.Expr): + r"""Represents a Lorentz boost matrix in the :math:`z`-direction. + + Args: + beta: Velocity in the :math:`z`-direction, :math:`\beta=p_z/E`. + n_events: Number of events :math:`n` for this matrix array of shape + :math:`n\times4\times4`. Defaults to the `len` of :code:`beta`. + """ + + beta: sp.Basic + n_events: sp.Basic + + def as_explicit(self) -> sp.MutableDenseMatrix: + beta = self.beta + gamma = 1 / ComplexSqrt(1 - beta**2) # type: ignore[operator] + return sp.Matrix([ + [gamma, 0, 0, -gamma * beta], + [0, 1, 0, 0], + [0, 0, 1, 0], + [-gamma * beta, 0, 0, gamma], + ]) + + def evaluate(self) -> _BoostZMatrixImplementation: + beta = self.beta + gamma = 1 / sp.sqrt(1 - beta**2) # type: ignore[operator] + n_events = self.n_events + return _BoostZMatrixImplementation( + beta=beta, + gamma=gamma, + gamma_beta=gamma * beta, + ones=_OnesArray(n_events), + zeros=_ZerosArray(n_events), + ) + + def _latex_repr_(self, printer: LatexPrinter, *args) -> str: + return printer._print(self.evaluate(), *args) + + +@unevaluated(implement_doit=False) +class _BoostZMatrixImplementation(NumPyPrintable): + beta: sp.Basic + gamma: sp.Basic + gamma_beta: sp.Basic + ones: _OnesArray + zeros: _ZerosArray + _latex_repr = R"\boldsymbol{{B_z}}\left({beta}\right)" + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + printer.module_imports[printer._module].add("array") + _, gamma, gamma_beta, ones, zeros = map(printer._print, self.args) + return f"""array( + [ + [{gamma}, {zeros}, {zeros}, -{gamma_beta}], + [{zeros}, {ones}, {zeros}, {zeros}], + [{zeros}, {zeros}, {ones}, {zeros}], + [-{gamma_beta}, {zeros}, {zeros}, {gamma}], + ] + ).transpose((2, 0, 1))""" + + +@unevaluated(commutative=False) +class BoostMatrix(sp.Expr): + r"""Compute a rank-3 Lorentz boost matrix from a `.FourMomentumSymbol`.""" + + momentum: sp.Basic + _latex_repr = R"\boldsymbol{{B}}\left({momentum}\right)" + + def as_explicit(self) -> sp.MutableDenseMatrix: + momentum = self.momentum + energy = Energy(momentum) + beta_sq = EuclideanNormSquared(ThreeMomentum(momentum)) / energy**2 + beta_x = FourMomentumX(momentum) / energy + beta_y = FourMomentumY(momentum) / energy + beta_z = FourMomentumZ(momentum) / energy + g = 1 / sp.sqrt(1 - beta_sq) + return sp.Matrix([ + [g, -g * beta_x, -g * beta_y, -g * beta_z], + [ + -g * beta_x, + 1 + (g - 1) * beta_x**2 / beta_sq, + (g - 1) * beta_y * beta_x / beta_sq, + (g - 1) * beta_z * beta_x / beta_sq, + ], + [ + -g * beta_y, + (g - 1) * beta_x * beta_y / beta_sq, + 1 + (g - 1) * beta_y**2 / beta_sq, + (g - 1) * beta_z * beta_y / beta_sq, + ], + [ + -g * beta_z, + (g - 1) * beta_x * beta_z / beta_sq, + (g - 1) * beta_y * beta_z / beta_sq, + 1 + (g - 1) * beta_z**2 / beta_sq, + ], + ]) + + def evaluate(self) -> _BoostMatrixImplementation: + p = self.momentum + energy = Energy(p) + beta_sq = EuclideanNormSquared(ThreeMomentum(p)) / energy**2 + beta_x = FourMomentumX(p) / energy + beta_y = FourMomentumY(p) / energy + beta_z = FourMomentumZ(p) / energy + gamma = 1 / sp.sqrt(1 - beta_sq) + return _BoostMatrixImplementation( + momentum=p, + b00=gamma, + b01=-gamma * beta_x, + b02=-gamma * beta_y, + b03=-gamma * beta_z, + b11=1 + (gamma - 1) * beta_x**2 / beta_sq, + b12=(gamma - 1) * beta_x * beta_y / beta_sq, + b13=(gamma - 1) * beta_x * beta_z / beta_sq, + b22=1 + (gamma - 1) * beta_y**2 / beta_sq, + b23=(gamma - 1) * beta_y * beta_z / beta_sq, + b33=1 + (gamma - 1) * beta_z**2 / beta_sq, + ) + + +@unevaluated(commutative=False, implement_doit=False) +class _BoostMatrixImplementation(NumPyPrintable): + momentum: sp.Basic + b00: sp.Basic + b01: sp.Basic + b02: sp.Basic + b03: sp.Basic + b11: sp.Basic + b12: sp.Basic + b13: sp.Basic + b22: sp.Basic + b23: sp.Basic + b33: sp.Basic + _latex_repr = R"\boldsymbol{{B}}\left({momentum}\right)" + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + _, b00, b01, b02, b03, b11, b12, b13, b22, b23, b33 = self.args + return f"""array( + [ + [{b00}, {b01}, {b02}, {b03}], + [{b01}, {b11}, {b12}, {b13}], + [{b02}, {b12}, {b22}, {b23}], + [{b03}, {b13}, {b23}, {b33}], + ] + ).transpose((2, 0, 1))""" + + +@unevaluated(commutative=False) +class RotationYMatrix(sp.Expr): + r"""Rotation matrix around the :math:`y`-axis for a `.FourMomentumSymbol`. + + Args: + angle: Angle with which to rotate, see e.g. `.Phi` and `.Theta`. + n_events: Number of events :math:`n` for this matrix array of shape + :math:`n\times4\times4`. Defaults to the `len` of :code:`angle`. + """ + + angle: sp.Basic + n_events: sp.Basic + + def as_explicit(self) -> sp.MutableDenseMatrix: + angle = self.angle + return sp.Matrix([ + [1, 0, 0, 0], + [0, sp.cos(angle), 0, sp.sin(angle)], + [0, 0, 1, 0], + [0, -sp.sin(angle), 0, sp.cos(angle)], + ]) + + def evaluate(self) -> _RotationYMatrixImplementation: + return _RotationYMatrixImplementation( + angle=self.angle, + cos_angle=sp.cos(self.angle), + sin_angle=sp.sin(self.angle), + ones=_OnesArray(self.n_events), + zeros=_ZerosArray(self.n_events), + ) + + def _latex_repr_(self, printer: LatexPrinter, *args) -> str: + return printer._print(self.evaluate(), *args) + + +@unevaluated(commutative=False, implement_doit=False) +class _RotationYMatrixImplementation(NumPyPrintable): + angle: sp.Basic + cos_angle: sp.Basic + sin_angle: sp.Basic + ones: _OnesArray + zeros: _ZerosArray + _latex_repr = R"\boldsymbol{{R_y}}\left({angle}\right)" + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + printer.module_imports[printer._module].add("array") + _, cos_angle, sin_angle, ones, zeros = map(printer._print, self.args) + return f"""array( + [ + [{ones}, {zeros}, {zeros}, {zeros}], + [{zeros}, {cos_angle}, {zeros}, {sin_angle}], + [{zeros}, {zeros}, {ones}, {zeros}], + [{zeros}, -{sin_angle}, {zeros}, {cos_angle}], + ] + ).transpose((2, 0, 1))""" + + +@unevaluated(commutative=False) +class RotationZMatrix(sp.Expr): + r"""Rotation matrix around the :math:`z`-axis for a `.FourMomentumSymbol`. + + Args: + angle: Angle with which to rotate, see e.g. `.Phi` and `.Theta`. + n_events: Number of events :math:`n` for this matrix array of shape + :math:`n\times4\times4`. Defaults to the `len` of :code:`angle`. + """ + + angle: sp.Basic + n_events: sp.Basic + + def as_explicit(self) -> sp.MutableDenseMatrix: + angle = self.angle + return sp.Matrix([ + [1, 0, 0, 0], + [0, sp.cos(angle), -sp.sin(angle), 0], + [0, sp.sin(angle), sp.cos(angle), 0], + [0, 0, 0, 1], + ]) + + def evaluate(self) -> _RotationZMatrixImplementation: + return _RotationZMatrixImplementation( + angle=self.angle, + cos_angle=sp.cos(self.angle), + sin_angle=sp.sin(self.angle), + ones=_OnesArray(self.n_events), + zeros=_ZerosArray(self.n_events), + ) + + def _latex_repr_(self, printer: LatexPrinter, *args) -> str: + return printer._print(self.evaluate(), *args) + + +@unevaluated(commutative=False, implement_doit=False) +class _RotationZMatrixImplementation(NumPyPrintable): + angle: sp.Basic + cos_angle: sp.Basic + sin_angle: sp.Basic + ones: _OnesArray + zeros: _ZerosArray + _latex_repr = R"\boldsymbol{{R_z}}\left({angle}\right)" + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + printer.module_imports[printer._module].add("array") + _, cos_angle, sin_angle, ones, zeros = map(printer._print, self.args) + return f"""array( + [ + [{ones}, {zeros}, {zeros}, {zeros}], + [{zeros}, {cos_angle}, -{sin_angle}, {zeros}], + [{zeros}, {sin_angle}, {cos_angle}, {zeros}], + [{zeros}, {zeros}, {zeros}, {ones}], + ] + ).transpose((2, 0, 1))""" + + +@unevaluated(implement_doit=False) +class _OnesArray(NumPyPrintable): + shape: Any + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + printer.module_imports[printer._module].add("ones") + shape = printer._print(self.shape) + return f"ones({shape})" + + +@unevaluated(implement_doit=False) +class _ZerosArray(NumPyPrintable): + shape: Any + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + printer.module_imports[printer._module].add("zeros") + shape = printer._print(self.shape) + return f"zeros({shape})" + + +@unevaluated(implement_doit=False) +class ArraySize(NumPyPrintable): + """Symbolic expression for getting the size of a numerical array.""" + + array: Any + _latex_repr_ = "N_{{{array}}}" + + def _numpycode(self, printer: NumPyPrinter, *args) -> str: + array = printer._print(self.array) + return f"len({array})" + + +def compute_boost_chain( + topology: Topology, momenta: FourMomenta, state_id: int +) -> list[BoostMatrix]: + boost_matrices = [] + decay_chain_state_ids = __get_boost_chain_ids(topology, state_id) + boosted_momenta: dict[int, sp.Expr] = { + i: get_four_momentum_sum(topology, momenta, i) for i in decay_chain_state_ids + } + for current_state_id in decay_chain_state_ids: + current_momentum = boosted_momenta[current_state_id] + boost = BoostMatrix(current_momentum) + boosted_momenta = { + i: ArrayMultiplication(boost, p) for i, p in boosted_momenta.items() + } + boost_matrices.append(boost) + return boost_matrices + + +def __get_boost_chain_ids(topology: Topology, state_id: int) -> list[int]: + """Get the state IDs from first resonance to this final state. + + >>> from qrules.topology import create_isobar_topologies + >>> topology = create_isobar_topologies(3)[0] + >>> __get_boost_chain_ids(topology, state_id=0) + [0] + >>> __get_boost_chain_ids(topology, state_id=1) + [3, 1] + >>> __get_boost_chain_ids(topology, state_id=2) + [3, 2] + """ + decay_chain_state_ids = list(reversed(list_decay_chain_ids(topology, state_id))) + initial_state_id = next(iter(topology.incoming_edge_ids)) + decay_chain_state_ids.remove(initial_state_id) + return decay_chain_state_ids + + +def get_four_momentum_sum( + topology: Topology, momenta: FourMomenta, state_id: int +) -> ArraySum | FourMomentumSymbol: + """Get the `.FourMomentumSymbol` or sum of momenta for **any** edge ID. + + If the edge ID is a final state ID, return its `.FourMomentumSymbol`. If it's an + intermediate edge ID, return the sum of the momenta of the final states to which it + decays. + + >>> from qrules.topology import create_isobar_topologies + >>> topology = create_isobar_topologies(3)[0] + >>> momenta = create_four_momentum_symbols(topology) + >>> get_four_momentum_sum(topology, momenta, state_id=0) + p0 + >>> get_four_momentum_sum(topology, momenta, state_id=3) + p1 + p2 + """ + if state_id in topology.outgoing_edge_ids: + return momenta[state_id] + sub_momenta_ids = determine_attached_final_state(topology, state_id) + return ArraySum(*[momenta[i] for i in sub_momenta_ids]) + + +def compute_invariant_masses( + four_momenta: FourMomenta, topology: Topology +) -> dict[sp.Symbol, sp.Expr]: + """Compute the invariant masses for all final state combinations.""" + if topology.outgoing_edge_ids != set(four_momenta): + msg = ( + f"Momentum IDs {set(four_momenta)} do not match final state edge IDs" + f" {set(topology.outgoing_edge_ids)}" + ) + raise ValueError(msg) + invariant_masses: dict[sp.Symbol, sp.Expr] = {} + for state_id in topology.edges: + attached_state_ids = determine_attached_final_state(topology, state_id) + total_momentum = ArraySum(*[four_momenta[i] for i in attached_state_ids]) + expr = InvariantMass(total_momentum) + symbol = get_invariant_mass_symbol(topology, state_id) + invariant_masses[symbol] = expr + return invariant_masses + + +def get_invariant_mass_symbol(topology: Topology, state_id: int) -> sp.Symbol: + """Generate an invariant mass label for a state (edge on a topology). + + Example + ------- + In the case shown in Figure :ref:`one-to-five-topology-0`, the invariant mass of + state :math:`5` is :math:`m_{034}`, because :math:`p_5=p_0+p_3+p_4`: + + >>> from qrules.topology import create_isobar_topologies + >>> topologies = create_isobar_topologies(5) + >>> get_invariant_mass_symbol(topologies[0], state_id=5) + m_034 + + Naturally, the 'invariant' mass label for a final state is just the mass of the + state itself: + + >>> get_invariant_mass_symbol(topologies[0], state_id=1) + m_1 + """ + final_state_ids = determine_attached_final_state(topology, state_id) + mass_name = f"m_{''.join(map(str, sorted(final_state_ids)))}" + return sp.Symbol(mass_name, nonnegative=True) diff --git a/src/ampform/sympy/__init__.py b/src/ampform/sympy/__init__.py index 601cbb20c..e0c9db2de 100644 --- a/src/ampform/sympy/__init__.py +++ b/src/ampform/sympy/__init__.py @@ -18,12 +18,14 @@ import logging import os import pickle +import re from abc import abstractmethod from os.path import abspath, dirname, expanduser from textwrap import dedent from typing import TYPE_CHECKING, Iterable, Sequence, SupportsFloat import sympy as sp +from sympy.printing.conventions import split_super_sub from sympy.printing.precedence import PRECEDENCE from ._decorator import ( @@ -247,6 +249,38 @@ def _is_regular_series(values: Sequence[SupportsFloat]) -> bool: return True +def determine_indices(symbol: sp.Basic) -> list[int]: + r"""Extract any indices if available from a `~sympy.core.symbol.Symbol`. + + >>> determine_indices(sp.Symbol("m1")) + [1] + >>> determine_indices(sp.Symbol("m_12")) + [12] + >>> determine_indices(sp.Symbol("m_a2")) + [2] + >>> determine_indices(sp.Symbol(R"\alpha_{i2, 5}")) + [2, 5] + >>> determine_indices(sp.Symbol("m")) + [] + + `~sympy.tensor.indexed.Indexed` instances can also be handled: + >>> m_a = sp.IndexedBase("m_a") + >>> determine_indices(m_a[0]) + [0] + """ + _, _, subscripts = split_super_sub(sp.latex(symbol)) + if not subscripts: + return [] + subscript: str = subscripts[-1] + subscript = re.sub(r"[^0-9^\,]", "", subscript) + subscript = f"[{subscript}]" + try: + indices = eval(subscript) # noqa: PGH001, S307 + except SyntaxError: + return [] + return list(indices) + + def perform_cached_doit( unevaluated_expr: sp.Expr, directory: str | None = None ) -> sp.Expr: diff --git a/tests/external/test_sympy.py b/tests/external/test_sympy.py index 2ca661de4..3b025c212 100644 --- a/tests/external/test_sympy.py +++ b/tests/external/test_sympy.py @@ -1,5 +1,4 @@ import operator -from copy import deepcopy from functools import reduce import sympy as sp @@ -39,22 +38,6 @@ def test_name(self): x.name = "x" assert x.name == "x" - def test_name_change(self): - x = sp.Symbol("a") - y = sp.Symbol("a") - assert x.name == y.name - assert x == y - assert x is y - x.name = "I am x" - assert y.name == "I am x" - z = deepcopy(x) - assert z == x - assert z is not x - assert z.name == "I am x" - z.name = "z" - assert x.name == "I am x" - assert z.name == "z" - def test_product(self): symbols = [ sp.Symbol("x"), diff --git a/tests/helicity/align/__init__.py b/tests/helicity/align/__init__.py new file mode 100644 index 000000000..948df262f --- /dev/null +++ b/tests/helicity/align/__init__.py @@ -0,0 +1 @@ +"""Required to set mypy options for the tests folder.""" diff --git a/tests/kinematics/conftest.py b/tests/kinematics/conftest.py new file mode 100644 index 000000000..385c10807 --- /dev/null +++ b/tests/kinematics/conftest.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +from qrules.topology import Topology, create_isobar_topologies + +from ampform.kinematics.lorentz import FourMomenta, create_four_momentum_symbols + +if TYPE_CHECKING: + import numpy as np + + +@pytest.fixture(scope="session") +def topology_and_momentum_symbols( + data_sample: dict[int, np.ndarray] +) -> tuple[Topology, FourMomenta]: + n = len(data_sample) + assert n == 4 + topologies = create_isobar_topologies(n) + topology = topologies[1] + momentum_symbols = create_four_momentum_symbols(topology) + return topology, momentum_symbols diff --git a/tests/kinematics/test_angles.py b/tests/kinematics/test_angles.py new file mode 100644 index 000000000..46470a679 --- /dev/null +++ b/tests/kinematics/test_angles.py @@ -0,0 +1,263 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest +import sympy as sp +from sympy.printing.numpy import NumPyPrinter + +from ampform.helicity.decay import get_parent_id +from ampform.kinematics.angles import ( + Phi, + Theta, + compute_helicity_angles, + compute_wigner_rotation_matrix, +) +from ampform.kinematics.lorentz import FourMomenta, FourMomentumSymbol + +if TYPE_CHECKING: + from qrules.topology import Topology + +m0, m1, m2, m3 = sp.symbols("m_:4", nonnegative=True) +s1: sp.Pow = sp.Symbol("m_23", nonnegative=True) ** 2 # type: ignore[assignment] +s2: sp.Pow = sp.Symbol("m_13", nonnegative=True) ** 2 # type: ignore[assignment] +s3: sp.Pow = sp.Symbol("m_12", nonnegative=True) ** 2 # type: ignore[assignment] + + +@pytest.fixture(scope="session") +def helicity_angles( + topology_and_momentum_symbols: tuple[Topology, FourMomenta] +) -> dict[sp.Symbol, sp.Expr]: + topology, momentum_symbols = topology_and_momentum_symbols + return compute_helicity_angles(momentum_symbols, topology) + + +class TestPhi: + @property + def phi(self): + p = FourMomentumSymbol("p", shape=[]) + return Phi(p) + + def test_latex(self): + latex = sp.latex(self.phi) + assert latex == R"\phi\left(p\right)" + + def test_numpy(self): + phi = self.phi.doit() + numpy_code = _generate_numpy_code(phi) + assert numpy_code == "numpy.arctan2(p[:, 2], p[:, 1])" + + +class TestTheta: + @property + def theta(self): + p = FourMomentumSymbol("p", shape=[]) + return Theta(p) + + def test_latex(self): + latex = sp.latex(self.theta) + assert latex == R"\theta\left(p\right)" + + def test_numpy(self): + theta = self.theta.doit() + numpy_code = _generate_numpy_code(theta) + assert ( + numpy_code == "numpy.arccos(p[:, 3]/numpy.sqrt(sum(p[:, 1:]**2, axis=1)))" + ) + + +@pytest.mark.parametrize("use_cse", [False, True]) +@pytest.mark.parametrize( + ("angle_name", "expected_values"), + [ + ( + "phi_0", + np.array([ + 2.79758, + 2.51292, + -1.07396, + -1.88051, + 1.06433, + -2.30129, + 2.36878, + -2.46888, + 0.568649, + -2.8792, + ]), + ), + ( + "theta_0", + np.arccos([ + -0.914298, + -0.994127, + 0.769715, + -0.918418, + 0.462214, + 0.958535, + 0.496489, + -0.674376, + 0.614968, + -0.0330843, + ]), + ), + ( + "phi_1^123", + np.array([ + 1.04362, + 1.87349, + 0.160733, + -2.81088, + 2.84379, + 2.29128, + 2.24539, + -1.20272, + 0.615838, + 2.98067, + ]), + ), + ( + "theta_1^123", + np.arccos([ + -0.772533, + 0.163659, + 0.556365, + 0.133251, + -0.0264361, + 0.227188, + -0.166924, + 0.652761, + 0.443122, + 0.503577, + ]), + ), + ( + "phi_2^23,123", + np.array([ # WARNING: subsystem solution (ComPWA) results in pi differences + -2.77203 + np.pi, + 1.45339 - np.pi, + -2.51096 + np.pi, + 2.71085 - np.pi, + -1.12706 + np.pi, + -3.01323 + np.pi, + 2.07305 - np.pi, + 0.502648 - np.pi, + -1.23689 + np.pi, + 1.7605 - np.pi, + ]), + ), + ( + "theta_2^23,123", + np.arccos([ + 0.460324, + -0.410464, + 0.248566, + -0.301959, + -0.522502, + 0.787267, + 0.488066, + 0.954167, + -0.553114, + 0.00256349, + ]), + ), + ], +) +def test_compute_helicity_angles( + use_cse: bool, + data_sample: dict[int, np.ndarray], + topology_and_momentum_symbols: tuple[Topology, FourMomenta], + angle_name: str, + expected_values: np.ndarray, + helicity_angles: dict[sp.Symbol, sp.Expr], +): + _, momentum_symbols = topology_and_momentum_symbols + four_momenta = data_sample.values() + angle_symbol = sp.Symbol(angle_name, real=True) + expr = helicity_angles[angle_symbol] + np_angle = sp.lambdify(momentum_symbols.values(), expr.doit(), cse=use_cse) + computed = np_angle(*four_momenta) + # cspell:ignore atol + np.testing.assert_allclose(computed, expected_values, atol=1e-5) + + +@pytest.mark.parametrize( + ("state_id", "expected"), + [ + ( + 0, + "MatrixMultiplication(BoostMatrix(NegativeMomentum(p0)), BoostMatrix(p0))", + ), + ( + 1, + ( + "MatrixMultiplication(BoostMatrix(NegativeMomentum(p1))," + " BoostMatrix(p1 + p2 + p3)," + " BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 + p3)," + " p1)))" + ), + ), + ( + 2, + ( + "MatrixMultiplication(BoostMatrix(NegativeMomentum(p2)), BoostMatrix(p1" + " + p2 + p3), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 +" + " p3), p2 + p3))," + " BoostMatrix(ArrayMultiplication(BoostMatrix(ArrayMultiplication(BoostMatrix(p1" + " + p2 + p3), p2 + p3)), ArrayMultiplication(BoostMatrix(p1 + p2 + p3)," + " p2))))" + ), + ), + ( + 3, + ( + "MatrixMultiplication(BoostMatrix(NegativeMomentum(p3)), BoostMatrix(p1" + " + p2 + p3), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 +" + " p3), p2 + p3))," + " BoostMatrix(ArrayMultiplication(BoostMatrix(ArrayMultiplication(BoostMatrix(p1" + " + p2 + p3), p2 + p3)), ArrayMultiplication(BoostMatrix(p1 + p2 + p3)," + " p3))))" + ), + ), + ], +) +def test_compute_wigner_rotation_matrix( + state_id: int, + expected: str, + topology_and_momentum_symbols: tuple[Topology, FourMomenta], +): + topology, momenta = topology_and_momentum_symbols + expr = compute_wigner_rotation_matrix(topology, momenta, state_id) + assert str(expr) == expected + + +@pytest.mark.parametrize( + "state_id", + [ + 0, + pytest.param(2, marks=pytest.mark.slow), + pytest.param(3, marks=pytest.mark.slow), + ], +) +def test_compute_wigner_rotation_matrix_numpy( + state_id: int, + data_sample: dict[int, np.ndarray], + topology_and_momentum_symbols: tuple[Topology, FourMomenta], +): + topology, momenta = topology_and_momentum_symbols + expr = compute_wigner_rotation_matrix(topology, momenta, state_id) + func = sp.lambdify(momenta.values(), expr.doit(), cse=True) + momentum_array = data_sample[state_id] + wigner_matrix_array = func(*data_sample.values()) + assert wigner_matrix_array.shape == (len(momentum_array), 4, 4) + if get_parent_id(topology, state_id) == -1: + product = np.einsum("...ij,...j->...j", wigner_matrix_array, momentum_array) + assert pytest.approx(product) == momentum_array + matrix_column_norms = np.linalg.norm(wigner_matrix_array, axis=1) + assert pytest.approx(matrix_column_norms) == 1 + + +def _generate_numpy_code(expr: sp.Expr) -> str: + # cspell:ignore doprint + printer = NumPyPrinter() + return printer.doprint(expr) diff --git a/tests/test_kinematics.py b/tests/kinematics/test_lorentz.py similarity index 68% rename from tests/test_kinematics.py rename to tests/kinematics/test_lorentz.py index e02a500eb..096ad2ae5 100644 --- a/tests/test_kinematics.py +++ b/tests/kinematics/test_lorentz.py @@ -1,4 +1,3 @@ -# cspell:ignore atol doprint from __future__ import annotations import inspect @@ -9,11 +8,9 @@ import pytest import sympy as sp from numpy.lib.scimath import sqrt as complex_sqrt -from qrules.topology import Topology, create_isobar_topologies from sympy.printing.numpy import NumPyPrinter -from ampform.helicity.decay import get_parent_id -from ampform.kinematics import ( +from ampform.kinematics.lorentz import ( ArraySize, BoostMatrix, BoostZMatrix, @@ -25,18 +22,13 @@ FourMomentumZ, InvariantMass, NegativeMomentum, - Phi, RotationYMatrix, RotationZMatrix, - Theta, ThreeMomentum, _OnesArray, _ZerosArray, compute_boost_chain, - compute_helicity_angles, compute_invariant_masses, - compute_wigner_rotation_matrix, - create_four_momentum_symbols, three_momentum_norm, ) from ampform.sympy._array_expressions import ( @@ -46,27 +38,9 @@ ) if TYPE_CHECKING: - from ampform.sympy import NumPyPrintable - - -@pytest.fixture(scope="session") -def topology_and_momentum_symbols( - data_sample: dict[int, np.ndarray] -) -> tuple[Topology, FourMomenta]: - n = len(data_sample) - assert n == 4 - topologies = create_isobar_topologies(n) - topology = topologies[1] - momentum_symbols = create_four_momentum_symbols(topology) - return topology, momentum_symbols + from qrules.topology import Topology - -@pytest.fixture(scope="session") -def helicity_angles( - topology_and_momentum_symbols: tuple[Topology, FourMomenta] -) -> dict[sp.Symbol, sp.Expr]: - topology, momentum_symbols = topology_and_momentum_symbols - return compute_helicity_angles(momentum_symbols, topology) + from ampform.sympy import NumPyPrintable class TestBoostMatrix: @@ -258,40 +232,6 @@ def test_numpy(self): assert numpy_code == "p[:, 1:]" -class TestPhi: - @property - def phi(self): - p = FourMomentumSymbol("p", shape=[]) - return Phi(p) - - def test_latex(self): - latex = sp.latex(self.phi) - assert latex == R"\phi\left(p\right)" - - def test_numpy(self): - phi = self.phi.doit() - numpy_code = _generate_numpy_code(phi) - assert numpy_code == "numpy.arctan2(p[:, 2], p[:, 1])" - - -class TestTheta: - @property - def theta(self): - p = FourMomentumSymbol("p", shape=[]) - return Theta(p) - - def test_latex(self): - latex = sp.latex(self.theta) - assert latex == R"\theta\left(p\right)" - - def test_numpy(self): - theta = self.theta.doit() - numpy_code = _generate_numpy_code(theta) - assert ( - numpy_code == "numpy.arccos(p[:, 3]/numpy.sqrt(sum(p[:, 1:]**2, axis=1)))" - ) - - class TestNegativeMomentum: def test_same_as_inverse(self, data_sample: dict[int, np.ndarray]): p = FourMomentumSymbol("p", shape=[]) @@ -421,119 +361,6 @@ def test_numpycode(self, array_type, shape): np.testing.assert_array_equal(array, array_func(shape)) -@pytest.mark.parametrize("use_cse", [False, True]) -@pytest.mark.parametrize( - ("angle_name", "expected_values"), - [ - ( - "phi_0", - np.array([ - 2.79758, - 2.51292, - -1.07396, - -1.88051, - 1.06433, - -2.30129, - 2.36878, - -2.46888, - 0.568649, - -2.8792, - ]), - ), - ( - "theta_0", - np.arccos([ - -0.914298, - -0.994127, - 0.769715, - -0.918418, - 0.462214, - 0.958535, - 0.496489, - -0.674376, - 0.614968, - -0.0330843, - ]), - ), - ( - "phi_1^123", - np.array([ - 1.04362, - 1.87349, - 0.160733, - -2.81088, - 2.84379, - 2.29128, - 2.24539, - -1.20272, - 0.615838, - 2.98067, - ]), - ), - ( - "theta_1^123", - np.arccos([ - -0.772533, - 0.163659, - 0.556365, - 0.133251, - -0.0264361, - 0.227188, - -0.166924, - 0.652761, - 0.443122, - 0.503577, - ]), - ), - ( - "phi_2^23,123", - np.array([ # WARNING: subsystem solution (ComPWA) results in pi differences - -2.77203 + np.pi, - 1.45339 - np.pi, - -2.51096 + np.pi, - 2.71085 - np.pi, - -1.12706 + np.pi, - -3.01323 + np.pi, - 2.07305 - np.pi, - 0.502648 - np.pi, - -1.23689 + np.pi, - 1.7605 - np.pi, - ]), - ), - ( - "theta_2^23,123", - np.arccos([ - 0.460324, - -0.410464, - 0.248566, - -0.301959, - -0.522502, - 0.787267, - 0.488066, - 0.954167, - -0.553114, - 0.00256349, - ]), - ), - ], -) -def test_compute_helicity_angles( - use_cse: bool, - data_sample: dict[int, np.ndarray], - topology_and_momentum_symbols: tuple[Topology, FourMomenta], - angle_name: str, - expected_values: np.ndarray, - helicity_angles: dict[sp.Symbol, sp.Expr], -): - _, momentum_symbols = topology_and_momentum_symbols - four_momenta = data_sample.values() - angle_symbol = sp.Symbol(angle_name, real=True) - expr = helicity_angles[angle_symbol] - np_angle = sp.lambdify(momentum_symbols.values(), expr.doit(), cse=use_cse) - computed = np_angle(*four_momenta) - np.testing.assert_allclose(computed, expected_values, atol=1e-5) - - def test_compute_invariant_masses_names( topology_and_momentum_symbols: tuple[Topology, FourMomenta] ): @@ -595,11 +422,6 @@ def __compute_mass(array: np.ndarray) -> np.ndarray: return complex_sqrt(mass_squared) -def _generate_numpy_code(expr: sp.Expr) -> str: - printer = NumPyPrinter() - return printer.doprint(expr) - - @pytest.mark.parametrize( ("state_id", "expected"), [ @@ -649,77 +471,7 @@ def test_compute_boost_chain( assert boost_chain_str == expected -@pytest.mark.parametrize( - ("state_id", "expected"), - [ - ( - 0, - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p0)), BoostMatrix(p0))", - ), - ( - 1, - ( - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p1))," - " BoostMatrix(p1 + p2 + p3)," - " BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 + p3)," - " p1)))" - ), - ), - ( - 2, - ( - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p2)), BoostMatrix(p1" - " + p2 + p3), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 +" - " p3), p2 + p3))," - " BoostMatrix(ArrayMultiplication(BoostMatrix(ArrayMultiplication(BoostMatrix(p1" - " + p2 + p3), p2 + p3)), ArrayMultiplication(BoostMatrix(p1 + p2 + p3)," - " p2))))" - ), - ), - ( - 3, - ( - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p3)), BoostMatrix(p1" - " + p2 + p3), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 +" - " p3), p2 + p3))," - " BoostMatrix(ArrayMultiplication(BoostMatrix(ArrayMultiplication(BoostMatrix(p1" - " + p2 + p3), p2 + p3)), ArrayMultiplication(BoostMatrix(p1 + p2 + p3)," - " p3))))" - ), - ), - ], -) -def test_compute_wigner_rotation_matrix( - state_id: int, - expected: str, - topology_and_momentum_symbols: tuple[Topology, FourMomenta], -): - topology, momenta = topology_and_momentum_symbols - expr = compute_wigner_rotation_matrix(topology, momenta, state_id) - assert str(expr) == expected - - -@pytest.mark.parametrize( - "state_id", - [ - 0, - pytest.param(2, marks=pytest.mark.slow), - pytest.param(3, marks=pytest.mark.slow), - ], -) -def test_compute_wigner_rotation_matrix_numpy( - state_id: int, - data_sample: dict[int, np.ndarray], - topology_and_momentum_symbols: tuple[Topology, FourMomenta], -): - topology, momenta = topology_and_momentum_symbols - expr = compute_wigner_rotation_matrix(topology, momenta, state_id) - func = sp.lambdify(momenta.values(), expr.doit(), cse=True) - momentum_array = data_sample[state_id] - wigner_matrix_array = func(*data_sample.values()) - assert wigner_matrix_array.shape == (len(momentum_array), 4, 4) - if get_parent_id(topology, state_id) == -1: - product = np.einsum("...ij,...j->...j", wigner_matrix_array, momentum_array) - assert pytest.approx(product) == momentum_array - matrix_column_norms = np.linalg.norm(wigner_matrix_array, axis=1) - assert pytest.approx(matrix_column_norms) == 1 +def _generate_numpy_code(expr: sp.Expr) -> str: + # cspell:ignore doprint + printer = NumPyPrinter() + return printer.doprint(expr)