DOC: show how to collect Wigner D-functions

.cspell.json

@@ -93,6 +93,7 @@
     "maxdepth",
     "meshgrid",
     "mypy",
+    "mystnb",
     "nansum",
     "nbconvert",
     "nbformat",
@@ -119,6 +120,7 @@
     "tqdm",
     "unsrt",
     "venv",
+    "wigners",
     "wspace",
     "xlabel",
     "xreplace",

docs/collect-wigners.ipynb

@@ -0,0 +1,361 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Collect angular distribution symbols\n",
+    "\n",
+    "```{autolink-concat}\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Import Python libraries"
+    },
+    "tags": [
+     "hide-cell",
+     "scroll-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations\n",
+    "\n",
+    "import itertools\n",
+    "from typing import Iterable\n",
+    "\n",
+    "import qrules\n",
+    "import sympy as sp\n",
+    "from IPython.display import Latex, Markdown\n",
+    "\n",
+    "from ampform_dpd import DalitzPlotDecompositionBuilder, simplify_latex_rendering\n",
+    "from ampform_dpd.decay import (\n",
+    "    IsobarNode,\n",
+    "    Particle,\n",
+    "    ThreeBodyDecay,\n",
+    "    ThreeBodyDecayChain,\n",
+    ")\n",
+    "from ampform_dpd.io import as_markdown_table, aslatex, perform_cached_doit\n",
+    "from ampform_dpd.spin import filter_parity_violating_ls, generate_ls_couplings\n",
+    "\n",
+    "simplify_latex_rendering()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Define particles"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "PDG = qrules.load_pdg()\n",
+    "PARTICLE_DB = {\n",
+    "    p.name: Particle(\n",
+    "        name=p.name,\n",
+    "        latex=p.latex,\n",
+    "        spin=p.spin,\n",
+    "        parity=int(p.parity),\n",
+    "        mass=p.mass,\n",
+    "        width=p.width,\n",
+    "    )\n",
+    "    for p in PDG\n",
+    "    if p.parity is not None\n",
+    "}\n",
+    "PARTICLE_DB[\"Lambda(2000)\"] = Particle(\n",
+    "    name=\"Lambda(2000)\",\n",
+    "    latex=R\"\\Lambda(2000)\",\n",
+    "    spin=0.5,\n",
+    "    parity=-1,\n",
+    "    mass=2.0,\n",
+    "    width=(0.020 + 0.400) / 2,\n",
+    ")\n",
+    "Λc = PARTICLE_DB[\"Lambda(c)+\"]\n",
+    "p = PARTICLE_DB[\"p\"]\n",
+    "π = PARTICLE_DB[\"pi+\"]\n",
+    "K = PARTICLE_DB[\"K-\"]\n",
+    "PARTICLE_TO_ID = {Λc: 0, p: 1, π: 2, K: 3}\n",
+    "Markdown(as_markdown_table(list(PARTICLE_TO_ID)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Select resonances"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "resonance_names = [\n",
+    "    \"Lambda(1405)\",\n",
+    "    \"Lambda(1520)\",\n",
+    "    \"Lambda(1600)\",\n",
+    "    \"Lambda(1670)\",\n",
+    "    \"Lambda(1690)\",\n",
+    "    \"Lambda(2000)\",\n",
+    "    \"Delta(1232)+\",\n",
+    "    \"Delta(1600)+\",\n",
+    "    \"Delta(1700)+\",\n",
+    "    \"K(0)*(700)+\",\n",
+    "    \"K*(892)0\",\n",
+    "    \"K(2)*(1430)0\",\n",
+    "]\n",
+    "resonances = [PARTICLE_DB[name] for name in resonance_names]\n",
+    "Markdown(as_markdown_table(resonances))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Generate decay"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "def load_three_body_decay(\n",
+    "    resonance_names: Iterable[str],\n",
+    "    particle_definitions: dict[str, Particle],\n",
+    "    min_ls: bool = True,\n",
+    ") -> ThreeBodyDecay:\n",
+    "    resonances = [particle_definitions[name] for name in resonance_names]\n",
+    "    chains: list[ThreeBodyDecayChain] = []\n",
+    "    for res in resonances:\n",
+    "        chains.extend(_create_isobar(res, min_ls))\n",
+    "    return ThreeBodyDecay(\n",
+    "        states={state_id: particle for particle, state_id in PARTICLE_TO_ID.items()},\n",
+    "        chains=tuple(chains),\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def _create_isobar(resonance: Particle, min_ls: bool) -> list[ThreeBodyDecayChain]:\n",
+    "    if resonance.name.startswith(\"K\"):\n",
+    "        child1, child2, spectator = π, K, p\n",
+    "    elif resonance.name.startswith(\"L\"):\n",
+    "        child1, child2, spectator = K, p, π\n",
+    "    elif resonance.name.startswith(\"D\"):\n",
+    "        child1, child2, spectator = p, π, K\n",
+    "    else:\n",
+    "        raise NotImplementedError\n",
+    "    prod_ls_couplings = _generate_ls(Λc, resonance, spectator, conserve_parity=False)\n",
+    "    dec_ls_couplings = _generate_ls(resonance, child1, child2, conserve_parity=True)\n",
+    "    if min_ls:\n",
+    "        decay = IsobarNode(\n",
+    "            parent=Λc,\n",
+    "            child1=IsobarNode(\n",
+    "                parent=resonance,\n",
+    "                child1=child1,\n",
+    "                child2=child2,\n",
+    "                interaction=min(dec_ls_couplings),\n",
+    "            ),\n",
+    "            child2=spectator,\n",
+    "            interaction=min(prod_ls_couplings),\n",
+    "        )\n",
+    "        return [ThreeBodyDecayChain(decay)]\n",
+    "    chains = []\n",
+    "    for dec_ls, prod_ls in itertools.product(dec_ls_couplings, prod_ls_couplings):\n",
+    "        decay = IsobarNode(\n",
+    "            parent=Λc,\n",
+    "            child1=IsobarNode(\n",
+    "                parent=resonance,\n",
+    "                child1=child1,\n",
+    "                child2=child2,\n",
+    "                interaction=dec_ls,\n",
+    "            ),\n",
+    "            child2=spectator,\n",
+    "            interaction=prod_ls,\n",
+    "        )\n",
+    "        chains.append(ThreeBodyDecayChain(decay))\n",
+    "    return chains\n",
+    "\n",
+    "\n",
+    "def _generate_ls(\n",
+    "    parent: Particle, child1: Particle, child2: Particle, conserve_parity: bool\n",
+    ") -> list[tuple[int, sp.Rational]]:\n",
+    "    ls = generate_ls_couplings(parent.spin, child1.spin, child2.spin)\n",
+    "    if conserve_parity:\n",
+    "        return filter_parity_violating_ls(\n",
+    "            ls, parent.parity, child1.parity, child2.parity\n",
+    "        )\n",
+    "    return ls\n",
+    "\n",
+    "\n",
+    "DECAY = load_three_body_decay(\n",
+    "    resonance_names,\n",
+    "    particle_definitions=PARTICLE_DB,\n",
+    "    min_ls=True,\n",
+    ")\n",
+    "Latex(aslatex(DECAY, with_jp=True))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "mystnb": {
+     "code_prompt_show": "Define dynamics builder that returns a symbol only"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "def formulate_dynamics_symbol(\n",
+    "    decay_chain: ThreeBodyDecayChain,\n",
+    ") -> tuple[sp.Symbol, dict[sp.Symbol, float]]:\n",
+    "    s = sp.latex(_get_mandelstam_s(decay_chain))\n",
+    "    l_dec = sp.Rational(decay_chain.outgoing_ls.L)\n",
+    "    l_prod = sp.Rational(decay_chain.incoming_ls.L)\n",
+    "    s_dec = sp.Rational(decay_chain.outgoing_ls.S)\n",
+    "    s_prod = sp.Rational(decay_chain.incoming_ls.S)\n",
+    "    parameter_defaults = {}\n",
+    "    dynamics = sp.Symbol(\n",
+    "        Rf\"X_{{{decay_chain.resonance.spin}}}^{{{l_prod},{s_prod};{l_dec},{s_dec}}}\\left({s}\\right)\"\n",
+    "    )\n",
+    "    return dynamics, parameter_defaults\n",
+    "\n",
+    "\n",
+    "def _get_mandelstam_s(decay: ThreeBodyDecayChain) -> sp.Symbol:\n",
+    "    s1, s2, s3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n",
+    "    m1, m2, m3 = map(_to_mass_symbol, [p, π, K])\n",
+    "    decay_masses = {_to_mass_symbol(p) for p in decay.decay_products}\n",
+    "    if decay_masses == {m2, m3}:\n",
+    "        return s1\n",
+    "    if decay_masses == {m1, m3}:\n",
+    "        return s2\n",
+    "    if decay_masses == {m1, m2}:\n",
+    "        return s3\n",
+    "    raise NotImplementedError(\n",
+    "        f\"Cannot find Mandelstam variable for {''.join(decay_masses)}\"\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def _to_mass_symbol(particle: Particle) -> sp.Symbol:\n",
+    "    state_id = PARTICLE_TO_ID.get(particle)\n",
+    "    if state_id is not None:\n",
+    "        return sp.Symbol(f\"m{state_id}\", nonnegative=True)\n",
+    "    return sp.Symbol(f\"m_{{{particle.latex}}}\", nonnegative=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "full-width",
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "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_dynamics_symbol)\n",
+    "model = model_builder.formulate(reference_subsystem=1)\n",
+    "model.intensity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "mystnb": {
+     "code_prompt_show": "Set all couplings to one"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "couplings_to_one = {s: 1 for s in model.parameter_defaults if \"mathcal\" in str(s)}\n",
+    "\n",
+    "full_amplitudes = {k: perform_cached_doit(v) for k, v in model.amplitudes.items()}\n",
+    "full_expression = perform_cached_doit(model.intensity).xreplace(full_amplitudes)\n",
+    "full_expression = full_expression.doit().xreplace(couplings_to_one)\n",
+    "dynamic_symbols = sorted(\n",
+    "    (s for s in full_expression.free_symbols if str(s).startswith(\"X\")),\n",
+    "    key=lambda s: (str(s)[-10:], str(s)),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "mystnb": {
+     "code_prompt_show": "Collect Wigner-D functions"
+    },
+    "tags": [
+     "hide-input",
+     "full-width"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "src = R\"\"\"\n",
+    "\\begin{eqnarray}\n",
+    "\"\"\"\n",
+    "for x in dynamic_symbols:\n",
+    "    filter_substitutions = {s: 1 if s == x else 0 for s in dynamic_symbols}\n",
+    "    factor = full_expression.xreplace(filter_substitutions).simplify()\n",
+    "    src += Rf\"{sp.latex(x)} &:& {sp.latex(factor)} \\\\\" \"\\n\"\n",
+    "src += R\"\\end{eqnarray}\"\n",
+    "Latex(src)"
+   ]
+  }
+ ],
+ "metadata": {
+  "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.8.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}