From fa7c91370da2a5e660afca34a75038b0ef1aaf37 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Sat, 30 Apr 2022 21:28:25 +0200
Subject: [PATCH] implement first version of decay data structures

---
 pyproject.toml             |  1 +
 pyrightconfig.json         |  1 +
 src/polarization/_attrs.py | 34 ++++++++++++++++++++++++++
 src/polarization/decay.py  | 50 ++++++++++++++++++++++++++++++++++++++
 src/polarization/io.py     | 28 +++++++++++++++++++++
 tests/test_decay.py        | 21 ++++++++++++++++
 tests/test_io.py           | 28 +++++++++++++++++++++
 7 files changed, 163 insertions(+)
 create mode 100644 src/polarization/_attrs.py
 create mode 100644 src/polarization/decay.py
 create mode 100644 tests/test_decay.py
 create mode 100644 tests/test_io.py

diff --git a/pyproject.toml b/pyproject.toml
index 96e8d84a..ca8d7f27 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -58,4 +58,5 @@ norecursedirs = [
 ]
 testpaths = [
     "src",
+    "tests",
 ]
diff --git a/pyrightconfig.json b/pyrightconfig.json
index 85d1aad0..04198d69 100644
--- a/pyrightconfig.json
+++ b/pyrightconfig.json
@@ -1,6 +1,7 @@
 {
   "include": ["src", "tests"],
   "reportGeneralTypeIssues": false,
+  "reportImportCycles": false,
   "reportMissingParameterType": false,
   "reportMissingTypeArgument": false,
   "reportMissingTypeStubs": false,
diff --git a/src/polarization/_attrs.py b/src/polarization/_attrs.py
new file mode 100644
index 00000000..dc6c4243
--- /dev/null
+++ b/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)
diff --git a/src/polarization/decay.py b/src/polarization/decay.py
new file mode 100644
index 00000000..2190b8cb
--- /dev/null
+++ b/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)
diff --git a/src/polarization/io.py b/src/polarization/io.py
index 82eff70e..ea6ee8a8 100644
--- a/src/polarization/io.py
+++ b/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
diff --git a/tests/test_decay.py b/tests/test_decay.py
new file mode 100644
index 00000000..b41097be
--- /dev/null
+++ b/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
diff --git a/tests/test_io.py b/tests/test_io.py
new file mode 100644
index 00000000..7da9f202
--- /dev/null
+++ b/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^-"