implement first version of decay data structures

pyproject.toml

@@ -58,4 +58,5 @@ norecursedirs = [
 ]
 testpaths = [
     "src",
+    "tests",
 ]

pyrightconfig.json

@@ -1,6 +1,7 @@
 {
   "include": ["src", "tests"],
   "reportGeneralTypeIssues": false,
+  "reportImportCycles": false,
   "reportMissingParameterType": false,
   "reportMissingTypeArgument": false,
   "reportMissingTypeStubs": false,

src/polarization/_attrs.py

@@ -0,0 +1,34 @@
+"""Helper functions for constructing `attrs` decorated classes."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, SupportsFloat
+
+import sympy as sp
+from attrs import Attribute
+
+if TYPE_CHECKING:
+    from polarization.decay import LSCoupling
+
+
+def assert_spin_value(instance, attribute: Attribute, value: sp.Rational) -> None:
+    if value.denominator not in {1, 2}:
+        raise ValueError(
+            f"{attribute.name} value should be integer or half-integer, not {value}"
+        )
+
+
+def to_ls(obj: LSCoupling | tuple[int, SupportsFloat] | None) -> LSCoupling:
+    from polarization.decay import LSCoupling
+
+    if obj is None:
+        return None
+    if isinstance(obj, LSCoupling):
+        return obj
+    if isinstance(obj, tuple):
+        L, S = obj
+        return LSCoupling(L, S)
+    raise TypeError(f"Cannot convert {type(obj).__name__} to {LSCoupling.__name__}")
+
+
+def to_rational(obj: SupportsFloat) -> sp.Rational:
+    return sp.Rational(obj)

src/polarization/decay.py

@@ -0,0 +1,50 @@
+"""Data structures that describe a three-body decay."""
+from __future__ import annotations
+
+import sympy as sp
+from attrs import field, frozen
+from attrs.validators import instance_of
+
+from polarization._attrs import assert_spin_value, to_ls, to_rational
+
+
+@frozen
+class ThreeBodyDecay:
+    initial_state: Particle
+    final_state: tuple[Particle, Particle, Particle]
+    resonances: tuple[IsobarNode, ...]
+
+    def __attrs_post_init__(self) -> None:
+        for resonance in self.resonances:
+            if self.final_state != resonance.children:
+                final_state = ", ".join(p.name for p in self.final_state)
+                raise ValueError(
+                    f"Resonance {resonance.parent.name} →"
+                    f" {resonance.child1.name} {resonance.child2.name} does not decay"
+                    f" to {final_state}"
+                )
+
+
+@frozen
+class Particle:
+    name: str
+    spin: sp.Rational = field(converter=to_rational, validator=assert_spin_value)
+    parity: int
+
+
+@frozen
+class IsobarNode:
+    parent: Particle = field(validator=instance_of(Particle))
+    child1: Particle = field(validator=instance_of(Particle))
+    child2: Particle = field(validator=instance_of(Particle))
+    interaction: LSCoupling | None = field(default=None, converter=to_ls)
+
+    @property
+    def children(self) -> tuple[Particle, Particle]:
+        return self.child1, self.child2
+
+
+@frozen
+class LSCoupling:
+    L: int
+    S: sp.Rational = field(converter=to_rational, validator=assert_spin_value)

src/polarization/io.py

@@ -23,12 +23,18 @@
 
 import sympy as sp
 
+from polarization.decay import IsobarNode, Particle
+
 
 @singledispatch
 def as_latex(obj) -> str:
     """Render objects as a LaTeX `str`.
 
     The resulting `str` can for instance be given to `IPython.display.Math`.
+
+    Optional keywords:
+
+    - `render_jp`: Render a `Particle` as :math:`J^P` (spin-parity).
     """
     return str(obj)
 
@@ -73,3 +79,25 @@ def _(obj: Iterable) -> str:
         latex += Rf"  {item} \\" + "\n"
     latex += R"\end{array}"
     return latex
+
+
+@as_latex.register(IsobarNode)
+def _(obj: IsobarNode, render_jp: bool = False) -> str:
+    def render_arrow(node: IsobarNode) -> str:
+        if node.interaction is None:
+            return R"\to"
+        return Rf"\xrightarrow[S={node.interaction.S}]{{L={node.interaction.L}}}"
+
+    parent = as_latex(obj.parent, render_jp)
+    to = render_arrow(obj)
+    child1 = as_latex(obj.child1, render_jp)
+    child2 = as_latex(obj.child2, render_jp)
+    return Rf"{parent} {to} {child1} {child2}"
+
+
+@as_latex.register(Particle)
+def _(obj: Particle, render_jp: bool = False) -> str:
+    if render_jp:
+        parity = "-1" if obj.parity < 0 else "+1"
+        return f"{{{obj.spin}}}^{{{parity}}}"
+    return obj.name

tests/test_decay.py

@@ -0,0 +1,21 @@
+from polarization.decay import IsobarNode, Particle
+
+# https://compwa-org--129.org.readthedocs.build/report/018.html#resonances-and-ls-scheme
+Λc = Particle(R"\Lambda_c^+", spin=0.5, parity=+1)
+p = Particle("p", spin=0.5, parity=+1)
+π = Particle(R"\pi^+", spin=0, parity=-1)
+K = Particle("K^-", spin=0, parity=-1)
+Λ1520 = Particle(R"\Lambda(1520)", spin=1.5, parity=-1)
+
+
+class TestIsobarNode:
+    def test_children(self):
+        decay = IsobarNode(Λ1520, p, K)
+        assert decay.children == (p, K)
+
+    def test_ls(self):
+        L, S = 2, 1
+        node = IsobarNode(Λ1520, p, K, interaction=(L, S))
+        assert node.interaction is not None
+        assert node.interaction.L == L
+        assert node.interaction.S == S

tests/test_io.py

@@ -0,0 +1,28 @@
+from polarization.decay import IsobarNode, Particle
+from polarization.io import as_latex
+
+# https://compwa-org--129.org.readthedocs.build/report/018.html#resonances-and-ls-scheme
+Λc = Particle(R"\Lambda_c^+", spin=0.5, parity=+1)
+p = Particle("p", spin=0.5, parity=+1)
+π = Particle(R"\pi^+", spin=0, parity=-1)
+K = Particle("K^-", spin=0, parity=-1)
+Λ1520 = Particle(R"\Lambda(1520)", spin=1.5, parity=-1)
+
+
+def test_as_latex_particle():
+    latex = as_latex(Λ1520)
+    assert latex == Λ1520.name
+    latex = as_latex(Λ1520, render_jp=True)
+    assert latex == R"{3/2}^{-1}"
+
+
+def test_as_latex_isobar_node():
+    node = IsobarNode(Λ1520, p, K)
+    latex = as_latex(node)
+    assert latex == R"\Lambda(1520) \to p K^-"
+    latex = as_latex(node, render_jp=True)
+    assert latex == R"{3/2}^{-1} \to {1/2}^{+1} {0}^{-1}"
+
+    node = IsobarNode(Λ1520, p, K, interaction=(2, 1))
+    latex = as_latex(node)
+    assert latex == R"\Lambda(1520) \xrightarrow[S=1]{L=2} p K^-"