feat: extract DynamicsSelector class

src/ampform/helicity/__init__.py

@@ -11,8 +11,9 @@
 import operator
 from collections import OrderedDict, abc
 from difflib import get_close_matches
-from functools import reduce
+from functools import reduce, singledispatchmethod
 from typing import (
+    Any,
     DefaultDict,
     Dict,
     ItemsView,
@@ -40,6 +41,7 @@
 from ampform.dynamics.builder import (
     ResonanceDynamicsBuilder,
     TwoBodyKinematicVariableSet,
+    create_non_dynamic,
 )
 from ampform.kinematics import HelicityAdapter, get_invariant_mass_label
 
@@ -251,6 +253,89 @@ def reset(self) -> None:
         self.kinematic_variables = {}
 
 
+class DynamicsSelector(abc.Mapping):
+    """Configure which `.ResonanceDynamicsBuilder` to use for each node."""
+
+    def __init__(
+        self, transitions: Union[ReactionInfo, Iterable[StateTransition]]
+    ) -> None:
+        if isinstance(transitions, ReactionInfo):
+            transitions = transitions.transitions
+        self.__choices: Dict[TwoBodyDecay, ResonanceDynamicsBuilder] = {}
+        for transition in transitions:
+            for node_id in transition.topology.nodes:
+                decay = TwoBodyDecay.from_transition(transition, node_id)
+                self.__choices[decay] = create_non_dynamic
+
+    @singledispatchmethod
+    def assign(
+        self, selection: Any, builder: ResonanceDynamicsBuilder
+    ) -> None:
+        """Assign a `.ResonanceDynamicsBuilder` to a selection of nodes.
+
+        Currently, the following types of selections are implements:
+
+        - `str`: Select transition nodes by the name of the
+          `~TwoBodyDecay.parent` `~qrules.particle.Particle`.
+        - `TwoBodyDecay` or `tuple` of a `StateTransition` with a node ID: set
+          dynamics for one specific transition node.
+        """
+        raise NotImplementedError(
+            "Cannot set dynamics builder for selection type"
+            f" {type(selection).__name__}"
+        )
+
+    @assign.register(TwoBodyDecay)
+    def _(
+        self, selection: TwoBodyDecay, builder: ResonanceDynamicsBuilder
+    ) -> None:
+        self.__choices[selection] = builder
+
+    @assign.register(tuple)
+    def _(
+        self,
+        selection: Tuple[StateTransition, int],
+        builder: ResonanceDynamicsBuilder,
+    ) -> None:
+        decay = TwoBodyDecay.create(selection)
+        return self.assign(decay, builder)
+
+    @assign.register(str)
+    def _(self, selection: str, builder: ResonanceDynamicsBuilder) -> None:
+        particle_name = selection
+        found_particle = False
+        for decay in self.__choices:
+            decaying_particle = decay.parent.particle
+            if decaying_particle.name == particle_name:
+                self.__choices[decay] = builder
+                found_particle = True
+        if not found_particle:
+            logging.warning(
+                f'Model contains no resonance with name "{particle_name}"'
+            )
+
+    def __getitem__(
+        self, __k: Union[TwoBodyDecay, Tuple[StateTransition, int]]
+    ) -> ResonanceDynamicsBuilder:
+        __k = TwoBodyDecay.create(__k)
+        return self.__choices[__k]
+
+    def __len__(self) -> int:
+        return len(self.__choices)
+
+    def __iter__(self) -> Iterator[TwoBodyDecay]:
+        return iter(self.__choices)
+
+    def items(self) -> ItemsView[TwoBodyDecay, ResonanceDynamicsBuilder]:
+        return self.__choices.items()
+
+    def keys(self) -> KeysView[TwoBodyDecay]:
+        return self.__choices.keys()
+
+    def values(self) -> ValuesView[ResonanceDynamicsBuilder]:
+        return self.__choices.values()
+
+
 class HelicityAmplitudeBuilder:  # pylint: disable=too-many-instance-attributes
     r"""Amplitude model generator for the helicity formalism.
 
@@ -279,18 +364,15 @@ def __init__(
         stable_final_state_ids: Optional[Iterable[int]] = None,
         scalar_initial_state_mass: bool = False,
     ) -> None:
-        self._name_generator = HelicityAmplitudeNameGenerator()
-        self.__reaction = reaction
-        self.__ingredients = _HelicityModelIngredients()
-        self.__dynamics_choices: Dict[
-            TwoBodyDecay, ResonanceDynamicsBuilder
-        ] = {}
-
         if len(reaction.transitions) < 1:
             raise ValueError(
                 f"At least one {StateTransition.__name__} required to"
                 " genenerate an amplitude model!"
             )
+        self._name_generator = HelicityAmplitudeNameGenerator()
+        self.__reaction = reaction
+        self.__ingredients = _HelicityModelIngredients()
+        self.__dynamics_choices = DynamicsSelector(reaction)
         self.__adapter = HelicityAdapter(reaction)
         self.stable_final_state_ids = stable_final_state_ids  # type: ignore[assignment]
         self.scalar_initial_state_mass = scalar_initial_state_mass  # type: ignore[assignment]
@@ -302,6 +384,10 @@ def adapter(self) -> HelicityAdapter:
         """Converter for computing kinematic variables from four-momenta."""
         return self.__adapter
 
+    @property
+    def dynamics_choices(self) -> DynamicsSelector:
+        return self.__dynamics_choices
+
     @property
     def stable_final_state_ids(self) -> Optional[Set[int]]:
         # noqa: D403
@@ -344,18 +430,7 @@ def scalar_initial_state_mass(self, value: bool) -> None:
     def set_dynamics(
         self, particle_name: str, dynamics_builder: ResonanceDynamicsBuilder
     ) -> None:
-        found_particle = False
-        for transition in self.__reaction.transitions:
-            for node_id in transition.topology.nodes:
-                decay = TwoBodyDecay.from_transition(transition, node_id)
-                decaying_particle = decay.parent.particle
-                if decaying_particle.name == particle_name:
-                    self.__dynamics_choices[decay] = dynamics_builder
-                    found_particle = True
-        if not found_particle:
-            logging.warning(
-                f'Model contains no resonance with name "{particle_name}"'
-            )
+        self.__dynamics_choices.assign(particle_name, dynamics_builder)
 
     def formulate(self) -> HelicityModel:
         self.__ingredients.reset()

src/ampform/helicity/decay.py

@@ -1,6 +1,7 @@
 """Extract two-body decay info from a `~qrules.transition.StateTransition`."""
 
-from typing import Iterable, List, Tuple
+from functools import singledispatch
+from typing import Any, Iterable, List, Tuple
 
 from attrs import frozen
 from qrules.quantum_numbers import InteractionProperties
@@ -47,6 +48,16 @@ class TwoBodyDecay:
     children: Tuple[StateWithID, StateWithID]
     interaction: InteractionProperties
 
+    @staticmethod
+    def create(obj: Any) -> "TwoBodyDecay":
+        """Create a `TwoBodyDecay` instance from an arbitrary object.
+
+        More implementations of :meth:`create` can be implemented with
+        :func:`@ampform.helicity.decay._create_two_body_decay.register(TYPE)
+        <functools.singledispatch>`.
+        """
+        return _create_two_body_decay(obj)
+
     @classmethod
     def from_transition(
         cls, transition: StateTransition, node_id: int
@@ -80,6 +91,28 @@ def from_transition(
         )
 
 
+@singledispatch
+def _create_two_body_decay(obj: Any) -> TwoBodyDecay:
+    raise NotImplementedError(
+        f"Cannot create a {TwoBodyDecay.__name__} from a {type(obj).__name__}"
+    )
+
+
+@_create_two_body_decay.register(TwoBodyDecay)
+def _(obj: TwoBodyDecay) -> TwoBodyDecay:
+    return obj
+
+
+@_create_two_body_decay.register(tuple)
+def _(obj: tuple) -> TwoBodyDecay:
+    if len(obj) == 2:
+        if isinstance(obj[0], StateTransition) and isinstance(obj[1], int):
+            return TwoBodyDecay.from_transition(*obj)
+    raise NotImplementedError(
+        f"Cannot create a {TwoBodyDecay.__name__} from {obj}"
+    )
+
+
 def get_helicity_info(
     transition: StateTransition, node_id: int
 ) -> Tuple[State, Tuple[State, State]]: