2021.12.13/ampform.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AmpForm demo PANDA Seminar December 2021"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook accompanies [these slides](https://docs.google.com/presentation/d/e/2PACX-1vSymz5AjdhPw4Kz1pKhdFMnFGYuQvVaC8WbV_HTg770x6RDYoP-Anv9tn88DSuzvSiiQ9F4pcDGVExv/pub). They were presented during a PANDA Seminar on 13 December 2021.\n",
    "\n",
    "Related notebooks for this presentation:\n",
    "- [QRules demo](./qrules.ipynb)\n",
    "- [TensorWaves demo](./tensorwaves.ipynb)\n",
    "\n",
    "For more extensive examples, see **[ampform.rtfd.io](https://ampform.readthedocs.io)**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "from functools import cache\n",
    "\n",
    "import ampform\n",
    "import graphviz\n",
    "import ipywidgets\n",
    "import qrules\n",
    "import symplot\n",
    "import sympy as sp\n",
    "from ampform.dynamics import BlattWeisskopfSquared, EnergyDependentWidth\n",
    "from ampform.dynamics.builder import (\n",
    "    create_analytic_breit_wigner,\n",
    "    create_relativistic_breit_wigner_with_ff,\n",
    ")\n",
    "from ampform.dynamics.kmatrix import NonRelativisticKMatrix, RelativisticKMatrix\n",
    "from ampform.sympy import UnevaluatedExpression\n",
    "from IPython.display import Math, display\n",
    "from ipywidgets import Checkbox, SelectionSlider, SelectMultiple, ToggleButtons\n",
    "from sympy.physics.quantum.spin import WignerD\n",
    "\n",
    "\n",
    "class Symbol(sp.Symbol):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## $K$-matrix expressions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = sp.Symbol(\"n_R\")\n",
    "matrix = RelativisticKMatrix.formulate(\n",
    "    n_channels=1,\n",
    "    n_poles=n,\n",
    ")\n",
    "matrix[0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix = NonRelativisticKMatrix.formulate(\n",
    "    n_channels=2,\n",
    "    n_poles=1,\n",
    ").doit()\n",
    "matrix[0, 0].simplify()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example $D^0 \\to K^0 K^- K^+$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate transitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=\"D0\",\n",
    "    final_state=[\"K0\", \"K-\", \"K+\"],\n",
    "    allowed_intermediate_particles=[\"a(0)(980)0\", \"a(0)(1450)0\", \"phi(1020)\"],\n",
    "    formalism=\"helicity\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(reaction, collapse_graphs=True, render_final_state_id=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that one of the resonances, $a_0(980)$, lies **below threshold**, which means we should parametrize it with an analytic continuation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PDG = qrules.load_pdg()\n",
    "PDG[\"a(0)(980)0\"].mass < PDG[\"K-\"].mass + PDG[\"K-\"].mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = ampform.get_builder(reaction)\n",
    "resonances = reaction.get_intermediate_particles()\n",
    "for p in resonances:\n",
    "    builder.set_dynamics(p.name, create_relativistic_breit_wigner_with_ff)\n",
    "builder.set_dynamics(\"a(0)(980)0\", create_analytic_breit_wigner)\n",
    "model = builder.formulate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "amplitude_expressions = {\n",
    "    expr: sp.Symbol(name)\n",
    "    for name, expr in model.components.items()\n",
    "    if name.startswith(\"A\")\n",
    "}\n",
    "top_expr = model.expression.xreplace(amplitude_expressions)\n",
    "top_expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examine one of the amplitudes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "some_amplitude = model.components[\n",
    "    R\"A_{D^{0}_{0} \\to K^{0}_{0} \\phi(1020)_{0}; \\phi(1020)_{0} \\to K^{+}_{0} K^{-}_{0}}\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "def round_nested(expression, n_decimals):\n",
    "    for node in sp.preorder_traversal(expression):\n",
    "        if node.free_symbols:\n",
    "            continue\n",
    "        if isinstance(node, float | sp.Float):\n",
    "            expression = expression.xreplace({node: node.n(n_decimals)})\n",
    "    return expression\n",
    "\n",
    "\n",
    "A_step1 = some_amplitude\n",
    "A_step2 = symplot.partial_doit(A_step1, (BlattWeisskopfSquared, WignerD))\n",
    "A_step2 = symplot.partial_doit(A_step2, EnergyDependentWidth)\n",
    "mass_substitutions = {\n",
    "    sp.Symbol(f\"m_{i}\", real=True): p.mass for i, p in reaction.final_state.items()\n",
    "}\n",
    "substitutions = {\n",
    "    sp.sqrt(2): sp.sqrt(2).n(),\n",
    "    **mass_substitutions,\n",
    "    **model.parameter_defaults,\n",
    "}\n",
    "A_step3 = some_amplitude.doit().xreplace(substitutions)\n",
    "A_step3 = round_nested(A_step3, n_decimals=2)\n",
    "with sp.assuming(*map(sp.Q.positive, A_step3.free_symbols)):\n",
    "    A_step3 = A_step3.refine()\n",
    "display(\n",
    "    A_step1,\n",
    "    Math(\"=\" + sp.latex(A_step2)),\n",
    "    Math(\"=\" + sp.latex(A_step3)),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize expression tree\n",
    "<!-- cspell:ignore bgcolor dodgerblue fillcolor indianred -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "dot_style = [\n",
    "    (sp.Basic, {\"style\": \"filled\", \"fillcolor\": \"white\"}),\n",
    "    (sp.Atom, {\"color\": \"gray\", \"style\": \"filled\", \"fillcolor\": \"white\"}),\n",
    "    (sp.Symbol, {\"color\": \"dodgerblue1\"}),\n",
    "    ((UnevaluatedExpression, WignerD), {\"color\": \"indianred2\"}),\n",
    "]\n",
    "dot = sp.dotprint(A_step3, styles=dot_style, size=10)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "aliases = [\"phi\", \"a(0)(980)\", \"a(0)(1450)\", \"full model\"]\n",
    "amplitudes = dict(zip(aliases, model.components.values(), strict=True))\n",
    "\n",
    "\n",
    "def partial_doit(expr, classes):\n",
    "    for sub_expr in sp.preorder_traversal(expr):\n",
    "        if type(sub_expr).__name__ in classes:\n",
    "            new_sub_expr = sub_expr.doit(deep=False)\n",
    "            new_sub_expr = partial_doit(new_sub_expr, classes)\n",
    "            expr = expr.xreplace({sub_expr: new_sub_expr})\n",
    "    return expr\n",
    "\n",
    "\n",
    "@cache\n",
    "def get_amplitude(\n",
    "    name,\n",
    "    doit,\n",
    "    substitute_masses,\n",
    "    substitute_parameters,\n",
    "    round_floats,\n",
    "    unfolded_classes,\n",
    "):\n",
    "    expr = amplitudes[name]\n",
    "    if len(unfolded_classes):\n",
    "        expr = partial_doit(expr, unfolded_classes)\n",
    "    if doit:\n",
    "        expr = expr.doit()\n",
    "    if substitute_masses:\n",
    "        expr = expr.subs(mass_substitutions)\n",
    "    if substitute_parameters:\n",
    "        expr = expr.subs(model.parameter_defaults)\n",
    "    if round_floats:\n",
    "        expr = round_nested(expr, n_decimals=2)\n",
    "    return expr\n",
    "\n",
    "\n",
    "@cache\n",
    "def create_graph(expression, fig_size, visualize_cse):\n",
    "    if fig_size == \"full\":\n",
    "        fig_size = None\n",
    "    dot = sp.dotprint(\n",
    "        expression, styles=dot_style, size=fig_size, repeat=not visualize_cse\n",
    "    )\n",
    "    return graphviz.Source(dot)\n",
    "\n",
    "\n",
    "@ipywidgets.interact(\n",
    "    amplitude=ToggleButtons(\n",
    "        description=\"Amplitude:\", value=aliases[2], options=list(amplitudes)\n",
    "    ),\n",
    "    rendering=ToggleButtons(\n",
    "        description=\"Rendering:\", options=[\"tree\", \"math\", \"unicode\"]\n",
    "    ),\n",
    "    doit=Checkbox(description=\"Unfold all nodes\"),\n",
    "    substitute_parameters=Checkbox(description=\"Substitute parameters\"),\n",
    "    substitute_masses=Checkbox(description=\"Substitute stable masses\"),\n",
    "    visualize_cse=Checkbox(description=\"Common sub-expressions\"),\n",
    "    round_floats=Checkbox(description=\"Round floats\", value=True),\n",
    "    fig_size=SelectionSlider(\n",
    "        description=\"Figure size\", options=[\"full\", 5, 7, 10, 15, 20, 50]\n",
    "    ),\n",
    "    unfolded_classes=SelectMultiple(\n",
    "        description=\"Unfold:\",\n",
    "        options=[\n",
    "            \"BlattWeisskopfSquared\",\n",
    "            \"BreakupMomentumSquared\",\n",
    "            \"EnergyDependentWidth\",\n",
    "            \"PhaseSpaceFactor\",\n",
    "            \"PhaseSpaceFactorAbs\",\n",
    "            \"PhaseSpaceFactorAnalytic\",\n",
    "            \"WignerD\",\n",
    "        ],\n",
    "        rows=7,\n",
    "    ),\n",
    ")\n",
    "def visualize(\n",
    "    rendering,\n",
    "    amplitude,\n",
    "    substitute_masses,\n",
    "    substitute_parameters,\n",
    "    round_floats,\n",
    "    doit,\n",
    "    visualize_cse,\n",
    "    unfolded_classes,\n",
    "    fig_size,\n",
    "):\n",
    "    expression = get_amplitude(\n",
    "        amplitude,\n",
    "        doit,\n",
    "        substitute_masses,\n",
    "        substitute_parameters,\n",
    "        round_floats,\n",
    "        unfolded_classes,\n",
    "    )\n",
    "    if rendering == \"unicode\":\n",
    "        print(sp.pretty(expression))\n",
    "    elif rendering == \"math\":\n",
    "        display(expression)\n",
    "    else:\n",
    "        graph = create_graph(expression, fig_size, visualize_cse)\n",
    "        display(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Expression manipulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "expr = amplitudes[\"a(0)(1450)\"].doit()\n",
    "graph_style = dict(\n",
    "    bgcolor=\"transparent\",\n",
    "    size=15,\n",
    "    styles=[\n",
    "        *dot_style,\n",
    "        (Symbol, {\"fillcolor\": \"lightblue\", \"style\": \"filled\"}),\n",
    "    ],\n",
    ")\n",
    "dot = sp.dotprint(expr, **graph_style)\n",
    "display(expr)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "mass_symbols = {s: Symbol(s.name, **s.assumptions0) for s in mass_substitutions}\n",
    "colored_symbols = {\n",
    "    s: sp.Symbol(Rf\"\\color{{RoyalBlue}}{{{s.name}}}\", **s.assumptions0)\n",
    "    for s in mass_substitutions\n",
    "}\n",
    "dot = sp.dotprint(expr.xreplace(mass_symbols), **graph_style)\n",
    "display(expr.xreplace(colored_symbols))\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "expr_subs = expr.subs(mass_substitutions)\n",
    "expr_subs = round_nested(expr_subs, n_decimals=2)\n",
    "node = next(s for s in sp.preorder_traversal(expr_subs) if isinstance(s, sp.Float))\n",
    "dot = sp.dotprint(expr_subs.subs(node, Symbol(f\"{node}\")), **graph_style)\n",
    "display(expr_subs.subs(node, sp.Symbol(Rf\"\\color{{RoyalBlue}}{node}\")))\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Symbolic kinematic variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kinematic_variables = sorted(A_step3.free_symbols, key=str)\n",
    "kinematic_variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_1_12 = kinematic_variables[1]\n",
    "model.kinematic_variables[theta_1_12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.kinematic_variables[theta_1_12].doit()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Pyproject Local",
   "language": "python",
   "name": "pyproject_local_kernel"
  },
  "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}