{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "remove-cell" ] }, "outputs": [], "source": [ "%config Completer.use_jedi = False\n", "%config InlineBackend.figure_formats = ['svg']" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# SymPy" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Some useful SymPy pages:\n", "- [Tutorial page](https://docs.sympy.org/latest/tutorial/index.html)\n", "- [SymPy API](https://docs.sympy.org/latest/modules/index.html)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Defining the amplitude model in terms of SymPY" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Parameters would become a mapping of `Symbol`s to initial, suggested values and dynamics would be a mapping of `Symbol`s to 'suggested' expressions. Intensity will be the eventual combined expression." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "hide-cell", "keep_output" ] }, "outputs": [], "source": [ "# !pip install matplotlib==3.3.3 sympy==1.7.1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [], "source": [ "from typing import Dict\n", "\n", "import attr\n", "import sympy as sp\n", "\n", "\n", "@attr.s\n", "class AmplitudeModel:\n", " initial_values: Dict[sp.Symbol, float] = attr.ib(default={})\n", " dynamics: Dict[sp.Symbol, sp.Function] = attr.ib(default={})\n", " intensity: sp.Expr = attr.ib(default=None)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "There needs to be one symbol $x$ that represents the four-momentum input:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [], "source": [ "x = sp.Symbol(\"x\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "As an example, let's create an `AmplitudeModel` with an `intensity` that is a sum of Gaussians. Each Gaussian here takes the rôle of a dynamics function:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [], "source": [ "model = AmplitudeModel()\n", "\n", "N_COMPONENTS = 3\n", "for i in range(1, N_COMPONENTS + 1):\n", " mu = sp.Symbol(fR\"\\mu_{i}\")\n", " sigma = sp.Symbol(fR\"\\sigma_{i}\")\n", " model.initial_values.update(\n", " {\n", " mu: float(i),\n", " sigma: 1 / (2 * i),\n", " }\n", " )\n", " gauss = sp.exp(-((x - mu) ** 2) / (sigma ** 2)) / (\n", " sigma * sp.sqrt(2 * sp.pi)\n", " )\n", " dyn_symbol = sp.Symbol(fR\"\\mathrm{{dyn}}_{i}\")\n", " model.dynamics[dyn_symbol] = gauss\n", "\n", "coherent_sum = sum(model.dynamics)\n", "model.intensity = coherent_sum" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/plain": [ "{\\mu_1: 1.0,\n", " \\sigma_1: 0.5,\n", " \\mu_2: 2.0,\n", " \\sigma_2: 0.25,\n", " \\mu_3: 3.0,\n", " \\sigma_3: 0.16666666666666666}" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "model.initial_values" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\mathrm{dyn}_1 + \\mathrm{dyn}_2 + \\mathrm{dyn}_3$" ], "text/plain": [ "\\mathrm{dyn}_1 + \\mathrm{dyn}_2 + \\mathrm{dyn}_3" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "model.intensity" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Dynamics are inserted into the intensity expression of the model:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\frac{\\sqrt{2} e^{- \\frac{\\left(- \\mu_{3} + x\\right)^{2}}{\\sigma_{3}^{2}}}}{2 \\sqrt{\\pi} \\sigma_{3}} + \\frac{\\sqrt{2} e^{- \\frac{\\left(- \\mu_{2} + x\\right)^{2}}{\\sigma_{2}^{2}}}}{2 \\sqrt{\\pi} \\sigma_{2}} + \\frac{\\sqrt{2} e^{- \\frac{\\left(- \\mu_{1} + x\\right)^{2}}{\\sigma_{1}^{2}}}}{2 \\sqrt{\\pi} \\sigma_{1}}$" ], "text/plain": [ "sqrt(2)*exp(-(-\\mu_3 + x)**2/\\sigma_3**2)/(2*sqrt(pi)*\\sigma_3) + sqrt(2)*exp(-(-\\mu_2 + x)**2/\\sigma_2**2)/(2*sqrt(pi)*\\sigma_2) + sqrt(2)*exp(-(-\\mu_1 + x)**2/\\sigma_1**2)/(2*sqrt(pi)*\\sigma_1)" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "model.intensity.subs(model.dynamics)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "And, for evaluating, the 'suggested' initial parameter values are inserted:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\frac{1.0 \\sqrt{2} e^{- 4.0 \\left(x - 1.0\\right)^{2}}}{\\sqrt{\\pi}} + \\frac{2.0 \\sqrt{2} e^{- 64.0 \\left(0.5 x - 1\\right)^{2}}}{\\sqrt{\\pi}} + \\frac{3.0 \\sqrt{2} e^{- 324.0 \\left(0.333333333333333 x - 1\\right)^{2}}}{\\sqrt{\\pi}}$" ], "text/plain": [ "1.0*sqrt(2)*exp(-4.0*(x - 1.0)**2)/sqrt(pi) + 2.0*sqrt(2)*exp(-64.0*(0.5*x - 1)**2)/sqrt(pi) + 3.0*sqrt(2)*exp(-324.0*(0.333333333333333*x - 1)**2)/sqrt(pi)" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "model.intensity.subs(model.dynamics).subs(model.initial_values)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Here's a small helper function to plot this model:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "hide-input", "keep_output" ] }, "outputs": [], "source": [ "def plot_model(model: AmplitudeModel) -> None:\n", " total_plot = sp.plotting.plot(\n", " model.intensity.subs(model.dynamics).subs(model.initial_values),\n", " (x, 0, 4),\n", " show=False,\n", " line_color=\"black\",\n", " )\n", " p1 = sp.plotting.plot(\n", " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_1\")].subs(\n", " model.initial_values\n", " ),\n", " (x, 0, 4),\n", " line_color=\"red\",\n", " show=False,\n", " )\n", " p2 = sp.plotting.plot(\n", " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_2\")].subs(\n", " model.initial_values\n", " ),\n", " (x, 0, 4),\n", " line_color=\"blue\",\n", " show=False,\n", " )\n", " p3 = sp.plotting.plot(\n", " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_3\")].subs(\n", " model.initial_values\n", " ),\n", " (x, 0, 4),\n", " line_color=\"green\",\n", " show=False,\n", " )\n", " total_plot.extend(p1)\n", " total_plot.extend(p2)\n", " total_plot.extend(p3)\n", " total_plot.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "image/svg+xml": [ "\n", "\n", "\n", "\n", " \n", " \n", " \n", " \n", " 2021-03-23T09:07:26.470246\n", " image/svg+xml\n", " \n", " \n", " Matplotlib v3.3.4, https://matplotlib.org/\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "plot_model(model)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we can couple parameters like this:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "image/svg+xml": [ "\n", "\n", "\n", "\n", " \n", " \n", " \n", " \n", " 2021-03-23T09:07:26.790497\n", " image/svg+xml\n", " \n", " \n", " Matplotlib v3.3.4, https://matplotlib.org/\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" }, { "data": { "image/svg+xml": [ "\n", "\n", "\n", "\n", " \n", " \n", " \n", " \n", " 2021-03-23T09:07:27.044751\n", " image/svg+xml\n", " \n", " \n", " Matplotlib v3.3.4, https://matplotlib.org/\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "model.initial_values[sp.Symbol(R\"\\sigma_1\")] = sp.Symbol(R\"\\sigma_3\")\n", "plot_model(model)\n", "model.initial_values[sp.Symbol(R\"\\sigma_3\")] = 1\n", "plot_model(model)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "And it's also possible to insert custom dynamics:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "image/svg+xml": [ "\n", "\n", "\n", "\n", " \n", " \n", " \n", " \n", " 2021-03-23T09:07:27.265598\n", " image/svg+xml\n", " \n", " \n", " Matplotlib v3.3.4, https://matplotlib.org/\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_3\")] = sp.sqrt(x)\n", "plot_model(model)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Implementation in TensorWaves" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Credits [@spflueger](https://github.com/spflueger)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-cell", "keep_output" ] }, "outputs": [], "source": [ "# !pip install jax==0.2.8 jaxlib==0.1.59 numpy==1.19.5 tensorflow==2.4.0" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 1. Create a double gaussian amp with SymPy" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "When building the model, we should be careful to pass the parameters as arguments as well, otherwise frameworks like jax can't determine the gradient." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\frac{0.398942280401433 A_{1} e^{- \\frac{\\left(- \\mu_{1} + x\\right)^{2}}{2 \\sigma_{1}}}}{\\sigma_{1}} + \\frac{0.398942280401433 A_{2} e^{- \\frac{\\left(- \\mu_{2} + x\\right)^{2}}{2 \\sigma_{2}}}}{\\sigma_{2}}$" ], "text/plain": [ "0.398942280401433*A1*exp(-(-mu1 + x)**2/(2*sigma1))/sigma1 + 0.398942280401433*A2*exp(-(-mu2 + x)**2/(2*sigma2))/sigma2" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import math\n", "\n", "import sympy as sp\n", "\n", "x, A1, mu1, sigma1, A2, mu2, sigma2 = sp.symbols(\n", " \"x, A1, mu1, sigma1, A2, mu2, sigma2\"\n", ")\n", "gaussian1 = (\n", " A1\n", " / (sigma1 * sp.sqrt(2.0 * math.pi))\n", " * sp.exp(-((x - mu1) ** 2) / (2 * sigma1))\n", ")\n", "gaussian2 = (\n", " A2\n", " / (sigma2 * sp.sqrt(2.0 * math.pi))\n", " * sp.exp(-((x - mu2) ** 2) / (2 * sigma2))\n", ")\n", "\n", "gauss_sum = gaussian1 + gaussian2\n", "gauss_sum" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 2. Convert this expression to a function using `lambdify`" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "[TensorFlow](https://www.tensorflow.org) as backend:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "def _lambdifygenerated(x, A1, mu1, sigma1, A2, mu2, sigma2):\n", " return (0.398942280401433*A1*exp(-1/2*pow(-mu1 + x, 2)/sigma1)/sigma1 + 0.398942280401433*A2*exp(-1/2*pow(-mu2 + x, 2)/sigma2)/sigma2)\n", "\n" ] } ], "source": [ "import inspect\n", "\n", "tf_gauss_sum = sp.lambdify(\n", " (x, A1, mu1, sigma1, A2, mu2, sigma2), gauss_sum, \"tensorflow\"\n", ")\n", "print(inspect.getsource(tf_gauss_sum))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "[NumPy](https://numpy.org) as backend:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "def _lambdifygenerated(x, A1, mu1, sigma1, A2, mu2, sigma2):\n", " return (0.398942280401433*A1*exp(-1/2*(-mu1 + x)**2/sigma1)/sigma1 + 0.398942280401433*A2*exp(-1/2*(-mu2 + x)**2/sigma2)/sigma2)\n", "\n" ] } ], "source": [ "numpy_gauss_sum = sp.lambdify(\n", " (x, A1, mu1, sigma1, A2, mu2, sigma2), gauss_sum, \"numpy\"\n", ")\n", "print(inspect.getsource(numpy_gauss_sum))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "[Jax](https://github.com/google/jax) as backend:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "def _lambdifygenerated(x, A1, mu1, sigma1, A2, mu2, sigma2):\n", " return (0.398942280401433*A1*exp(-1/2*(-mu1 + x)**2/sigma1)/sigma1 + 0.398942280401433*A2*exp(-1/2*(-mu2 + x)**2/sigma2)/sigma2)\n", "\n" ] } ], "source": [ "from jax import numpy as jnp\n", "from jax import scipy as jsp\n", "\n", "jax_gauss_sum = sp.lambdify(\n", " (x, A1, mu1, sigma1, A2, mu2, sigma2),\n", " gauss_sum,\n", " modules=(jnp, jsp.special),\n", ")\n", "print(inspect.getsource(jax_gauss_sum))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 3. Natively create the respective packages" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [], "source": [ "import math\n", "\n", "import tensorflow as tf\n", "\n", "\n", "def gaussian(x, A, mu, sigma):\n", " return (\n", " A\n", " / (sigma * tf.sqrt(tf.constant(2.0, dtype=tf.float64) * math.pi))\n", " * tf.exp(\n", " -tf.pow(-tf.constant(0.5, dtype=tf.float64) * (x - mu) / sigma, 2)\n", " )\n", " )\n", "\n", "\n", "def native_tf_gauss_sum(x_, A1_, mu1_, sigma1_, A2_, mu2_, sigma2_):\n", " return gaussian(x_, A1_, mu1_, sigma1_) + gaussian(x_, A2_, mu2_, sigma2_)\n", "\n", "\n", "# @jx.pmap\n", "def jax_gaussian(x, A, mu, sigma):\n", " return (\n", " A\n", " / (sigma * jnp.sqrt(2.0 * math.pi))\n", " * jnp.exp(-((-0.5 * (x - mu) / sigma) ** 2))\n", " )\n", "\n", "\n", "def native_jax_gauss_sum(x_, A1_, mu1_, sigma1_, A2_, mu2_, sigma2_):\n", " return jax_gaussian(x_, A1_, mu1_, sigma1_) + jax_gaussian(\n", " x_, A2_, mu2_, sigma2_\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 4. Compare performance" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [], "source": [ "import numpy as np\n", "\n", "parameter_values = (1.0, 0.0, 0.1, 2.0, 2.0, 0.2)\n", "np_x = np.random.uniform(-1, 3, 10000)\n", "tf_x = tf.constant(np_x)\n", "\n", "\n", "def evaluate_with_parameters(function):\n", " def wrapper():\n", " return function(np_x, *(parameter_values))\n", "\n", " return wrapper\n", "\n", "\n", "def call_native_tf():\n", " func = native_tf_gauss_sum\n", " params = tuple(tf.Variable(v, dtype=tf.float64) for v in parameter_values)\n", "\n", " def wrapper():\n", " return func(tf_x, *params)\n", "\n", " return wrapper" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "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" ] }, { "name": "stdout", "output_type": "stream", "text": [ "sympy tf lambdify 0.0923923499999546\n", "sympy numpy lambdify 0.021591910000097414\n", "sympy jax lambdify 0.15498641800013502\n", "native tf 0.1038757030000852\n", "native jax 0.19923505700012356\n" ] } ], "source": [ "import timeit\n", "\n", "from jax.config import config\n", "\n", "config.update(\"jax_enable_x64\", True)\n", "\n", "print(\n", " \"sympy tf lambdify\",\n", " timeit.timeit(evaluate_with_parameters(tf_gauss_sum), number=100),\n", ")\n", "print(\n", " \"sympy numpy lambdify\",\n", " timeit.timeit(evaluate_with_parameters(numpy_gauss_sum), number=100),\n", ")\n", "print(\n", " \"sympy jax lambdify\",\n", " timeit.timeit(evaluate_with_parameters(jax_gauss_sum), number=100),\n", ")\n", "print(\"native tf\", timeit.timeit(call_native_tf(), number=100))\n", "\n", "print(\n", " \"native jax\",\n", " timeit.timeit(evaluate_with_parameters(native_jax_gauss_sum), number=100),\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 5. Handling parameters" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Some options:\n", "\n", "#### 5.1 Changing parameter values\n", "\n", "Can be done in the model itself...\n", "\n", "But how can the values be propagated to the `AmplitudeModel`?\n", "\n", "Well, if an amplitude model only defines parameters with a name and the values are supplied in the function evaluation, then everything is decoupled and there are no problems.\n", "\n", "#### 5.2 Changing parameter names\n", "\n", "Names can be changed in the sympy `AmplitudeModel`. Since this sympy model serves as the source of truth for the `Function`, all things generated from this model will reflect the name changes as well.\n", "\n", "But going even further, since the Parameters are passed into the functions as arguments, the whole naming becomes irrelevant anyways.\n", "\n", "`tf_var_A1 = tf.Variable(1.0, dtype=tf.float64)` <- does not carry a name!!\n", "\n", "#### 5.3 Coupling parameters\n", "This means that one parameter is just assigned to another one?" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/plain": [ "DeviceArray([2.8856688 , 1.81293941, 0.8762776 , ..., 1.57889689,\n", " 3.45055759, 3.30175027], dtype=float64)" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "result = evaluate_with_parameters(jax_gauss_sum)()\n", "result" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/plain": [ "array([1.64006008, 1.43829976, 2.77864469, ..., 1.39104059, 2.24092382,\n", " 1.72490419])" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "np_x" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "image/svg+xml": [ "\n", "\n", "\n", "\n", " \n", " \n", " \n", " \n", " 2021-03-23T09:07:29.505627\n", " image/svg+xml\n", " \n", " \n", " Matplotlib v3.3.4, https://matplotlib.org/\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "import matplotlib.pyplot as plt\n", "\n", "plt.hist(np_x, bins=100, weights=result);" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "image/svg+xml": [ "\n", "\n", "\n", "\n", " \n", " \n", " \n", " \n", " 2021-03-23T09:07:29.809335\n", " image/svg+xml\n", " \n", " \n", " Matplotlib v3.3.4, https://matplotlib.org/\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "parameter_values = (1.0, 0.0, 0.1, 2.0, 2.0, 0.1)\n", "result = evaluate_with_parameters(jax_gauss_sum)()\n", "plt.hist(np_x, bins=100, weights=result);" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 6. Exchange a gaussian with some other function" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This should be easy if you know the exact expression that you want to replace:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\sin{\\left(a x \\right)} + \\cos{\\left(b x \\right)}$" ], "text/plain": [ "sin(a*x) + cos(b*x)" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from sympy.abc import C, a, b, x\n", "\n", "expr = sp.sin(a * x) + sp.cos(b * x)\n", "expr" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle C + \\cos{\\left(b x \\right)}$" ], "text/plain": [ "C + cos(b*x)" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "expr.subs(sp.sin(a * x), C)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 7. Matrix operations?" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\operatorname{re}{\\left(\\left[\\begin{matrix}1 - i & 2 - i & 3 - i\\end{matrix}\\right] \\rho \\left[\\begin{matrix}1 + i\\\\2 + i\\\\3 + i\\end{matrix}\\right]\\right)}$" ], "text/plain": [ "re(Matrix([[1 - I, 2 - I, 3 - I]])*rho*Matrix([\n", "[1 + I],\n", "[2 + I],\n", "[3 + I]]))" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from sympy.physics.quantum.dagger import Dagger\n", "\n", "spin_density = sp.MatrixSymbol(\"rho\", 3, 3)\n", "amplitudes = sp.Matrix([[1 + sp.I], [2 + sp.I], [3 + sp.I]])\n", "\n", "dummy_intensity = sp.re(\n", " Dagger(amplitudes) * spin_density * amplitudes,\n", " evaluate=False\n", " # evaluate=False is important otherwise it generates some function that cant\n", " # be lambdified anymore\n", ")\n", "dummy_intensity" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "def _lambdifygenerated(rho):\n", " return (real(matmul(matmul(constant([[1 - 1j, 2 - 1j, 3 - 1j]]), rho), constant([[1 + 1j], [2 + 1j], [3 + 1j]]))))\n", "\n" ] } ], "source": [ "tf_intensity = sp.lambdify(\n", " spin_density,\n", " dummy_intensity,\n", " modules=(tf,),\n", ")\n", "print(inspect.getsource(tf_intensity))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "keep_output" ] }, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "real0 = tf.constant(0, dtype=tf.float64)\n", "real1 = tf.constant(1, dtype=tf.float64)\n", "intensity_result = tf_intensity(\n", " np.array(\n", " [\n", " [\n", " tf.complex(real1, real0),\n", " tf.complex(real0, real0),\n", " -tf.complex(real0, real1),\n", " ],\n", " [\n", " tf.complex(real0, real0),\n", " tf.complex(real1, real0),\n", " tf.complex(real0, real0),\n", " ],\n", " [\n", " tf.complex(real0, real1),\n", " tf.complex(real0, real0),\n", " tf.complex(real1, real0),\n", " ],\n", " ]\n", " ),\n", ")\n", "intensity_result" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.8" }, "toc-autonumbering": false, "toc-showcode": true, "toc-showmarkdowntxt": false, "toc-showtags": true }, "nbformat": 4, "nbformat_minor": 4 }