From b4f6b33be7b7dce740dcf834aa10f2f536f06c1d Mon Sep 17 00:00:00 2001
From: Remco de Boer <redeboer@gmx.com>
Date: Thu, 24 Jun 2021 12:39:33 +0200
Subject: [PATCH] feat: implement HelicityModel.sum_components (#90)

---
 docs/usage/amplitude.ipynb       | 11 +++++++++
 src/ampform/helicity/__init__.py | 41 ++++++++++++++++++++++++++++++++
 tests/helicity/test_helicity.py  | 33 +++++++++++++++++++++++++
 3 files changed, 85 insertions(+)

diff --git a/docs/usage/amplitude.ipynb b/docs/usage/amplitude.ipynb
index ff3b89003..d891f4ebe 100644
--- a/docs/usage/amplitude.ipynb
+++ b/docs/usage/amplitude.ipynb
@@ -417,6 +417,17 @@
     "plots[0].show()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    ":::{tip}\n",
+    "\n",
+    "Use {meth}`.HelicityModel.sum_components` for adding up separate components of the model.\n",
+    "\n",
+    ":::"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/src/ampform/helicity/__init__.py b/src/ampform/helicity/__init__.py
index 35650c7f3..599a4cb79 100644
--- a/src/ampform/helicity/__init__.py
+++ b/src/ampform/helicity/__init__.py
@@ -3,6 +3,7 @@
 import logging
 import operator
 from collections import defaultdict
+from difflib import get_close_matches
 from functools import reduce
 from typing import (
     Any,
@@ -83,6 +84,46 @@ def parameter_defaults(self) -> Dict[sp.Symbol, ParameterValue]:
     def adapter(self) -> HelicityAdapter:
         return self._adapter
 
+    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"'
+        )
+
 
 class HelicityAmplitudeBuilder:
     """Amplitude model generator for the helicity formalism."""
diff --git a/tests/helicity/test_helicity.py b/tests/helicity/test_helicity.py
index b2a547d6f..a8fb4b713 100644
--- a/tests/helicity/test_helicity.py
+++ b/tests/helicity/test_helicity.py
@@ -1,4 +1,6 @@
 # pylint: disable=no-member, no-self-use
+from typing import Tuple
+
 import pytest
 import sympy as sp
 from qrules import ReactionInfo
@@ -6,6 +8,7 @@
 
 from ampform import get_builder
 from ampform.helicity import (
+    HelicityModel,
     _generate_kinematic_variables,
     formulate_wigner_d,
     group_transitions,
@@ -48,6 +51,36 @@ def test_formulate(self, reaction: ReactionInfo):
             assert no_dynamics == 8.0 - 4.0 * sin(theta) ** 2
 
 
+class TestHelicityModel:
+    def test_sum_components(self, amplitude_model: Tuple[str, HelicityModel]):
+        # pylint: disable=cell-var-from-loop, line-too-long
+        _, model = amplitude_model
+        from_intensities = model.sum_components(
+            components=filter(lambda c: c.startswith("I"), model.components),
+        )
+        assert from_intensities == model.expression
+        for spin_jpsi in ["-1", "+1"]:
+            for spin_gamma in ["-1", "+1"]:
+                jpsi_with_spin = fR"J/\psi(1S)_{{{spin_jpsi}}}"
+                gamma_with_spin = fR"\gamma_{{{spin_gamma}}}"
+                from_amplitudes = model.sum_components(
+                    components=filter(
+                        lambda c: c.startswith("A")
+                        and jpsi_with_spin in c
+                        and gamma_with_spin in c,
+                        model.components,
+                    )
+                )
+                selected_intensities = filter(
+                    lambda c: c.startswith("I")
+                    and jpsi_with_spin in c
+                    and gamma_with_spin in c,
+                    model.components,
+                )
+                selected_intensity = next(selected_intensities)
+                assert from_amplitudes == model.components[selected_intensity]
+
+
 @pytest.mark.parametrize(
     ("node_id", "mass", "phi", "theta"),
     [