Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/main' into create-pull-request/p…
Browse files Browse the repository at this point in the history 
…atch-1709724303
  • Loading branch information
@redeboer
redeboer committed Mar 7, 2024
2 parents 42e4ec1 + 00f79fa commit 604f9fa
Show file tree
Hide file tree
Showing 6 changed files with 113 additions and 166 deletions.
2 changes: 1 addition & 1 deletion benchmarks/ampform.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()


Expand Down
255 changes: 105 additions & 150 deletions docs/amplitude-analysis.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down Expand Up @@ -173,6 +193,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"(build-amplitude-model)=\n",
"### 1.2 Build amplitude model"
]
},
Expand Down Expand Up @@ -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."
]
},
{
Expand All @@ -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()"
]
},
Expand Down Expand Up @@ -778,7 +799,7 @@
"source": [
":::{seealso}\n",
"\n",
"[](#extract-intensity-components)\n",
"[](#fit-fractions)\n",
"\n",
":::"
]
Expand Down Expand Up @@ -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 <sympy.core.expr.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:"
]
},
{
Expand All @@ -1863,139 +1813,144 @@
"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"
]
},
"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()"
]
},
{
"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"
]
},
{
Expand Down Expand Up @@ -2032,7 +1987,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.18"
"version": "3.10.13"
}
},
"nbformat": 4,
Expand Down
Loading

0 comments on commit 604f9fa

Please sign in to comment.