docs/usage/sympy.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell",
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "# WARNING: advised to install a specific version, e.g. ampform==0.1.2\n",
    "%pip install -q ampform[doc,viz] IPython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SymPy helper functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {mod}`ampform.sympy` module contains a few classes that make it easier to construct larger expressions that consist of several mathematical definitions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unevaluated expressions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {func}`.unevaluated` decorator makes it easier to write classes that represent a mathematical function definition. It makes a class that derives from {class}`sp.Expr <sympy.core.expr.Expr>` behave more like a {func}`~.dataclasses.dataclass` (see [PEP&nbsp;861](https://peps.python.org/pep-0681)). All you have to do is:\n",
    "\n",
    "1. Specify the arguments the function requires.\n",
    "2. Specify how to render the 'unevaluated' or 'folded' form of the expression with a `_latex_repr_` string or method.\n",
    "3. Specify how to unfold the expression using an `evaluate()` method.\n",
    "\n",
    "In the example below, we define a phase space factor $\\rho^\\text{CM}$ using the Chew-Mandelstam function (see PDG Resonances section, [Eq.&nbsp;(50.44)](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15)). For this, you need to define a break-up momentum $q$ as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "\n",
    "from ampform.sympy import unevaluated\n",
    "\n",
    "\n",
    "@unevaluated(real=False)\n",
    "class PhspFactorSWave(sp.Expr):\n",
    "    s: sp.Symbol\n",
    "    m1: sp.Symbol\n",
    "    m2: sp.Symbol\n",
    "    _latex_repr_ = R\"\\rho^\\text{{CM}}\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2 = self.args\n",
    "        q = BreakupMomentum(s, m1, m2)\n",
    "        cm = (\n",
    "            (2 * q / sp.sqrt(s))\n",
    "            * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n",
    "            - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n",
    "        ) / (16 * sp.pi**2)\n",
    "        return 16 * sp.pi * sp.I * cm\n",
    "\n",
    "\n",
    "@unevaluated(real=False)\n",
    "class BreakupMomentum(sp.Expr):\n",
    "    s: sp.Symbol\n",
    "    m1: sp.Symbol\n",
    "    m2: sp.Symbol\n",
    "    _latex_repr_ = R\"q\\left({s}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        s, m1, m2 = self.args\n",
    "        return sp.sqrt((s - (m1 + m2) ** 2) * (s - (m1 - m2) ** 2) / (s * 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, the LaTeX rendering of these classes makes them ideal for mathematically defining and building up larger amplitude models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from IPython.display import Math\n",
    "\n",
    "from ampform.io import aslatex\n",
    "\n",
    "s, m1, m2 = sp.symbols(\"s m1 m2\")\n",
    "q_expr = BreakupMomentum(s, m1, m2)\n",
    "rho_expr = PhspFactorSWave(s, m1, m2)\n",
    "Math(aslatex({e: e.evaluate() for e in [rho_expr, q_expr]}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Class variables and default arguments to instance arguments are also supported. They can either be indicated with {class}`typing.ClassVar` or by not providing a type hint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from typing import Any, ClassVar\n",
    "\n",
    "\n",
    "@unevaluated\n",
    "class FunkyPower(sp.Expr):\n",
    "    x: Any\n",
    "    m: int = 1\n",
    "    default_return: ClassVar[sp.Expr | None] = None\n",
    "    class_name = \"my name\"\n",
    "    _latex_repr_ = R\"f_{{{m}}}\\left({x}\\right)\"\n",
    "\n",
    "    def evaluate(self) -> sp.Expr | None:\n",
    "        if self.default_return is None:\n",
    "            return self.x**self.m\n",
    "        return self.default_return\n",
    "\n",
    "\n",
    "x = sp.Symbol(\"x\")\n",
    "exprs = (\n",
    "    FunkyPower(x),\n",
    "    FunkyPower(x, 2),\n",
    "    FunkyPower(x, m=3),\n",
    ")\n",
    "Math(aslatex({e: e.doit() for e in exprs}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FunkyPower.default_return = sp.Rational(0.5)\n",
    "Math(aslatex({e: e.doit() for e in exprs}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, instance attributes are converted ['sympified'](https://docs.sympy.org/latest/modules/core.html#module-sympy.core.sympify). To avoid this behavior, use the {func}`.argument` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Callable\n",
    "\n",
    "from ampform.sympy import argument\n",
    "\n",
    "\n",
    "class Transformation:\n",
    "    def __init__(self, power: int) -> None:\n",
    "        self.power = power\n",
    "\n",
    "    def __call__(self, x: sp.Basic, y: sp.Basic) -> sp.Expr:\n",
    "        return x + y**self.power\n",
    "\n",
    "\n",
    "@unevaluated\n",
    "class MyExpr(sp.Expr):\n",
    "    x: Any\n",
    "    y: Any\n",
    "    functor: Callable = argument(sympify=False)\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        return self.functor(self.x, self.y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the `functor` attribute has not been sympified (there is no SymPy equivalent for a callable object), but the `functor` can be called in the `evaluate()`/`doit()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a, b, k = sp.symbols(\"a b k\")\n",
    "expr = MyExpr(a, y=b, functor=Transformation(power=k))\n",
    "assert expr.x is a\n",
    "assert expr.y is b\n",
    "assert not isinstance(expr.functor, sp.Basic)\n",
    "Math(aslatex({expr: expr.doit()}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip}\n",
    "An example where this is used, is in the {class}`.EnergyDependentWidth` class, where we do not want to sympify the {attr}`~.EnergyDependentWidth.phsp_factor` protocol.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numerical integrals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In hadron physics and high-energy physics, it often happens that models contain integrals that do not have an analytical solution.. They can arise in theoretical models, complex scattering problems, or in the analysis of experimental data. In such cases, we need to resort to numerical integrations.\n",
    "\n",
    "SymPy provides the [`sympy.Integral`](https://docs.sympy.org/latest/modules/integrals/integrals.html#sympy.integrals.integrals.Integral) class, but this does not give us control over whether or not we want to avoid integrating the class analytically. An example of such an analytically unsolvable integral is shown below. Note that the integral does not evaluate despite the `doit()` call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "\n",
    "x, a, b = sp.symbols(\"x a b\")\n",
    "p = sp.Symbol(\"p\", positive=True)\n",
    "integral_expr = sp.Integral(sp.exp(x) / (x**p + 1), (x, a, b))\n",
    "integral_expr.doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For amplitude models that contain such integrals that should not be solved analytically, AmpForm provides the {class}`.UnevaluatableIntegral` class. It functions in the same way as [`sympy.Integral`](https://docs.sympy.org/latest/modules/integrals/integrals.html#sympy.integrals.integrals.Integral), but prevents the class from evaluating at all, even if the integral can be solved analytically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ampform.sympy import UnevaluatableIntegral\n",
    "\n",
    "UnevaluatableIntegral(x**p, (x, a, b)).doit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.Integral(x**p, (x, a, b)).doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This allows {class}`.UnevaluatableIntegral` to serve as a placeholder in expression trees that we call `doit` on when lambdifying to a numerical function. The resulting numerical function takes **complex-valued** and **multidimensional arrays** as function arguments.\n",
    "\n",
    "In the following, we see an example where the parameter $p$ inside the integral gets an array as input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "integral_expr = UnevaluatableIntegral(sp.exp(x) / (x**p + 1), (x, a, b))\n",
    "integral_func = sp.lambdify(args=[p, a, b], expr=integral_expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "a_val = 1.2\n",
    "b_val = 3.6\n",
    "p_array = np.array([0.4, 0.6, 0.8])\n",
    "\n",
    "areas = integral_func(p_array, a_val, b_val)\n",
    "areas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "x_area = np.linspace(a_val, b_val, num=100)\n",
    "x_line = np.linspace(0, 4, num=100)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_xlabel(\"$x$\")\n",
    "ax.set_ylabel(\"$x^p$\")\n",
    "\n",
    "for i, p_val in enumerate(p_array):\n",
    "    ax.plot(x_line, x_line**p_val, label=f\"$p={p_val}$\", c=f\"C{i}\")\n",
    "    ax.fill_between(x_area, x_area**p_val, alpha=(0.7 - i * 0.2), color=\"C0\")\n",
    "\n",
    "ax.text(\n",
    "    x=(a_val + b_val) / 2,\n",
    "    y=((a_val ** p_array[0] + b_val ** p_array[0]) / 2) * 0.5,\n",
    "    s=\"Area\",\n",
    "    horizontalalignment=\"center\",\n",
    "    verticalalignment=\"center\",\n",
    ")\n",
    "text_kwargs = dict(ha=\"center\", textcoords=\"offset points\", xytext=(0, -15))\n",
    "ax.annotate(\"a\", (a_val, 0.08), **text_kwargs)\n",
    "ax.annotate(\"b\", (b_val, 0.08), **text_kwargs)\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The arrays can be complex-valued as well. This is particularly useful when calculating dispersion integrals (see **[TR-003](https://compwa.github.io/report/003#general-dispersion-integral)**)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "integral_func(\n",
    "    p=np.array([1.5 - 8.6j, -4.6 + 5.5j]),\n",
    "    a=a_val,\n",
    "    b=b_val,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Summations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {class}`.PoolSum` class makes it possible to write sums over non-integer ranges. This is for instance useful when summing over allowed helicities. Here are some examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ampform.sympy import PoolSum\n",
    "\n",
    "i, j, m, n = sp.symbols(\"i j m n\")\n",
    "expr = PoolSum(i**m + j**n, (i, (-1, 0, +1)), (j, (2, 4, 5)))\n",
    "Math(aslatex({expr: expr.doit()}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "A = sp.IndexedBase(\"A\")\n",
    "λ, μ = sp.symbols(\"lambda mu\")\n",
    "to_range = lambda a, b: tuple(sp.Rational(i) for i in np.arange(a, b + 0.5))\n",
    "expr = abs(PoolSum(A[λ, μ], (λ, to_range(-0.5, +0.5)), (μ, to_range(-1, +1)))) ** 2\n",
    "Math(aslatex({expr: expr.doit()}))"
   ]
  }
 ],
 "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}