diff --git a/benchmarks/ampform.py b/benchmarks/ampform.py index 54f76c16..981688e4 100644 --- a/benchmarks/ampform.py +++ b/benchmarks/ampform.py @@ -49,7 +49,7 @@ def formulate_amplitude_model( builder = ampform.get_builder(reaction) for name in reaction.get_intermediate_particles().names: - builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff) + builder.dynamics.assign(name, create_relativistic_breit_wigner_with_ff) return builder.formulate() diff --git a/docs/amplitude-analysis.ipynb b/docs/amplitude-analysis.ipynb index 18eca084..19847fad 100644 --- a/docs/amplitude-analysis.ipynb +++ b/docs/amplitude-analysis.ipynb @@ -84,6 +84,26 @@ ":::" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "editable": true, + "jupyter": { + "source_hidden": true + }, + "slideshow": { + "slide_type": "" + }, + "tags": [ + "remove-cell" + ] + }, + "outputs": [], + "source": [ + "from __future__ import annotations" + ] + }, { "cell_type": "markdown", "metadata": { @@ -173,6 +193,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ + "(build-amplitude-model)=\n", "### 1.2 Build amplitude model" ] }, @@ -242,7 +263,7 @@ "1. The coefficients for the different amplitudes are **complex** valued.\n", "2. By default there is no dynamics in the model, so it still has to be specified.\n", "\n", - "We choose to use {func}`~ampform.dynamics.relativistic_breit_wigner_with_ff` as the lineshape for all resonances and use a Blatt-Weisskopf form factor ({func}`~ampform.dynamics.builder.create_non_dynamic_with_ff`) for the production decay. The {meth}`~ampform.helicity.HelicityAmplitudeBuilder.set_dynamics` is a convenience interface for replacing the dynamics for intermediate states." + "We choose to use {func}`~ampform.dynamics.relativistic_breit_wigner_with_ff` as the lineshape for all resonances and use a Blatt-Weisskopf form factor ({func}`~ampform.dynamics.builder.create_non_dynamic_with_ff`) for the production decay. The {meth}`~ampform.helicity.DynamicsSelector.assign` method of the {attr}`~ampform.helicity.HelicityAmplitudeBuilder.dynamics` attribute is a convenience interface for replacing the dynamics for intermediate states." ] }, { @@ -258,9 +279,9 @@ " create_relativistic_breit_wigner_with_ff,\n", ")\n", "\n", - "model_builder.set_dynamics(\"J/psi(1S)\", create_non_dynamic_with_ff)\n", + "model_builder.dynamics.assign(\"J/psi(1S)\", create_non_dynamic_with_ff)\n", "for name in reaction.get_intermediate_particles().names:\n", - " model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", + " model_builder.dynamics.assign(name, create_relativistic_breit_wigner_with_ff)\n", "model = model_builder.formulate()" ] }, @@ -778,7 +799,7 @@ "source": [ ":::{seealso}\n", "\n", - "[](#extract-intensity-components)\n", + "[](#fit-fractions)\n", "\n", ":::" ] @@ -1769,92 +1790,21 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "#### Extract intensity components" + "#### Fit fractions" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Notice that the {class}`~ampform.helicity.HelicityModel` contains {attr}`~ampform.helicity.HelicityModel.components` attribute. This is {obj}`dict` of component names to {class}`sympy.Expr `s helps us to identify amplitudes and intensities in the total amplitude." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "sorted(model.components)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "model.components[\n", - " R\"A_{J/\\psi(1S)_{+1} \\to {f_{0}(1370)}_{0} \\gamma_{+1};\"\n", - " R\" {f_{0}(1370)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}\"\n", - "].subs(model.parameter_defaults).doit()" + "As we have seen [when formulating the amplitude model](#build-amplitude-model), the helicity model is built up of an incoherent sum of real-valued intensities (one for each spin combination), which are each built up of a coherent sum of complex-valued amplitudes (one for each resonance). If we want to compute what each of these amplitudes contribute to the main intensity, we should use the optimized intensity and set coefficients or helicity couplings of amplitudes that were are not interested in to zero." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "We can also compute separate sums of these components with the method {meth}`~ampform.helicity.HelicityModel.sum_components` from the original {class}`~ampform.helicity.HelicityModel`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "added_components = model.sum_components(\n", - " components=[\n", - " R\"A_{J/\\psi(1S)_{+1} \\to {f_{0}(500)}_{0} \\gamma_{+1};\"\n", - " R\" {f_{0}(500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}\",\n", - " R\"A_{J/\\psi(1S)_{+1} \\to {f_{0}(980)}_{0} \\gamma_{+1};\"\n", - " R\" {f_{0}(980)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}\",\n", - " R\"A_{J/\\psi(1S)_{+1} \\to {f_{0}(1370)}_{0} \\gamma_{+1};\"\n", - " R\" {f_{0}(1370)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}\",\n", - " R\"A_{J/\\psi(1S)_{+1} \\to {f_{0}(1500)}_{0} \\gamma_{+1};\"\n", - " R\" {f_{0}(1500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}\",\n", - " R\"A_{J/\\psi(1S)_{+1} \\to {f_{0}(1710)}_{0} \\gamma_{+1};\"\n", - " R\" {f_{0}(1710)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}\",\n", - " ]\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Just like in [](compwa-step-2.2), these _intensity components_ can each be expressed in a computational backend. We do not have to parametrize this function now that we have already optimized the parameters, so we can just substitute the {class}`~sympy.core.symbol.Symbol`s in all expression beforehand and use {func}`.create_function` instead:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "parameter_substitutions = model.parameter_defaults\n", - "for symbol in unfolded_expression.free_symbols:\n", - " optimized_value = fit_result.parameter_values.get(symbol.name)\n", - " if optimized_value is not None:\n", - " parameter_substitutions[symbol] = optimized_value" + "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:" ] }, { @@ -1863,54 +1813,90 @@ "metadata": {}, "outputs": [], "source": [ - "from tensorwaves.function.sympy import create_function\n", + "import re\n", "\n", - "function_from_amplitude_sum = create_function(\n", - " added_components.doit().subs(parameter_substitutions).subs(mass_substitutions),\n", - " backend=\"numpy\",\n", - ")\n", - "function_from_intensity = create_function(\n", - " model.components[R\"I_{J/\\psi(1S)_{+1} \\to \\gamma_{+1} \\pi^{0}_{0} \\pi^{0}_{0}}\"]\n", - " .doit()\n", - " .subs(parameter_substitutions)\n", - " .subs(mass_substitutions),\n", - " backend=\"numpy\",\n", - ")" + "from tensorwaves.interface import ParametrizedFunction\n", + "\n", + "\n", + "def compute_sub_intensity(\n", + " func: ParametrizedFunction,\n", + " 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", + "\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)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { + "editable": true, "jupyter": { "source_hidden": true }, + "slideshow": { + "slide_type": "" + }, "tags": [ - "remove-cell" + "remove-input" ] }, "outputs": [], "source": [ - "difference = np.average(\n", - " function_from_amplitude_sum(phsp) - function_from_intensity(phsp)\n", + "all_resonances_latex = [p.latex for p in resonances]\n", + "np.testing.assert_array_equal(\n", + " compute_sub_intensity(intensity_func, phsp, all_resonances_latex),\n", + " intensity_func(phsp),\n", ")\n", - "assert np.round(difference, decimals=0) == 0" + "del all_resonances_latex" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "The result is a {class}`.PositionalArgumentFunction` that can be plotted just like in [](#plot-optimized-model):" + "These functions can be used to compute and visualize the sub-intensity distributions over the phase space." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "total_intensities = intensity_func(phsp)\n", + "sub_intensities = {\n", + " p: compute_sub_intensity(intensity_func, phsp, resonances=[p.latex])\n", + " for p in resonances\n", + "}" ] }, { "cell_type": "code", "execution_count": null, "metadata": { + "editable": true, "jupyter": { "source_hidden": true }, + "slideshow": { + "slide_type": "" + }, "tags": [ "hide-input" ] @@ -1918,26 +1904,31 @@ "outputs": [], "source": [ "fig, ax = plt.subplots(figsize=(8, 5))\n", + "ax.set_xlim(0.25, 2.5)\n", + "ax.set_xlabel(R\"$m_{\\pi^0\\pi^0}$ [GeV]\")\n", + "ax.set_yticks([])\n", + "\n", "bins = 150\n", - "phsp_projection = np.real(phsp[\"m_12\"])\n", + "phsp_projection = phsp[\"m_12\"].real\n", "ax.hist(\n", " phsp_projection,\n", - " weights=np.array(optimized_function(phsp)),\n", + " weights=total_intensities,\n", " bins=bins,\n", - " alpha=0.2,\n", + " color=\"red\",\n", + " histtype=\"step\",\n", " label=\"full intensity\",\n", ")\n", "ax.hist(\n", - " phsp_projection,\n", - " weights=np.array(function_from_intensity(phsp)),\n", + " len(sub_intensities) * [phsp_projection],\n", + " weights=list(sub_intensities.values()),\n", " bins=bins,\n", - " histtype=\"step\",\n", - " label=R\"$J/\\psi(1S)_{-1} \\to \\gamma_{-1} \\pi^0 \\pi^0$\",\n", + " alpha=0.6,\n", + " label=[Rf\"${p.latex}$\" for p in sub_intensities],\n", + " stacked=True,\n", ")\n", - "ax.set_xlim(0.25, 2.5)\n", - "ax.set_xlabel(R\"$m_{\\pi^0\\pi^0}$ [GeV]\")\n", - "ax.set_yticks([])\n", - "plt.legend()\n", + "\n", + "fig.legend()\n", + "plt.tight_layout()\n", "plt.show()" ] }, @@ -1945,57 +1936,21 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Or generically, so that we can stack all the sub-intensities:" + "We can also use these functions to compute the decay rates for each resonance. Notice how the sum of the decay rates does not add up to a 100%. This is because of the strong constructive interference between the resonances in this model." ] }, { "cell_type": "code", "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, + "metadata": {}, "outputs": [], "source": [ - "import logging\n", - "import warnings\n", - "\n", - "logging.getLogger(\"tensorwaves.data\").setLevel(logging.ERROR) # hide progress bars\n", - "warnings.filterwarnings(\"ignore\") # hide negative sqrt warning\n", - "\n", - "sub_intensity_functions = {\n", - " name: create_function(\n", - " sub_expression.doit().subs(parameter_substitutions).subs(mass_substitutions),\n", - " backend=\"jax\",\n", - " )\n", - " for name, sub_expression in model.components.items()\n", - " if name.startswith(\"I\")\n", + "total_intensity = intensity_func(phsp).sum()\n", + "fit_fractions = {\n", + " resonance.name: f\"{sub_intensity.sum() / total_intensity:.1%}\"\n", + " for resonance, sub_intensity in sub_intensities.items()\n", "}\n", - "initial_state_mass = reaction.initial_state[-1].mass\n", - "final_state_masses = {i: p.mass for i, p in reaction.final_state.items()}\n", - "\n", - "masses = []\n", - "for sub_function in sub_intensity_functions.values():\n", - " sub_data_generator = IntensityDistributionGenerator(\n", - " phsp_generator,\n", - " function=sub_function,\n", - " domain_transformer=helicity_transformer,\n", - " )\n", - " sub_events = sub_data_generator.generate(5_000, rng)\n", - " sub_dataset = helicity_transformer(sub_events)\n", - " masses.append(np.real(sub_dataset[\"m_12\"]))\n", - "\n", - "fig, ax = plt.subplots(figsize=(8, 5))\n", - "plt.hist(masses, bins=100, stacked=True, alpha=0.6)\n", - "ax.set_xlim(0.25, 2.5)\n", - "ax.set_xlabel(R\"$m_{\\pi^0\\pi^0}$ [GeV]\")\n", - "ax.set_yticks([])\n", - "ax.legend(labels=[f\"${name[3:-1]}$\" for name in sub_intensity_functions])\n", - "plt.show()" + "fit_fractions" ] }, { @@ -2032,7 +1987,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.18" + "version": "3.10.13" } }, "nbformat": 4, diff --git a/docs/amplitude-analysis/analytic-continuation.ipynb b/docs/amplitude-analysis/analytic-continuation.ipynb index d7c06870..1f5f529f 100644 --- a/docs/amplitude-analysis/analytic-continuation.ipynb +++ b/docs/amplitude-analysis/analytic-continuation.ipynb @@ -188,18 +188,9 @@ " energy_dependent_width=True,\n", " phsp_factor=PhaseSpaceFactorSWave,\n", ")\n", - "model_builder.set_dynamics(\n", - " \"D0\",\n", - " create_non_dynamic_with_ff,\n", - ")\n", - "model_builder.set_dynamics(\n", - " \"a(0)(980)0\",\n", - " analytic_breit_wigner_builder,\n", - ")\n", - "model_builder.set_dynamics(\n", - " \"a(2)(1320)0\",\n", - " create_relativistic_breit_wigner_with_ff,\n", - ")\n", + "model_builder.dynamics.assign(\"D0\", create_non_dynamic_with_ff)\n", + "model_builder.dynamics.assign(\"a(0)(980)0\", analytic_breit_wigner_builder)\n", + "model_builder.dynamics.assign(\"a(2)(1320)0\", create_relativistic_breit_wigner_with_ff)\n", "model = model_builder.formulate()" ] }, diff --git a/docs/conf.py b/docs/conf.py index 3ba765a6..ae25c4bf 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -195,7 +195,7 @@ def get_tensorflow_url() -> str: } html_title = REPO_TITLE intersphinx_mapping = { - "ampform": (f"https://ampform.readthedocs.io/en/{pin('ampform')}", None), + "ampform": (f"https://ampform.readthedocs.io/{pin('ampform')}", None), "compwa": ("https://compwa.github.io", None), "graphviz": ("https://graphviz.readthedocs.io/en/stable", None), "iminuit": ("https://scikit-hep.org/iminuit", None), @@ -205,7 +205,7 @@ def get_tensorflow_url() -> str: "pandas": (f"https://pandas.pydata.org/pandas-docs/version/{pin('pandas')}", None), "pwa": ("https://pwa.readthedocs.io", None), "python": ("https://docs.python.org/3", None), - "qrules": (f"https://qrules.readthedocs.io/en/{pin('qrules')}", None), + "qrules": (f"https://qrules.readthedocs.io/{pin('qrules')}", None), "scipy": (get_scipy_url(), None), "sympy": ("https://docs.sympy.org/latest", None), "tensorflow": (get_tensorflow_url(), "tensorflow.inv"), diff --git a/docs/usage/faster-lambdify.ipynb b/docs/usage/faster-lambdify.ipynb index 8f55f4be..42266d3e 100644 --- a/docs/usage/faster-lambdify.ipynb +++ b/docs/usage/faster-lambdify.ipynb @@ -281,7 +281,7 @@ ")\n", "model_builder = ampform.get_builder(reaction)\n", "for name in reaction.get_intermediate_particles().names:\n", - " model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", + " model_builder.dynamics.assign(name, create_relativistic_breit_wigner_with_ff)\n", "model = model_builder.formulate()" ] }, diff --git a/pyproject.toml b/pyproject.toml index dd61a5a5..f0ba36f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -370,6 +370,7 @@ split-on-trailing-comma = false "D", "E402", "E703", + "F404", "N806", "N816", "PLR09",