Skip to content

Commit

Permalink
FEAT: implement unevaluated_expression decorator (#365)
Browse files Browse the repository at this point in the history
* 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
redeboer committed Dec 15, 2023
1 parent a2c7524 commit ce5d6ae
Show file tree
Hide file tree
Showing 17 changed files with 694 additions and 39 deletions.
1 change: 1 addition & 0 deletions .cspell.json
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@
"pyright",
"pytestconfig",
"rankdir",
"repr",
"richman",
"rightarrow",
"risch",
Expand Down
3 changes: 3 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ jobs:
id-token: write
with:
apt-packages: graphviz
python-version: "3.11"
specific-pip-packages: ${{ inputs.specific-pip-packages }}
pytest:
uses: ComPWA/actions/.github/workflows/pytest.yml@v1
Expand All @@ -45,3 +46,5 @@ jobs:
secrets:
token: ${{ secrets.PAT }}
uses: ComPWA/actions/.github/workflows/pre-commit.yml@v1
with:
python-version: "3.11"
4 changes: 2 additions & 2 deletions .gitpod.yml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
tasks:
- init: pyenv local 3.8
- init: pip install -c .constraints/py3.8.txt -e .[dev]
- init: pyenv local 3.11
- init: pip install -c .constraints/py3.11.txt -e .[dev]

github:
prebuilds:
Expand Down
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ repos:
- id: check-dev-files
args:
- --doc-apt-packages=graphviz
- --dev-python-version=3.11
- --no-prettierrc
- --pin-requirements=monthly
- --repo-name=ampform
Expand Down
4 changes: 2 additions & 2 deletions .readthedocs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,9 @@ formats:
build:
os: ubuntu-22.04
tools:
python: "3.8"
python: "3.11"
apt_packages:
- graphviz
jobs:
post_install:
- pip install -c .constraints/py3.8.txt -e .[doc]
- pip install -c .constraints/py3.11.txt -e .[doc]
1 change: 1 addition & 0 deletions .vscode/settings.json
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
"ruff.enable": true,
"ruff.organizeImports": true,
"search.exclude": {
"typings/**": true,
"**/tests/**/__init__.py": true,
".constraints/*.txt": true
},
Expand Down
6 changes: 3 additions & 3 deletions codecov.yml
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,13 @@ coverage:
threshold: 1% # allow drops by this percentage
base: auto
# advanced
branches: null
branches: []
if_no_uploads: error
if_not_found: success
if_ci_failed: error
only_pulls: false
flags: null
paths: null
flags: []
paths: []
patch:
default:
# basic
Expand Down
5 changes: 4 additions & 1 deletion docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,11 +52,13 @@
"BuilderReturnType": ("obj", "ampform.dynamics.builder.BuilderReturnType"),
"DecoratedClass": ("obj", "ampform.sympy.DecoratedClass"),
"DecoratedExpr": ("obj", "ampform.sympy.DecoratedExpr"),
"ExprClass": "ampform.sympy.ExprClass",
"FourMomenta": ("obj", "ampform.kinematics.FourMomenta"),
"FourMomentumSymbol": ("obj", "ampform.kinematics.FourMomentumSymbol"),
"InteractionProperties": "qrules.quantum_numbers.InteractionProperties",
"LatexPrinter": "sympy.printing.printer.Printer",
"Literal[(-1, 1)]": "typing.Literal",
"Literal[-1, 1]": "typing.Literal",
"NumPyPrinter": "sympy.printing.printer.Printer",
"ParameterValue": ("obj", "ampform.helicity.ParameterValue"),
"Particle": "qrules.particle.Particle",
Expand All @@ -68,6 +70,7 @@
"WignerD": "sympy.physics.quantum.spin.WignerD",
"ampform.helicity._T": "typing.TypeVar",
"ampform.sympy._decorator.ExprClass": ("obj", "ampform.sympy.ExprClass"),
"ampform.sympy._decorator.SymPyAssumptions": "ampform.sympy.SymPyAssumptions",
"an object providing a view on D's values": "typing.ValuesView",
"sp.Basic": "sympy.core.basic.Basic",
"sp.Expr": "sympy.core.expr.Expr",
Expand Down Expand Up @@ -283,7 +286,7 @@
nb_execution_timeout = -1
nb_output_stderr = "remove"
nitpick_ignore = [
("py:class", "ArraySum"),
("py:class", "ampform.sympy._array_expressions.ArraySum"),
("py:class", "ampform.sympy._array_expressions.MatrixMultiplication"),
]
nitpicky = True
Expand Down
1 change: 1 addition & 0 deletions docs/usage.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,7 @@
"usage/helicity/formalism\n",
"usage/helicity/spin-alignment\n",
"usage/kinematics\n",
"usage/sympy\n",
"```"
]
}
Expand Down
232 changes: 232 additions & 0 deletions docs/usage/sympy.ipynb
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&nbsp;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.&nbsp;(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
}
4 changes: 2 additions & 2 deletions environment.yml
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@ name: ampform
channels:
- defaults
dependencies:
- python==3.8.*
- python==3.11.*
- pip
- graphviz # for binder
- pip:
- -c .constraints/py3.8.txt -e .[dev]
- -c .constraints/py3.11.txt -e .[dev]
variables:
PRETTIER_LEGACY_CLI: "1"
PYTHONHASHSEED: 0
4 changes: 4 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,7 @@ branch = true
source = ["src"]

[tool.mypy]
enable_incomplete_feature = "Unpack"
exclude = "_build"
show_error_codes = true
warn_unused_configs = true
Expand Down Expand Up @@ -207,6 +208,7 @@ exclude = [
"**/.tox",
"**/__pycache__",
"**/_build",
"**/typings",
]
reportGeneralTypeIssues = false
reportIncompatibleMethodOverride = false
Expand All @@ -224,6 +226,7 @@ reportUnknownVariableType = false
reportUnnecessaryComparison = false
reportUnnecessaryContains = false
reportUnnecessaryIsInstance = false
reportUntypedClassDecorator = false
reportUntypedFunctionDecorator = false
reportUnusedClass = true
reportUnusedFunction = true
Expand Down Expand Up @@ -324,6 +327,7 @@ task-tags = ["cspell"]
known-third-party = ["sympy"]

[tool.ruff.per-file-ignores]
"**/docs/usage/sympy.ipynb" = ["E731"]
"*.ipynb" = [
"B018",
"C408",
Expand Down
Loading

0 comments on commit ce5d6ae

Please sign in to comment.