diff --git a/docs/amplitude-analysis.ipynb b/docs/amplitude-analysis.ipynb index e5889ba9..1ed6a1bc 100644 --- a/docs/amplitude-analysis.ipynb +++ b/docs/amplitude-analysis.ipynb @@ -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", @@ -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", @@ -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." ] }, { @@ -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)" ] }, { @@ -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", @@ -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." ] }, { @@ -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)" ] }, { @@ -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:" ] }, { @@ -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", + " }" ] }, { @@ -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." ] }, { @@ -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", "}" ] @@ -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", @@ -1908,7 +1910,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.13" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/amplitude-analysis/analytic-continuation.ipynb b/docs/amplitude-analysis/analytic-continuation.ipynb index fbc23a55..faa9c349 100644 --- a/docs/amplitude-analysis/analytic-continuation.ipynb +++ b/docs/amplitude-analysis/analytic-continuation.ipynb @@ -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", @@ -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\")" ] }, @@ -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", @@ -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()" ] @@ -236,7 +236,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.13" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/conf.py b/docs/conf.py index 2df67045..571459b0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import os import subprocess import warnings diff --git a/docs/usage.ipynb b/docs/usage.ipynb index 4ebebbb3..65354725 100644 --- a/docs/usage.ipynb +++ b/docs/usage.ipynb @@ -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", @@ -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()" @@ -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", @@ -201,15 +209,16 @@ "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", @@ -217,12 +226,12 @@ "\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", @@ -230,33 +239,35 @@ " 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", @@ -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", @@ -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", @@ -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)" ] }, { @@ -793,7 +806,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/usage/basics.ipynb b/docs/usage/basics.ipynb index ae16c143..48272fcb 100644 --- a/docs/usage/basics.ipynb +++ b/docs/usage/basics.ipynb @@ -46,6 +46,7 @@ "import sympy as sp\n", "from IPython.display import display\n", "from matplotlib import MatplotlibDeprecationWarning\n", + "from matplotlib_inline.backend_inline import set_matplotlib_formats\n", "from sympy.plotting import plot3d\n", "\n", "from tensorwaves.estimator import UnbinnedNLL\n", @@ -60,6 +61,7 @@ "from tensorwaves.optimizer.scipy import ScipyMinimizer\n", "\n", "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", + "set_matplotlib_formats(\"svg\")\n", "warnings.filterwarnings(\"ignore\", category=MatplotlibDeprecationWarning)" ] }, @@ -151,8 +153,6 @@ "metadata": {}, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", "parameter_defaults = {\n", " a: 0.15,\n", " b: 0.05,\n", @@ -255,7 +255,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The {meth}`.ParametrizedFunction.__call__` takes a {class}`dict` of variable names (here, `\"x\"` only) to the value(s) that should be used in their place." + "The {meth}`.ParametrizedFunction.__call__` takes a {class}`dict` of variable names (here, `\"x\"` only) to the value(s) that should be used in their place.\n", + "\n", + "A {class}`.ParametrizedFunction` is **immutable**: evaluating it never changes the function itself. There are therefore two ways of using other parameter values than the defaults: pass them as second argument to {meth}`~.ParametrizedFunction.__call__` if you need them for one evaluation only, or use {meth}`~.ParametrizedFunction.with_parameters` to create a new function that carries them as its defaults. Both are cheap, because the lambdified backend function is shared between the two." ] }, { @@ -296,8 +298,6 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", "fig, ax = plt.subplots()\n", "ax.scatter(x_values, y_values)\n", "ax.set_xlabel(\"$x$\")\n", @@ -353,8 +353,6 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", "x_domain = np.linspace(0, 5, num=200)\n", "y_values = function_1d({\"x\": x_domain})\n", "fig, ax = plt.subplots(figsize=(8, 5))\n", @@ -497,14 +495,13 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", "plt.hist(\n", " domain[\"x\"],\n", " bins=200,\n", " density=True,\n", " alpha=0.5,\n", " label=\"uniform\",\n", + " rasterized=True,\n", ")\n", "plt.hist(\n", " domain[\"x\"],\n", @@ -513,6 +510,7 @@ " alpha=0.5,\n", " density=True,\n", " label=\"weighted with $f$\",\n", + " rasterized=True,\n", ")\n", "plt.legend()\n", "plt.show()" @@ -570,9 +568,7 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", - "plt.hist(data[\"x\"], bins=200)\n", + "plt.hist(data[\"x\"], bins=200, rasterized=True)\n", "plt.show()" ] }, @@ -589,7 +585,7 @@ "source": [ "For the rest, the procedure is really just the same as that sketched in {ref}`compwa-step-3`.\n", "\n", - "We tweak the parameters a bit, then use {meth}`.ParametrizedBackendFunction.update_parameters` to change the function..." + "We tweak the parameters a bit, then use {meth}`.ParametrizedBackendFunction.with_parameters` to create a new function with these parameter values, leaving `function_1d` itself untouched..." ] }, { @@ -606,7 +602,7 @@ " \"sigma_0\": 0.4,\n", " \"sigma_1\": 0.4,\n", "}\n", - "function_1d.update_parameters(initial_parameters)" + "initial_function_1d = function_1d.with_parameters(initial_parameters)" ] }, { @@ -629,16 +625,15 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", - "plt.hist(data[\"x\"], bins=200, density=True)\n", + "plt.hist(data[\"x\"], bins=200, density=True, rasterized=True)\n", "plt.hist(\n", " domain[\"x\"],\n", - " weights=np.array(function_1d(domain)),\n", + " weights=np.array(initial_function_1d(domain)),\n", " bins=200,\n", " histtype=\"step\",\n", " color=\"red\",\n", " density=True,\n", + " rasterized=True,\n", ")\n", "plt.show()" ] @@ -723,16 +718,6 @@ "And again, we have a look at the resulting fit, as well as what happened during the optimization." ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "optimized_parameters = fit_result.parameter_values\n", - "function_1d.update_parameters(optimized_parameters)" - ] - }, { "cell_type": "code", "execution_count": null, @@ -746,16 +731,17 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", + "optimized_function_1d = function_1d.with_parameters(fit_result.parameter_values)\n", "\n", - "plt.hist(data[\"x\"], bins=200, density=True)\n", + "plt.hist(data[\"x\"], bins=200, density=True, rasterized=True)\n", "plt.hist(\n", " domain[\"x\"],\n", - " weights=np.array(function_1d(domain)),\n", + " weights=np.array(optimized_function_1d(domain)),\n", " bins=200,\n", " histtype=\"step\",\n", " color=\"red\",\n", " density=True,\n", + " rasterized=True,\n", ")\n", "plt.show()" ] @@ -816,8 +802,6 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", "fit_traceback = pd.read_csv(\"traceback-1D.csv\")\n", "fig, (ax1, ax2) = plt.subplots(\n", " nrows=2, figsize=(7, 9), sharex=True, gridspec_kw={\"height_ratios\": [1, 2]}\n", @@ -913,11 +897,12 @@ "metadata": {}, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", "y_range = (y, -sp.pi, +sp.pi)\n", "substituted_expr_2d = expression_2d.subs(parameter_defaults)\n", - "plot3d(substituted_expr_2d, x_range, y_range)\n", + "plot_2d = plot3d(substituted_expr_2d, x_range, y_range, show=False)\n", + "plot_2d.process_series()\n", + "for surface in plot_2d.ax.collections:\n", + " surface.set_rasterized(True)\n", "plt.show()" ] }, @@ -967,14 +952,13 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", "fig, axes = plt.subplots(1, 2, figsize=(8, 3))\n", "intensities = np.array(function_2d(domain_2d))\n", "kwargs = {\n", " \"weights\": intensities,\n", " \"bins\": 100,\n", " \"density\": True,\n", + " \"rasterized\": True,\n", "}\n", "axes[0].hist(domain_2d[\"x\"], **kwargs)\n", "axes[1].hist(domain_2d[\"y\"], **kwargs)\n", @@ -1007,8 +991,7 @@ " \"omega\": 0.35,\n", " \"sigma_0\": 0.4,\n", " \"sigma_1\": 0.4,\n", - "}\n", - "function_2d.update_parameters(initial_parameters)" + "}" ] }, { @@ -1024,17 +1007,20 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", "fig, axes = plt.subplots(ncols=2, figsize=(9, 4), sharey=True, tight_layout=True)\n", - "axes[0].hist2d(**data_2d, bins=50)\n", - "axes[1].hist2d(**domain_2d, weights=function_2d(domain_2d), bins=50)\n", + "axes[0].hist2d(**data_2d, bins=50, rasterized=True)\n", + "axes[1].hist2d(\n", + " **domain_2d,\n", + " weights=function_2d(domain_2d, initial_parameters),\n", + " bins=50,\n", + " rasterized=True,\n", + ")\n", "axes[0].set_xlabel(\"$x$\")\n", "axes[0].set_ylim([-3, +3])\n", "axes[1].set_xlabel(\"$x$\")\n", "axes[0].set_ylabel(\"$y$\")\n", "axes[0].set_title(\"Data sample\")\n", - "axes[1].set_title(\"Function with optimized parameters\")\n", + "axes[1].set_title(\"Function with initial parameters\")\n", "plt.show()" ] }, @@ -1148,17 +1134,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "If we update the parameters in the {class}`.ParametrizedFunction` with the optimized parameter values found by the {class}`.Optimizer`, we can compare the data distribution with the function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "optimized_parameters = fit_result.parameter_values\n", - "function_2d.update_parameters(optimized_parameters)" + "Here, we do not need a new function: passing the optimized parameter values found by the {class}`.Optimizer` to {meth}`~.ParametrizedFunction.__call__` is enough to compare the data distribution with the function." ] }, { @@ -1174,12 +1150,15 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", "fig, axes = plt.subplots(1, 2, figsize=(9, 4), sharey=True, tight_layout=True)\n", "fig.suptitle(\"Final fit result\")\n", - "axes[0].hist2d(**data_2d, bins=50)\n", - "axes[1].hist2d(**domain_2d, weights=function_2d(domain_2d), bins=50)\n", + "axes[0].hist2d(**data_2d, bins=50, rasterized=True)\n", + "axes[1].hist2d(\n", + " **domain_2d,\n", + " weights=function_2d(domain_2d, fit_result.parameter_values),\n", + " bins=50,\n", + " rasterized=True,\n", + ")\n", "axes[0].set_xlabel(\"$x$\")\n", "axes[0].set_ylim([-3, +3])\n", "axes[1].set_xlabel(\"$x$\")\n", @@ -1210,8 +1189,6 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", "minuit_traceback = pd.read_csv(\"traceback.csv\")\n", "scipy_traceback = pd.read_csv(\"traceback-scipy.csv\")\n", "fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(\n", @@ -1263,7 +1240,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.13" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/usage/binned-fit.ipynb b/docs/usage/binned-fit.ipynb index 7d858ff0..e4540d03 100644 --- a/docs/usage/binned-fit.ipynb +++ b/docs/usage/binned-fit.ipynb @@ -49,14 +49,15 @@ "metadata": {}, "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=(8, 5))\n", "ax.set_xlabel(\"$x$\")\n", "n_bins = 50\n", - "bin_values, bin_edges, _ = ax.hist(x_distribution, bins=n_bins)\n", + "bin_values, bin_edges, _ = ax.hist(x_distribution, bins=n_bins, rasterized=True)\n", "plt.show()" ] }, @@ -126,7 +127,6 @@ }, "outputs": [], "source": [ - "function.update_parameters(initial_parameters)\n", "fig, ax = plt.subplots(figsize=(8, 5))\n", "ax.set_xlabel(\"$x$\")\n", "ax.hist(x_distribution, bins=n_bins, label=\"Data distribution\")\n", @@ -198,11 +198,15 @@ }, "outputs": [], "source": [ - "function.update_parameters(fit_result.parameter_values)\n", "fig, ax = plt.subplots(figsize=(8, 5))\n", "ax.set_xlabel(\"$x$\")\n", "ax.hist(x_distribution, bins=n_bins, label=\"Data distribution\")\n", - "ax.plot(x_values, function({\"x\": x_values}), label=\"Optimized model\", c=\"red\")\n", + "ax.plot(\n", + " x_values,\n", + " function({\"x\": x_values}, fit_result.parameter_values),\n", + " color=\"red\",\n", + " label=\"Optimized model\",\n", + ")\n", "ax.legend()\n", "plt.show()" ] @@ -270,7 +274,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/usage/chi-squared.ipynb b/docs/usage/chi-squared.ipynb index 72fa45f0..5749b814 100644 --- a/docs/usage/chi-squared.ipynb +++ b/docs/usage/chi-squared.ipynb @@ -87,7 +87,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "To make the fit a bit more interesting, we give the {attr}`~.ParametrizedFunction.parameters` a different value than the ones we used to generate the $\\mathbf{y}$-values with." + "To make the fit a bit more interesting, we create a new function with different values for the {attr}`~.ParametrizedFunction.parameters` than the ones we used to generate the $\\mathbf{y}$-values with. A {class}`.ParametrizedFunction` is immutable, so the original `function` still carries the original parameter values." ] }, { @@ -96,9 +96,8 @@ "metadata": {}, "outputs": [], "source": [ - "original_parameters = function.parameters\n", "initial_parameters = {\"a\": -25, \"b\": 1.5, \"c\": 2.6}\n", - "function.update_parameters(initial_parameters)" + "initial_function = function.with_parameters(initial_parameters)" ] }, { @@ -114,9 +113,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", "\n", "def compare_model(function, x_values, observed_y):\n", @@ -125,11 +125,11 @@ " ax.plot(\n", " linear_domain[\"x\"],\n", " function(linear_domain),\n", - " c=\"red\",\n", + " color=\"red\",\n", " linewidth=3,\n", " label=\"initial fit model\",\n", " )\n", - " ax.scatter(x_values, observed_y, s=2, label=\"generated data\")\n", + " ax.scatter(x_values, observed_y, s=2, label=\"generated data\", rasterized=True)\n", " ax.set_xlabel(\"$x$\")\n", " ax.set_ylabel(\"$y$\")\n", " ax.set_ylim((-30, 50))\n", @@ -137,7 +137,7 @@ " plt.show()\n", "\n", "\n", - "compare_model(function, x_values, observed_y)" + "compare_model(initial_function, x_values, observed_y)" ] }, { @@ -191,7 +191,7 @@ "metadata": {}, "outputs": [], "source": [ - "original_parameters" + "function.parameters" ] }, { @@ -207,7 +207,8 @@ }, "outputs": [], "source": [ - "compare_model(function, x_values, observed_y)" + "optimized_function = function.with_parameters(fit_result.parameter_values)\n", + "compare_model(optimized_function, x_values, observed_y)" ] } ], @@ -230,7 +231,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/docs/usage/unbinned-fit.ipynb b/docs/usage/unbinned-fit.ipynb index a59b8126..ede3e914 100644 --- a/docs/usage/unbinned-fit.ipynb +++ b/docs/usage/unbinned-fit.ipynb @@ -55,12 +55,12 @@ }, "outputs": [], "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", "import matplotlib as mpl\n", "import matplotlib.pyplot as plt\n", + "from matplotlib_inline.backend_inline import set_matplotlib_formats\n", "\n", "mpl.rcParams[\"figure.dpi\"] = 300\n", + "set_matplotlib_formats(\"svg\")\n", "\n", "fig, ((ax_x, empty), (ax2d, ax_y)) = plt.subplots(\n", " ncols=2,\n", @@ -85,7 +85,11 @@ "\n", "bins_x, bins_y = 80, 50\n", "bin_values, bin_edges_x, bin_edges_y, _ = ax2d.hist2d(\n", - " data[\"x\"], data[\"y\"], bins=(bins_x, bins_y), cmap=\"coolwarm\"\n", + " data[\"x\"],\n", + " data[\"y\"],\n", + " bins=(bins_x, bins_y),\n", + " cmap=\"coolwarm\",\n", + " rasterized=True,\n", ")\n", "xlim = 0, 4\n", "ylim = -3, +3\n", @@ -182,7 +186,6 @@ "Y = np.linspace(*ylim, bins_y)\n", "X, Y = np.meshgrid(X, Y)\n", "\n", - "function.update_parameters(initial_parameters)\n", "Z = function({\"x\": X, \"y\": Y})\n", "\n", "fig, (ax1, ax2) = plt.subplots(figsize=(8, 7), nrows=2, sharex=True, tight_layout=True)\n", @@ -258,8 +261,7 @@ }, "outputs": [], "source": [ - "function.update_parameters(fit_result.parameter_values)\n", - "Z = function({\"x\": X, \"y\": Y})\n", + "Z = function({\"x\": X, \"y\": Y}, fit_result.parameter_values)\n", "\n", "fig, (ax1, ax2) = plt.subplots(figsize=(8, 7), nrows=2, sharex=True, tight_layout=True)\n", "ax1.set_title(\"Data distribution\")\n", @@ -292,7 +294,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.12" + "version": "3.13.15" } }, "nbformat": 4, diff --git a/src/tensorwaves/estimator.py b/src/tensorwaves/estimator.py index 3e3ef68e..dbc3f28c 100644 --- a/src/tensorwaves/estimator.py +++ b/src/tensorwaves/estimator.py @@ -146,8 +146,7 @@ def __init__( self.__sum = find_function("sum", backend) def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: - self.__function.update_parameters(parameters) - computed_values = self.__function(self.__domain) + computed_values = self.__function(self.__domain, parameters) chi_squared = self.__weights * (computed_values - self.__observed_values) ** 2 return self.__sum(chi_squared) @@ -213,9 +212,8 @@ def __init__( self.__phsp_volume = phsp_volume def __call__(self, parameters: Mapping[str, ParameterValue]) -> float: - self.__function.update_parameters(parameters) - data_intensities = self.__function(self.__data) - phsp_intensities = self.__function(self.__phsp) + data_intensities = self.__function(self.__data, parameters) + phsp_intensities = self.__function(self.__phsp, parameters) if self.__phsp_weights is not None: phsp_intensities *= self.__phsp_weights normalization_integral = self.__phsp_volume * self.__mean(phsp_intensities) diff --git a/src/tensorwaves/function/__init__.py b/src/tensorwaves/function/__init__.py index 4b008013..325e2e4b 100644 --- a/src/tensorwaves/function/__init__.py +++ b/src/tensorwaves/function/__init__.py @@ -106,8 +106,12 @@ def __init__( self.__function = PositionalArgumentFunction(function, argument_order) self.__parameters = dict(parameters) - def __call__(self, data: DataSample) -> np.ndarray: - extended_data = {**data, **self.__parameters} + def __call__( + self, + data: DataSample, + parameters: Mapping[str, ParameterValue] | None = None, + ) -> np.ndarray: + extended_data = {**data, **self.__merge_parameters(parameters)} return self.__function(extended_data) # ty:ignore[invalid-argument-type] @property @@ -122,8 +126,21 @@ def argument_order(self) -> tuple[str, ...]: def parameters(self) -> dict[str, ParameterValue]: return dict(self.__parameters) - def update_parameters(self, new_parameters: Mapping[str, ParameterValue]) -> None: - over_defined = set(new_parameters) - set(self.__parameters) + def with_parameters( + self, parameters: Mapping[str, ParameterValue] + ) -> ParametrizedBackendFunction: + return ParametrizedBackendFunction( + function=self.function, + argument_order=self.argument_order, + parameters=self.__merge_parameters(parameters), + ) + + def __merge_parameters( + self, parameters: Mapping[str, ParameterValue] | None + ) -> dict[str, ParameterValue]: + if parameters is None: + return self.__parameters + over_defined = set(parameters) - set(self.__parameters) if over_defined: sep = "\n " parameter_listing = f"{sep}".join(sorted(self.__parameters)) @@ -132,7 +149,7 @@ def update_parameters(self, new_parameters: Mapping[str, ParameterValue]) -> Non f" Expecting one of:{sep}{parameter_listing}" ) raise ValueError(msg) - self.__parameters.update(new_parameters) + return {**self.__parameters, **parameters} def get_source_code(function: Function) -> str: diff --git a/src/tensorwaves/function/sympy/__init__.py b/src/tensorwaves/function/sympy/__init__.py index 28bcad82..868b6daa 100644 --- a/src/tensorwaves/function/sympy/__init__.py +++ b/src/tensorwaves/function/sympy/__init__.py @@ -119,8 +119,10 @@ def create_parametrized_function( # ruff:ignore[too-many-arguments] ... ) >>> array = np.linspace(0, 1, num=5) >>> data = {"x": array, "y": array} - >>> function.update_parameters({"b": 1}) - >>> function(data).tolist() + >>> function(data, parameters={"b": 1}).tolist() + [0.0, 0.0, 0.0, 0.0, 0.0] + >>> function_b1 = function.with_parameters({"b": 1}) + >>> function_b1(data).tolist() [0.0, 0.0, 0.0, 0.0, 0.0] """ expression = _substitute_matrix_elements(expression) diff --git a/src/tensorwaves/interface.py b/src/tensorwaves/interface.py index 6a97b784..18fa37ac 100644 --- a/src/tensorwaves/interface.py +++ b/src/tensorwaves/interface.py @@ -43,26 +43,43 @@ def __call__(self, data: InputType) -> OutputType: ... class ParametrizedFunction(Function[InputType, OutputType]): - """Interface of a callable function. + """Interface of a callable function with parameters. A `ParametrizedFunction` identifies certain variables in a mathematical expression as **parameters**. Remaining variables are considered **domain variables**. Domain - variables are the argument of the evaluation (see - :func:`~ParametrizedFunction.__call__`), while the parameters are controlled via - :attr:`parameters` (getter) and :meth:`update_parameters` (setter). This mechanism - is especially important for an `Estimator`. + variables are the first argument of the evaluation (see + :func:`~ParametrizedFunction.__call__`), while parameter values can be passed as the + second argument. Parameter values that are not provided at the call fall back to the + default values in :attr:`parameters`. A `ParametrizedFunction` is immutable: a call + never affects later calls, which makes it thread-safe and safe to trace for JIT + compilers like :code:`jax.jit`. Use :meth:`with_parameters` to create a new function + with different default parameter values. .. automethod:: __call__ """ + @abstractmethod + def __call__( + self, + data: InputType, + parameters: Mapping[str, ParameterValue] | None = None, + ) -> OutputType: + """Evaluate the function over :code:`data` for these parameter values. + + Given parameter values are merged with the defaults in :attr:`parameters` for + this evaluation only. + """ + @property @abstractmethod def parameters(self) -> dict[str, ParameterValue]: - """`dict` of parameters.""" + """`dict` of default parameter values.""" @abstractmethod - def update_parameters(self, new_parameters: Mapping[str, ParameterValue]) -> None: - """Update the collection of parameters.""" + def with_parameters( + self, parameters: Mapping[str, ParameterValue] + ) -> ParametrizedFunction[InputType, OutputType]: + """Create a new function with updated default parameter values.""" class DataTransformer(Function[DataSample, DataSample]): diff --git a/tests/function/test_function.py b/tests/function/test_function.py index d25e66f1..bac5cf3c 100644 --- a/tests/function/test_function.py +++ b/tests/function/test_function.py @@ -58,23 +58,45 @@ def test_call( def test_function(self, function: ParametrizedBackendFunction): assert callable(function.function) - def test_update_parameter(self): - initial_parameter_values = {"a": 1, "b": 1} + def test_call_with_parameters(self): + initial_parameter_values = {"a": 1.0, "b": 2.0} func = ParametrizedBackendFunction( lambda a, b, x: a * x + b, argument_order=("a", "b", "x"), parameters=initial_parameter_values, ) + data: DataSample = {"x": np.array([0.0, 1.0, 2.0])} + np.testing.assert_array_equal(func(data), [2.0, 3.0, 4.0]) + np.testing.assert_array_equal( + func(data, parameters={"a": -1.0}), [2.0, 1.0, 0.0] + ) with pytest.raises( ValueError, match=r"^Parameters {'c'} do not exist in function arguments\.", ): - func.update_parameters({"a": 2, "c": 1}) + func(data, parameters={"a": 2.0, "c": 1.0}) + assert func.parameters == initial_parameter_values + np.testing.assert_array_equal(func(data), [2.0, 3.0, 4.0]) + + def test_with_parameters(self): + initial_parameter_values = {"a": 1.0, "b": 2.0} + func = ParametrizedBackendFunction( + lambda a, b, x: a * x + b, + argument_order=("a", "b", "x"), + parameters=initial_parameter_values, + ) + new_func = func.with_parameters({"a": 2.0}) + assert new_func is not func + assert new_func.parameters == {"a": 2.0, "b": 2.0} + assert new_func.function is func.function assert func.parameters == initial_parameter_values - new_parameter_values = {"a": 2, "b": 2} - func.update_parameters(new_parameter_values) - assert func.parameters == new_parameter_values - assert new_parameter_values != initial_parameter_values + data: DataSample = {"x": np.array([0.0, 1.0, 2.0])} + np.testing.assert_array_equal(new_func(data), [2.0, 4.0, 6.0]) + with pytest.raises( + ValueError, + match=r"^Parameters {'c'} do not exist in function arguments\.", + ): + func.with_parameters({"c": 1.0}) class TestPositionalArgumentFunction: diff --git a/tests/test_estimator.py b/tests/test_estimator.py index 6f41f77f..ef21a52e 100644 --- a/tests/test_estimator.py +++ b/tests/test_estimator.py @@ -32,6 +32,7 @@ def test_call(self, backend): assert estimator({}) == 0 assert estimator({"b": 2}) == 5.0 assert estimator({"a": 1, "b": 2}) == 14.0 + assert function.parameters == {"a": 0, "b": 1}, "estimator call is not pure" estimator = ChiSquared( function, x_data, @@ -213,6 +214,7 @@ def test_sympy_unbinned_nll( true_params: dict[str, ParameterValue], phsp: DataSample, ): + original_parameters = function.parameters estimator = UnbinnedNLL( function, data, @@ -224,6 +226,7 @@ def test_sympy_unbinned_nll( estimator, initial_parameters=true_params, ) + assert function.parameters == original_parameters, "optimize() is not pure" par_values = fit_result.parameter_values par_errors = fit_result.parameter_errors