Skip to content

BREAK: make parametrized functions pure #579

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
grayson-helmholz wants to merge 11 commits into from
Open

BREAK: make parametrized functions pure #579

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
2895169
BREAK: replace update_parameters() with pure evaluation API
redeboer Aug 5, 2026
d1955d7
MAINT: implement updates from formatters
web-flow Aug 7, 2026
54e0436
MAINT: implement updates from formatters
web-flow Sep 1, 2026
4b9a41d
DOC: name parameters keyword in call examples
redeboer Sep 1, 2026
0c20c71
MAINT: minor formatting improvements
redeboer Sep 1, 2026
96a3f99
MAINT: remove redundant `from __future__import annotations`
redeboer Sep 1, 2026
3eed813
DOC: use immutable function interface idiomatically
redeboer Sep 1, 2026
8aea5de
BEHAVIOR: compute fit fractions with optimized parameters
redeboer Sep 1, 2026
8481a7f
DOC: distinguish initial and optimized functions
redeboer Sep 1, 2026
44582ce
DX: replace inline backend magic with set_matplotlib_formats
redeboer Sep 1, 2026
2900121
MAINT: implement updates from formatters
web-flow Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 28 additions & 26 deletions docs/amplitude-analysis.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -689,9 +689,10 @@
},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['svg']\n",
"\n",
"import matplotlib.pyplot as plt\n",
"from matplotlib_inline.backend_inline import set_matplotlib_formats\n",
"\n",
"set_matplotlib_formats(\"svg\")\n",
"\n",
"resonances = sorted(\n",
" reaction.get_intermediate_particles(),\n",
Expand All @@ -705,6 +706,7 @@
" bins=100,\n",
" alpha=0.5,\n",
" density=True,\n",
" rasterized=True,\n",
")\n",
"ax.set_xlabel(\"$m$ [GeV]\")\n",
"for p, color in zip(resonances, colors, strict=True):\n",
Expand Down Expand Up @@ -1395,7 +1397,7 @@
"\n",
"Let's have a look at our [first guess for the parameter values](#determine-free-parameters). Recall that a {class}`.ParametrizedFunction` object computes the intensity for a certain {obj}`.DataSample`. This can be seen nicely when we use these intensities as weights on the phase space sample and plot it together with the original data sample. Here, we look at the invariant mass distribution projection of the final states `1` and `2`, which, [as we saw before](compwa-step-2.3), is the final state particle pair $\\pi^0\\pi^0$.\n",
"\n",
"Don't forget to use {meth}`~.ParametrizedFunction.update_parameters` first!"
"Note that we do not modify `optimized_function` itself: {meth}`~.ParametrizedFunction.with_parameters` returns a *new* function with these initial parameter values."
]
},
{
Expand Down Expand Up @@ -1472,9 +1474,8 @@
"metadata": {},
"outputs": [],
"source": [
"original_parameters = optimized_function.parameters\n",
"optimized_function.update_parameters(initial_parameters)\n",
"compare_model(\"m_12\", data_real, phsp_real, optimized_function)"
"initial_function = optimized_function.with_parameters(initial_parameters)\n",
"compare_model(\"m_12\", data_real, phsp_real, initial_function)"
]
},
{
Expand Down Expand Up @@ -1558,6 +1559,7 @@
"outputs": [],
"source": [
"optimized_parameters = fit_result.parameter_values\n",
"original_parameters = optimized_function.parameters\n",
"for p in optimized_parameters:\n",
" print(p)\n",
" print(f\" initial: {initial_parameters[p]:.3}\")\n",
Expand Down Expand Up @@ -1622,7 +1624,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Using the same method as above, we renew the parameters of the {class}`.ParametrizedFunction` and plot it again over the phase space sample."
"Using the same method as above, we create a new {class}`.ParametrizedFunction` with the optimized parameter values and plot it again over the phase space sample."
]
},
{
Expand All @@ -1631,8 +1633,8 @@
"metadata": {},
"outputs": [],
"source": [
"optimized_function.update_parameters(fit_result.parameter_values)\n",
"compare_model(\"m_12\", data_real, phsp_real, optimized_function)"
"fitted_function = optimized_function.with_parameters(fit_result.parameter_values)\n",
"compare_model(\"m_12\", data_real, phsp_real, fitted_function)"
]
},
{
Expand Down Expand Up @@ -1733,7 +1735,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Here's an example function that can do this. Using regular expressions, we set all coefficients in the intensity function to zero if they do not contain a certain resonance $\\LaTeX$ name:"
"Here's an example function that can do this. Using regular expressions, we select all coefficients that do not contain a certain resonance $\\LaTeX$ name and evaluate the intensity function with those coefficients set to zero:"
]
},
{
Expand All @@ -1752,22 +1754,21 @@
" input_data: DataSample,\n",
" resonances: list[str],\n",
"):\n",
" original_parameters = dict(func.parameters)\n",
" negative_lookahead = f\"(?!{'|'.join(map(re.escape, resonances))})\"\n",
" # https://regex101.com/r/WrgGyD/1\n",
" pattern = rf\"^(\\\\mathcal{{H}}|C_)({negative_lookahead}.)*$\"\n",
" set_parameters_to_zero(func, pattern)\n",
" array = func(input_data)\n",
" func.update_parameters(original_parameters)\n",
" return array\n",
" zeroed_parameters = get_zeroed_parameters(func, pattern)\n",
" return func(input_data, zeroed_parameters)\n",
"\n",
"\n",
"def set_parameters_to_zero(func: ParametrizedFunction, name_pattern: str) -> None:\n",
" new_parameters = dict(func.parameters)\n",
" for par_name in func.parameters:\n",
" if re.match(name_pattern, par_name) is not None:\n",
" new_parameters[par_name] = 0\n",
" func.update_parameters(new_parameters)"
"def get_zeroed_parameters(\n",
" func: ParametrizedFunction, name_pattern: str\n",
") -> dict[str, complex]:\n",
" return {\n",
" par_name: 0\n",
" for par_name in func.parameters\n",
" if re.match(name_pattern, par_name) is not None\n",
" }"
]
},
{
Expand Down Expand Up @@ -1795,7 +1796,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"These functions can be used to compute and visualize the sub-intensity distributions over the phase space."
"These functions can be used to compute and visualize the sub-intensity distributions over the phase space. Note that we have to compute them with the *optimized* parameter values, so we first create a function that carries the values found by the fit."
]
},
{
Expand All @@ -1804,9 +1805,10 @@
"metadata": {},
"outputs": [],
"source": [
"total_intensities = intensity_func(phsp)\n",
"fitted_intensity_func = intensity_func.with_parameters(fit_result.parameter_values)\n",
"total_intensities = fitted_intensity_func(phsp)\n",
"sub_intensities = {\n",
" p: compute_sub_intensity(intensity_func, phsp, resonances=[p.latex])\n",
" p: compute_sub_intensity(fitted_intensity_func, phsp, resonances=[p.latex])\n",
" for p in resonances\n",
"}"
]
Expand Down Expand Up @@ -1866,7 +1868,7 @@
"metadata": {},
"outputs": [],
"source": [
"total_intensity = intensity_func(phsp).sum()\n",
"total_intensity = total_intensities.sum()\n",
"fit_fractions = {\n",
" resonance.name: f\"{sub_intensity.sum() / total_intensity:.1%}\"\n",
" for resonance, sub_intensity in sub_intensities.items()\n",
Expand Down Expand Up @@ -1908,7 +1910,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.13"
"version": "3.13.15"
}
},
"nbformat": 4,
Expand Down
8 changes: 4 additions & 4 deletions docs/amplitude-analysis/analytic-continuation.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
"import matplotlib.pyplot as plt\n",
"import qrules\n",
"from IPython.display import Math, display\n",
"from matplotlib_inline.backend_inline import set_matplotlib_formats\n",
"\n",
"from tensorwaves.data import (\n",
" SympyDataTransformer,\n",
Expand All @@ -40,6 +41,7 @@
"from tensorwaves.function.sympy import create_parametrized_function\n",
"\n",
"logging.getLogger(\"tensorwaves.data\").setLevel(logging.ERROR) # hide progress bars\n",
"set_matplotlib_formats(\"svg\")\n",
"warnings.filterwarnings(\"ignore\")"
]
},
Expand Down Expand Up @@ -192,8 +194,6 @@
},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['svg']\n",
"\n",
"import numpy as np\n",
"\n",
"phsp = helicity_transformer(phsp_momenta)\n",
Expand All @@ -211,7 +211,7 @@
"for p, color in zip(resonances, colors, strict=True):\n",
" ax.axvline(x=p.mass, linestyle=\"dotted\", label=p.name, color=color)\n",
"ax.set_yticks([])\n",
"ax.hist(phsp[\"m_01\"], bins=100, alpha=0.5, weights=intensities)\n",
"ax.hist(phsp[\"m_01\"], bins=100, alpha=0.5, rasterized=True, weights=intensities)\n",
"ax.legend()\n",
"plt.show()"
]
Expand All @@ -236,7 +236,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.13"
"version": "3.13.15"
}
},
"nbformat": 4,
Expand Down
2 changes: 0 additions & 2 deletions docs/conf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
from __future__ import annotations

import os
import subprocess
import warnings
Expand Down
73 changes: 43 additions & 30 deletions docs/usage.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,9 +130,10 @@
},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['svg']\n",
"\n",
"import matplotlib.pyplot as plt\n",
"from matplotlib_inline.backend_inline import set_matplotlib_formats\n",
"\n",
"set_matplotlib_formats(\"svg\")\n",
"\n",
"fig, ax = plt.subplots(figsize=(5, 3))\n",
"fig.canvas.toolbar_visible = False\n",
Expand All @@ -141,12 +142,21 @@
"ax.set_title(\"First parameter guess\")\n",
"ax.set_xlabel(\"$x$\")\n",
"ax.set_yticks([])\n",
"bin_values, bin_edges, _ = ax.hist(data[\"x\"], bins=50, alpha=0.7, label=\"data\")\n",
"bin_values, bin_edges, _ = ax.hist(\n",
" data[\"x\"],\n",
" bins=50,\n",
" alpha=0.7,\n",
" label=\"data\",\n",
" rasterized=True,\n",
")\n",
"x_values = (bin_edges[1:] + bin_edges[:-1]) / 2\n",
"y_values = bin_values\n",
"function.update_parameters(initial_parameters)\n",
"lines = ax.plot(\n",
" x_values, function({\"x\": x_values}), c=\"red\", linewidth=2, label=\"model\"\n",
" x_values,\n",
" function({\"x\": x_values}),\n",
" color=\"red\",\n",
" linewidth=2,\n",
" label=\"model\",\n",
")\n",
"ax.legend(loc=\"upper right\")\n",
"plt.show()"
Expand Down Expand Up @@ -187,8 +197,6 @@
},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['svg']\n",
"\n",
"%matplotlib widget\n",
"import matplotlib.pyplot as plt\n",
"from matplotlib.animation import PillowWriter\n",
Expand All @@ -201,62 +209,65 @@
"class FitAnimation(Callback):\n",
" def __init__(self, data, function, x_values, output_file, estimated_iterations=140):\n",
" self.__function = function\n",
" self.__parameters = dict(function.parameters)\n",
" self.__fig, (self.__ax1, self.__ax2) = plt.subplots(\n",
" nrows=2, figsize=(7, 7), tight_layout=True\n",
" )\n",
" self.__ax2.set_yticks(np.arange(-30, 80, 10))\n",
" self.__ax1.hist(data[\"x\"], bins=50, alpha=0.7, label=\"data\")\n",
" self.__ax1.hist(data[\"x\"], bins=50, alpha=0.7, label=\"data\", rasterized=True)\n",
" self.__line = self.__ax1.plot(\n",
" x_values,\n",
" function({\"x\": x_values}),\n",
" c=\"red\",\n",
" color=\"red\",\n",
" linewidth=2,\n",
" label=\"model\",\n",
" )[0]\n",
" self.__ax1.legend(loc=\"upper right\")\n",
"\n",
" self.__par_lines = [\n",
" self.__ax2.plot(0, value, label=par)[0]\n",
" for par, value in function.parameters.items()\n",
" for par, value in self.__parameters.items()\n",
" ]\n",
" self.__ax2.set_xlim(0, estimated_iterations)\n",
" self.__ax2.set_title(\"Parameter values\")\n",
" self.__ax2.legend(\n",
" [f\"${sp.latex(sp.Symbol(par_name))}$\" for par_name in function.parameters],\n",
" [f\"${sp.latex(sp.Symbol(par_name))}$\" for par_name in self.__parameters],\n",
" loc=\"upper right\",\n",
" )\n",
"\n",
" self.__writer = PillowWriter(fps=15)\n",
" self.__writer.setup(self.__fig, outfile=output_file)\n",
"\n",
" def on_optimize_start(self, logs):\n",
" self._update_plot()\n",
" self._update_plot(logs)\n",
"\n",
" def on_optimize_end(self, logs):\n",
" self._update_plot()\n",
" self._update_plot(logs)\n",
" self.__writer.finish()\n",
"\n",
" def on_iteration_end(self, iteration, logs):\n",
" self._update_plot()\n",
" self._update_plot(logs)\n",
" self.__writer.finish()\n",
"\n",
" def on_function_call_end(self, function_call, logs):\n",
" self._update_plot()\n",
" self._update_plot(logs)\n",
"\n",
" def _update_plot(self):\n",
" def _update_plot(self, logs):\n",
" if logs is not None:\n",
" self.__parameters.update(logs[\"parameters\"])\n",
" self._update_parametrization_plot()\n",
" self._update_traceback()\n",
" self.__writer.grab_frame()\n",
"\n",
" def _update_parametrization_plot(self):\n",
" title = self._render_parameters(self.__function.parameters)\n",
" title = self._render_parameters(self.__parameters)\n",
" self.__ax1.set_title(title)\n",
" self.__line.set_ydata(self.__function({\"x\": x_values}))\n",
" self.__line.set_ydata(self.__function({\"x\": x_values}, self.__parameters))\n",
"\n",
" def _update_traceback(self):\n",
" for line in self.__par_lines:\n",
" par_name = line.get_label()\n",
" new_value = function.parameters[par_name]\n",
" new_value = self.__parameters[par_name]\n",
" x = line.get_xdata()\n",
" x = [*x, x[-1] + 1]\n",
" y = [*line.get_ydata(), new_value]\n",
Expand Down Expand Up @@ -632,14 +643,18 @@
},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['png']\n",
"\n",
"fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(7, 4.3))\n",
"fig.canvas.toolbar_visible = False\n",
"fig.canvas.header_visible = False\n",
"fig.canvas.footer_visible = False\n",
"ax1.hist2d(*cartesian_data.values(), bins=100, cmap=\"coolwarm\")\n",
"ax2.hist2d(polar_data[\"phi\"], polar_data[\"r\"], bins=100, cmap=\"coolwarm\")\n",
"ax1.hist2d(*cartesian_data.values(), bins=100, cmap=\"coolwarm\", rasterized=True)\n",
"ax2.hist2d(\n",
" polar_data[\"phi\"],\n",
" polar_data[\"r\"],\n",
" bins=100,\n",
" cmap=\"coolwarm\",\n",
" rasterized=True,\n",
")\n",
"fig.suptitle(\"Hit-and-miss intensity distribution\")\n",
"ax1.set_title(\"cartesian\")\n",
"ax2.set_title(\"polar\")\n",
Expand Down Expand Up @@ -696,8 +711,6 @@
},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['png']\n",
"\n",
"%matplotlib widget\n",
"import ipywidgets\n",
"import matplotlib.pyplot as plt\n",
Expand Down Expand Up @@ -730,16 +743,16 @@
")\n",
"def plot(dphi, k_r, k_phi, sigma):\n",
" global color_mesh, X, Y\n",
" polar_function.update_parameters({\n",
" parameters = {\n",
" R\"\\Delta\\phi\": dphi,\n",
" \"k_r\": k_r,\n",
" \"k_phi\": k_phi,\n",
" \"sigma\": sigma,\n",
" })\n",
" Z = polar_function(polar_domain)\n",
" }\n",
" Z = polar_function(polar_domain, parameters)\n",
" if color_mesh is not None:\n",
" color_mesh.remove()\n",
" color_mesh = ax_interactive.pcolormesh(X, Y, Z, cmap=\"coolwarm\")"
" color_mesh = ax_interactive.pcolormesh(X, Y, Z, cmap=\"coolwarm\", rasterized=True)"
]
},
{
Expand Down Expand Up @@ -793,7 +806,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.12"
"version": "3.13.15"
}
},
"nbformat": 4,
Expand Down
Loading