refactor: remove decay/naming functions from kinematics (#227)

* refactor: move Topology functions to .decay module
* refactor: move label functions to helicity.naming

src/ampform/helicity/__init__.py

@@ -4,7 +4,6 @@
 import collections
 import logging
 import operator
-import re
 from collections import OrderedDict
 from difflib import get_close_matches
 from functools import reduce
@@ -32,17 +31,15 @@
     ResonanceDynamicsBuilder,
     TwoBodyKinematicVariableSet,
 )
-from ampform.kinematics import (
-    HelicityAdapter,
-    get_helicity_angle_label,
-    get_invariant_mass_label,
-)
+from ampform.kinematics import HelicityAdapter, get_invariant_mass_label
 
 from .decay import TwoBodyDecay
 from .naming import (
     CanonicalAmplitudeNameGenerator,
     HelicityAmplitudeNameGenerator,
     generate_transition_label,
+    get_helicity_angle_label,
+    natural_sorting,
 )
 
 ParameterValue = Union[float, complex, int]
@@ -53,7 +50,7 @@ def _order_component_mapping(
     mapping: Mapping[str, ParameterValue]
 ) -> "OrderedDict[str, ParameterValue]":
     return collections.OrderedDict(
-        [(key, mapping[key]) for key in sorted(mapping, key=_natural_sorting)]
+        [(key, mapping[key]) for key in sorted(mapping, key=natural_sorting)]
     )
 
 
@@ -64,27 +61,12 @@ def _order_symbol_mapping(
         [
             (symbol, mapping[symbol])
             for symbol in sorted(
-                mapping, key=lambda s: _natural_sorting(s.name)
+                mapping, key=lambda s: natural_sorting(s.name)
             )
         ]
     )
 
 
-def _natural_sorting(text: str) -> List[Union[float, str]]:
-    # https://stackoverflow.com/a/5967539/13219025
-    return [
-        __attempt_number_cast(c)
-        for c in re.split(r"[+-]?([0-9]+(?:[.][0-9]*)?|[.][0-9]+)", text)
-    ]
-
-
-def __attempt_number_cast(text: str) -> Union[float, str]:
-    try:
-        return float(text)
-    except ValueError:
-        return text
-
-
 @attr.frozen
 class HelicityModel:  # noqa: R701
     expression: sp.Expr = attr.ib(

src/ampform/helicity/decay.py

@@ -4,10 +4,9 @@
 
 import attr
 from qrules.quantum_numbers import InteractionProperties
+from qrules.topology import Topology
 from qrules.transition import State, StateTransition
 
-from ampform.kinematics import _assert_two_body_decay
-
 
 @attr.frozen
 class StateWithID(State):
@@ -97,7 +96,7 @@ def get_helicity_info(
     transition: StateTransition, node_id: int
 ) -> Tuple[State, Tuple[State, State]]:
     """Extract in- and outgoing states for a two-body decay node."""
-    _assert_two_body_decay(transition.topology, node_id)
+    assert_two_body_decay(transition.topology, node_id)
     in_edge_ids = transition.topology.get_edge_ids_ingoing_to_node(node_id)
     out_edge_ids = transition.topology.get_edge_ids_outgoing_from_node(node_id)
     in_helicity_list = get_sorted_states(transition, in_edge_ids)
@@ -120,3 +119,48 @@ def get_sorted_states(
     """
     states = [transition.states[i] for i in state_ids]
     return sorted(states, key=lambda s: s.particle.name)
+
+
+def assert_isobar_topology(topology: Topology) -> None:
+    for node_id in topology.nodes:
+        assert_two_body_decay(topology, node_id)
+
+
+def assert_two_body_decay(topology: Topology, node_id: int) -> None:
+    parent_state_ids = topology.get_edge_ids_ingoing_to_node(node_id)
+    if len(parent_state_ids) != 1:
+        raise ValueError(
+            f"Node {node_id} has {len(parent_state_ids)} parent states,"
+            " so this is not an isobar decay"
+        )
+    child_state_ids = topology.get_edge_ids_outgoing_from_node(node_id)
+    if len(child_state_ids) != 2:
+        raise ValueError(
+            f"Node {node_id} decays to {len(child_state_ids)} states,"
+            " so this is not an isobar decay"
+        )
+
+
+def determine_attached_final_state(
+    topology: Topology, state_id: int
+) -> List[int]:
+    """Determine all final state particles of a transition.
+
+    These are attached downward (forward in time) for a given edge (resembling
+    the root).
+
+    Example
+    -------
+    For **edge 5** in Figure :ref:`one-to-five-topology-0`, we get:
+
+    >>> from qrules.topology import create_isobar_topologies
+    >>> topologies = create_isobar_topologies(5)
+    >>> determine_attached_final_state(topologies[0], state_id=5)
+    [0, 3, 4]
+    """
+    edge = topology.edges[state_id]
+    if edge.ending_node_id is None:
+        return [state_id]
+    return sorted(
+        topology.get_originating_final_state_edge_ids(edge.ending_node_id)
+    )

src/ampform/helicity/naming.py

@@ -1,11 +1,19 @@
 """Generate descriptions used in the `~ampform.helicity` formalism."""
 
-from typing import Dict, List, Optional, Tuple
+import re
+from functools import lru_cache
+from typing import Dict, List, Optional, Tuple, Union
 
 import sympy as sp
+from qrules.topology import Topology
 from qrules.transition import State, StateTransition
 
-from .decay import get_helicity_info, get_sorted_states
+from .decay import (
+    assert_isobar_topology,
+    determine_attached_final_state,
+    get_helicity_info,
+    get_sorted_states,
+)
 
 
 class HelicityAmplitudeNameGenerator:
@@ -201,6 +209,130 @@ def generate_transition_label(transition: StateTransition) -> str:
     )
 
 
+def get_helicity_angle_label(
+    topology: Topology, state_id: int
+) -> Tuple[str, str]:
+    r"""Generate a nested helicity angle label for :math:`\phi,\theta`.
+
+    See :func:`get_boost_chain_suffix` for the meaning of the suffix.
+    """
+    suffix = get_boost_chain_suffix(topology, state_id)
+    return f"phi{suffix}", f"theta{suffix}"
+
+
+@lru_cache(maxsize=None)
+def get_boost_chain_suffix(topology: Topology, state_id: int) -> str:
+    """Generate a subscript-superscript to identify a chain of Lorentz boosts.
+
+    The generated subscripts describe the decay sequence from the right to the
+    left, separated by commas. Resonance edge IDs are expressed as a sum of the
+    final state IDs that lie below them (see
+    :func:`.determine_attached_final_state`). The generated label does not
+    state the top-most edge (the initial state).
+
+    Example
+    -------
+    The following two allowed isobar topologies for a **1-to-5-body** decay
+    illustrates how the naming scheme results in a unique label for each of the
+    **eight edges** in the decay topology. Note that label only uses final
+    state IDs, but still reflects the internal decay topology.
+
+    >>> from qrules.topology import create_isobar_topologies
+    >>> topologies = create_isobar_topologies(5)
+    >>> topology = topologies[0]
+    >>> for i in topology.intermediate_edge_ids | topology.outgoing_edge_ids:
+    ...     suffix = get_boost_chain_suffix(topology, i)
+    ...     print(f"{i}: 'phi{suffix}'")
+    0: 'phi_0^034'
+    1: 'phi_1^12'
+    2: 'phi_2^12'
+    3: 'phi_3^34,034'
+    4: 'phi_4^34,034'
+    5: 'phi_034'
+    6: 'phi_12'
+    7: 'phi_34^034'
+    >>> topology = topologies[1]
+    >>> for i in topology.intermediate_edge_ids | topology.outgoing_edge_ids:
+    ...     suffix = get_boost_chain_suffix(topology, i)
+    ...     print(f"{i}: 'phi{suffix}'")
+    0: 'phi_0^01'
+    1: 'phi_1^01'
+    2: 'phi_2^234'
+    3: 'phi_3^34,234'
+    4: 'phi_4^34,234'
+    5: 'phi_01'
+    6: 'phi_234'
+    7: 'phi_34^234'
+
+    Some labels explained:
+
+    - :code:`phi_12`: **edge 6** on the *left* topology, because for this
+      topology, we have :math:`p_6=p_1+p_2`.
+    - :code:`phi_234`: **edge 6** *right*, because for this topology,
+      :math:`p_6=p_2+p_3+p_4`.
+    - :code:`phi_1^12`: **edge 1** *left*, because 1 decays from
+      :math:`p_6=p_1+p_2`.
+    - :code:`phi_1^01`: **edge 1** *right*, because it decays from
+      :math:`p_5=p_0+p_1`.
+    - :code:`phi_4^34,234`: **edge 4** *right*, because it decays from edge 7
+      (:math:`p_7=p_3+p_4`), which comes from edge 6 (:math:`p_7=p_2+p_3+p_4`).
+
+    As noted, the top-most parent (initial state) is not listed in the label.
+    """
+    assert_isobar_topology(topology)
+
+    def recursive_label(topology: Topology, state_id: int) -> str:
+        edge = topology.edges[state_id]
+        if edge.ending_node_id is None:
+            label = f"{state_id}"
+        else:
+            attached_final_state_ids = determine_attached_final_state(
+                topology, state_id
+            )
+            label = "".join(map(str, attached_final_state_ids))
+        if edge.originating_node_id is not None:
+            incoming_state_ids = topology.get_edge_ids_ingoing_to_node(
+                edge.originating_node_id
+            )
+            state_id = next(iter(incoming_state_ids))
+            if state_id not in topology.incoming_edge_ids:
+                label += f",{recursive_label(topology, state_id)}"
+        return label
+
+    label = recursive_label(topology, state_id)
+
+    index_groups = label.split(",")
+    subscript = index_groups[0]
+    suffix = f"_{subscript}"
+    if len(index_groups) > 1:
+        superscript = ",".join(index_groups[1:])
+        suffix += f"^{superscript}"
+    return suffix
+
+
+def natural_sorting(text: str) -> List[Union[float, str]]:
+    """Function that can be used for natural sort order in :func:`sorted`.
+
+    See `natural sort order
+    <https://en.wikipedia.org/wiki/Natural_sort_order>`_.
+
+    >>> sorted(["z11", "z2"], key=natural_sorting)
+    ['z2', 'z11']
+    """
+    # https://stackoverflow.com/a/5967539/13219025
+    return [
+        __attempt_number_cast(c)
+        for c in re.split(r"[+-]?([0-9]+(?:[.][0-9]*)?|[.][0-9]+)", text)
+    ]
+
+
+def __attempt_number_cast(text: str) -> Union[float, str]:
+    try:
+        return float(text)
+    except ValueError:
+        return text
+
+
 def _state_to_str(
     state: State,
     use_helicity: bool = True,