{
"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": [
"# Subclassing `sympy.Expr`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-cell",
"remove-cell"
]
},
"outputs": [],
"source": [
"# !pip install matplotlib==3.3.3 sympy==1.7.1"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"jupyter": {
"source_hidden": true
},
"tags": [
"hide-cell"
]
},
"outputs": [],
"source": [
"from abc import abstractmethod\n",
"from typing import Any, Callable, Set, Type\n",
"\n",
"import sympy as sp\n",
"from helpers import (\n",
" StateTransitionGraph,\n",
" blatt_weisskopf,\n",
" determine_attached_final_state,\n",
" two_body_momentum_squared,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The major disadvantage of {doc}`function`, is that there is no way to identify which `Symbol`s are variables and which are parameters. This can be solved by sub-classing from {class}`sympy.core.expr.Expr`.\n",
"\n",
"An example of a class that does this is {class}`~sympy.physics.quantum.spin.WignerD`. There, the implementation of the dynamics expression can be evaluated through a {meth}`~sympy.core.basic.Basic.doit` call. This method can call anything, but {mod}`sympy` seems to follow the convention that it returns an 'evaluated' version of the class itself, where 'evaluated' means that any randomly named method of the class has been called on the `*args` that are implemented through the `__new__` method (the examples below make this clearer).\n",
"\n",
"For our purposes, the follow `DynamicsExpr` base class illustrates the interface that we expect. Here, `evaluate` is where expression is implemented and (just as in {doc}`function`) `from_graph` is the builder method."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class DynamicsExpr(sp.Expr):\n",
" @classmethod\n",
" @abstractmethod\n",
" def __new__(cls, *args: sp.Symbol, **hints: Any) -> sp.Expr:\n",
" pass\n",
"\n",
" @abstractmethod\n",
" def doit(self, **hints: Any) -> sp.Expr:\n",
" pass\n",
"\n",
" @abstractmethod\n",
" def evaluate(self) -> sp.Expr:\n",
" pass\n",
"\n",
" @classmethod\n",
" @abstractmethod\n",
" def from_graph(cls, graph: StateTransitionGraph, edge_id: int) -> sp.Basic:\n",
" pass"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `__new__` and `doit` methods split the construction from the evaluation of the expression. This allows one to distinguish `variables` and `parameters` and present them as properties:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class RelativisticBreitWigner(DynamicsExpr):\n",
" def __new__(cls, *args: sp.Symbol, **hints: Any) -> sp.Expr:\n",
" if len(args) != 3:\n",
" raise ValueError(f\"3 parameters expected, got {len(args)}\")\n",
" args = sp.sympify(args)\n",
" evaluate = hints.get(\"evaluate\", False)\n",
" if evaluate:\n",
" return sp.Expr.__new__(cls, *args).evaluate() # type: ignore # pylint: disable=no-member\n",
" return sp.Expr.__new__(cls, *args)\n",
"\n",
" @property\n",
" def mass(self) -> sp.Symbol:\n",
" return self.args[0]\n",
"\n",
" @property\n",
" def mass0(self) -> sp.Symbol:\n",
" return self.args[1]\n",
"\n",
" @property\n",
" def gamma0(self) -> sp.Symbol:\n",
" return self.args[2]\n",
"\n",
" @property\n",
" def variables(self) -> Set[sp.Symbol]:\n",
" return {self.mass}\n",
"\n",
" @property\n",
" def parameters(self) -> Set[sp.Symbol]:\n",
" return {self.mass0, self.gamma0}\n",
"\n",
" def doit(self, **hints: Any) -> sp.Expr:\n",
" return RelativisticBreitWigner(*self.args, **hints, evaluate=True)\n",
"\n",
" def evaluate(self) -> sp.Expr:\n",
" return (\n",
" self.gamma0\n",
" * self.mass0\n",
" / (\n",
" self.mass0 ** 2\n",
" - self.mass ** 2\n",
" - self.gamma0 * self.mass0 * sp.I\n",
" )\n",
" )\n",
"\n",
" @classmethod\n",
" def from_graph(\n",
" cls, graph: StateTransitionGraph, edge_id: int\n",
" ) -> \"RelativisticBreitWigner\":\n",
" edge_ids = determine_attached_final_state(graph, edge_id)\n",
" final_state_ids = map(str, edge_ids)\n",
" mass = sp.Symbol(f\"m_{{{'+'.join(final_state_ids)}}}\")\n",
" particle, _ = graph.get_edge_props(edge_id)\n",
" mass0 = sp.Symbol(f\"m_{{{particle.latex}}}\")\n",
" gamma0 = sp.Symbol(fR\"\\Gamma_{{{particle.latex}}}\")\n",
" return cls(mass, mass0, gamma0)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"jupyter": {
"source_hidden": true
},
"tags": [
"hide-cell",
"keep_output"
]
},
"outputs": [],
"source": [
"class RelativisticBreitWignerWithFF(DynamicsExpr):\n",
" def __new__(cls, *args: sp.Symbol, **hints: Any) -> sp.Expr:\n",
" if len(args) != 7:\n",
" raise ValueError(f\"7 parameters expected, got {len(args)}\")\n",
" args = sp.sympify(args)\n",
" evaluate = hints.get(\"evaluate\", False)\n",
" if evaluate:\n",
" return sp.Expr.__new__(cls, *args).evaluate() # type: ignore # pylint: disable=no-member\n",
" return sp.Expr.__new__(cls, *args)\n",
"\n",
" def doit(self, **hints: Any) -> sp.Expr:\n",
" return RelativisticBreitWignerWithFF(\n",
" *self.args, **hints, evaluate=True\n",
" )\n",
"\n",
" @property\n",
" def mass(self) -> sp.Symbol:\n",
" return self.args[0]\n",
"\n",
" @property\n",
" def mass0(self) -> sp.Symbol:\n",
" return self.args[1]\n",
"\n",
" @property\n",
" def gamma0(self) -> sp.Symbol:\n",
" return self.args[2]\n",
"\n",
" @property\n",
" def m_a(self) -> sp.Symbol:\n",
" return self.args[3]\n",
"\n",
" @property\n",
" def m_b(self) -> sp.Symbol:\n",
" return self.args[4]\n",
"\n",
" @property\n",
" def angular_momentum(self) -> sp.Symbol:\n",
" return self.args[5]\n",
"\n",
" @property\n",
" def meson_radius(self) -> sp.Symbol:\n",
" return self.args[6]\n",
"\n",
" def evaluate(self) -> sp.Expr:\n",
" # Computed variables\n",
" q_squared = two_body_momentum_squared(self.mass, self.m_a, self.m_b)\n",
" q0_squared = two_body_momentum_squared(self.mass0, self.m_a, self.m_b)\n",
" ff2 = blatt_weisskopf(\n",
" q_squared, self.meson_radius, self.angular_momentum\n",
" )\n",
" ff02 = blatt_weisskopf(\n",
" q0_squared, self.meson_radius, self.angular_momentum\n",
" )\n",
" width = self.gamma0 * (self.mass0 / self.mass) * (ff2 / ff02)\n",
" width = width * sp.sqrt(q_squared / q0_squared)\n",
" # Expression\n",
" return (\n",
" RelativisticBreitWigner(self.mass, self.mass0, width)\n",
" * self.mass0\n",
" * self.gamma0\n",
" * sp.sqrt(ff2)\n",
" )\n",
"\n",
" @classmethod\n",
" def from_graph(\n",
" cls, graph: StateTransitionGraph, edge_id: int\n",
" ) -> \"RelativisticBreitWignerWithFF\":\n",
" edge_ids = determine_attached_final_state(graph, edge_id)\n",
" final_state_ids = map(str, edge_ids)\n",
" mass = sp.Symbol(f\"m_{{{'+'.join(final_state_ids)}}}\")\n",
" particle, _ = graph.get_edge_props(edge_id)\n",
" mass0 = sp.Symbol(f\"m_{{{particle.latex}}}\")\n",
" gamma0 = sp.Symbol(fR\"\\Gamma_{{{particle.latex}}}\")\n",
" m_a = sp.Symbol(f\"m_{edge_ids[0]}\")\n",
" m_b = sp.Symbol(f\"m_{edge_ids[1]}\")\n",
" angular_momentum = particle.spin # helicity formalism only!\n",
" meson_radius = sp.Symbol(fR\"R_{{{particle.latex}}}\")\n",
" return cls(\n",
" mass, mass0, gamma0, m_a, m_b, angular_momentum, meson_radius\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The following illustrates the difference with {doc}`function`. First, notice that a class derived from `DynamicsExpr` is still **identifiable** upon construction:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"keep_output"
]
},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle RelativisticBreitWigner\\left(m, m_{0}, \\Gamma\\right)$"
],
"text/plain": [
"RelativisticBreitWigner(m, m_0, \\Gamma)"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"m, m0, w0 = sp.symbols(R\"m m_0 \\Gamma\")\n",
"rel_bw = RelativisticBreitWigner(m, m0, w0)\n",
"rel_bw"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```{toggle}\n",
"The way in which this expression is rendered in a Jupyter notebook can be changed by overwriting the `_pretty` and/or `_latex` methods.\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Only once `doit()` is called, is the `DynamicsExpr` converted into a mathematical expression:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"keep_output"
]
},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\frac{\\Gamma m_{0}}{- i \\Gamma m_{0} - m^{2} + m_{0}^{2}}$"
],
"text/plain": [
"\\Gamma*m_0/(-I*\\Gamma*m_0 - m**2 + m_0**2)"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"evaluated_bw = rel_bw.doit()\n",
"evaluated_bw"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"keep_output"
]
},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {
"needs_background": "light"
},
"output_type": "display_data"
}
],
"source": [
"sp.plot(\n",
" sp.Abs(evaluated_bw.subs({m0: 1, w0: 0.2})),\n",
" (m, 0, 2),\n",
" axis_center=(0, 0),\n",
" ylim=(0, 1),\n",
");"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Decorator"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"There are a lot of implicit conventions that need to be followed to provide a correct implementation of a `DynamicsExpr`. Some of this may be mitigated by proving some class decorator that can easily construct the `__new__()` and `doit()` methods for you."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def dynamics_expression(n_args: int) -> Callable[[type], Type[DynamicsExpr]]:\n",
" def decorator(decorated_class: type) -> Type[DynamicsExpr]:\n",
" def __new__(cls, *args: sp.Symbol, **hints: Any) -> sp.Expr:\n",
" if len(args) != n_args:\n",
" raise ValueError(\n",
" f\"{n_args} parameters expected, got {len(args)}\"\n",
" )\n",
" args = sp.sympify(args)\n",
" evaluate = hints.get(\"evaluate\", False)\n",
" if evaluate:\n",
" return sp.Expr.__new__(cls, *args).evaluate()\n",
" return sp.Expr.__new__(cls, *args)\n",
"\n",
" def doit(self, **hints: Any) -> sp.Expr:\n",
" return decorated_class(*self.args, **hints, evaluate=True)\n",
"\n",
" decorated_class.__new__ = __new__\n",
" decorated_class.doit = doit\n",
" return decorated_class\n",
"\n",
" return decorator"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This saves some lines of code:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@dynamics_expression(n_args=3)\n",
"class Gauss(DynamicsExpr):\n",
" @property\n",
" def mass(self) -> sp.Symbol:\n",
" return self.args[0]\n",
"\n",
" @property\n",
" def mu(self) -> sp.Symbol:\n",
" return self.args[1]\n",
"\n",
" @property\n",
" def sigma(self) -> sp.Symbol:\n",
" return self.args[2]\n",
"\n",
" @property\n",
" def variables(self) -> Set[sp.Symbol]:\n",
" return {self.mass}\n",
"\n",
" @property\n",
" def parameters(self) -> Set[sp.Symbol]:\n",
" return {self.mu, self.sigma}\n",
"\n",
" def evaluate(self) -> sp.Expr:\n",
" return sp.exp(-((self.mass - self.mu) ** 2) / self.sigma ** 2)\n",
"\n",
" @classmethod\n",
" def from_graph(\n",
" cls, graph: StateTransitionGraph, edge_id: int\n",
" ) -> \"RelativisticBreitWigner\":\n",
" edge_ids = determine_attached_final_state(graph, edge_id)\n",
" final_state_ids = map(str, edge_ids)\n",
" mass = sp.Symbol(f\"m_{{{'+'.join(final_state_ids)}}}\")\n",
" particle, _ = graph.get_edge_props(edge_id)\n",
" mass0 = sp.Symbol(fR\"\\mu_{{{particle.latex}}}\")\n",
" gamma0 = sp.Symbol(fR\"\\sigma_{{{particle.latex}}}\")\n",
" return cls(mass, mass0, gamma0)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"keep_output"
]
},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle Gauss\\left(x, \\mu, \\Gamma\\right)$"
],
"text/plain": [
"Gauss(x, \\mu, \\Gamma)"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"x, mu, sigma = sp.symbols(R\"x \\mu \\sigma\")\n",
"Gauss(x, mu, w0)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"keep_output"
]
},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle e^{- \\frac{\\left(- \\mu + x\\right)^{2}}{\\sigma^{2}}}$"
],
"text/plain": [
"exp(-(-\\mu + x)**2/\\sigma**2)"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"Gauss(x, mu, sigma).doit()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Issue with lambdify"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"It's not possible to plot a `DynamicsExpr` directly as long as no `lambdify` hook has been provided: **`doit()` has to be executed first**."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"keep_output",
"hide-cell"
]
},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {
"needs_background": "light"
},
"output_type": "display_data"
}
],
"source": [
"sp.plot(sp.Abs(rel_bw.subs({m0: 1, w0: 0.2})).doit(), (m, 0, 2));"
]
}
],
"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"
}
},
"nbformat": 4,
"nbformat_minor": 4
}