ENH: automatically select best reference subsystem (#178)

* ENH: document `formulate(reference_subsystem=None)`

Author: redeboer

docs/lc2pkpi.ipynb

@@ -254,31 +254,44 @@
     "model_builder = DalitzPlotDecompositionBuilder(DECAY, min_ls=(False, True))\n",
     "for chain in model_builder.decay.chains:\n",
     "    model_builder.dynamics_choices.register_builder(chain, formulate_breit_wigner)\n",
-    "model = model_builder.formulate(reference_subsystem=1)\n",
+    "model = model_builder.formulate()\n",
     "model.intensity"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Notice that subsystem `3` (with the $\\Lambda^*$ resonances and recoil $\\pi^+$) has been selected as reference subsystem for the Wigner‑$d$ functions. This is computationally more efficient, because that subsystem contains most resonances."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
     "tags": [
-     "hide-input",
-     "full-width",
-     "hide-output",
      "scroll-output"
     ]
    },
    "outputs": [],
    "source": [
-    "Latex(aslatex(model.amplitudes).replace(R\"\\sum_\", R\"\\sum\\limits_\"))"
+    "Latex(aslatex(model.amplitudes, terms_per_line=1))"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "By default, the aligned amplitudes are built up of Wigner-$d$ functions, Clebsch–Gordan coefficients ($C$), a resonance parametrization ($\\mathcal{R}(\\sigma)$), and two coupling symbols $\\mathcal{H}^\\text{prod}_\\dots,\\mathcal{H}^\\text{dec}_\\dots$. In some cases, you want to combine the couplings into one scaling factor. That can be done with the `use_coefficients` flag:"
+    "By default, the aligned amplitudes are built up of Wigner-$d$ functions, Clebsch–Gordan coefficients ($C$), a resonance parametrization ($\\mathcal{R}(\\sigma)$), and two coupling symbols $\\mathcal{H}^\\text{prod}_{\\dots}, \\mathcal{H}^\\text{dec}_{\\dots}$. In some cases, you want to combine the couplings into one scaling factor. That can be done with the `use_coefficients` flag. The following example also explicitly specifies the reference subsystem (which is otherwise automatically chosen)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = model_builder.formulate(reference_subsystem=1, use_coefficients=True)"
    ]
   },
   {
@@ -289,14 +302,13 @@
      "source_hidden": true
     },
     "tags": [
-     "full-width"
+     "hide-input"
     ]
    },
    "outputs": [],
    "source": [
-    "model = model_builder.formulate(reference_subsystem=1, use_coefficients=True)\n",
     "(symbol, expr), *_ = model.amplitudes.items()\n",
-    "Latex(aslatex({symbol: expr}).replace(R\"\\sum_\", R\"\\sum\\limits_\"))"
+    "Latex(aslatex({symbol: expr}, terms_per_line=1))"
    ]
   }
  ],
@@ -319,7 +331,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.8"
+   "version": "3.12.11"
   }
  },
  "nbformat": 4,

src/ampform_dpd/__init__.py

@@ -91,12 +91,25 @@ def __init__(
 
     def formulate(
         self,
-        reference_subsystem: FinalStateID = 1,
+        reference_subsystem: FinalStateID | None = None,
         *,
         cleanup_summations: bool = False,
         use_coefficients: bool = False,
     ) -> AmplitudeModel:
-        _check_reference_subsystems(self.decay, reference_subsystem)
+        """Formulate the amplitude model given the configuration of this builder.
+
+        Args:
+            reference_subsystem: The subsystem to use as reference for the alignment of
+                helicities. If `None`, the subsystem with the most resonances is chosen.
+            cleanup_summations: Whether to remove helicity indices in the summations if
+                their corresponding state is spinless.
+            use_coefficients: Whether to use a single complex coefficient per decay
+                chain, instead of separate coefficients for each helicity coupling.
+        """
+        if reference_subsystem is None:
+            reference_subsystem = _get_best_reference_subsystems(self.decay)
+        else:
+            _check_reference_subsystems(self.decay, reference_subsystem)
         helicity_symbols: tuple[sp.Symbol, sp.Symbol, sp.Symbol, sp.Symbol] = (
             sp.symbols("lambda:4", rational=True)
         )
@@ -244,9 +257,12 @@ def formulate_aligned_amplitude(
         λ1: sp.Rational | sp.Symbol,
         λ2: sp.Rational | sp.Symbol,
         λ3: sp.Rational | sp.Symbol,
-        reference_subsystem: FinalStateID = 1,
+        reference_subsystem: FinalStateID | None = None,
     ) -> tuple[PoolSum, dict[sp.Symbol, sp.Expr]]:
-        _check_reference_subsystems(self.decay, reference_subsystem)
+        if reference_subsystem is None:
+            reference_subsystem = _get_best_reference_subsystems(self.decay)
+        else:
+            _check_reference_subsystems(self.decay, reference_subsystem)
         wigner_generator = _AlignmentWignerGenerator(reference_subsystem)
         _λ0, _λ1, _λ2, _λ3 = sp.symbols(R"\lambda_(0:4)^{\prime}", rational=True)
         j0, j1, j2, j3 = (self.decay.states[i].spin for i in sorted(self.decay.states))
@@ -274,6 +290,17 @@ def _product(obj: Any | Iterable):
     return obj
 
 
+def _get_best_reference_subsystems(decay: ThreeBodyDecay) -> FinalStateID:
+    subsystem_ids = _get_subsystem_ids(decay)
+    if not subsystem_ids:
+        msg = f"Decay {_get_decay_description(decay)} has no subsystems"
+        raise ValueError(msg)
+    resonances_per_subsystem = [
+        (k, len(decay.get_subsystem(k).chains)) for k in subsystem_ids
+    ]
+    return max(resonances_per_subsystem, key=operator.itemgetter(1))[0]
+
+
 def _check_reference_subsystems(
     decay: ThreeBodyDecay, reference_subsystem: FinalStateID
 ) -> None:

tests/test_decay.py

@@ -1,5 +1,18 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from qrules.transition import ReactionInfo
+
+from ampform_dpd import (
+    _get_best_reference_subsystems,  # pyright: ignore[reportPrivateUsage]
+)
+from ampform_dpd.adapter.qrules import to_three_body_decay
 from ampform_dpd.decay import IsobarNode, Particle
 
+if TYPE_CHECKING:
+    from qrules.transition import ReactionInfo
+
 # https://compwa-org--129.org.readthedocs.build/report/018.html#resonances-and-ls-scheme
 dummy_args = dict(mass=0, width=0)
 Λc = Particle("Λc", latex=R"\Lambda_c^+", spin=0.5, parity=+1, **dummy_args)
@@ -20,3 +33,16 @@ def test_ls(self):
         assert node.interaction is not None
         assert node.interaction.L == L
         assert node.interaction.S == S
+
+
+def test_get_best_reference_subsystems(
+    a2pipipi_reaction: ReactionInfo,
+    xib2pkk_reaction: ReactionInfo,
+    jpsi2pksigma_reaction: ReactionInfo,
+):
+    decay = to_three_body_decay(a2pipipi_reaction.transitions)
+    assert _get_best_reference_subsystems(decay) == 1
+    decay = to_three_body_decay(xib2pkk_reaction.transitions)
+    assert _get_best_reference_subsystems(decay) == 2
+    decay = to_three_body_decay(jpsi2pksigma_reaction.transitions)
+    assert _get_best_reference_subsystems(decay) == 2