docs/report/001.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "lambdification",
     "sympy"
    ]
   },
   "source": [
    "::::{margin}\n",
    ":::{card} Custom lambdification\n",
    "TR-001\n",
    "^^^\n",
    "See also {doc}`SymPy's tutorial page on the printing modules <sympy:modules/printing>`.\n",
    "+++\n",
    "✅&nbsp;[ampform#72](https://github.com/ComPWA/ampform/issues/72), [tensorwaves#284](https://github.com/ComPWA/tensorwaves/issues/284)\n",
    ":::\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom lambdification\n",
    "<!-- cspell:disable -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q black==21.5b2 jax==0.4.28 jaxlib==0.4.28 matplotlib==3.4.2 numpy==1.26.4 sympy==1.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import inspect\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from black import FileMode, format_str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overwrite printer methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As noted in [TR-000](000.ipynb), it's hard to lambdify a {func}`sympy.sqrt <sympy.functions.elementary.miscellaneous.sqrt>` to {doc}`JAX <jax:index>`. One possible way out is to define a custom class that derives from {class}`sympy.Expr <sympy.core.expr.Expr>` and {doc}`overwrite its printer methods <sympy:modules/printing>`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy.printing.printer import Printer\n",
    "\n",
    "\n",
    "class ComplexSqrt(sp.Expr):\n",
    "    def __new__(cls, x, *args, **kwargs):\n",
    "        x = sp.sympify(x)\n",
    "        expr = sp.Expr.__new__(cls, x, *args, **kwargs)\n",
    "        if hasattr(x, \"free_symbols\") and not x.free_symbols:\n",
    "            return expr.evaluate()\n",
    "        return expr\n",
    "\n",
    "    def evaluate(self):\n",
    "        x = self.args[0]\n",
    "        if not x.is_real:\n",
    "            return sp.sqrt(x)\n",
    "        return sp.Piecewise(\n",
    "            (sp.I * sp.sqrt(-x), x < 0),\n",
    "            (sp.sqrt(x), True),\n",
    "        )\n",
    "\n",
    "    def _latex(self, printer: Printer, *args) -> str:\n",
    "        x = printer._print(self.args[0])\n",
    "        return Rf\"\\sqrt[\\mathrm{{c}}]{{{x}}}\"\n",
    "\n",
    "    def _numpycode(self, printer: Printer, *args) -> str:\n",
    "        printer.module_imports[\"numpy.lib\"].add(\"scimath\")\n",
    "        x = printer._print(self.args[0])\n",
    "        return f\"scimath.sqrt({x})\"\n",
    "\n",
    "    def _pythoncode(self, printer: Printer, *args) -> str:\n",
    "        printer.module_imports[\"cmath\"].add(\"sqrt as csqrt\")\n",
    "        x = printer._print(self.args[0])\n",
    "        return f\"csqrt({x})\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As opposed to the {doc}`derivation of a sympy.Expr </adr/002/expr>`, this class evaluates directly, because the `evaluate` key-word argument is not used processed by the `__new__` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 2 i$"
      ],
      "text/plain": [
       "2*I"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ComplexSqrt(-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `_latex()` method ensures that `ComplexSqrt` renders nicely in notebooks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sqrt[\\mathrm{c}]{x}$"
      ],
      "text/plain": [
       "ComplexSqrt(x)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = sp.Symbol(\"x\")\n",
    "ComplexSqrt(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot custom class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, one may modify this `Lambdifier` class, so that {func}`sympy.plot() <sympy.plotting.plot.plot>` also works on this custom class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy.plotting.experimental_lambdify import Lambdifier\n",
    "\n",
    "Lambdifier.builtin_functions_different[\"ComplexSqrt\"] = \"sqrt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "\n",
    "x = sp.Symbol(\"x\")\n",
    "expr = ComplexSqrt(x)\n",
    "p1 = sp.plot(sp.re(expr), (x, -1, 2), show=False, line_color=\"red\")\n",
    "p2 = sp.plot(sp.im(expr), (x, -1, 2), show=False)\n",
    "p1.append(p2[0])\n",
    "p1.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164990419-9cd03001-d6f3-44b1-a8f9-beed2c6bf69b.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lambdifying"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The important part, lambdifying to {mod}`numpy` or {mod}`math` works well as well now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return (csqrt(x))\n",
      "\n"
     ]
    }
   ],
   "source": [
    "lambdified_py = sp.lambdify(x, ComplexSqrt(x), \"math\")\n",
    "source = inspect.getsource(lambdified_py)\n",
    "print(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return (scimath.sqrt(x))\n",
      "\n"
     ]
    }
   ],
   "source": [
    "numpy_lambdified = sp.lambdify(x, ComplexSqrt(x), \"numpy\")\n",
    "source = inspect.getsource(numpy_lambdified)\n",
    "print(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.        +1.j        , 0.        +0.70710678j,\n",
       "       0.        +0.j        , 0.70710678+0.j        ,\n",
       "       1.        +0.j        ])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample = np.linspace(-1, +1, 5)\n",
    "numpy_lambdified(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just as noted in {ref}`report/000:Complex square root` though, {mod}`numpy.emath` is not provided by the NumPy API of {doc}`JAX <jax:index>`. As discussed there, we can at most decorate the {mod}`numpy.emath` version with {func}`jax.jit` and work with static arguments only:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.+1.j, dtype=complex64)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax_lambdified = jax.jit(numpy_lambdified, backend=\"cpu\", static_argnums=0)\n",
    "jax_lambdified(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In that case, unhashable (non-static) input samples are still not accepted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class 'numpy.ndarray'>, [-1.  -0.5  0.   0.5  1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-94f797bd7204>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjax_lambdified\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class 'numpy.ndarray'>, [-1.  -0.5  0.   0.5  1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n"
     ]
    }
   ],
   "source": [
    "jax_lambdified(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Handle for JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As concluded in {ref}`report/000:Conditional square root`, the alternative to lambdify to {mod}`numpy.emath` is to lambdify to {func}`numpy.select`. This has some caveats, though, like that you should not use `__dict__`. Worse, JAX is not immediately supported as backend. Fortunately, we now know how to overwrite lambdify methods.\n",
    "\n",
    "An additional tool we need now is to {doc}`define a new printer class <sympy:modules/printing>` for JAX, so that we can also define a special rendering method for `ComplexSqrt` in the case of JAX. Most of its printing methods should be the same as that of SymPy's `NumPyPrinter`, the rest we can overwrite:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "Alternative would be to add a method `_jaxcode` to the `ComplexSqrt` class above. See {doc}`sympy:modules/printing`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sympy.printing.numpy import NumPyPrinter\n",
    "\n",
    "\n",
    "class JaxPrinter(NumPyPrinter):\n",
    "    _module = \"jax\"\n",
    "\n",
    "    def _print_ComplexSqrt(self, expr: sp.Expr) -> str:\n",
    "        arg = expr.args[0]\n",
    "        x = self._print(arg)\n",
    "        return (\n",
    "            f\"select([less({x}, 0), True], [1j * sqrt(-{x}), sqrt({x})],\"\n",
    "            \" default=nan,)\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numpy_expr = sp.lambdify(x, ComplexSqrt(x), modules=np, printer=JaxPrinter)\n",
    "source = inspect.getsource(numpy_expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return select(\n",
      "        [less(x, 0), True],\n",
      "        [1j * sqrt(-x), sqrt(x)],\n",
      "        default=nan,\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(format_str(source, mode=FileMode()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jax_expr = sp.lambdify(x, ComplexSqrt(x), modules=jnp, printer=JaxPrinter)\n",
    "source = inspect.getsource(jax_expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return select(\n",
      "        [less(x, 0), True],\n",
      "        [1j * sqrt(-x), sqrt(x)],\n",
      "        default=nan,\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(format_str(source, mode=FileMode()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-stderr",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray([0.        +1.j        , 0.        +0.70710677j,\n",
       "             0.        +0.j        , 0.70710677+0.j        ,\n",
       "             1.        +0.j        ], dtype=complex64)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax_expr(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The lambdified function can of course also be decorated with {func}`jax.jit`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jit_expr = jax.jit(jax_expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performance check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng()\n",
    "sample = rng.normal(size=1_000_000)\n",
    "jax_sample = jnp.array(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip} section\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.91 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit jit_expr(jax_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.31 ms ± 42.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit jax_expr(jax_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output",
     "remove-stderr"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<lambdifygenerated-3>:2: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return (select([less(x, 0), True], [1j * sqrt(-x), sqrt(x)], default=nan,))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16.9 ms ± 614 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit numpy_expr(sample)"
   ]
  }
 ],
 "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"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}