{
"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.Function`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"jupyter": {
"source_hidden": true
},
"tags": [
"hide-cell",
"keep_output",
"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",
"\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": [
"One way to address the Cons of {doc}`composition`, is to sub-class {class}`~sympy.core.function.Function`. The expression is implemented by overwriting the inherited `eval()` method and the builder is provided through the class through an additional `from_graph` class method. The interface would look like this:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class DynamicsFunction(sp.Function):\n",
" @classmethod\n",
" @abstractmethod\n",
" def eval(cls, *args: sp.Symbol) -> sp.Expr:\n",
" \"\"\"Implementation of the dynamics function.\"\"\"\n",
"\n",
" @classmethod\n",
" @abstractmethod\n",
" def from_graph(cls, graph: StateTransitionGraph, edge_id: int) -> sp.Basic:\n",
" pass"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As can be seen from the implementation of a relativistic Breit-Wigner, the implementation of the expression is nicely kept together with the implementation of the expression:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class RelativisticBreitWigner(DynamicsFunction):\n",
" @classmethod\n",
" def eval(cls, *args: sp.Symbol) -> sp.Expr:\n",
" mass = args[0]\n",
" mass0 = args[1]\n",
" gamma0 = args[2]\n",
" return (\n",
" gamma0 * mass0 / (mass0 ** 2 - mass ** 2 - gamma0 * mass0 * sp.I)\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": "markdown",
"metadata": {},
"source": [
"It becomes a bit less clear when using a form factor, but the `DynamicsFunction` base class enforces a correct interfaces:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class RelativisticBreitWignerWithFF(DynamicsFunction):\n",
" @classmethod\n",
" def eval(cls, *args: sp.Symbol) -> sp.Expr:\n",
" # Arguments\n",
" mass = args[0]\n",
" mass0 = args[1]\n",
" gamma0 = args[2]\n",
" m_a = args[3]\n",
" m_b = args[4]\n",
" angular_momentum = args[5]\n",
" meson_radius = args[6]\n",
" # Computed variables\n",
" q_squared = two_body_momentum_squared(mass, m_a, m_b)\n",
" q0_squared = two_body_momentum_squared(mass0, m_a, m_b)\n",
" ff2 = blatt_weisskopf(q_squared, meson_radius, angular_momentum)\n",
" ff02 = blatt_weisskopf(q0_squared, meson_radius, angular_momentum)\n",
" width = gamma0 * (mass0 / mass) * (ff2 / ff02)\n",
" width = width * sp.sqrt(q_squared / q0_squared)\n",
" # Expression\n",
" return (\n",
" RelativisticBreitWigner(mass, mass0, width)\n",
" * mass0\n",
" * 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 `expression_builder` used in the syntax proposed at {ref}`adr/002:Considered solutions`, would now just be a class that is derived of `DynamicsFunction`."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `sympy.Function` class provides mixin methods, so that the derived class behaves as a `sympy` expression. So the expression can be inspected with the usual `sympy` tools (compare the Pros of {doc}`composition`):"
]
},
{
"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"
},
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {
"needs_background": "light"
},
"output_type": "display_data"
},
{
"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": [
"m, m0, w0 = sp.symbols(R\"m m_0 \\Gamma\")\n",
"evaluated_bw = RelativisticBreitWigner(m, 1.0, 0.3)\n",
"sp.plot(sp.Abs(evaluated_bw), (m, 0, 2), axis_center=(0, 0), ylim=(0, 1))\n",
"sp.plot(sp.arg(evaluated_bw), (m, 0, 2), axis_center=(0, 0), ylim=(0, sp.pi))\n",
"RelativisticBreitWigner(m, m0, w0)"
]
}
],
"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
}