docs/report/014.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "sympy"
    ]
   },
   "source": [
    "::::{margin}\n",
    ":::{card} Amplitude model with sum notation\n",
    "TR-014\n",
    "^^^\n",
    "See also [TR-013](013.ipynb) and [TR-015](015.ipynb).\n",
    "+++\n",
    "✅&nbsp;[ampform#245](https://github.com/ComPWA/ampform/issues/245)\n",
    ":::\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amplitude model with sum notation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q qrules[viz]==0.9.7 sympy==1.9 git+https://github.com/ComPWA/ampform@3ed3ed5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import inspect\n",
    "import itertools\n",
    "import logging\n",
    "from functools import lru_cache\n",
    "from typing import TYPE_CHECKING, Iterable, Sequence\n",
    "\n",
    "import ampform\n",
    "import attrs\n",
    "import graphviz\n",
    "import qrules\n",
    "import symplot\n",
    "import sympy as sp\n",
    "from ampform.dynamics.builder import (\n",
    "    ResonanceDynamicsBuilder,\n",
    "    create_non_dynamic,\n",
    "    create_relativistic_breit_wigner,\n",
    ")\n",
    "from ampform.helicity import (\n",
    "    _generate_kinematic_variable_set,\n",
    "    _generate_kinematic_variables,\n",
    ")\n",
    "from ampform.helicity.decay import TwoBodyDecay\n",
    "from ampform.sympy import (\n",
    "    UnevaluatedExpression,\n",
    "    create_expression,\n",
    "    implement_doit_method,\n",
    ")\n",
    "from IPython.display import Math, display\n",
    "from sympy.core.symbol import Str\n",
    "from sympy.physics.quantum.spin import Rotation as Wigner\n",
    "from sympy.printing.precedence import PRECEDENCE\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    import sys\n",
    "\n",
    "    from qrules.topology import Topology\n",
    "    from qrules.transition import StateTransition\n",
    "    from sympy.printing.latex import LatexPrinter\n",
    "\n",
    "LOGGER = logging.getLogger()\n",
    "LOGGER.setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(tr-014-problem)=\n",
    "## Problem description"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{admonition} Challenge\n",
    ":class: purple\n",
    "\n",
    "Formulate the [`HelicityModel`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel) in such a way that:\n",
    "\n",
    "1. The sum over the amplitudes is concise and expresses that the sum depends only on helicities.\n",
    "2. It is the {class}`sympy.Expr <sympy.core.expr.Expr>` of the amplitude model can easily and uniquely be constructed from the data in the [`HelicityModel`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel). (Currently, this is as simple as {meth}`HelicityModel.expression.doit <sympy.core.basic.Basic.doit>`).\n",
    "3. All parameters under [`parameter_defaults`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel.parameter_defaults) are of type {class}`sympy.Symbol <sympy.core.symbol.Symbol>`. This is important for a correct lambdification of the arguments with {func}`sympy.utilities.lambdify.lambdify`.\n",
    "\n",
    "This report presents two solutions:\n",
    "- {ref}`tr-014-solution-1`\n",
    "- {ref}`tr-014-solution-2`\n",
    "\n",
    ":::\n",
    "\n",
    "[ampform#245](https://github.com/ComPWA/ampform/pull/245) implements spin alignment, which results in large sum combinatorics for all helicity combinations. The result is an amplitude model expression that is too large to be rendered as LaTeX.\n",
    "\n",
    "To some extend, this is already the case with the [current implementation](https://ampform.readthedocs.io/en/0.12.3/usage/formalism.html) of the 'standard' helicity formalism {cite}`jacobGeneralTheoryCollisions1959, richmanExperimenterGuideHelicity1984, kutschkeAngularDistributionCookbook1996, chungSpinFormalismsUpdated2014`: many of the terms in the total intensity expression differ only by the helicities of the final and initial state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Simplify resonance notation\n",
    "PDG = qrules.load_pdg()\n",
    "delta_res = PDG[\"Delta(1600)++\"]\n",
    "lambda_res = PDG[\"Lambda(1405)\"]\n",
    "particles = set(PDG)\n",
    "particles.remove(delta_res)\n",
    "particles.remove(lambda_res)\n",
    "particles.add(attrs.evolve(delta_res, latex=R\"\\Delta\"))\n",
    "particles.add(attrs.evolve(lambda_res, latex=R\"\\Lambda\"))\n",
    "MODIFIED_PDG = qrules.ParticleCollection(particles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=\"Lambda(c)+\",\n",
    "    final_state=[\"K-\", \"p\", \"pi+\"],\n",
    "    formalism=\"helicity\",\n",
    "    allowed_intermediate_particles=[\"Delta(1600)++\"],\n",
    "    particle_db=MODIFIED_PDG,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "display(*(graphviz.Source(qrules.io.asdot(t, size=3)) for t in reaction.transitions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{container} full-width\n",
    "![](https://user-images.githubusercontent.com/29308176/164992090-84c38c7d-9c1c-4e57-abef-1ad46d557eda.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992091-e24324cc-8e5c-46b6-a5ec-df1bd2b98cb3.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992092-f2b2d578-be50-4102-aa8f-780eea7cf8c7.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992095-75f21354-3e50-4919-b07c-9bb68ee71929.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992096-d18f2f49-d317-4874-b728-d14192cb876d.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992097-be7b64a7-0fa5-44a6-863a-120ccf648a5a.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992099-ccd9b237-27e7-46ac-88b5-69417f4102cc.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992101-99d6d3e4-5a30-49dc-888a-e9f6ee580eff.svg)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "builder = ampform.get_builder(reaction)\n",
    "model = builder.formulate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "I = & \\left|{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \\\\\n",
       "& + \\left|{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \\\\\n",
       "& + \\left|{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \\\\\n",
       "& + \\left|{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def remove_coefficients(expr: sp.Expr) -> sp.Expr:\n",
    "    coefficients = {s: 1 for s in expr.free_symbols if s.name.startswith(\"C_\")}\n",
    "    return expr.subs(coefficients)\n",
    "\n",
    "\n",
    "model = builder.formulate()\n",
    "full_expression = remove_coefficients(model.expression)\n",
    "I = sp.Symbol(\"I\")\n",
    "latex = sp.multiline_latex(I, full_expression)\n",
    "Math(latex)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we did not insert any dynamics, but it is unusual that dynamics expressions depend on helicity or spin projection (see [`TwoBodyKinematicVariableSet`](https://ampform.readthedocs.io/en/stable/api/ampform.dynamics.builder.html#ampform.dynamics.builder.TwoBodyKinematicVariableSet))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simplified notation with `PoolSum`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both {ref}`tr-014-solution-1` and {ref}`tr-014-solution-2` require the definition of a special \"`PoolSum`\" class to simplify the summation over the amplitudes. The class mimics {class}`sympy.Sum <sympy.concrete.summations.Sum>` in that it substitutes certain {class}`~sympy.core.symbol.Symbol`s in an expression over which we symbol a range of values. The range of values in a `PoolSum` does not have to be a sequential range, but can be a collection of arbitrary items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "@implement_doit_method\n",
    "class PoolSum(UnevaluatedExpression):\n",
    "    precedence = PRECEDENCE[\"Add\"]\n",
    "\n",
    "    def __new__(\n",
    "        cls,\n",
    "        expression: sp.Expr,\n",
    "        *indices: tuple[sp.Symbol, Iterable[sp.Float]],\n",
    "        **hints,\n",
    "    ) -> PoolSum:\n",
    "        indices = tuple((s, tuple(v)) for s, v in indices)\n",
    "        return create_expression(cls, expression, *indices, **hints)\n",
    "\n",
    "    @property\n",
    "    def expression(self) -> sp.Expr:\n",
    "        return self.args[0]\n",
    "\n",
    "    @property\n",
    "    def indices(self) -> list[tuple[sp.Symbol, tuple[sp.Float, ...]]]:\n",
    "        return self.args[1:]\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        indices = dict(self.indices)\n",
    "        return sum(\n",
    "            self.expression.subs(zip(indices, combi))\n",
    "            for combi in itertools.product(*indices.values())\n",
    "        )\n",
    "\n",
    "    def _latex(self, printer: LatexPrinter, *args) -> str:\n",
    "        indices = dict(self.indices)\n",
    "        sum_symbols: list[str] = []\n",
    "        for idx, values in indices.items():\n",
    "            sum_symbols.append(_render_sum_symbol(printer, idx, values))\n",
    "        expression = printer._print(self.expression)\n",
    "        return R\" \".join(sum_symbols) + f\"{{{expression}}}\"\n",
    "\n",
    "    def cleanup(self) -> sp.Expr | PoolSum:\n",
    "        substitutions = {}\n",
    "        new_indices = []\n",
    "        for idx, values in self.indices:\n",
    "            if idx not in self.expression.free_symbols:\n",
    "                continue\n",
    "            if len(values) == 0:\n",
    "                continue\n",
    "            if len(values) == 1:\n",
    "                substitutions[idx] = values[0]\n",
    "            else:\n",
    "                new_indices.append((idx, values))\n",
    "        new_expression = self.expression.xreplace(substitutions)\n",
    "        if len(new_indices) == 0:\n",
    "            return new_expression\n",
    "        return PoolSum(new_expression, *new_indices)\n",
    "\n",
    "\n",
    "def _render_sum_symbol(\n",
    "    printer: LatexPrinter, idx: sp.Symbol, values: Sequence[float]\n",
    ") -> str:\n",
    "    if len(values) == 0:\n",
    "        return \"\"\n",
    "    idx = printer._print(idx)\n",
    "    if len(values) == 1:\n",
    "        value = values[0]\n",
    "        return Rf\"\\sum_{{{idx}={value}}}\"\n",
    "    if _is_regular_series(values):\n",
    "        sorted_values = sorted(values)\n",
    "        first_value = sorted_values[0]\n",
    "        last_value = sorted_values[-1]\n",
    "        return Rf\"\\sum_{{{idx}={first_value}}}^{{{last_value}}}\"\n",
    "    idx_values = \",\".join(map(printer._print, values))\n",
    "    return Rf\"\\sum_{{{idx}\\in\\left\\{{{idx_values}\\right\\}}}}\"\n",
    "\n",
    "\n",
    "def _is_regular_series(values: Sequence[float]) -> bool:\n",
    "    if len(values) <= 1:\n",
    "        return False\n",
    "    sorted_values = sorted(values)\n",
    "    for val, next_val in zip(sorted_values, sorted_values[1:]):\n",
    "        difference = float(next_val - val)\n",
    "        if difference != 1.0:\n",
    "            return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a sketch of how to construct the amplitude model with a `PoolSum`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left. \\sum_{\\lambda_{\\Lambda_c}=-1/2}^{1/2} \\sum_{\\lambda_{p}=-1/2}^{1/2} \\sum_{\\lambda_{\\pi}=0} \\sum_{\\lambda_{K}=0}{\\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{s_{\\Delta}}_{\\lambda_{\\Delta},- \\lambda_{\\pi} + \\lambda_{p}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{s_{\\Lambda_c}}_{\\lambda_{\\Lambda_c},\\lambda_{K} - \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right)}}\\right|^{2}} \\right|_{\\substack{ s_{\\Lambda_c}=\\frac{1}{2}\\\\ s_{\\Delta}=\\frac{3}{2} }}$"
      ],
      "text/plain": [
       "Subs(PoolSum(Abs(PoolSum(WignerD(s_\\Delta, \\lambda_\\Delta, -\\lambda_\\pi + \\lambda_p, phi_1^12, theta_1^12, 0)*WignerD(s_{\\Lambda_c}, \\lambda_{\\Lambda_c}, \\lambda_K - \\lambda_\\Delta, phi_12, theta_12, 0), (\\lambda_\\Delta, (-1/2, 1/2))))**2, (\\lambda_{\\Lambda_c}, (-1/2, 1/2)), (\\lambda_p, (-1/2, 1/2)), (\\lambda_\\pi, (0,)), (\\lambda_K, (0,))), (s_{\\Lambda_c}, s_\\Delta), (1/2, 3/2))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{\\lambda_{\\Lambda_c}=-1/2}^{1/2} \\sum_{\\lambda_{p}=-1/2}^{1/2}{\\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{s_{\\Delta}}_{\\lambda_{\\Delta},\\lambda_{p}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{s_{\\Lambda_c}}_{\\lambda_{\\Lambda_c},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right)}}\\right|^{2}}$"
      ],
      "text/plain": [
       "PoolSum(Abs(PoolSum(WignerD(s_\\Delta, \\lambda_\\Delta, \\lambda_p, phi_1^12, theta_1^12, 0)*WignerD(s_{\\Lambda_c}, \\lambda_{\\Lambda_c}, -\\lambda_\\Delta, phi_12, theta_12, 0), (\\lambda_\\Delta, (-1/2, 1/2))))**2, (\\lambda_{\\Lambda_c}, (-1/2, 1/2)), (\\lambda_p, (-1/2, 1/2)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "I = & \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},- \\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n",
       "& + \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},\\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n",
       "& + \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},- \\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n",
       "& + \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},\\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\sin^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)}}{8} - \\frac{3 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\sin{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\sin{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{4} + \\frac{9 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\sin^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{\\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{3 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\cos{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{4} + \\frac{9 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{\\sin^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)}}{8} - \\frac{3 \\sin{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\sin{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)}}{4} + \\frac{9 \\sin^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)}}{8} + \\frac{\\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{3 \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\cos{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{4} + \\frac{9 \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{8}$"
      ],
      "text/plain": [
       "sin(theta_12/2)**2*sin(theta_1^12/2)**2/8 - 3*sin(theta_12/2)**2*sin(theta_1^12/2)*sin(3*theta_1^12/2)/4 + 9*sin(theta_12/2)**2*sin(3*theta_1^12/2)**2/8 + sin(theta_12/2)**2*cos(theta_1^12/2)**2/8 + 3*sin(theta_12/2)**2*cos(theta_1^12/2)*cos(3*theta_1^12/2)/4 + 9*sin(theta_12/2)**2*cos(3*theta_1^12/2)**2/8 + sin(theta_1^12/2)**2*cos(theta_12/2)**2/8 - 3*sin(theta_1^12/2)*sin(3*theta_1^12/2)*cos(theta_12/2)**2/4 + 9*sin(3*theta_1^12/2)**2*cos(theta_12/2)**2/8 + cos(theta_12/2)**2*cos(theta_1^12/2)**2/8 + 3*cos(theta_12/2)**2*cos(theta_1^12/2)*cos(3*theta_1^12/2)/4 + 9*cos(theta_12/2)**2*cos(3*theta_1^12/2)**2/8"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "half = sp.S.Half\n",
    "\n",
    "spin_parent = sp.Symbol(R\"s_{\\Lambda_c}\", real=True)\n",
    "spin_resonance = sp.Symbol(R\"s_\\Delta\", real=True)\n",
    "\n",
    "phi_12, theta_12 = sp.symbols(\"phi_12 theta_12\", real=True)\n",
    "phi_1_12, theta_1_12 = sp.symbols(R\"phi_1^12 theta_1^12\", real=True)\n",
    "\n",
    "lambda_parent = sp.Symbol(R\"\\lambda_{\\Lambda_c}\", real=True)\n",
    "lambda_resonance = sp.Symbol(R\"\\lambda_\\Delta\", real=True)\n",
    "lambda_p = sp.Symbol(R\"\\lambda_p\", real=True)\n",
    "lambda_k = sp.Symbol(R\"\\lambda_K\", real=True)\n",
    "lambda_pi = sp.Symbol(R\"\\lambda_\\pi\", real=True)\n",
    "sum_expr = sp.Subs(\n",
    "    PoolSum(\n",
    "        sp.Abs(\n",
    "            PoolSum(\n",
    "                Wigner.D(\n",
    "                    spin_parent,\n",
    "                    lambda_parent,\n",
    "                    lambda_k - lambda_resonance,\n",
    "                    phi_12,\n",
    "                    theta_12,\n",
    "                    0,\n",
    "                )\n",
    "                * Wigner.D(\n",
    "                    spin_resonance,\n",
    "                    lambda_resonance,\n",
    "                    lambda_p - lambda_pi,\n",
    "                    phi_1_12,\n",
    "                    theta_1_12,\n",
    "                    0,\n",
    "                ),\n",
    "                (lambda_resonance, (-half, +half)),\n",
    "            )\n",
    "        )\n",
    "        ** 2,\n",
    "        (lambda_parent, (-half, +half)),\n",
    "        (lambda_p, (-half, +half)),\n",
    "        (lambda_pi, (0,)),\n",
    "        (lambda_k, (0,)),\n",
    "    ),\n",
    "    (spin_parent, spin_resonance),\n",
    "    (half, 3 * half),\n",
    ")\n",
    "display(\n",
    "    sum_expr,\n",
    "    sum_expr.expr.cleanup(),\n",
    "    Math(sp.multiline_latex(I, sum_expr.doit(deep=False))),\n",
    "    sum_expr.doit(deep=False).doit(deep=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert sum_expr.doit(deep=False).doit() == full_expression.doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(tr-014-solution-1)=\n",
    "## Solution 1: `Indexed` coefficients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{margin}\n",
    "\n",
    "When introducing spin alignment ([ampform#245](https://github.com/ComPWA/ampform/pull/245)), we have to distinguish the helicity symbols between different topologies.\n",
    "\n",
    ":::\n",
    "\n",
    "The current implementation of the [`HelicityAmplitudeBuilder`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityAmplitudeBuilder) has to be changed quite a bit to produce an amplitude model with `PoolSum`s. First of all, we have to introduce special {class}`~sympy.core.symbol.Symbol`s for the helicities, $\\lambda_i$, with $i$ the state ID (taking a sum of attached final state IDs in case of a resonance ID). Next, [`formulate_wigner_d()`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.formulate_wigner_d) has to be modified to insert these {class}`~sympy.core.symbol.Symbol`s into the {class}`~sympy.physics.quantum.spin.WignerD`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def formulate_wigner_d(transition: StateTransition, node_id: int) -> sp.Expr:\n",
    "    from sympy.physics.quantum.spin import Rotation as Wigner\n",
    "\n",
    "    decay = TwoBodyDecay.from_transition(transition, node_id)\n",
    "    topology = transition.topology\n",
    "    parent_helicity = create_helicity_symbol(topology, decay.parent.id)\n",
    "    child1_helicity = create_helicity_symbol(topology, decay.children[0].id)\n",
    "    child2_helicity = create_helicity_symbol(topology, decay.children[1].id)\n",
    "    _, phi, theta = _generate_kinematic_variables(transition, node_id)\n",
    "    return Wigner.D(\n",
    "        j=sp.Rational(decay.parent.particle.spin),\n",
    "        m=parent_helicity,\n",
    "        mp=child1_helicity - child2_helicity,\n",
    "        alpha=-phi,\n",
    "        beta=theta,\n",
    "        gamma=0,\n",
    "    )\n",
    "\n",
    "\n",
    "def create_helicity_symbol(topology: Topology, state_id: int) -> sp.Symbol:\n",
    "    if state_id in topology.incoming_edge_ids:\n",
    "        suffix = \"\"\n",
    "    else:\n",
    "        suffix = f\"_{state_id}\"\n",
    "    return sp.Symbol(f\"lambda{suffix}\", rational=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)$"
      ],
      "text/plain": [
       "WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wigner_functions = {\n",
    "    sp.Mul(*[\n",
    "        formulate_wigner_d(transition, node_id) for node_id in transition.topology.nodes\n",
    "    ])\n",
    "    for transition in reaction.transitions\n",
    "}\n",
    "display(*wigner_functions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also have to collect the allowed helicity values for each of these helicity symbols."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "from qrules import ReactionInfo\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    if sys.version_info >= (3, 8):\n",
    "        from typing import Literal\n",
    "    else:\n",
    "        from typing_extensions import Literal\n",
    "\n",
    "\n",
    "@lru_cache(maxsize=None)\n",
    "def get_helicities(\n",
    "    reaction: ReactionInfo, which: Literal[\"inner\", \"outer\"]\n",
    ") -> dict[int, set[sp.Rational]]:\n",
    "    helicities = defaultdict(set)\n",
    "    initial_state_ids = set(reaction.initial_state)\n",
    "    final_state_ids = set(reaction.final_state)\n",
    "    intermediate_state_ids = (\n",
    "        set(reaction.transitions[0].states) - initial_state_ids - final_state_ids\n",
    "    )\n",
    "    if which == \"inner\":\n",
    "        state_ids = sorted(intermediate_state_ids)\n",
    "    elif which == \"outer\":\n",
    "        state_ids = sorted(initial_state_ids | final_state_ids)\n",
    "    for transition in reaction.transitions:\n",
    "        for state_id in state_ids:\n",
    "            state = transition.states[state_id]\n",
    "            helicity = sp.Rational(state.spin_projection)\n",
    "            symbol = create_helicity_symbol(transition.topology, state_id)\n",
    "            helicities[symbol].add(helicity)\n",
    "    return dict(helicities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{lambda_3: {-1/2, 1/2}}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{lambda: {-1/2, 1/2}, lambda_0: {0}, lambda_1: {-1/2, 1/2}, lambda_2: {0}}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "inner_helicities = get_helicities(reaction, which=\"inner\")\n",
    "outer_helicities = get_helicities(reaction, which=\"outer\")\n",
    "display(inner_helicities, outer_helicities)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These collected helicity values can then be combined with the Wigner-$D$ expressions through a `PoolSum`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2}}$"
      ],
      "text/plain": [
       "PoolSum(Abs(PoolSum(WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def formulate_intensity(reaction: ReactionInfo):\n",
    "    wigner_functions = {\n",
    "        sp.Mul(*[\n",
    "            formulate_wigner_d(transition, node_id)\n",
    "            for node_id in transition.topology.nodes\n",
    "        ])\n",
    "        for transition in reaction.transitions\n",
    "    }\n",
    "    inner_helicities = get_helicities(reaction, which=\"inner\")\n",
    "    outer_helicities = get_helicities(reaction, which=\"outer\")\n",
    "    return PoolSum(\n",
    "        sp.Abs(\n",
    "            PoolSum(\n",
    "                sum(wigner_functions),\n",
    "                *inner_helicities.items(),\n",
    "            )\n",
    "        )\n",
    "        ** 2,\n",
    "        *outer_helicities.items(),\n",
    "    )\n",
    "\n",
    "\n",
    "formulate_intensity(reaction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is indeed identical to the model as formulated with the existing implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert formulate_intensity(reaction).doit() == full_expression.doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how this approach also works in case there are two decay topologies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "reaction_two_resonances = qrules.generate_transitions(\n",
    "    initial_state=\"Lambda(c)+\",\n",
    "    final_state=[\"K-\", \"p\", \"pi+\"],\n",
    "    formalism=\"helicity\",\n",
    "    allowed_intermediate_particles=[\"Lambda(1405)\", \"Delta(1600)++\"],\n",
    "    particle_db=MODIFIED_PDG,\n",
    ")\n",
    "assert len(reaction_two_resonances.transition_groups) == 2\n",
    "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n",
    "display(*[\n",
    "    graphviz.Source(\n",
    "        qrules.io.asdot(g, collapse_graphs=True, size=4, render_resonance_id=True)\n",
    "    )\n",
    "    for g in reaction_two_resonances.transition_groups\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{container} full-width\n",
    "![](https://user-images.githubusercontent.com/29308176/164992102-fcac3af8-285a-49e8-b830-58fc947fef30.svg)\n",
    "![](https://user-images.githubusercontent.com/29308176/164992104-f37f1a89-3cf0-43bb-a013-9f20e3064aed.svg)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{1}=-1/2}^{1/2}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\lambda,\\lambda_{3}}\\left(- \\phi_{01},\\theta_{01},0\\right) D^{\\frac{1}{2}}_{\\lambda_{3},- \\lambda_{1}}\\left(- \\phi^{01}_{0},\\theta^{01}_{0},0\\right) + D^{\\frac{1}{2}}_{\\lambda,\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2}}$"
      ],
      "text/plain": [
       "PoolSum(Abs(PoolSum(WignerD(1/2, lambda, lambda_3, -phi_01, theta_01, 0)*WignerD(1/2, lambda_3, -lambda_1, -phi_0^01, theta_0^01, 0) + WignerD(1/2, lambda, lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1, -phi_1^12, theta_1^12, 0), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_1, (-1/2, 1/2)))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formulate_intensity(reaction_two_resonances).cleanup()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "ampform_expr = ampform.get_builder(reaction_two_resonances).formulate().expression\n",
    "ampform_expr = remove_coefficients(ampform_expr)\n",
    "sum_expr = formulate_intensity(reaction_two_resonances)\n",
    "assert ampform_expr.doit() == sum_expr.doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inserting coefficients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is a problem, though: the sums above could be written in sum form because the values over which we sum appear as arguments ({attr}`~sympy.core.basic.Basic.args` in the {class}`~sympy.physics.quantum.spin.WignerD` functions). This is only true because we previously set all coefficients to $1$. The coefficient _names_ are, however, also 'dependent' on the helicities in the final state over which we sum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[C_{\\Lambda_{c}^{+} \\to \\Delta_{-1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}},\n",
       " C_{\\Lambda_{c}^{+} \\to \\Delta_{+1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sorted(model.parameter_defaults, key=str)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We therefore have to somehow introduce a dependence in these {class}`~sympy.core.symbol.Symbol`s on the helicity values.  An idea may be to use {class}`~sympy.tensor.indexed.IndexedBase`. Modifying the function introduced in {ref}`tr-014-solution-1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2}} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2}}$"
      ],
      "text/plain": [
       "PoolSum(Abs(PoolSum(C[lambda_0, lambda_1, lambda_2]*WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C = sp.IndexedBase(\"C\")\n",
    "\n",
    "\n",
    "@lru_cache(maxsize=None)\n",
    "def formulate_intensity_with_coefficient(reaction: ReactionInfo):\n",
    "    amplitudes = {\n",
    "        sp.Mul(\n",
    "            C[\n",
    "                [\n",
    "                    create_helicity_symbol(transition.topology, state_id)\n",
    "                    for state_id in transition.final_states\n",
    "                ]\n",
    "            ],\n",
    "            *[\n",
    "                formulate_wigner_d(transition, node_id)\n",
    "                for node_id in transition.topology.nodes\n",
    "            ],\n",
    "        )\n",
    "        for transition in reaction.transitions\n",
    "    }\n",
    "    inner_helicities = get_helicities(reaction, which=\"inner\")\n",
    "    outer_helicities = get_helicities(reaction, which=\"outer\")\n",
    "    return PoolSum(\n",
    "        sp.Abs(\n",
    "            PoolSum(\n",
    "                sum(amplitudes),\n",
    "                *inner_helicities.items(),\n",
    "            )\n",
    "        )\n",
    "        ** 2,\n",
    "        *outer_helicities.items(),\n",
    "    )\n",
    "\n",
    "\n",
    "indexed_coefficient_expr = formulate_intensity_with_coefficient(reaction)\n",
    "indexed_coefficient_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "I = & \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,- \\frac{1}{2},0} D^{\\frac{1}{2}}_{- \\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n",
       "& + \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,- \\frac{1}{2},0} D^{\\frac{1}{2}}_{\\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n",
       "& + \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,\\frac{1}{2},0} D^{\\frac{1}{2}}_{- \\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n",
       "& + \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,\\frac{1}{2},0} D^{\\frac{1}{2}}_{\\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "latex = sp.multiline_latex(I, indexed_coefficient_expr.doit(deep=False))\n",
    "Math(latex)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{admonition} Caveat\n",
    "\n",
    "Using {class}`~sympy.tensor.indexed.IndexedBase` makes the coefficient names concise, but harder to understand.\n",
    "\n",
    ":::\n",
    "\n",
    "This seems to work rather well, but there is a subtle problems introduced by writing the coefficients as a {class}`~sympy.tensor.indexed.IndexedBase`: the {class}`~sympy.tensor.indexed.IndexedBase` itself is considered listed under the {attr}`~sympy.core.basic.Basic.free_symbols` of the expression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[C, C[0, -1/2, 0], C[0, 1/2, 0], phi_12, phi_1^12, theta_12, theta_1^12]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "free_symbols = sorted(indexed_coefficient_expr.doit().free_symbols, key=str)\n",
    "free_symbols"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, not all symbols in the expression are of type {class}`~sympy.core.symbol.Symbol` anymore:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{C: (sympy.core.symbol.Symbol, True),\n",
       " C[0, -1/2, 0]: (sympy.tensor.indexed.Indexed, False),\n",
       " C[0, 1/2, 0]: (sympy.tensor.indexed.Indexed, False),\n",
       " phi_12: (sympy.core.symbol.Symbol, True),\n",
       " phi_1^12: (sympy.core.symbol.Symbol, True),\n",
       " theta_12: (sympy.core.symbol.Symbol, True),\n",
       " theta_1^12: (sympy.core.symbol.Symbol, True)}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{s: (type(s), isinstance(s, sp.Symbol)) for s in free_symbols}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "This will become problematic when lambdifying, because it results in an additional argument in the signature of the generated function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (C, Dummy_24, Dummy_23, phi_12, Dummy_26, theta_12, Dummy_25)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "func = sp.lambdify(free_symbols, indexed_coefficient_expr.doit())\n",
    "inspect.signature(func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A solution may be to use [`symplot.substitute_indexed_symbols()`](https://ampform.readthedocs.io/en/0.12.3/usage/symplot.html#symplot.substitute_indexed_symbols):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{C_{0,-1/2,0}, C_{0,1/2,0}, phi_12, phi_1^12, theta_12, theta_1^12}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indexed_coefficient_expr_symbols_only = symplot.substitute_indexed_symbols(\n",
    "    indexed_coefficient_expr.doit()\n",
    ")\n",
    "indexed_coefficient_expr_symbols_only.free_symbols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Signature (Dummy_30, Dummy_29, phi_12, Dummy_28, theta_12, Dummy_27)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = sorted(indexed_coefficient_expr_symbols_only.free_symbols, key=str)\n",
    "func = sp.lambdify(args, indexed_coefficient_expr_symbols_only)\n",
    "inspect.signature(func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{admonition} Caveat\n",
    "\n",
    "This seems clumsy, because [`substitute_indexed_symbols()`](https://ampform.readthedocs.io/en/0.12.3/usage/symplot.html#symplot.substitute_indexed_symbols) would have to be actively called before creating a computational function with [TensorWaves](https://tensorwaves.rtfd.io). It also becomes a hassle to keep track of the correct {class}`~sympy.core.symbol.Symbol` names in [`HelicityModel.parameter_defaults`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityModel.parameter_defaults).\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# One topology\n",
    "expr = ampform.get_builder(reaction).formulate().expression\n",
    "expr = remove_coefficients(expr.doit())\n",
    "sum_expr = formulate_intensity_with_coefficient(reaction)\n",
    "sum_expr = symplot.substitute_indexed_symbols(sum_expr.doit())\n",
    "sum_expr = remove_coefficients(sum_expr)\n",
    "assert sum_expr == expr\n",
    "\n",
    "# Two topologies\n",
    "expr = ampform.get_builder(reaction_two_resonances).formulate().expression\n",
    "expr = remove_coefficients(expr.doit())\n",
    "sum_expr = formulate_intensity_with_coefficient(reaction_two_resonances)\n",
    "sum_expr = symplot.substitute_indexed_symbols(sum_expr.doit())\n",
    "sum_expr = remove_coefficients(sum_expr)\n",
    "assert sum_expr == expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Inserting dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dynamics pose a challenge that is similar to {ref}`report/014:Inserting coefficients` in that we have to introduce expressions that are dependent on spin. Still, as can be seen from the available attributes on a [`TwoBodyKinematicVariableSet`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.dynamics.builder.html#ampform.dynamics.builder.TwoBodyKinematicVariableSet) (which serves as input to [`ResonanceDynamicsBuilder`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.dynamics.builder.html#ampform.dynamics.builder.ResonanceDynamicsBuilder)s), dynamics (currently) cannot depend on helicities.\n",
    "\n",
    "What may become a problem are $LS$-combinations. So far we have only considered a [`ReactionInfo`](https://qrules.readthedocs.io/en/0.9.6/api/qrules.transition.html#qrules.transition.ReactionInfo) that was created with `formalism=\"helicity\"`, but we also have to sum over $LS$-combinations when using `formalism=\"canonical-helicity\"`. This is particularly important when using dynamics with form factors, which depend on [`angular_momentum`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.dynamics.builder.html#ampform.dynamics.builder.TwoBodyKinematicVariableSet.angular_momentum).\n",
    "\n",
    "\n",
    ":::{note}\n",
    "\n",
    "The {class}`sympy.tensor.indexed.Indexed` now also contains the names of the resonances.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def formulate_intensity_with_dynamics(\n",
    "    reaction: ReactionInfo,\n",
    "    dynamics_choices: dict[str, ResonanceDynamicsBuilder],\n",
    "):\n",
    "    amplitudes = set()\n",
    "    for transition in reaction.transitions:\n",
    "        final_state_helicities = [\n",
    "            create_helicity_symbol(transition.topology, state_id)\n",
    "            for state_id in transition.final_states\n",
    "        ]\n",
    "        resonances = [\n",
    "            Str(s.particle.latex) for s in transition.intermediate_states.values()\n",
    "        ]\n",
    "        indices = [*final_state_helicities, *resonances]\n",
    "        coefficient = C[indices]\n",
    "        expr: sp.Expr = coefficient\n",
    "        for node_id in sorted(transition.topology.nodes):\n",
    "            expr *= formulate_wigner_d(transition, node_id)\n",
    "            decay = TwoBodyDecay.from_transition(transition, node_id)\n",
    "            parent_particle = decay.parent.particle\n",
    "            dynamics_builder = dynamics_choices.get(\n",
    "                parent_particle.name, create_non_dynamic\n",
    "            )\n",
    "            variables = _generate_kinematic_variable_set(transition, node_id)\n",
    "            dynamics, _ = dynamics_builder(parent_particle, variables)\n",
    "            expr *= dynamics\n",
    "        amplitudes.add(expr)\n",
    "    inner_helicities = get_helicities(reaction, which=\"inner\")\n",
    "    outer_helicities = get_helicities(reaction, which=\"outer\")\n",
    "    return PoolSum(\n",
    "        sp.Abs(\n",
    "            PoolSum(sum(amplitudes), *inner_helicities.items()),\n",
    "            evaluate=False,\n",
    "        )\n",
    "        ** 2,\n",
    "        *outer_helicities.items(),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output",
     "full-width"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{\\frac{\\Gamma_{\\Delta} m_{\\Delta} {C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2},\\Delta} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}}}}\\right|^{2}}$"
      ],
      "text/plain": [
       "PoolSum(Abs(PoolSum(\\Gamma_{\\Delta}*m_{\\Delta}*C[lambda_0, lambda_1, lambda_2, \\Delta]*WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0)/(-I*\\Gamma_{\\Delta}*m_{\\Delta} - m_12**2 + m_{\\Delta}**2), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formulate_intensity_with_dynamics(\n",
    "    reaction,\n",
    "    dynamics_choices={\n",
    "        resonance.name: create_relativistic_breit_wigner\n",
    "        for resonance in reaction.get_intermediate_particles()\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{\\frac{\\Gamma_{\\Delta} m_{\\Delta} {C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2},\\Delta} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} + \\frac{\\Gamma_{\\Lambda} m_{\\Lambda} {C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2},\\Lambda} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{2} + \\lambda_{3}}\\left(- \\phi_{01},\\theta_{01},0\\right) D^{\\frac{1}{2}}_{\\lambda_{3},\\lambda_{0} - \\lambda_{1}}\\left(- \\phi^{01}_{0},\\theta^{01}_{0},0\\right)}{- i \\Gamma_{\\Lambda} m_{\\Lambda} - m_{01}^{2} + m_{\\Lambda}^{2}}}}\\right|^{2}}$"
      ],
      "text/plain": [
       "PoolSum(Abs(PoolSum(\\Gamma_{\\Delta}*m_{\\Delta}*C[lambda_0, lambda_1, lambda_2, \\Delta]*WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0)/(-I*\\Gamma_{\\Delta}*m_{\\Delta} - m_12**2 + m_{\\Delta}**2) + \\Gamma_{\\Lambda}*m_{\\Lambda}*C[lambda_0, lambda_1, lambda_2, \\Lambda]*WignerD(1/2, lambda, -lambda_2 + lambda_3, -phi_01, theta_01, 0)*WignerD(1/2, lambda_3, lambda_0 - lambda_1, -phi_0^01, theta_0^01, 0)/(-I*\\Gamma_{\\Lambda}*m_{\\Lambda} - m_01**2 + m_{\\Lambda}**2), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formulate_intensity_with_dynamics(\n",
    "    reaction_two_resonances,\n",
    "    dynamics_choices={\n",
    "        resonance.name: create_relativistic_breit_wigner\n",
    "        for resonance in reaction_two_resonances.get_intermediate_particles()\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting amplitude is again identical to the original [`HelicityModel.expression`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel.expression):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# One topology\n",
    "b = ampform.get_builder(reaction)\n",
    "for p in reaction.get_intermediate_particles():\n",
    "    b.set_dynamics(p.name, create_relativistic_breit_wigner)\n",
    "expr = b.formulate().expression\n",
    "expr = remove_coefficients(expr.doit())\n",
    "sum_expr = formulate_intensity_with_dynamics(\n",
    "    reaction,\n",
    "    dynamics_choices={\n",
    "        resonance.name: create_relativistic_breit_wigner\n",
    "        for resonance in reaction.get_intermediate_particles()\n",
    "    },\n",
    ")\n",
    "sum_expr = symplot.substitute_indexed_symbols(sum_expr.doit())\n",
    "sum_expr = remove_coefficients(sum_expr)\n",
    "assert sum_expr == expr\n",
    "\n",
    "# Two topologies\n",
    "b = ampform.get_builder(reaction_two_resonances)\n",
    "for p in reaction_two_resonances.get_intermediate_particles():\n",
    "    b.set_dynamics(p.name, create_relativistic_breit_wigner)\n",
    "expr = b.formulate().expression\n",
    "expr = remove_coefficients(expr)\n",
    "sum_expr = formulate_intensity_with_dynamics(\n",
    "    reaction_two_resonances,\n",
    "    dynamics_choices={\n",
    "        resonance.name: create_relativistic_breit_wigner\n",
    "        for resonance in reaction_two_resonances.get_intermediate_particles()\n",
    "    },\n",
    ")\n",
    "sum_expr = symplot.partial_doit(sum_expr, doit_classes=(PoolSum,))\n",
    "sum_expr = symplot.partial_doit(sum_expr, doit_classes=(PoolSum,))  # recurse\n",
    "sum_expr = symplot.substitute_indexed_symbols(sum_expr)\n",
    "sum_expr = remove_coefficients(sum_expr)\n",
    "assert sum_expr.free_symbols == expr.free_symbols\n",
    "for intensity1, intensity2 in zip(sum_expr.args, expr.args):\n",
    "    # Annoyingly, Abs is rewritten with conjugates when using PoolSum...\n",
    "    amp1 = intensity1.args[0]\n",
    "    amp2 = intensity2.rewrite(sp.conjugate).args[0]\n",
    "    amp1 = sp.factor(amp1, deep=True, fraction=False)\n",
    "    amp2 = sp.factor(amp2, deep=True, fraction=False)\n",
    "    assert amp1 == amp2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(tr-014-solution-2)=\n",
    "## Solution 2: `Indexed` amplitude components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main problem with {ref}`tr-014-solution-2` is that it requires changing coefficient {class}`~sympy.core.symbol.Symbol`s to instances of {class}`~sympy.tensor.indexed.Indexed`, which have to be substituted using [`substitute_indexed_symbols()`](https://ampform.readthedocs.io/en/0.12.3/usage/symplot.html#symplot.substitute_indexed_symbols) (**after** calling {meth}`~sympy.core.basic.Basic.doit`).\n",
    "\n",
    "An alternative would be insert dynamics (and coefficients) into the `PoolSum`s over the helicities is to index the amplitude itself. The [`HelicityModel.expression`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityModel.expression) would then contain {class}`~sympy.tensor.indexed.Indexed` symbols that represent specific amplitudes. A definition of these amplitudes can be provided through [`HelicityModel.components`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityModel.components) or an equivalent attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from ampform.helicity.naming import HelicityAmplitudeNameGenerator\n",
    "\n",
    "A = sp.IndexedBase(R\"\\mathcal{A}\")\n",
    "\n",
    "\n",
    "def formulate_intensity_indexed_amplitudes_only(\n",
    "    reaction: ReactionInfo,\n",
    "    dynamics_choices: dict[str, ResonanceDynamicsBuilder],\n",
    ") -> tuple[sp.Expr, dict[sp.Indexed, sp.Expr]]:\n",
    "    name_generator = HelicityAmplitudeNameGenerator()\n",
    "    amplitudes = set()\n",
    "    amplitude_definitions = {}\n",
    "    for transition in reaction.transitions:\n",
    "        name_generator.register_amplitude_coefficient_name(transition)\n",
    "    for transition in reaction.transitions:\n",
    "        suffix = name_generator.generate_sequential_amplitude_suffix(transition)\n",
    "        expr: sp.Expr = sp.Symbol(f\"C_{{{suffix}}}\")\n",
    "        for node_id in sorted(transition.topology.nodes):\n",
    "            expr *= ampform.helicity.formulate_wigner_d(transition, node_id)\n",
    "            decay = TwoBodyDecay.from_transition(transition, node_id)\n",
    "            parent_particle = decay.parent.particle\n",
    "            dynamics_builder = dynamics_choices.get(\n",
    "                parent_particle.name, create_non_dynamic\n",
    "            )\n",
    "            variables = _generate_kinematic_variable_set(transition, node_id)\n",
    "            dynamics, _ = dynamics_builder(parent_particle, variables)\n",
    "            expr *= dynamics\n",
    "        resonances = [\n",
    "            Str(s.particle.latex) for s in transition.intermediate_states.values()\n",
    "        ]\n",
    "        helicity_symbols = [\n",
    "            create_helicity_symbol(transition.topology, state_id)\n",
    "            for state_id in sorted(transition.states)\n",
    "        ]\n",
    "        helicities = [\n",
    "            sp.Rational(transition.states[state_id].spin_projection)\n",
    "            for state_id in sorted(transition.states)\n",
    "        ]\n",
    "        amplitudes.add(A[[*helicity_symbols, *resonances]])\n",
    "        amplitude_definitions[A[[*helicities, *resonances]]] = expr\n",
    "    inner_helicities = get_helicities(reaction, which=\"inner\")\n",
    "    outer_helicities = get_helicities(reaction, which=\"outer\")\n",
    "    expression = PoolSum(\n",
    "        sp.Abs(\n",
    "            PoolSum(sum(amplitudes), *inner_helicities.items()),\n",
    "            evaluate=False,\n",
    "        )\n",
    "        ** 2,\n",
    "        *outer_helicities.items(),\n",
    "    )\n",
    "    return expression, amplitude_definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expression, amplitudes = formulate_intensity_indexed_amplitudes_only(\n",
    "    reaction_two_resonances,\n",
    "    dynamics_choices={\n",
    "        resonance.name: create_relativistic_breit_wigner\n",
    "        for resonance in reaction_two_resonances.get_intermediate_particles()\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "I = & \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{\\mathcal{A}}_{\\lambda,\\lambda_{0},\\lambda_{1},\\lambda_{2},\\lambda_{3},\\Delta} + {\\mathcal{A}}_{\\lambda,\\lambda_{0},\\lambda_{1},\\lambda_{2},\\lambda_{3},\\Lambda}}}\\right|^{2}} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "I = & \\left|{{\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \\\\\n",
       "& + \\left|{{\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \\\\\n",
       "& + \\left|{{\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \\\\\n",
       "& + \\left|{{\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "{\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Delta} = & \\frac{C_{\\Lambda_{c}^{+} \\to \\Delta_{-1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}} \\Gamma_{\\Delta} m_{\\Delta} D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "{\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Delta} = & \\frac{C_{\\Lambda_{c}^{+} \\to \\Delta_{+1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}} \\Gamma_{\\Delta} m_{\\Delta} D^{\\frac{1}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align*}\n",
       "{\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Delta} = & \\frac{C_{\\Lambda_{c}^{+} \\to \\Delta_{-1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}} \\Gamma_{\\Delta} m_{\\Delta} D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} \n",
       "\\end{align*}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Math(sp.multiline_latex(I, expression)))\n",
    "display(Math(sp.multiline_latex(I, expression.doit())))\n",
    "for i, (symbol, expr) in enumerate(amplitudes.items(), 1):\n",
    "    latex = sp.multiline_latex(symbol, expr)\n",
    "    display(Math(latex))\n",
    "    if i == 3:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\dots$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting amplitude is indeed identical to the original [`HelicityModel.expression`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel.expression):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "b = ampform.get_builder(reaction_two_resonances)\n",
    "for resonance in reaction_two_resonances.get_intermediate_particles():\n",
    "    b.set_dynamics(resonance.name, create_relativistic_breit_wigner)\n",
    "model_two_res = b.formulate()\n",
    "\n",
    "assert model_two_res.expression == expression.doit().xreplace(amplitudes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip}\n",
    "\n",
    "Currently, amplitudes with different resonances are put under a different amplitude symbol, identified by that resonance. Such resonances can be combined, e.g. $\\mathcal{A}_{\\lambda_i} = \\mathcal{A}_{\\lambda_i,\\Delta} + \\mathcal{A}_{\\lambda_i,\\Lambda}$. This would also make it easier to introduce correct interference terms through the [$K$-matrix](https://ampform.readthedocs.io/en/0.12.3/usage/dynamics/k-matrix.html).\n",
    "\n",
    ":::\n",
    "\n",
    ":::{admonition} Question\n",
    "\n",
    "The helicity of the intermediate state is also passed to the indexed amplitude. This is required for the coefficient name, which has a helicity subscript for the intermediate state, e.g. $C_{\\Lambda_{c}^{+} \\to \\Lambda_{\\color{red}+1/2} \\pi^{+}_{0}; \\Lambda \\to K^{-}_{0} p_{+1/2}}$. Does it really make sense to distinguish coefficients for different helicities of intermediate states?\n",
    "\n",
    ":::"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}