feat: implement permutate_registered_topologies

.cspell.json

@@ -99,6 +99,7 @@
         "mypy",
         "nishijima",
         "numpy",
+        "permutate",
         "pydocstyle",
         "pylint",
         "pyplot",

docs/usage/amplitude.ipynb

@@ -313,6 +313,102 @@
     "model_builder.stable_final_state_ids = None"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Extend kinematic variables"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The {class}`.HelicityAmplitudeBuilder` by default only generates {attr}`.kinematic_variables` (helicity angles and invariant masses) for the topologies that are available in the {class}`~qrules.transition.ReactionInfo` object that it was created with. If you want to calculate more kinematic variables, you can use the method {meth}`.register_topology` of its helicity {attr}`.HelicityAmplitudeBuilder.adapter` to register more topologies and generate more kinematic variables. This is especially useful when generating data later on with [TensorWaves](https://tensorwaves.rtfd.io).\n",
+    "\n",
+    "To make this a bit easier, there is {meth}`.permutate_registered_topologies`, which is a small convenience function makes it possible to generate permutations of all {attr}`.registered_topologies` and register them as well. Note that initially, only one {class}`~qrules.topology.Topology` is registered in the {attr}`.HelicityAmplitudeBuilder.adapter`, namely the one for the decay $J/\\psi \\to \\gamma f_0, f_0 \\to \\pi^0\\pi^0$:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dot = qrules.io.asdot(model_builder.adapter.registered_topologies)\n",
+    "graphviz.Source(dot)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We now {meth}`.permutate_registered_topologies` to register permutations of this {class}`~qrules.topology.Topology`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_builder.adapter.permutate_registered_topologies()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "There are now **three** {attr}`.registered_topologies`―one for each permutation:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(model_builder.adapter.registered_topologies)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert len(model_builder.adapter.registered_topologies) == 3\n",
+    "dot = qrules.io.asdot(model_builder.adapter.registered_topologies)\n",
+    "graphviz.Source(dot)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "And if we {meth}`~.HelicityAmplitudeBuilder.formulate` a new {class}`.HelicityModel`, we see that it has many more {attr}`.kinematic_variables`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "set(model_builder.formulate().kinematic_variables)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    ":::{tip}\n",
+    "\n",
+    "To register even more topologies, use e.g. {func}`~qrules.topology.create_isobar_topologies` to generate other, non-isomorphic topologies that cannot be created with permutations. This is relevant for more than three final states.\n",
+    "\n",
+    ":::"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

src/ampform/helicity/__init__.py

@@ -165,7 +165,7 @@ class HelicityAmplitudeBuilder:  # pylint: disable=too-many-instance-attributes
     Args:
         reaction: The `~qrules.transition.ReactionInfo` from which to
             :meth:`formulate` an amplitude model.
-        use_scalar_masses: Put final state 'invariant' masses
+        stable_final_state_ids: Put final state 'invariant' masses
             (:math:`m_0, m_1, \dots`) under `.HelicityModel.parameter_defaults`
             (with a *scalar* suggested value) instead of
             `~.HelicityModel.kinematic_variables` (which are expressions to
@@ -198,8 +198,19 @@ def __init__(
         for grouping in reaction.transition_groups:
             self.__adapter.register_topology(grouping.topology)
 
+    @property
+    def adapter(self) -> HelicityAdapter:
+        """Converter for computing kinematic variables from four-momenta."""
+        return self.__adapter
+
     @property
     def stable_final_state_ids(self) -> Optional[Set[int]]:
+        # noqa: D403
+        """IDs of the final states that should be considered stable.
+
+        The 'invariant' mass symbols for these final states will be inserted as
+        **scalar** values into the `.parameter_defaults`.
+        """
         return self.__stable_final_state_ids
 
     @stable_final_state_ids.setter

src/ampform/kinematics.py

@@ -3,6 +3,7 @@
 """Kinematics of an amplitude model in the helicity formalism."""
 
 import functools
+import itertools
 from typing import (
     Any,
     Callable,
@@ -111,6 +112,24 @@ def register_topology(self, topology: Topology) -> None:
                 raise ValueError("Edge or node IDs of topology do not match")
         self.registered_topologies.add(topology)
 
+    def permutate_registered_topologies(self) -> None:
+        """Register permutations of all `registered_topologies`.
+
+        See :ref:`usage/amplitude:Extend kinematic variables`.
+        """
+        for topology in set(self.registered_topologies):
+            final_state_ids = topology.outgoing_edge_ids
+            for permutation in itertools.permutations(final_state_ids):
+                id_mapping = dict(zip(topology.outgoing_edge_ids, permutation))
+                permuted_topology = attr.evolve(
+                    topology,
+                    edges={
+                        id_mapping.get(i, i): edge
+                        for i, edge in topology.edges.items()
+                    },
+                )
+                self.register_topology(permuted_topology)
+
     def create_expressions(self) -> Dict[str, sp.Expr]:
         output = {}
         for topology in self.registered_topologies:

tests/helicity/test_helicity.py

@@ -16,10 +16,16 @@
 
 
 class TestHelicityAmplitudeBuilder:
+    @pytest.mark.parametrize("permutate_topologies", [False, True])
     @pytest.mark.parametrize(
         "stable_final_state_ids", [None, (1, 2), (0, 1, 2)]
     )
-    def test_formulate(self, reaction: ReactionInfo, stable_final_state_ids):
+    def test_formulate(
+        self,
+        permutate_topologies,
+        reaction: ReactionInfo,
+        stable_final_state_ids,
+    ):
         # pylint: disable=too-many-locals
         if reaction.formalism == "canonical-helicity":
             n_amplitudes = 16
@@ -35,6 +41,10 @@ def test_formulate(self, reaction: ReactionInfo, stable_final_state_ids):
 
         model_builder: HelicityAmplitudeBuilder = get_builder(reaction)
         model_builder.stable_final_state_ids = stable_final_state_ids
+        if permutate_topologies:
+            model_builder.adapter.permutate_registered_topologies()
+            n_kinematic_variables += 10
+
         model = model_builder.formulate()
         assert len(model.parameter_defaults) == n_parameters
         assert len(model.components) == 4 + n_amplitudes