# pylint: disable=import-outside-toplevel
"""Generate an amplitude model with the helicity formalism."""
import collections
import logging
import operator
import re
from collections import OrderedDict
from difflib import get_close_matches
from functools import reduce
from typing import (
DefaultDict,
Dict,
Iterable,
List,
Mapping,
Optional,
Tuple,
Union,
)
import attr
import sympy as sp
from attr.validators import instance_of
from qrules.combinatorics import (
perform_external_edge_identical_particle_combinatorics,
)
from qrules.particle import ParticleCollection
from qrules.transition import ReactionInfo, StateTransition
from ampform.dynamics.builder import (
ResonanceDynamicsBuilder,
TwoBodyKinematicVariableSet,
)
from ampform.kinematics import (
HelicityAdapter,
get_helicity_angle_label,
get_invariant_mass_label,
)
from .decay import TwoBodyDecay
from .naming import (
CanonicalAmplitudeNameGenerator,
HelicityAmplitudeNameGenerator,
generate_transition_label,
)
ParameterValue = Union[float, complex, int]
"""Allowed value types for parameters."""
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)]
)
def _order_symbol_mapping(
mapping: Mapping[sp.Symbol, sp.Expr]
) -> "OrderedDict[sp.Symbol, sp.Expr]":
return collections.OrderedDict(
[
(symbol, mapping[symbol])
for symbol in sorted(
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
[docs]@attr.s(frozen=True)
class HelicityModel:
_expression: sp.Expr = attr.ib(
validator=attr.validators.instance_of(sp.Expr)
)
_parameter_defaults: "OrderedDict[sp.Symbol, ParameterValue]" = attr.ib(
converter=_order_symbol_mapping
)
_components: "OrderedDict[str, sp.Expr]" = attr.ib(
converter=_order_component_mapping
)
_adapter: HelicityAdapter = attr.ib(
validator=attr.validators.instance_of(HelicityAdapter)
)
particles: ParticleCollection = attr.ib(
validator=instance_of(ParticleCollection)
)
@property
def expression(self) -> sp.Expr:
return self._expression
@property
def components(self) -> "OrderedDict[str, sp.Expr]":
"""A mapping for identifying main components in the :attr:`expression`.
Keys are the component names (`str`), formatted as LaTeX, and values
are sub-expressions in the main :attr:`expression`. The mapping is an
`~collections.OrderedDict` that orders the component names
alphabetically with `natural sort order
<https://en.wikipedia.org/wiki/Natural_sort_order>`_.
"""
return self._components
@property
def parameter_defaults(self) -> "OrderedDict[sp.Symbol, ParameterValue]":
"""A mapping of suggested parameter values.
Keys are `~sympy.core.symbol.Symbol` instances from the main
:attr:`expression` that should be interpreted as parameters (as opposed
to variables). The symbols are ordered alphabetically by name with
`natural sort order
<https://en.wikipedia.org/wiki/Natural_sort_order>`_. Values have been
extracted from the input `~qrules.transition.ReactionInfo`.
"""
return self._parameter_defaults
@property
def adapter(self) -> HelicityAdapter:
"""Adapter for converting four-momenta to kinematic variables."""
return self._adapter
[docs] def sum_components( # noqa: R701
self, components: Iterable[str]
) -> sp.Expr:
"""Coherently or incoherently add components of a helicity model."""
components = list(components) # copy
for component in components:
if component not in self.components:
first_letter = component[0]
candidates = get_close_matches(
component,
filter(
lambda c: c.startswith(
first_letter # pylint: disable=cell-var-from-loop
),
self.components,
),
)
raise KeyError(
f'Component "{component}" not in model components. '
f"Did you mean any of these?",
candidates,
)
if any(map(lambda c: c.startswith("I"), components)) and any(
map(lambda c: c.startswith("A"), components)
):
intensity_sum = self.sum_components(
components=filter(lambda c: c.startswith("I"), components),
)
amplitude_sum = self.sum_components(
components=filter(lambda c: c.startswith("A"), components),
)
return intensity_sum + amplitude_sum
if all(map(lambda c: c.startswith("I"), components)):
return sum(self.components[c] for c in components)
if all(map(lambda c: c.startswith("A"), components)):
return abs(sum(self.components[c] for c in components)) ** 2
raise ValueError(
'Not all component names started with either "A" or "I"'
)
[docs]class HelicityAmplitudeBuilder:
"""Amplitude model generator for the helicity formalism."""
def __init__(self, reaction: ReactionInfo) -> None:
self._name_generator = HelicityAmplitudeNameGenerator()
self.__reaction = reaction
self.__parameter_defaults: Dict[sp.Symbol, ParameterValue] = {}
self.__components: Dict[str, sp.Expr] = {}
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.__adapter = HelicityAdapter(reaction)
for grouping in reaction.transition_groups:
self.__adapter.register_topology(grouping.topology)
self.__particles = extract_particle_collection(reaction.transitions)
[docs] 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}"'
)
def __formulate_top_expression(self) -> sp.Expr:
transition_groups = group_transitions(self.__reaction.transitions)
self.__register_parameter_couplings(transition_groups)
coherent_intensities = [
self.__formulate_coherent_intensity(group)
for group in transition_groups
]
return sum(coherent_intensities)
def __register_parameter_couplings(
self, transition_groups: List[List[StateTransition]]
) -> None:
for graph_group in transition_groups:
for transition in graph_group:
self._name_generator.register_amplitude_coefficient_name(
transition
)
def __formulate_coherent_intensity(
self, transition_group: List[StateTransition]
) -> sp.Expr:
graph_group_label = generate_transition_label(transition_group[0])
sequential_expressions: List[sp.Expr] = []
for transition in transition_group:
sequential_graphs = (
perform_external_edge_identical_particle_combinatorics(
transition.to_graph()
)
)
for graph in sequential_graphs:
transition = StateTransition.from_graph(graph)
expression = self.__formulate_sequential_decay(transition)
sequential_expressions.append(expression)
amplitude_sum = sum(sequential_expressions)
coherent_intensity = abs(amplitude_sum) ** 2
self.__components[fR"I_{{{graph_group_label}}}"] = coherent_intensity
return coherent_intensity
def __formulate_sequential_decay(
self, transition: StateTransition
) -> sp.Expr:
partial_decays: List[sp.Symbol] = [
self._formulate_partial_decay(transition, node_id)
for node_id in transition.topology.nodes
]
sequential_amplitudes = reduce(operator.mul, partial_decays)
coefficient = self.__generate_amplitude_coefficient(transition)
prefactor = self.__generate_amplitude_prefactor(transition)
expression = coefficient * sequential_amplitudes
if prefactor is not None:
expression = prefactor * expression
self.__components[
f"A_{{{self._name_generator.generate_amplitude_name(transition)}}}"
] = expression
return expression
def _formulate_partial_decay(
self, transition: StateTransition, node_id: int
) -> sp.Expr:
wigner_d = formulate_wigner_d(transition, node_id)
dynamics = self.__formulate_dynamics(transition, node_id)
return wigner_d * dynamics
def __formulate_dynamics(
self, transition: StateTransition, node_id: int
) -> sp.Expr:
decay = TwoBodyDecay.from_transition(transition, node_id)
if decay not in self.__dynamics_choices:
return 1
builder = self.__dynamics_choices[decay]
variable_set = _generate_kinematic_variable_set(transition, node_id)
expression, parameters = builder(decay.parent.particle, variable_set)
for par, value in parameters.items():
if par in self.__parameter_defaults:
previous_value = self.__parameter_defaults[par]
if value != previous_value:
logging.warning(
f'New default value {value} for parameter "{par.name}"'
f" is inconsistent with existing value {previous_value}"
)
self.__parameter_defaults[par] = value
return expression
def __generate_amplitude_coefficient(
self, transition: StateTransition
) -> sp.Symbol:
"""Generate coefficient parameter for a sequential amplitude.
Generally, each partial amplitude of a sequential amplitude transition
should check itself if it or a parity partner is already defined. If so
a coupled coefficient is introduced.
"""
suffix = self._name_generator.generate_sequential_amplitude_suffix(
transition
)
coefficient_symbol = sp.Symbol(f"C_{{{suffix}}}")
self.__parameter_defaults[coefficient_symbol] = complex(1, 0)
return coefficient_symbol
def __generate_amplitude_prefactor(
self, transition: StateTransition
) -> Optional[float]:
prefactor = get_prefactor(transition)
if prefactor != 1.0:
for node_id in transition.topology.nodes:
raw_suffix = self._name_generator.generate_coefficient_name(
transition, node_id
)
if (
raw_suffix
in self._name_generator.parity_partner_coefficient_mapping
):
coefficient_suffix = self._name_generator.parity_partner_coefficient_mapping[
raw_suffix
]
if coefficient_suffix != raw_suffix:
return prefactor
return None
[docs]class CanonicalAmplitudeBuilder(HelicityAmplitudeBuilder):
r"""Amplitude model generator for the canonical helicity formalism.
This class defines a full amplitude in the canonical formalism, using the
helicity formalism as a foundation. The key here is that we take the full
helicity intensity as a template, and just exchange the helicity amplitudes
:math:`F` as a sum of canonical amplitudes :math:`A`:
.. math::
F^J_{\lambda_1,\lambda_2} = \sum_{LS} \mathrm{norm}(A^J_{LS})C^2.
Here, :math:`C` stands for `Clebsch-Gordan factor
<https://en.wikipedia.org/wiki/Clebsch%E2%80%93Gordan_coefficients>`_.
"""
def __init__(self, reaction_result: ReactionInfo) -> None:
super().__init__(reaction_result)
self._name_generator = CanonicalAmplitudeNameGenerator()
def _formulate_partial_decay(
self, transition: StateTransition, node_id: int
) -> sp.Expr:
amplitude = super()._formulate_partial_decay(transition, node_id)
cg_coefficients = formulate_clebsch_gordan_coefficients(
transition, node_id
)
return cg_coefficients * amplitude
[docs]def extract_particle_collection(
transitions: Iterable[StateTransition],
) -> ParticleCollection:
"""Collect all particles from a collection of state transitions."""
particles = ParticleCollection()
for transition in transitions:
for state in transition.states.values():
if state.particle not in particles:
particles.add(state.particle)
return particles
[docs]def get_prefactor(transition: StateTransition) -> float:
"""Calculate the product of all prefactors defined in this transition.
.. seealso:: `qrules.quantum_numbers.InteractionProperties.parity_prefactor`
"""
prefactor = 1.0
for node_id in transition.topology.nodes:
interaction = transition.interactions[node_id]
if interaction and interaction.parity_prefactor is not None:
prefactor *= interaction.parity_prefactor
return prefactor
[docs]def group_transitions(
transitions: Iterable[StateTransition],
) -> List[List[StateTransition]]:
"""Match final and initial states in groups.
Each `~qrules.transition.StateTransition` corresponds to a specific state
transition amplitude. This function groups together transitions, which have
the same initial and final state (including spin). This is needed to
determine the coherency of the individual amplitude parts.
"""
transition_groups: DefaultDict[
Tuple[
Tuple[Tuple[str, float], ...],
Tuple[Tuple[str, float], ...],
],
List[StateTransition],
] = collections.defaultdict(list)
for transition in transitions:
initial_state = sorted(
(
transition.states[i].particle.name,
transition.states[i].spin_projection,
)
for i in transition.topology.incoming_edge_ids
)
final_state = sorted(
(
transition.states[i].particle.name,
transition.states[i].spin_projection,
)
for i in transition.topology.outgoing_edge_ids
)
group_key = (tuple(initial_state), tuple(final_state))
transition_groups[group_key].append(transition)
return list(transition_groups.values())
def _generate_kinematic_variable_set(
transition: StateTransition, node_id: int
) -> TwoBodyKinematicVariableSet:
decay = TwoBodyDecay.from_transition(transition, node_id)
inv_mass, phi, theta = _generate_kinematic_variables(transition, node_id)
child1_mass = sp.Symbol(
get_invariant_mass_label(transition.topology, decay.children[0].id),
real=True,
)
child2_mass = sp.Symbol(
get_invariant_mass_label(transition.topology, decay.children[1].id),
real=True,
)
return TwoBodyKinematicVariableSet(
incoming_state_mass=inv_mass,
outgoing_state_mass1=child1_mass,
outgoing_state_mass2=child2_mass,
helicity_theta=theta,
helicity_phi=phi,
angular_momentum=decay.extract_angular_momentum(),
)
def _generate_kinematic_variables(
transition: StateTransition, node_id: int
) -> Tuple[sp.Symbol, sp.Symbol, sp.Symbol]:
"""Generate symbol for invariant mass, phi angle, and theta angle."""
decay = TwoBodyDecay.from_transition(transition, node_id)
phi_label, theta_label = get_helicity_angle_label(
transition.topology, decay.children[0].id
)
inv_mass_label = get_invariant_mass_label(
transition.topology, decay.parent.id
)
return (
sp.Symbol(inv_mass_label, real=True),
sp.Symbol(phi_label, real=True),
sp.Symbol(theta_label, real=True),
)
