-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
FEAT: implement
unevaluated_expression
decorator (#365)
* DOC: add notebook for SymPy helper functions * DOC: illustrate usage of `PoolSum` * DX: skip ipywidgets tests that copy widgets * DX: upgrade developer environment to Python 3.11 * ENH: make implementation method public as `evaluate()` * FEAT: implement `unevaluated_expression` decorator * MAINT: update Codecov config style
- Loading branch information
Showing
17 changed files
with
694 additions
and
39 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -168,6 +168,7 @@ | |
"pyright", | ||
"pytestconfig", | ||
"rankdir", | ||
"repr", | ||
"richman", | ||
"rightarrow", | ||
"risch", | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,232 @@ | ||
{ | ||
"cells": [ | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"hideCode": true, | ||
"hideOutput": true, | ||
"hidePrompt": true, | ||
"jupyter": { | ||
"source_hidden": true | ||
}, | ||
"slideshow": { | ||
"slide_type": "skip" | ||
}, | ||
"tags": [ | ||
"remove-cell", | ||
"skip-execution" | ||
] | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"# WARNING: advised to install a specific version, e.g. ampform==0.1.2\n", | ||
"%pip install -q ampform[doc,viz] IPython" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"hideCode": true, | ||
"hideOutput": true, | ||
"hidePrompt": true, | ||
"jupyter": { | ||
"source_hidden": true | ||
}, | ||
"slideshow": { | ||
"slide_type": "skip" | ||
}, | ||
"tags": [ | ||
"remove-cell" | ||
] | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"import os\n", | ||
"\n", | ||
"STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"```{autolink-concat}\n", | ||
"```" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"# SymPy helper functions" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"The {mod}`ampform.sympy` module contains a few classes that make it easier to construct larger expressions that consist of several mathematical definitions." | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"## Unevaluated expressions" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"The {func}`.unevaluated_expression` decorator makes it easier to write classes that represent a mathematical function definition. It makes a class that derives from {class}`sp.Expr <sympy.core.expr.Expr>` behave more like a {func}`~.dataclasses.dataclass` (see [PEP 861](https://peps.python.org/pep-0681)). All you have to do is:\n", | ||
"\n", | ||
"1. Specify the arguments the function requires.\n", | ||
"2. Specify how to render the 'unevaluated' or 'folded' form of the expression with a `_latex_repr_` string or method.\n", | ||
"3. Specify how to unfold the expression using an `evaluate()` method.\n", | ||
"\n", | ||
"In the example below, we define a phase space factor $\\rho^\\text{CM}$ using the Chew-Mandelstam function (see PDG Resonances section, [Eq. (50.44)](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15)). For this, you need to define a break-up momentum $q$ as well." | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"import sympy as sp\n", | ||
"\n", | ||
"from ampform.sympy import unevaluated_expression\n", | ||
"\n", | ||
"\n", | ||
"@unevaluated_expression(real=False)\n", | ||
"class PhspFactorSWave(sp.Expr):\n", | ||
" s: sp.Symbol\n", | ||
" m1: sp.Symbol\n", | ||
" m2: sp.Symbol\n", | ||
" _latex_repr_ = R\"\\rho^\\text{{CM}}\\left({s}\\right)\"\n", | ||
"\n", | ||
" def evaluate(self) -> sp.Expr:\n", | ||
" s, m1, m2 = self.args\n", | ||
" q = BreakupMomentum(s, m1, m2)\n", | ||
" cm = (\n", | ||
" (2 * q / sp.sqrt(s))\n", | ||
" * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n", | ||
" - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", | ||
" ) / (16 * sp.pi**2)\n", | ||
" return 16 * sp.pi * sp.I * cm\n", | ||
"\n", | ||
"\n", | ||
"@unevaluated_expression(real=False)\n", | ||
"class BreakupMomentum(sp.Expr):\n", | ||
" s: sp.Symbol\n", | ||
" m1: sp.Symbol\n", | ||
" m2: sp.Symbol\n", | ||
" _latex_repr_ = R\"q\\left({s}\\right)\"\n", | ||
"\n", | ||
" def evaluate(self) -> sp.Expr:\n", | ||
" s, m1, m2 = self.args\n", | ||
" return sp.sqrt((s - (m1 + m2) ** 2) * (s - (m1 - m2) ** 2) / (s * 4))" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"As can be seen, the LaTeX rendering of these classes makes them ideal for mathematically defining and building up larger amplitude models:" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"editable": true, | ||
"slideshow": { | ||
"slide_type": "" | ||
}, | ||
"tags": [ | ||
"hide-input" | ||
] | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"from IPython.display import Math\n", | ||
"\n", | ||
"from ampform.io import aslatex\n", | ||
"\n", | ||
"s, m1, m2 = sp.symbols(\"s m1 m2\")\n", | ||
"q_expr = BreakupMomentum(s, m1, m2)\n", | ||
"rho_expr = PhspFactorSWave(s, m1, m2)\n", | ||
"Math(aslatex({e: e.evaluate() for e in [rho_expr, q_expr]}))" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"## Summations" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"The {class}`.PoolSum` class makes it possible to write sums over non-integer ranges. This is for instance useful when summing over allowed helicities. Here are some examples:" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"from ampform.sympy import PoolSum\n", | ||
"\n", | ||
"i, j, m, n = sp.symbols(\"i j m n\")\n", | ||
"expr = PoolSum(i**m + j**n, (i, (-1, 0, +1)), (j, (2, 4, 5)))\n", | ||
"Math(aslatex({expr: expr.doit()}))" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"import numpy as np\n", | ||
"\n", | ||
"A = sp.IndexedBase(\"A\")\n", | ||
"λ, μ = sp.symbols(\"lambda mu\")\n", | ||
"to_range = lambda a, b: tuple(sp.Rational(i) for i in np.arange(a, b + 0.5))\n", | ||
"expr = abs(PoolSum(A[λ, μ], (λ, to_range(-0.5, +0.5)), (μ, to_range(-1, +1)))) ** 2\n", | ||
"Math(aslatex({expr: expr.doit()}))" | ||
] | ||
} | ||
], | ||
"metadata": { | ||
"colab": { | ||
"toc_visible": true | ||
}, | ||
"kernelspec": { | ||
"display_name": "Python 3 (ipykernel)", | ||
"language": "python", | ||
"name": "python3" | ||
}, | ||
"language_info": { | ||
"codemirror_mode": { | ||
"name": "ipython", | ||
"version": 3 | ||
}, | ||
"file_extension": ".py", | ||
"mimetype": "text/x-python", | ||
"name": "python", | ||
"nbconvert_exporter": "python", | ||
"pygments_lexer": "ipython3", | ||
"version": "3.11.5" | ||
} | ||
}, | ||
"nbformat": 4, | ||
"nbformat_minor": 4 | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.