diff --git a/docs/usage/helicity/formalism.ipynb b/docs/usage/helicity/formalism.ipynb index 1ff6b58be..5677a4206 100644 --- a/docs/usage/helicity/formalism.ipynb +++ b/docs/usage/helicity/formalism.ipynb @@ -142,6 +142,12 @@ "cell_type": "markdown", "metadata": {}, "source": [ + ":::{margin}\n", + "\n", + "{ref}`usage/helicity/formalism:Coefficient names` shows how to generate different coefficient names.\n", + "\n", + ":::\n", + "\n", "From {attr}`.components` and {attr}`.parameter_defaults`, we can see that the canonical formalism has a larger number of amplitudes." ] }, @@ -365,6 +371,90 @@ " render_selection(cano_model),\n", ")" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Coefficient names" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the previous section, we saw that the {class}`.HelicityAmplitudeBuilder` by default generates coefficient names that only contain **helicities of the decay products**, while coefficients generated by the {class}`.CanonicalAmplitudeBuilder` contain only **$LS$-combinations**. It's possible to tweak this behavior with the {attr}`~.HelicityAmplitudeBuilder.naming` attribute. Here are two extreme examples, where we generate coefficient names that contain $LS$-combinations, the helicities of each parent state, and the helicity of each decay product, as well as a {class}`.HelicityModel` of which the coefficient names only contain information about the resonances:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "reaction = qrules.generate_transitions(\n", + " initial_state=(\"D(1)(2420)0\", [+1]),\n", + " final_state=[\"K+\", \"K-\", \"K~0\"],\n", + " allowed_intermediate_particles=[\"a(1)(1260)+\"],\n", + " formalism=\"canonical-helicity\",\n", + ")\n", + "builder = ampform.get_builder(reaction)\n", + "builder.naming.insert_parent_helicities = True\n", + "builder.naming.insert_child_helicities = True\n", + "builder.naming.insert_ls_combinations = True\n", + "model = builder.formulate()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "amplitudes = [c for c in model.components if c.startswith(\"A\")]\n", + "assert len(model.parameter_defaults) == len(amplitudes)\n", + "sp.Matrix(model.parameter_defaults)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "builder.naming.insert_parent_helicities = False\n", + "builder.naming.insert_child_helicities = False\n", + "builder.naming.insert_ls_combinations = False\n", + "model = builder.formulate()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "assert len(model.parameter_defaults) == 1\n", + "display(*model.parameter_defaults)" + ] } ], "metadata": { @@ -372,6 +462,10 @@ "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.8.12" } }, "nbformat": 4, diff --git a/src/ampform/helicity/__init__.py b/src/ampform/helicity/__init__.py index 51a2af35c..bd091d14d 100644 --- a/src/ampform/helicity/__init__.py +++ b/src/ampform/helicity/__init__.py @@ -63,6 +63,7 @@ from .naming import ( CanonicalAmplitudeNameGenerator, HelicityAmplitudeNameGenerator, + NameGenerator, generate_transition_label, get_helicity_angle_label, get_helicity_suffix, @@ -509,7 +510,11 @@ def __init__( " genenerate an amplitude model!" ) self.__reaction = reaction - self._name_generator = HelicityAmplitudeNameGenerator(reaction) + self.naming: NameGenerator = HelicityAmplitudeNameGenerator(reaction) + """Name generator for amplitude names and coefficient names. + + .. seealso:: :ref:`usage/helicity/formalism:Coefficient names`. + """ self.__ingredients = _HelicityModelIngredients() self.__dynamics_choices = DynamicsSelector(reaction) self.__adapter = HelicityAdapter(reaction) @@ -708,7 +713,7 @@ def __formulate_sequential_decay( expression = coefficient * sequential_amplitudes if prefactor is not None: expression = prefactor * expression - subscript = self._name_generator.generate_amplitude_name(transition) + subscript = self.naming.generate_amplitude_name(transition) self.__ingredients.components[f"A_{{{subscript}}}"] = expression return expression @@ -751,9 +756,7 @@ def __generate_amplitude_coefficient( 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 - ) + suffix = self.naming.generate_sequential_amplitude_suffix(transition) symbol = sp.Symbol(f"C_{{{suffix}}}") value = complex(1, 0) self.__ingredients.parameter_defaults[symbol] = value @@ -765,16 +768,18 @@ def __generate_amplitude_prefactor( prefactor = get_prefactor(transition) if prefactor != 1.0: for node_id in transition.topology.nodes: - raw_suffix = self._name_generator.generate_coefficient_name( + raw_suffix = self.naming.generate_coefficient_suffix( transition, node_id ) if ( raw_suffix - in self._name_generator.parity_partner_coefficient_mapping + in self.naming.parity_partner_coefficient_mapping ): - coefficient_suffix = self._name_generator.parity_partner_coefficient_mapping[ - raw_suffix - ] + coefficient_suffix = ( + self.naming.parity_partner_coefficient_mapping[ + raw_suffix + ] + ) if coefficient_suffix != raw_suffix: return prefactor return None @@ -861,7 +866,7 @@ class CanonicalAmplitudeBuilder(HelicityAmplitudeBuilder): def __init__(self, reaction: ReactionInfo) -> None: super().__init__(reaction) - self._name_generator = CanonicalAmplitudeNameGenerator(reaction) + self.naming = CanonicalAmplitudeNameGenerator(reaction) def _formulate_partial_decay( self, transition: StateTransition, node_id: int diff --git a/src/ampform/helicity/naming.py b/src/ampform/helicity/naming.py index 5b21130c2..472908047 100644 --- a/src/ampform/helicity/naming.py +++ b/src/ampform/helicity/naming.py @@ -2,6 +2,7 @@ from __future__ import annotations import re +from abc import ABC, abstractmethod from functools import lru_cache from typing import Iterable @@ -17,14 +18,83 @@ ) -class HelicityAmplitudeNameGenerator: +class NameGenerator(ABC): + """Name generator for amplitudes and coefficients in a `.HelicityModel`. + + .. seealso:: :ref:`usage/helicity/formalism:Coefficient names` + """ + + @abstractmethod + def generate_amplitude_name( + self, transition: StateTransition, node_id: int | None = None + ) -> str: + """Generates a unique name for the amplitude corresponding. + + That is, corresponging to the given + `~qrules.transition.StateTransition`. If ``node_id`` is given, it + generates a unique name for the partial amplitude corresponding to the + interaction node of the given `~qrules.transition.StateTransition`. + """ + + @abstractmethod + def generate_sequential_amplitude_suffix( + self, transition: StateTransition + ) -> str: + """Generate unique suffix for a sequential amplitude transition.""" + + @abstractmethod + def generate_coefficient_suffix( # pylint: disable=no-self-use + self, transition: StateTransition, node_id: int + ) -> str: + """Generate partial amplitude coefficient name suffix.""" + + @property + @abstractmethod + def parity_partner_coefficient_mapping(self) -> dict[str, str]: + ... + + +class HelicityAmplitudeNameGenerator(NameGenerator): def __init__( - self, transitions: ReactionInfo | Iterable[StateTransition] + self, + transitions: ReactionInfo | Iterable[StateTransition], + insert_parent_helicities: bool = False, + insert_child_helicities: bool = True, ) -> None: if isinstance(transitions, ReactionInfo): transitions = transitions.transitions - self.parity_partner_coefficient_mapping: dict[str, str] = {} - for transition in transitions: + self.__transitions = transitions + self.__insert_parent_helicities = insert_parent_helicities + self.__insert_child_helicities = insert_child_helicities + self._register_amplitude_coefficients() + + @property + def parity_partner_coefficient_mapping(self) -> dict[str, str]: + return self.__parity_partner_coefficient_mapping + + @property + def insert_parent_helicities(self) -> bool: + """Insert helicities of each parent state in the coefficient names.""" + return self.__insert_parent_helicities + + @insert_parent_helicities.setter + def insert_parent_helicities(self, value: bool) -> None: + self.__insert_parent_helicities = value + self._register_amplitude_coefficients() + + @property + def insert_child_helicities(self) -> bool: + """Embed the helicity of each decay product in the coefficients.""" + return self.__insert_child_helicities + + @insert_child_helicities.setter + def insert_child_helicities(self, value: bool) -> None: + self.__insert_child_helicities = value + self._register_amplitude_coefficients() + + def _register_amplitude_coefficients(self) -> None: + self.__parity_partner_coefficient_mapping: dict[str, str] = {} + for transition in self.__transitions: self.__register_amplitude_coefficient_name(transition) def __register_amplitude_coefficient_name( @@ -44,30 +114,30 @@ def __register_amplitude_coefficient_name( if ( coefficient_suffix - not in self.parity_partner_coefficient_mapping + not in self.__parity_partner_coefficient_mapping ): if ( parity_partner_coefficient_suffix - in self.parity_partner_coefficient_mapping + in self.__parity_partner_coefficient_mapping ): if ( parity_partner_coefficient_suffix == priority_partner_coefficient_suffix ): - self.parity_partner_coefficient_mapping[ + self.__parity_partner_coefficient_mapping[ coefficient_suffix ] = parity_partner_coefficient_suffix else: - self.parity_partner_coefficient_mapping[ + self.__parity_partner_coefficient_mapping[ parity_partner_coefficient_suffix ] = coefficient_suffix - self.parity_partner_coefficient_mapping[ + self.__parity_partner_coefficient_mapping[ coefficient_suffix ] = coefficient_suffix else: # if neither this coefficient nor its partner are registered just add it - self.parity_partner_coefficient_mapping[ + self.__parity_partner_coefficient_mapping[ coefficient_suffix ] = coefficient_suffix @@ -77,7 +147,7 @@ def __generate_amplitude_coefficient_couple( incoming_state, outgoing_states = get_helicity_info( transition, node_id ) - par_name_suffix = self.generate_coefficient_name(transition, node_id) + par_name_suffix = self.generate_coefficient_suffix(transition, node_id) pp_par_name_suffix = ( _state_to_str(incoming_state, use_helicity=False) @@ -102,13 +172,6 @@ def generate_amplitude_name( # pylint: disable=no-self-use transition: StateTransition, node_id: int | None = None, ) -> str: - """Generates a unique name for the amplitude corresponding. - - That is, corresponging to the given - `~qrules.transition.StateTransition`. If ``node_id`` is given, it - generates a unique name for the partial amplitude corresponding to the - interaction node of the given `~qrules.transition.StateTransition`. - """ name = "" if node_id is None: node_ids = transition.topology.nodes @@ -125,24 +188,37 @@ def generate_amplitude_name( # pylint: disable=no-self-use names.append(name) return "; ".join(names) - def generate_coefficient_name( # pylint: disable=no-self-use + def generate_coefficient_suffix( self, transition: StateTransition, node_id: int ) -> str: - """Generate partial amplitude coefficient name suffix.""" + components = self._get_coefficient_components(transition, node_id) + return "".join(components) + + def _get_coefficient_components( + self, transition: StateTransition, node_id: int + ) -> tuple[str, str, str]: in_hel_info, out_hel_info = get_helicity_info(transition, node_id) return ( - _state_to_str(in_hel_info, use_helicity=False) - + R" \to " - + " ".join(_state_to_str(s) for s in out_hel_info) + _state_to_str( + in_hel_info, + use_helicity=self.insert_parent_helicities, + ), + R" \to ", + " ".join( + _state_to_str( + state, + use_helicity=self.insert_child_helicities, + ) + for state in out_hel_info + ), ) def generate_sequential_amplitude_suffix( self, transition: StateTransition ) -> str: - """Generate unique suffix for a sequential amplitude transition.""" coefficient_names: list[str] = [] for node_id in transition.topology.nodes: - suffix = self.generate_coefficient_name(transition, node_id) + suffix = self.generate_coefficient_suffix(transition, node_id) if suffix in self.parity_partner_coefficient_mapping: suffix = self.parity_partner_coefficient_mapping[suffix] coefficient_names.append(suffix) @@ -150,6 +226,30 @@ def generate_sequential_amplitude_suffix( class CanonicalAmplitudeNameGenerator(HelicityAmplitudeNameGenerator): + def __init__( + self, + transitions: ReactionInfo | Iterable[StateTransition], + insert_parent_helicities: bool = False, + insert_child_helicities: bool = False, + insert_ls_combinations: bool = True, + ) -> None: + self.__insert_ls_combinations = insert_ls_combinations + super().__init__( + transitions, + insert_parent_helicities=insert_parent_helicities, + insert_child_helicities=insert_child_helicities, + ) + + @property + def insert_ls_combinations(self) -> bool: + """Embed each :math:`LS`-combination in the coefficient names.""" + return self.__insert_ls_combinations + + @insert_ls_combinations.setter + def insert_ls_combinations(self, value: bool) -> None: + self.__insert_ls_combinations = value + self._register_amplitude_coefficients() + def generate_amplitude_name( self, transition: StateTransition, @@ -169,18 +269,16 @@ def generate_amplitude_name( names.append(canonical_name) return "; ".join(names) - def generate_coefficient_name( + def _get_coefficient_components( self, transition: StateTransition, node_id: int - ) -> str: - incoming_state, outgoing_states = get_helicity_info( - transition, node_id - ) + ) -> tuple[str, str, str]: + components = super()._get_coefficient_components(transition, node_id) + if not self.insert_ls_combinations: + return components return ( - _state_to_str(incoming_state, use_helicity=False) - + self.__generate_ls_arrow(transition, node_id) - + " ".join( - _state_to_str(s, use_helicity=False) for s in outgoing_states - ) + components[0], + self.__generate_ls_arrow(transition, node_id), + components[2], ) @staticmethod diff --git a/tests/helicity/test_naming.py b/tests/helicity/test_naming.py index 818b0d6db..e20353b2e 100644 --- a/tests/helicity/test_naming.py +++ b/tests/helicity/test_naming.py @@ -1,6 +1,16 @@ +from __future__ import annotations + +import pytest from qrules import ReactionInfo -from ampform.helicity.naming import _render_float, generate_transition_label +from ampform import get_builder +from ampform.helicity import HelicityModel +from ampform.helicity.naming import ( + CanonicalAmplitudeNameGenerator, + HelicityAmplitudeNameGenerator, + _render_float, + generate_transition_label, +) def test_generate_transition_label(reaction: ReactionInfo): @@ -12,3 +22,83 @@ def test_generate_transition_label(reaction: ReactionInfo): Rf"J/\psi(1S)_{{{jpsi_spin}}} \to \gamma_{{{gamma_spin}}}" R" \pi^{0}_{0} \pi^{0}_{0}" ) + + +@pytest.mark.parametrize("parent_helicities", [False, True]) +@pytest.mark.parametrize("child_helicities", [False, True]) +@pytest.mark.parametrize("ls_combinations", [False, True]) +def test_coefficient_names( # noqa: R701 + reaction: ReactionInfo, + parent_helicities, + child_helicities, + ls_combinations, +): + # pylint: disable=too-many-branches, too-many-statements + builder = get_builder(reaction) + assert isinstance(builder.naming, HelicityAmplitudeNameGenerator) + builder.naming.insert_parent_helicities = parent_helicities + builder.naming.insert_child_helicities = child_helicities + if ls_combinations: + if reaction.formalism == "helicity": + pytest.skip("No LS-combinations if using helicity formalism") + if isinstance(builder.naming, CanonicalAmplitudeNameGenerator): + builder.naming.insert_ls_combinations = ls_combinations + model = builder.formulate() + + coefficients = get_coefficients(model) + n_resonances = len(reaction.get_intermediate_particles()) + if reaction.formalism == "helicity": + if parent_helicities: + if child_helicities: + assert len(coefficients) == 4 * n_resonances + else: + assert len(coefficients) == 2 * n_resonances + else: + if child_helicities: + assert len(coefficients) == n_resonances + else: + assert len(coefficients) == n_resonances + elif reaction.formalism == "canonical-helicity": + if ls_combinations: + if parent_helicities: + if child_helicities: + assert len(coefficients) == 8 * n_resonances + else: + assert len(coefficients) == 4 * n_resonances + else: + if child_helicities: + assert len(coefficients) == 4 * n_resonances + else: + assert len(coefficients) == 2 * n_resonances + else: + if parent_helicities: + if child_helicities: + assert len(coefficients) == 4 * n_resonances + else: + assert len(coefficients) == 2 * n_resonances + else: + assert len(coefficients) == n_resonances + + coefficient_name = coefficients[0] + if parent_helicities: + assert R"J/\psi(1S)_{-1}" in coefficient_name + else: + assert R"J/\psi(1S) " in coefficient_name + + if child_helicities: + assert R"\gamma_{" in coefficient_name + else: + assert R"\gamma;" in coefficient_name + + if ls_combinations: + assert R"\xrightarrow[S=1]{L=0}" in coefficient_name + else: + assert R"\to" in coefficient_name + + +def get_coefficients(model: HelicityModel) -> list[str]: + return [ + symbol.name + for symbol in model.parameter_defaults + if symbol.name.startswith("C_") + ]