diff --git a/src/ampform/helicity/__init__.py b/src/ampform/helicity/__init__.py index 158852bd2..d0a884d16 100644 --- a/src/ampform/helicity/__init__.py +++ b/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 `~qrules.transition.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() diff --git a/src/ampform/helicity/decay.py b/src/ampform/helicity/decay.py index 106fac668..308b62417 100644 --- a/src/ampform/helicity/decay.py +++ b/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) + `. + """ + 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]]: