diff --git a/.cspell.json b/.cspell.json index 2dbbd97b..ae092f87 100644 --- a/.cspell.json +++ b/.cspell.json @@ -159,6 +159,7 @@ "axhline", "axvline", "azim", + "bbox", "bdist", "bgcolor", "boldsymbol", @@ -214,6 +215,7 @@ "forall", "framealpha", "funcs", + "fvector", "getitem", "getsource", "graphviz", @@ -225,6 +227,7 @@ "heli", "hepstats", "histtype", + "hotpink", "hoverinfo", "hspace", "hypotests", @@ -242,12 +245,15 @@ "isinstance", "isnan", "isort", + "isospin", + "isrealobj", "jaxlib", "joinpath", "jpsi", "juliaup", "jupyterlab", "kernelspec", + "kmatrix", "kutschke", "lambdifier", "lambdifygenerated", @@ -295,6 +301,7 @@ "noreply", "nrows", "nsimplify", + "nstar", "numpycode", "operatorname", "pandoc", @@ -371,6 +378,7 @@ "toprettyxml", "tqdm", "treewise", + "twinx", "unevaluatable", "unsrt", "venv", diff --git a/docs/report/.gitignore b/docs/report/.gitignore index c2728bc9..98a133c7 100644 --- a/docs/report/.gitignore +++ b/docs/report/.gitignore @@ -6,3 +6,5 @@ 002-*-graph 013-graph? 018-graph + +!030/*.yml diff --git a/docs/report/000.ipynb b/docs/report/000.ipynb index f326c088..ccf08615 100644 --- a/docs/report/000.ipynb +++ b/docs/report/000.ipynb @@ -12,6 +12,7 @@ "cell_type": "markdown", "metadata": { "tags": [ + "dynamics", "lambdification", "sympy" ] @@ -917,7 +918,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.13" + "version": "3.10.14" }, "orphan": true }, diff --git a/docs/report/004.ipynb b/docs/report/004.ipynb index efea7e57..bb6546a3 100644 --- a/docs/report/004.ipynb +++ b/docs/report/004.ipynb @@ -12,6 +12,7 @@ "cell_type": "markdown", "metadata": { "tags": [ + "dynamics", "physics" ] }, diff --git a/docs/report/005.ipynb b/docs/report/005.ipynb index b3520591..c40485fe 100644 --- a/docs/report/005.ipynb +++ b/docs/report/005.ipynb @@ -12,6 +12,8 @@ "cell_type": "markdown", "metadata": { "tags": [ + "dynamics", + "K-matrix", "physics" ] }, diff --git a/docs/report/009.ipynb b/docs/report/009.ipynb index 99b0000c..6208f40e 100644 --- a/docs/report/009.ipynb +++ b/docs/report/009.ipynb @@ -12,6 +12,8 @@ "cell_type": "markdown", "metadata": { "tags": [ + "K-matrix", + "dynamics", "physics", "sympy" ] diff --git a/docs/report/010.ipynb b/docs/report/010.ipynb index b96335a8..e4c67457 100644 --- a/docs/report/010.ipynb +++ b/docs/report/010.ipynb @@ -12,6 +12,8 @@ "cell_type": "markdown", "metadata": { "tags": [ + "dynamics", + "K-matrix", "physics", "sympy" ] diff --git a/docs/report/025.ipynb b/docs/report/025.ipynb index 2bda35b4..58edec15 100644 --- a/docs/report/025.ipynb +++ b/docs/report/025.ipynb @@ -11,7 +11,9 @@ { "cell_type": "markdown", "metadata": { - "tags": [] + "tags": [ + "K-matrix" + ] }, "source": [ "::::{margin}\n", diff --git a/docs/report/026.ipynb b/docs/report/026.ipynb index 9fa8229d..e2942936 100755 --- a/docs/report/026.ipynb +++ b/docs/report/026.ipynb @@ -11,7 +11,10 @@ { "cell_type": "markdown", "metadata": { - "tags": [] + "tags": [ + "dynamics", + "K-matrix" + ] }, "source": [ "::::{margin}\n", diff --git a/docs/report/027.ipynb b/docs/report/027.ipynb index 95b40f11..688f228f 100755 --- a/docs/report/027.ipynb +++ b/docs/report/027.ipynb @@ -11,7 +11,10 @@ { "cell_type": "markdown", "metadata": { - "tags": [] + "tags": [ + "dynamics", + "K-matrix" + ] }, "source": [ "::::{margin}\n", diff --git a/docs/report/029.ipynb b/docs/report/029.ipynb index 67b17c2c..de528a8c 100644 --- a/docs/report/029.ipynb +++ b/docs/report/029.ipynb @@ -14,7 +14,8 @@ "cell_type": "markdown", "metadata": { "tags": [ - "PDG" + "dynamics", + "sympy" ] }, "source": [ diff --git a/docs/report/030.ipynb b/docs/report/030.ipynb new file mode 100644 index 00000000..bb0bdc35 --- /dev/null +++ b/docs/report/030.ipynb @@ -0,0 +1,1427 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "```{autolink-concat}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [ + "dynamics", + "K-matrix" + ] + }, + "source": [ + "::::{margin}\n", + ":::{card} Sub-intensities of P-vector amplitude model\n", + "TR-030\n", + "^^^\n", + "Sub-intensity plots for a model with $P$-vector dynamics. Also includes an investigation of phases in a $P$-vector lineshape.\n", + "+++\n", + "🚧 [compwa.github.io#278](https://github.com/ComPWA/compwa.github.io/pull/278)\n", + ":::\n", + "::::" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "# Sub-intensities of P vector" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-cell" + ] + }, + "outputs": [], + "source": [ + "%pip install -q 'qrules[viz]==0.10.2' 'tensorwaves[jax,phsp]==0.4.12' ampform==0.15.4 sympy==1.12" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import logging\n", + "import os\n", + "import re\n", + "from collections import defaultdict\n", + "from functools import lru_cache\n", + "from typing import Any\n", + "\n", + "import ampform\n", + "import attrs\n", + "import graphviz\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import qrules\n", + "import sympy as sp\n", + "from ampform.dynamics.builder import TwoBodyKinematicVariableSet\n", + "from ampform.helicity import ParameterValues\n", + "from ampform.io import aslatex, improve_latex_rendering\n", + "from ampform.kinematics.phasespace import Kallen\n", + "from ampform.sympy import perform_cached_doit, unevaluated\n", + "from attrs import define, field\n", + "from IPython.display import Math\n", + "from qrules.particle import Particle, ParticleCollection\n", + "from sympy import Abs\n", + "from tensorwaves.data import (\n", + " SympyDataTransformer,\n", + " TFPhaseSpaceGenerator,\n", + " TFUniformRealNumberGenerator,\n", + ")\n", + "from tensorwaves.function.sympy import create_parametrized_function\n", + "from tensorwaves.interface import DataSample, Function, ParametrizedFunction\n", + "\n", + "improve_latex_rendering()\n", + "logging.getLogger(\"absl\").setLevel(logging.ERROR)\n", + "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", + "plt.rc(\"font\", size=12)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "%config InlineBackend.figure_formats = ['svg']" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Studied decay" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Define N* resonances" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@lru_cache(maxsize=1)\n", + "def create_particle_database() -> ParticleCollection:\n", + " particles = qrules.load_default_particles()\n", + " for nstar in particles.filter(lambda p: p.name.startswith(\"N\")):\n", + " particles.remove(nstar)\n", + " particles += create_nstar(mass=1.82, width=0.6, parity=+1, spin=1.5, idx=1)\n", + " particles += create_nstar(mass=1.92, width=0.6, parity=+1, spin=1.5, idx=2)\n", + " return particles\n", + "\n", + "\n", + "def create_nstar(\n", + " mass: float, width: float, parity: int, spin: float, idx: int\n", + ") -> Particle:\n", + " spin = sp.Rational(spin)\n", + " parity_symbol = \"⁺\" if parity > 0 else \"⁻\"\n", + " unicode_subscripts = list(\"₀₁₂₃₄₅₆₇₈₉\")\n", + " return Particle(\n", + " name=f\"N{unicode_subscripts[idx]}({spin}{parity_symbol})\",\n", + " latex=Rf\"N_{idx}({spin.numerator}/{spin.denominator}^-)\",\n", + " pid=2024_05_00_00 + 100 * bool(parity + 1) + idx,\n", + " mass=mass,\n", + " width=width,\n", + " baryon_number=1,\n", + " charge=+1,\n", + " isospin=(0.5, +0.5),\n", + " parity=parity,\n", + " spin=1.5,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "reaction = qrules.generate_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"eta\", \"p\", \"p~\"],\n", + " allowed_intermediate_particles=[\"N\"],\n", + " allowed_interaction_types=[\"strong\"],\n", + " formalism=\"helicity\",\n", + " particle_db=create_particle_database(),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", + "graphviz.Source(dot)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Amplitude builder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Dynamics builder with X symbols of J^PC channels" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@define\n", + "class DynamicsSymbolBuilder:\n", + " collected_symbols: set[sp.Symbol, tuple[Particle, TwoBodyKinematicVariableSet]] = (\n", + " field(factory=lambda: defaultdict(set))\n", + " )\n", + "\n", + " def __call__(\n", + " self, resonance: Particle, variable_pool: TwoBodyKinematicVariableSet\n", + " ) -> tuple[sp.Expr, dict[sp.Symbol, float]]:\n", + " jp = render_jp(resonance)\n", + " charge = resonance.charge\n", + " if variable_pool.angular_momentum is not None:\n", + " L = sp.Rational(variable_pool.angular_momentum)\n", + " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}^{{l={L}}}\")\n", + " else:\n", + " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}\")\n", + " self.collected_symbols[X].add((resonance, variable_pool))\n", + " parameter_defaults = {}\n", + " return X, parameter_defaults\n", + "\n", + "\n", + "def render_jp(particle: Particle) -> str:\n", + " spin = sp.Rational(particle.spin)\n", + " j = (\n", + " str(spin)\n", + " if spin.denominator == 1\n", + " else Rf\"\\frac{{{spin.numerator}}}{{{spin.denominator}}}\"\n", + " )\n", + " if particle.parity is None:\n", + " return f\"J={j}\"\n", + " p = \"-\" if particle.parity < 0 else \"+\"\n", + " return f\"J^P={{{j}}}^{{{p}}}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "model_builder = ampform.get_builder(reaction)\n", + "model_builder.adapter.permutate_registered_topologies()\n", + "model_builder.config.scalar_initial_state_mass = True\n", + "model_builder.config.stable_final_state_ids = [0, 1, 2]\n", + "create_dynamics_symbol = DynamicsSymbolBuilder()\n", + "for resonance in reaction.get_intermediate_particles():\n", + " model_builder.set_dynamics(resonance.name, create_dynamics_symbol)\n", + "model = model_builder.formulate()\n", + "model.intensity.cleanup()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "selected_amplitudes = {\n", + " k: v for i, (k, v) in enumerate(model.amplitudes.items()) if i == 0\n", + "}\n", + "Math(aslatex(selected_amplitudes, terms_per_line=1))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "src = R\"\\begin{array}{cll}\" \"\\n\"\n", + "for symbol, resonances in create_dynamics_symbol.collected_symbols.items():\n", + " src += Rf\" {symbol} \\\\\" \"\\n\"\n", + " for p, _ in resonances:\n", + " src += Rf\" {p.latex} & m={p.mass:g}\\text{{ GeV}} & \\Gamma={p.width:g}\\text{{ GeV}} \\\\\"\n", + " src += \"\\n\"\n", + "src += R\"\\end{array}\"\n", + "Math(src)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Dynamics parametrization" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Phasespace factor" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + ":::{seealso}\n", + "**[TR-026](./026.ipynb)** and **[TR-027](./027.ipynb)** on analyticity and Riemann sheets.\n", + ":::" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Expression classes for phase space factors" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@unevaluated(real=False)\n", + "class PhaseSpaceCM(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\n", + " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " s, m1, m2 = self.args\n", + " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", + "\n", + "\n", + "@unevaluated(real=False)\n", + "class ChewMandelstam(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\n", + " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " s, m1, m2 = self.args\n", + " q = BreakupMomentum(s, m1, m2)\n", + " return (\n", + " (2 * q / sp.sqrt(s))\n", + " * sp.log(Abs((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", + "\n", + "\n", + "@unevaluated(real=False)\n", + "class BreakupMomentum(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\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(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "s, m1, m2 = sp.symbols(\"s m1 m2\", nonnegative=True)\n", + "exprs = [\n", + " PhaseSpaceCM(s, m1, m2),\n", + " ChewMandelstam(s, m1, m2),\n", + " BreakupMomentum(s, m1, m2),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Relativistic Breit-Wigner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "PARAMETERS_BW = dict(model.parameter_defaults)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_breit_wigner(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " m1 = variables.outgoing_state_mass1\n", + " m2 = variables.outgoing_state_mass2\n", + " ρ = PhaseSpaceCM(s, m1, m2)\n", + " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", + " Γ0 = [sp.Symbol(Rf\"\\Gamma_{{{p.latex}}}\") for p, _ in resonances]\n", + " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", + " expr = sum(\n", + " (β_ * m_ * Γ0_) / (m_**2 - s - m_ * Γ0_ * ρ) for m_, Γ0_, β_ in zip(m, Γ0, β)\n", + " )\n", + " for i, (resonance, _) in enumerate(resonances):\n", + " PARAMETERS_BW[β[i]] = 1 + 0j\n", + " PARAMETERS_BW[m[i]] = resonance.mass\n", + " PARAMETERS_BW[Γ0[i]] = resonance.width\n", + " return expr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "dynamics_expressions_bw = {\n", + " symbol: formulate_breit_wigner(resonances)\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + "}\n", + "model_bw = attrs.evolve(\n", + " model,\n", + " parameter_defaults=ParameterValues({\n", + " **model.parameter_defaults,\n", + " **PARAMETERS_BW,\n", + " }),\n", + ")\n", + "Math(aslatex(dynamics_expressions_bw))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### $P$ vector" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "PARAMETERS_F = dict(model.parameter_defaults)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_k_matrix(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", + " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", + " expr = sum((g_**2) / (m_**2 - s) for m_, g_ in zip(m, g))\n", + " for i, (resonance, _) in enumerate(resonances):\n", + " PARAMETERS_F[m[i]] = resonance.mass\n", + " PARAMETERS_F[g[i]] = 1\n", + " return expr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_p_vector(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", + " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", + " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", + " expr = sum((g_ * β_) / (m_**2 - s) for m_, g_, β_ in zip(m, g, β))\n", + " for i, (resonance, _) in enumerate(resonances):\n", + " PARAMETERS_F[β[i]] = 1 + 0j\n", + " PARAMETERS_F[m[i]] = resonance.mass\n", + " PARAMETERS_F[g[i]] = 1\n", + " return expr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_f_vector(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " m1 = variables.outgoing_state_mass1\n", + " m2 = variables.outgoing_state_mass2\n", + " rho = PhaseSpaceCM(s, m1, m2)\n", + " K = formulate_k_matrix(resonances)\n", + " P = formulate_p_vector(resonances)\n", + " return (1 / (1 - rho * K)) * P" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "dynamics_expressions_fvector = {\n", + " symbol: formulate_f_vector(resonances)\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + "}\n", + "model_fvector = attrs.evolve(\n", + " model,\n", + " parameter_defaults=ParameterValues({\n", + " **model.parameter_defaults,\n", + " **PARAMETERS_F,\n", + " }),\n", + ")\n", + "Math(aslatex(dynamics_expressions_fvector))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Create numerical functions" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Amplitude model function" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "full_expression_bw = perform_cached_doit(model_bw.expression).xreplace(\n", + " dynamics_expressions_bw\n", + ")\n", + "intensity_func_bw = create_parametrized_function(\n", + " expression=perform_cached_doit(full_expression_bw),\n", + " backend=\"jax\",\n", + " parameters=PARAMETERS_BW,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "full_expression_fvector = perform_cached_doit(model_fvector.expression).xreplace(\n", + " dynamics_expressions_fvector\n", + ")\n", + "intensity_func_fvector = create_parametrized_function(\n", + " expression=perform_cached_doit(full_expression_fvector),\n", + " backend=\"jax\",\n", + " parameters=PARAMETERS_F,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Dynamics function" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "mystnb": { + "code_prompt_show": "Breit–Wigner parametrization" + }, + "tags": [ + "full-width", + "hide-input" + ] + }, + "outputs": [], + "source": [ + "dynamics_expr_bw, *_ = dynamics_expressions_bw.values()\n", + "dynamics_expr_bw" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "mystnb": { + "code_prompt_show": "F-vector parametrization" + }, + "tags": [ + "full-width", + "hide-input" + ] + }, + "outputs": [], + "source": [ + "dynamics_expr_fvector, *_ = dynamics_expressions_fvector.values()\n", + "dynamics_expr_fvector.simplify(doit=False)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "dynamics_func_bw = create_parametrized_function(\n", + " expression=perform_cached_doit(dynamics_expr_bw),\n", + " backend=\"jax\",\n", + " parameters=model_bw.parameter_defaults,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "dynamics_func_fvector = create_parametrized_function(\n", + " expression=perform_cached_doit(dynamics_expr_fvector),\n", + " backend=\"jax\",\n", + " parameters=model_fvector.parameter_defaults,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Generate data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Generate phase space sample" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "rng = TFUniformRealNumberGenerator(seed=0)\n", + "phsp_generator = TFPhaseSpaceGenerator(\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", + "phsp_momenta = phsp_generator.generate(500_000, rng)\n", + "\n", + "ε = 1e-8\n", + "transformer = SympyDataTransformer.from_sympy(model.kinematic_variables, backend=\"jax\")\n", + "phsp = transformer(phsp_momenta)\n", + "phsp = {k: v + ε * 1j if re.match(r\"^m_\\d\\d$\", k) else v for k, v in phsp.items()}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Update function parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "m_res1 = 1.82\n", + "m_res2 = 1.92\n", + "g_res1 = 1\n", + "g_res2 = 1" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "toy_parameters_bw = {\n", + " R\"m_{N_1(3/2^-)}\": m_res1,\n", + " R\"m_{N_2(3/2^-)}\": m_res2,\n", + " R\"\\Gamma_{N_1(3/2^-)}\": g_res1 / m_res1,\n", + " R\"\\Gamma_{N_2(3/2^-)}\": g_res2 / m_res2,\n", + "}\n", + "dynamics_func_bw.update_parameters(toy_parameters_bw)\n", + "intensity_func_bw.update_parameters(toy_parameters_bw)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "toy_parameters_fvector = {\n", + " R\"\\beta_{N_1(3/2^-)}\": 1 + 0j,\n", + " R\"\\beta_{N_2(3/2^-)}\": 1 + 0j,\n", + " R\"m_{N_1(3/2^-)}\": m_res1,\n", + " R\"m_{N_2(3/2^-)}\": m_res2,\n", + " R\"g_{N_1(3/2^-)}\": g_res1,\n", + " R\"g_{N_2(3/2^-)}\": g_res2,\n", + "}\n", + "dynamics_func_fvector.update_parameters(toy_parameters_fvector)\n", + "intensity_func_fvector.update_parameters(toy_parameters_fvector)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Plots" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Sub-intensities" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Function for computing sub-intensities" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def compute_sub_intensity(\n", + " func: ParametrizedFunction,\n", + " input_data: DataSample,\n", + " resonances: list[str],\n", + " coupling_pattern: 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\"^{coupling_pattern}({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", + " toy_parameters = dict(func.parameters)\n", + " for par_name in func.parameters:\n", + " if re.match(name_pattern, par_name) is not None:\n", + " toy_parameters[par_name] = 0\n", + " func.update_parameters(toy_parameters)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "total_intensities_bw = intensity_func_bw(phsp)\n", + "sub_intensities_bw = {\n", + " p: compute_sub_intensity(\n", + " intensity_func_bw,\n", + " phsp,\n", + " resonances=[p.latex],\n", + " coupling_pattern=r\"\\\\beta\",\n", + " )\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + " for p, _ in resonances\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "total_intensities_fvector = intensity_func_fvector(phsp)\n", + "sub_intensities_fvector = {\n", + " p: compute_sub_intensity(\n", + " intensity_func_fvector,\n", + " phsp,\n", + " resonances=[p.latex],\n", + " coupling_pattern=r\"\\\\beta\",\n", + " )\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + " for p, _ in resonances\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Function for plotting histograms with JAX" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def fast_histogram(\n", + " data: jnp.ndarray,\n", + " weights: jnp.ndarray | None = None,\n", + " bins: int = 100,\n", + " density: bool | None = None,\n", + " fill: bool = True,\n", + " ax=plt,\n", + " **plot_kwargs,\n", + ") -> None:\n", + " bin_values, bin_edges = jnp.histogram(\n", + " data,\n", + " bins=bins,\n", + " density=density,\n", + " weights=weights,\n", + " )\n", + " if fill:\n", + " bin_rights = bin_edges[1:]\n", + " ax.fill_between(bin_rights, bin_values, step=\"pre\", **plot_kwargs)\n", + " else:\n", + " bin_mids = (bin_edges[:-1] + bin_edges[1:]) / 2\n", + " ax.step(bin_mids, bin_values, **plot_kwargs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(8, 5))\n", + "ax.set_xlim(2, 5)\n", + "ax.set_xlabel(R\"$m_{p\\eta}^{2}$ [GeV$^2$]\")\n", + "ax.set_ylabel(R\"Intensity [a. u.]\")\n", + "ax.set_yticks([])\n", + "\n", + "bins = 150\n", + "phsp_projection = np.real(phsp[\"m_01\"]) ** 2\n", + "fast_histogram(\n", + " phsp_projection,\n", + " weights=total_intensities_fvector,\n", + " alpha=0.2,\n", + " bins=bins,\n", + " color=\"hotpink\",\n", + " label=\"Full intensity $F$ vector\",\n", + " ax=ax,\n", + ")\n", + "fast_histogram(\n", + " phsp_projection,\n", + " weights=total_intensities_bw,\n", + " alpha=0.2,\n", + " bins=bins,\n", + " color=\"grey\",\n", + " label=\"Full intensity Breit-Wigner\",\n", + " ax=ax,\n", + ")\n", + "for i, (p, v) in enumerate(sub_intensities_fvector.items()):\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=v,\n", + " alpha=0.6,\n", + " bins=bins,\n", + " color=f\"C{i}\",\n", + " fill=False,\n", + " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV}}$ $F$ vector\",\n", + " linewidth=2,\n", + " ax=ax,\n", + " )\n", + "for i, (p, v) in enumerate(sub_intensities_bw.items()):\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=v,\n", + " alpha=0.6,\n", + " bins=bins,\n", + " color=f\"C{i}\",\n", + " fill=False,\n", + " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV^2}}$ Breit-Wigner\",\n", + " linestyle=\"dashed\",\n", + " ax=ax,\n", + " )\n", + "\n", + "ax.set_ylim(0, None)\n", + "fig.legend(loc=\"upper right\")\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Argand plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "ε = 1e-8\n", + "x = np.linspace(2, 5, num=400)\n", + "plot_data = {\"m_01\": np.sqrt(x) + ε * 1j}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "total_dynamics_bw = dynamics_func_bw(plot_data)\n", + "sub_dynamics_bw = {\n", + " p: compute_sub_intensity(\n", + " dynamics_func_bw,\n", + " plot_data,\n", + " resonances=[p.latex],\n", + " coupling_pattern=r\"\\\\beta\",\n", + " )\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + " for p, _ in resonances\n", + "}\n", + "total_dynamics_fvector = dynamics_func_fvector(plot_data)\n", + "sub_dynamics_fvector = {\n", + " p: compute_sub_intensity(\n", + " dynamics_func_fvector,\n", + " plot_data,\n", + " resonances=[p.latex],\n", + " coupling_pattern=r\"\\\\beta\",\n", + " )\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + " for p, _ in resonances\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "x1 = np.linspace(2.0, (m_res1**2 + m_res2**2) / 2, num=500)\n", + "x2 = np.linspace((m_res1**2 + m_res2**2) / 2, 5.0, num=500)\n", + "plot_data1 = {\"m_01\": np.sqrt(x1) + ε * 1j}\n", + "plot_data2 = {\"m_01\": np.sqrt(x2) + ε * 1j}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Function definition for plotting Argand plots" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def plot_argand(\n", + " total_func: Function, sub_funcs: dict[Particle, Function], title: str\n", + ") -> None:\n", + " fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharey=True)\n", + " fig.subplots_adjust(wspace=0.05)\n", + " fig.suptitle(title, y=0.99)\n", + " ax1, ax2 = axes\n", + " ax1.set_title(\"Total amplitude\")\n", + " ax2.set_title(\"Amplitude for resonance only\")\n", + " ax1.set_ylabel(R\"$\\text{Im}\\,F$\")\n", + " for ax in axes:\n", + " ax.axhline(0, color=\"black\", linewidth=0.5)\n", + " ax.axvline(0, color=\"black\", linewidth=0.5)\n", + " ax.set_xlabel(R\"$\\text{Re}\\,F$\")\n", + "\n", + " y1 = total_func(plot_data1)\n", + " ax1.plot(\n", + " y1.real,\n", + " y1.imag,\n", + " label=f\"Domain of {m_res1}-GeV resonance \",\n", + " color=\"C0\",\n", + " )\n", + " y2 = total_func(plot_data2)\n", + " ax1.plot(\n", + " y2.real,\n", + " y2.imag,\n", + " label=f\"Domain of {m_res2}-GeV resonance \",\n", + " color=\"C1\",\n", + " )\n", + " for i, (k, v) in enumerate(sub_funcs.items()):\n", + " ax2.plot(\n", + " v.real,\n", + " v.imag,\n", + " color=f\"C{i}\",\n", + " label=f\"Resonance at {k.mass} GeV $F$-vector\",\n", + " )\n", + "\n", + " ax1.legend(loc=\"upper left\")\n", + " fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "plot_argand(\n", + " dynamics_func_fvector,\n", + " sub_dynamics_fvector,\n", + " title=\"F vector\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "plot_argand(\n", + " dynamics_func_bw,\n", + " sub_dynamics_bw,\n", + " title=\"Breit-Wigner\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Phase" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "total_phase_bw = np.angle(total_dynamics_bw)\n", + "total_phase_fvector = np.angle(total_dynamics_fvector)\n", + "sub_phase_bw = {p: np.angle(v) for p, v in sub_dynamics_bw.items()}\n", + "sub_phase_fvector = {p: np.angle(v) for p, v in sub_dynamics_fvector.items()}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Function definition for plotting phases" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def plot_phases(\n", + " total_intensity_array: np.ndarray,\n", + " sub_intensity_arrays: dict[Particle, np.ndarray],\n", + " total_phase_array: np.ndarray,\n", + " sub_phase_arrays: dict[Particle, np.ndarray],\n", + " title: str,\n", + ") -> None:\n", + " fig, ax1 = plt.subplots(figsize=(10, 6))\n", + " ax1.set_title(title)\n", + " ax2 = ax1.twinx()\n", + " ax1.set_xlim(2.0, 5.0)\n", + " ax1.set_xlabel(R\"$m_{p\\eta}^{2}$ [GeV$^{2}$]\")\n", + " ax1.set_ylabel(\"Intensity [a. u.]\")\n", + " ax2.set_ylabel(\"Angle\")\n", + " ax1.set_yticks([])\n", + " ax2.set_ylim([-np.pi, +np.pi])\n", + " ax2.set_yticks([\n", + " -np.pi,\n", + " -np.pi / 2,\n", + " 0,\n", + " +np.pi / 2,\n", + " +np.pi,\n", + " ])\n", + " ax2.set_yticklabels([\n", + " R\"$-\\pi$\",\n", + " R\"$-\\frac{\\pi}{2}$\",\n", + " \"0\",\n", + " R\"$+\\frac{\\pi}{2}$\",\n", + " R\"$+\\pi$\",\n", + " ])\n", + " ax2.axhline(0, c=\"black\", lw=0.5)\n", + "\n", + " # Plot background histograms\n", + " phsp_projection = np.real(phsp[\"m_01\"]) ** 2\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=total_intensity_array,\n", + " bins=bins,\n", + " alpha=0.2,\n", + " color=\"gray\",\n", + " label=\"Full intensity\",\n", + " ax=ax1,\n", + " )\n", + " for i, (k, v) in enumerate(sub_intensity_arrays.items()):\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=v,\n", + " bins=bins,\n", + " alpha=0.2,\n", + " color=f\"C{i}\",\n", + " label=Rf\"Resonance at ${k.mass}\\,\\mathrm{{GeV}}$\",\n", + " ax=ax1,\n", + " )\n", + " ax1.set_ylim(0, None)\n", + "\n", + " # Plot phases\n", + " ax2.scatter(\n", + " x,\n", + " total_phase_array,\n", + " color=\"gray\",\n", + " label=\"Total Phase\",\n", + " s=18,\n", + " )\n", + " for i, (k, v) in enumerate(sub_phase_arrays.items()):\n", + " ax2.scatter(\n", + " x,\n", + " v,\n", + " alpha=0.5,\n", + " color=f\"C{i}\",\n", + " label=f\"Resonance at {k.mass} GeV\",\n", + " s=8,\n", + " )\n", + " ax2.axvline(k.mass**2, linestyle=\"dotted\", color=f\"C{i}\")\n", + "\n", + " # Add legends\n", + " fig.legend(bbox_to_anchor=(0.1, 0.9), loc=\"upper left\")\n", + " fig.tight_layout()\n", + " fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "plot_phases(\n", + " total_intensity_array=total_intensities_fvector,\n", + " sub_intensity_arrays=sub_intensities_fvector,\n", + " total_phase_array=total_phase_fvector,\n", + " sub_phase_arrays=sub_phase_fvector,\n", + " title=\"F vector\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "plot_phases(\n", + " total_intensity_array=total_intensities_bw,\n", + " sub_intensity_arrays=sub_intensities_bw,\n", + " total_phase_array=total_phase_bw,\n", + " sub_phase_arrays=sub_phase_bw,\n", + " title=\"Breit-Wigner\",\n", + ")" + ] + } + ], + "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.10.14" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/report/031.ipynb b/docs/report/031.ipynb new file mode 100644 index 00000000..57cb8c89 --- /dev/null +++ b/docs/report/031.ipynb @@ -0,0 +1,1473 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "```{autolink-concat}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [ + "dynamics", + "K-matrix" + ] + }, + "source": [ + "::::{margin}\n", + ":::{card} Single-channel amplitude model fit with $P$-vector dynamics\n", + "TR-031\n", + "^^^\n", + "Comparison between fit performance for an amplitude model with Breit–Wigner and $P$-vector dynamics. In both cases, data is generated with $P$-vector dynamics.\n", + "+++\n", + "🚧 [compwa.github.io#278](https://github.com/ComPWA/compwa.github.io/pull/278)\n", + ":::\n", + "::::" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "# P-vector model fit, single channel" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-cell" + ] + }, + "outputs": [], + "source": [ + "%pip install -q 'qrules[viz]==0.10.2' 'tensorwaves[jax,phsp]==0.4.12' ampform==0.15.4 pandas==2.2.2 sympy==1.12" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import logging\n", + "import os\n", + "import re\n", + "from collections import defaultdict\n", + "from functools import lru_cache\n", + "from typing import Any\n", + "\n", + "import ampform\n", + "import attrs\n", + "import graphviz\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pandas as pd\n", + "import qrules\n", + "import sympy as sp\n", + "from ampform.dynamics.builder import TwoBodyKinematicVariableSet\n", + "from ampform.helicity import ParameterValues\n", + "from ampform.io import aslatex, improve_latex_rendering\n", + "from ampform.kinematics.phasespace import Kallen\n", + "from ampform.sympy import perform_cached_doit, unevaluated\n", + "from attrs import define, field\n", + "from IPython.display import Math\n", + "from matplotlib import cm\n", + "from qrules.particle import Particle, ParticleCollection\n", + "from sympy import Abs\n", + "from tensorwaves.data import (\n", + " IntensityDistributionGenerator,\n", + " SympyDataTransformer,\n", + " TFPhaseSpaceGenerator,\n", + " TFUniformRealNumberGenerator,\n", + " TFWeightedPhaseSpaceGenerator,\n", + ")\n", + "from tensorwaves.estimator import UnbinnedNLL\n", + "from tensorwaves.function.sympy import create_parametrized_function\n", + "from tensorwaves.interface import DataSample, FitResult, ParametrizedFunction\n", + "from tensorwaves.optimizer import Minuit2\n", + "\n", + "improve_latex_rendering()\n", + "logging.getLogger(\"absl\").setLevel(logging.ERROR)\n", + "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", + "plt.rc(\"font\", size=12)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "%config InlineBackend.figure_formats = ['svg']" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Studied decay" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Define N* resonances" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@lru_cache(maxsize=1)\n", + "def create_particle_database() -> ParticleCollection:\n", + " particles = qrules.load_default_particles()\n", + " for nstar in particles.filter(lambda p: p.name.startswith(\"N\")):\n", + " particles.remove(nstar)\n", + " particles += create_nstar(mass=1.65, width=0.6, parity=-1, spin=0.5, idx=1)\n", + " particles += create_nstar(mass=1.75, width=0.6, parity=-1, spin=0.5, idx=2)\n", + " particles += create_nstar(mass=1.82, width=0.6, parity=+1, spin=1.5, idx=1)\n", + " particles += create_nstar(mass=1.92, width=0.6, parity=+1, spin=1.5, idx=2)\n", + " return particles\n", + "\n", + "\n", + "def create_nstar(\n", + " mass: float, width: float, parity: int, spin: float, idx: int\n", + ") -> Particle:\n", + " spin = sp.Rational(spin)\n", + " parity_symbol = \"⁺\" if parity > 0 else \"⁻\"\n", + " unicode_subscripts = list(\"₀₁₂₃₄₅₆₇₈₉\")\n", + " return Particle(\n", + " name=f\"N{unicode_subscripts[idx]}({spin}{parity_symbol})\",\n", + " latex=Rf\"N_{idx}({spin.numerator}/{spin.denominator}^-)\",\n", + " pid=2024_05_00_00 + 100 * bool(parity + 1) + idx,\n", + " mass=mass,\n", + " width=width,\n", + " baryon_number=1,\n", + " charge=+1,\n", + " isospin=(0.5, +0.5),\n", + " parity=parity,\n", + " spin=1.5,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "reaction = qrules.generate_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"eta\", \"p\", \"p~\"],\n", + " allowed_intermediate_particles=[\"N\"],\n", + " allowed_interaction_types=[\"strong\"],\n", + " formalism=\"helicity\",\n", + " particle_db=create_particle_database(),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", + "graphviz.Source(dot)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Amplitude builder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Dynamics builder with X symbols of J^PC channels" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@define\n", + "class DynamicsSymbolBuilder:\n", + " collected_symbols: set[sp.Symbol, tuple[Particle, TwoBodyKinematicVariableSet]] = (\n", + " field(factory=lambda: defaultdict(set))\n", + " )\n", + "\n", + " def __call__(\n", + " self, resonance: Particle, variable_pool: TwoBodyKinematicVariableSet\n", + " ) -> tuple[sp.Expr, dict[sp.Symbol, float]]:\n", + " jp = render_jp(resonance)\n", + " charge = resonance.charge\n", + " if variable_pool.angular_momentum is not None:\n", + " L = sp.Rational(variable_pool.angular_momentum)\n", + " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}^{{l={L}}}\")\n", + " else:\n", + " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}\")\n", + " self.collected_symbols[X].add((resonance, variable_pool))\n", + " parameter_defaults = {}\n", + " return X, parameter_defaults\n", + "\n", + "\n", + "def render_jp(particle: Particle) -> str:\n", + " spin = sp.Rational(particle.spin)\n", + " j = (\n", + " str(spin)\n", + " if spin.denominator == 1\n", + " else Rf\"\\frac{{{spin.numerator}}}{{{spin.denominator}}}\"\n", + " )\n", + " if particle.parity is None:\n", + " return f\"J={j}\"\n", + " p = \"-\" if particle.parity < 0 else \"+\"\n", + " return f\"J^P={{{j}}}^{{{p}}}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "model_builder = ampform.get_builder(reaction)\n", + "model_builder.adapter.permutate_registered_topologies()\n", + "model_builder.config.scalar_initial_state_mass = True\n", + "model_builder.config.stable_final_state_ids = [0, 1, 2]\n", + "create_dynamics_symbol = DynamicsSymbolBuilder()\n", + "for resonance in reaction.get_intermediate_particles():\n", + " model_builder.set_dynamics(resonance.name, create_dynamics_symbol)\n", + "model = model_builder.formulate()\n", + "model.intensity.cleanup()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "selected_amplitudes = {\n", + " k: v for i, (k, v) in enumerate(model.amplitudes.items()) if i == 0\n", + "}\n", + "Math(aslatex(selected_amplitudes, terms_per_line=1))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "src = R\"\\begin{array}{cll}\" \"\\n\"\n", + "for symbol, resonances in create_dynamics_symbol.collected_symbols.items():\n", + " src += Rf\" {symbol} \\\\\" \"\\n\"\n", + " for p, _ in resonances:\n", + " src += Rf\" {p.latex} & m={p.mass:g}\\text{{ GeV}} & \\Gamma={p.width:g}\\text{{ GeV}} \\\\\"\n", + " src += \"\\n\"\n", + "src += R\"\\end{array}\"\n", + "Math(src)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Dynamics parametrization" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Phasespace factor" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + ":::{seealso}\n", + "**[TR-026](./026.ipynb)** and **[TR-027](./027.ipynb)** on analyticity and Riemann sheets.\n", + ":::" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Expression classes for phase space factors" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@unevaluated(real=False)\n", + "class PhaseSpaceCM(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\n", + " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " s, m1, m2 = self.args\n", + " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", + "\n", + "\n", + "@unevaluated(real=False)\n", + "class ChewMandelstam(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\n", + " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " s, m1, m2 = self.args\n", + " q = BreakupMomentum(s, m1, m2)\n", + " return (\n", + " (2 * q / sp.sqrt(s))\n", + " * sp.log(Abs((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", + "\n", + "\n", + "@unevaluated(real=False)\n", + "class BreakupMomentum(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\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(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "s, m1, m2 = sp.symbols(\"s m1 m2\", nonnegative=True)\n", + "exprs = [\n", + " PhaseSpaceCM(s, m1, m2),\n", + " ChewMandelstam(s, m1, m2),\n", + " BreakupMomentum(s, m1, m2),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Relativistic Breit-Wigner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "PARAMETERS_BW = dict(model.parameter_defaults)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_breit_wigner(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " m1 = variables.outgoing_state_mass1\n", + " m2 = variables.outgoing_state_mass2\n", + " ρ = PhaseSpaceCM(s, m1, m2)\n", + " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", + " Γ0 = [sp.Symbol(Rf\"\\Gamma_{{{p.latex}}}\") for p, _ in resonances]\n", + " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", + " expr = sum(\n", + " (β_ * m_ * Γ0_) / (m_**2 - s - m_ * Γ0_ * ρ) for m_, Γ0_, β_ in zip(m, Γ0, β)\n", + " )\n", + " for i, (resonance, _) in enumerate(resonances):\n", + " PARAMETERS_BW[β[i]] = 1 + 0j\n", + " PARAMETERS_BW[m[i]] = resonance.mass\n", + " PARAMETERS_BW[Γ0[i]] = resonance.width\n", + " return expr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "dynamics_expressions_bw = {\n", + " symbol: formulate_breit_wigner(resonances)\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + "}\n", + "model_bw = attrs.evolve(\n", + " model,\n", + " parameter_defaults=ParameterValues({\n", + " **model.parameter_defaults,\n", + " **PARAMETERS_BW,\n", + " }),\n", + ")\n", + "Math(aslatex(dynamics_expressions_bw))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### $P$ vector" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "PARAMETERS_F = dict(model.parameter_defaults)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_k_matrix(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", + " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", + " expr = sum((g_**2) / (m_**2 - s) for m_, g_ in zip(m, g))\n", + " for i, (resonance, _) in enumerate(resonances):\n", + " PARAMETERS_F[m[i]] = resonance.mass\n", + " PARAMETERS_F[g[i]] = 1\n", + " return expr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_p_vector(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", + " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", + " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", + " expr = sum((g_ * β_) / (m_**2 - s) for m_, g_, β_ in zip(m, g, β))\n", + " for i, (resonance, _) in enumerate(resonances):\n", + " PARAMETERS_F[β[i]] = 1 + 0j\n", + " PARAMETERS_F[m[i]] = resonance.mass\n", + " PARAMETERS_F[g[i]] = 1\n", + " return expr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_f_vector(\n", + " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", + ") -> sp.Expr:\n", + " (_, variables), *_ = resonances\n", + " s = variables.incoming_state_mass**2\n", + " m1 = variables.outgoing_state_mass1\n", + " m2 = variables.outgoing_state_mass2\n", + " rho = PhaseSpaceCM(s, m1, m2)\n", + " K = formulate_k_matrix(resonances)\n", + " P = formulate_p_vector(resonances)\n", + " return (1 / (1 - rho * K)) * P" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "dynamics_expressions_fvector = {\n", + " symbol: formulate_f_vector(resonances)\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + "}\n", + "model_fvector = attrs.evolve(\n", + " model,\n", + " parameter_defaults=ParameterValues({\n", + " **model.parameter_defaults,\n", + " **PARAMETERS_F,\n", + " }),\n", + ")\n", + "Math(aslatex(dynamics_expressions_fvector))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Create numerical functions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "full_expression_bw = perform_cached_doit(model_bw.expression).xreplace(\n", + " dynamics_expressions_bw\n", + ")\n", + "intensity_func_bw = create_parametrized_function(\n", + " expression=perform_cached_doit(full_expression_bw),\n", + " backend=\"jax\",\n", + " parameters=PARAMETERS_BW,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "full_expression_fvector = perform_cached_doit(model_fvector.expression).xreplace(\n", + " dynamics_expressions_fvector\n", + ")\n", + "intensity_func_fvector = create_parametrized_function(\n", + " expression=perform_cached_doit(full_expression_fvector),\n", + " backend=\"jax\",\n", + " parameters=PARAMETERS_F,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Generate data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Generate phase space sample" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "rng = TFUniformRealNumberGenerator(seed=0)\n", + "phsp_generator = TFPhaseSpaceGenerator(\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", + "phsp_momenta = phsp_generator.generate(100_000, rng)\n", + "\n", + "ε = 1e-8\n", + "transformer = SympyDataTransformer.from_sympy(model.kinematic_variables, backend=\"jax\")\n", + "phsp = transformer(phsp_momenta)\n", + "phsp = {k: v + ε * 1j if re.match(r\"^m_\\d\\d$\", k) else v for k, v in phsp.items()}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Update function parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "toy_parameters_bw = {\n", + " R\"m_{N_1(1/2^-)}\": 1.65,\n", + " R\"m_{N_2(1/2^-)}\": 1.75,\n", + " R\"m_{N_1(3/2^-)}\": 1.85,\n", + " R\"m_{N_2(3/2^-)}\": 1.9,\n", + " R\"\\Gamma_{N_1(1/2^-)}\": 1 / 1.65,\n", + " R\"\\Gamma_{N_2(1/2^-)}\": 1 / 1.75,\n", + " R\"\\Gamma_{N_1(3/2^-)}\": 1 / 1.85,\n", + " R\"\\Gamma_{N_2(3/2^-)}\": 1 / 1.9,\n", + "}\n", + "intensity_func_bw.update_parameters(toy_parameters_bw)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "toy_parameters_fvector = {\n", + " R\"\\beta_{N_1(1/2^-)}\": 1 + 0j,\n", + " R\"\\beta_{N_2(1/2^-)}\": 1 + 0j,\n", + " R\"\\beta_{N_1(3/2^-)}\": 1 + 0j,\n", + " R\"\\beta_{N_2(3/2^-)}\": 1 + 0j,\n", + " R\"m_{N_1(1/2^-)}\": 1.65,\n", + " R\"m_{N_2(1/2^-)}\": 1.75,\n", + " R\"m_{N_1(3/2^-)}\": 1.95,\n", + " R\"m_{N_2(3/2^-)}\": 1.9,\n", + " R\"g_{N_1(1/2^-)}\": 1.65,\n", + " R\"g_{N_2(1/2^-)}\": 1,\n", + " R\"g_{N_1(3/2^-)}\": 1,\n", + " R\"g_{N_2(3/2^-)}\": 1,\n", + "}\n", + "intensity_func_fvector.update_parameters(toy_parameters_fvector)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Plot sub-intensities" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Function for computing sub-intensities" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def compute_sub_intensity(\n", + " func: ParametrizedFunction,\n", + " input_data: DataSample,\n", + " resonances: list[str],\n", + " coupling_pattern: 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\"^{coupling_pattern}({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", + " toy_parameters = dict(func.parameters)\n", + " for par_name in func.parameters:\n", + " if re.match(name_pattern, par_name) is not None:\n", + " toy_parameters[par_name] = 0\n", + " func.update_parameters(toy_parameters)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "total_intensities_bw = intensity_func_bw(phsp)\n", + "sub_intensities_bw = {\n", + " p: compute_sub_intensity(\n", + " intensity_func_bw,\n", + " phsp,\n", + " resonances=[p.latex],\n", + " coupling_pattern=r\"\\\\beta\",\n", + " )\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + " for p, _ in resonances\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "total_intensities_fvector = intensity_func_fvector(phsp)\n", + "sub_intensities_fvector = {\n", + " p: compute_sub_intensity(\n", + " intensity_func_fvector,\n", + " phsp,\n", + " resonances=[p.latex],\n", + " coupling_pattern=r\"\\\\beta\",\n", + " )\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + " for p, _ in resonances\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Function for plotting histograms with JAX" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def fast_histogram(\n", + " data: jnp.ndarray,\n", + " weights: jnp.ndarray | None = None,\n", + " bins: int = 100,\n", + " density: bool | None = None,\n", + " fill: bool = True,\n", + " ax=plt,\n", + " **plot_kwargs,\n", + ") -> None:\n", + " bin_values, bin_edges = jnp.histogram(\n", + " data,\n", + " bins=bins,\n", + " density=density,\n", + " weights=weights,\n", + " )\n", + " if fill:\n", + " bin_rights = bin_edges[1:]\n", + " ax.fill_between(bin_rights, bin_values, step=\"pre\", **plot_kwargs)\n", + " else:\n", + " bin_mids = (bin_edges[:-1] + bin_edges[1:]) / 2\n", + " ax.step(bin_mids, bin_values, **plot_kwargs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(8, 5))\n", + "ax.set_xlim(2, 5)\n", + "ax.set_xlabel(R\"$m_{p\\eta}^{2}$ [GeV$^2$]\")\n", + "ax.set_ylabel(R\"Intensity [a. u.]\")\n", + "ax.set_yticks([])\n", + "\n", + "bins = 150\n", + "phsp_projection = np.real(phsp[\"m_01\"]) ** 2\n", + "fast_histogram(\n", + " phsp_projection,\n", + " weights=total_intensities_fvector,\n", + " alpha=0.2,\n", + " bins=bins,\n", + " color=\"hotpink\",\n", + " label=\"Full intensity $F$ vector\",\n", + " ax=ax,\n", + ")\n", + "fast_histogram(\n", + " phsp_projection,\n", + " weights=total_intensities_bw,\n", + " alpha=0.2,\n", + " bins=bins,\n", + " color=\"grey\",\n", + " label=\"Full intensity Breit-Wigner\",\n", + " ax=ax,\n", + ")\n", + "for i, (p, v) in enumerate(sub_intensities_fvector.items()):\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=v,\n", + " alpha=0.6,\n", + " bins=bins,\n", + " color=f\"C{i}\",\n", + " fill=False,\n", + " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV}}$ $F$ vector\",\n", + " linewidth=2,\n", + " ax=ax,\n", + " )\n", + "for i, (p, v) in enumerate(sub_intensities_bw.items()):\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=v,\n", + " alpha=0.6,\n", + " bins=bins,\n", + " color=f\"C{i}\",\n", + " fill=False,\n", + " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV^2}}$ Breit-Wigner\",\n", + " linestyle=\"dashed\",\n", + " ax=ax,\n", + " )\n", + "\n", + "ax.set_ylim(0, None)\n", + "fig.legend(loc=\"upper right\")\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Weighted data with $F$ vector " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(6, 5))\n", + "fast_histogram(\n", + " phsp[\"m_01\"].real,\n", + " bins=100,\n", + " weights=np.real(intensity_func_fvector(phsp)),\n", + " ax=ax,\n", + ")\n", + "ax.set_xlabel(R\"$M^2\\left(\\eta p\\right)\\, \\mathrm{[(GeV/c)^2]}$\")\n", + "ax.set_ylabel(R\"Intensity [a.u.]\")\n", + "ax.set_ylim(0, None)\n", + "fig.tight_layout()\n", + "fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(\n", + " initial_state_mass=model.reaction_info.initial_state[-1].mass,\n", + " final_state_masses={i: p.mass for i, p in model.reaction_info.final_state.items()},\n", + ")\n", + "data_generator = IntensityDistributionGenerator(\n", + " domain_generator=weighted_phsp_generator,\n", + " function=intensity_func_fvector,\n", + " domain_transformer=transformer,\n", + ")\n", + "data_momenta = data_generator.generate(50_000, rng)\n", + "data = transformer(data_momenta)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(9, 4))\n", + "fast_histogram(\n", + " np.real(data[\"m_01\"]),\n", + " bins=200,\n", + " alpha=0.5,\n", + " density=True,\n", + " ax=ax,\n", + ")\n", + "mass_parameters = {p: v for p, v in toy_parameters_bw.items() if p.startswith(\"m_{\")}\n", + "evenly_spaced_interval = np.linspace(0, 1, num=len(mass_parameters))\n", + "colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n", + "ax.set_xlabel(\"$m$ [GeV]\")\n", + "for (k, v), color in zip(mass_parameters.items(), colors):\n", + " ax.axvline(v, c=color, label=f\"${k}$\", ls=\"dotted\")\n", + "ax.set_ylim(0, None)\n", + "ax.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Perform fit" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Estimator definition" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "estimator_bw = UnbinnedNLL(\n", + " intensity_func_bw,\n", + " data=data,\n", + " phsp=phsp,\n", + " backend=\"jax\",\n", + ")\n", + "estimator_fvector = UnbinnedNLL(\n", + " intensity_func_fvector,\n", + " data=data,\n", + " phsp=phsp,\n", + " backend=\"jax\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Initial parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Functions for comparing model to data" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def indicate_masses(ax, intensity_func, ls: str, lw: float, typ: str):\n", + " mass_pars = {\n", + " k: v for k, v in intensity_func.parameters.items() if k.startswith(\"m_{\")\n", + " }\n", + " for i, (k, v) in enumerate(mass_pars.items()):\n", + " ax.axvline(v, c=f\"C{i}\", label=f\"${k}$ ({typ})\", ls=ls, lw=lw)\n", + "\n", + "\n", + "def compare_model(\n", + " variable_name,\n", + " data,\n", + " phsp,\n", + " function1,\n", + " function2,\n", + " bins=100,\n", + ") -> None:\n", + " intensities1 = function1(phsp)\n", + " intensities2 = function2(phsp)\n", + " fig, ax = plt.subplots(figsize=(11, 4))\n", + " fig.subplots_adjust(right=0.85, top=0.95)\n", + " ax.set_xlabel(R\"$m_{p\\eta}$ [GeV]\")\n", + " ax.set_ylabel(\"Intensity [a. u.]\")\n", + " ax.set_yticks([])\n", + " data_projection = np.real(data[variable_name])\n", + " fast_histogram(\n", + " data_projection,\n", + " bins=bins,\n", + " alpha=0.5,\n", + " label=\"data\",\n", + " density=True,\n", + " ax=ax,\n", + " )\n", + " phsp_projection = np.real(phsp[variable_name])\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=np.array(intensities1),\n", + " bins=bins,\n", + " fill=False,\n", + " color=\"red\",\n", + " label=\"Fit model with F vector\",\n", + " density=True,\n", + " ax=ax,\n", + " )\n", + " fast_histogram(\n", + " phsp_projection,\n", + " weights=np.array(intensities2),\n", + " bins=bins,\n", + " fill=False,\n", + " color=\"blue\",\n", + " label=\"Fit model with Breit-Wigner\",\n", + " density=True,\n", + " ax=ax,\n", + " )\n", + " ax.set_ylim(0, None)\n", + " indicate_masses(ax, function1, ls=\"dashed\", lw=1, typ=\"F vector\")\n", + " indicate_masses(ax, function2, ls=\"dotted\", lw=1, typ=\"Breit-Wigner\")\n", + " fig.legend(loc=\"outside upper right\")\n", + " fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "initial_parameters_beta = {\n", + " R\"\\beta_{N_2(1/2^-)}\": 1 + 0j,\n", + " R\"\\beta_{N_2(3/2^-)}\": 1 + 0j,\n", + "}\n", + "initial_parameters_masses = {\n", + " R\"m_{N_1(1/2^-)}\": 1.6,\n", + " R\"m_{N_2(1/2^-)}\": 1.7,\n", + " R\"m_{N_1(3/2^-)}\": 1.8,\n", + " R\"m_{N_2(3/2^-)}\": 1.93,\n", + "}\n", + "initial_parameters_bw = {\n", + " **initial_parameters_beta,\n", + " **initial_parameters_masses,\n", + " R\"\\Gamma_{N_1(1/2^-)}\": 1 / 1.6,\n", + " R\"\\Gamma_{N_2(1/2^-)}\": 1 / 1.65,\n", + " R\"\\Gamma_{N_1(3/2^-)}\": 1 / 1.85,\n", + " R\"\\Gamma_{N_2(3/2^-)}\": 1 / 1.93,\n", + "}\n", + "initial_parameters_fvector = {\n", + " **initial_parameters_beta,\n", + " **initial_parameters_masses,\n", + " R\"g_{N_1(1/2^-)}\": 1.6,\n", + " R\"g_{N_2(1/2^-)}\": 1,\n", + " R\"g_{N_1(3/2^-)}\": 1.0,\n", + " R\"g_{N_2(3/2^-)}\": 1.0,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "original_parameters_bw = dict(intensity_func_bw.parameters)\n", + "intensity_func_bw.update_parameters(initial_parameters_bw)\n", + "original_parameters_fvector = dict(intensity_func_fvector.parameters)\n", + "intensity_func_fvector.update_parameters(initial_parameters_fvector)\n", + "compare_model(\"m_01\", data, phsp, intensity_func_fvector, intensity_func_bw)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Optimize parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "minuit2 = Minuit2()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "scroll-output" + ] + }, + "outputs": [], + "source": [ + "fit_result_bw = minuit2.optimize(estimator_bw, initial_parameters_bw)\n", + "assert fit_result_bw.minimum_valid\n", + "fit_result_bw" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "scroll-output" + ] + }, + "outputs": [], + "source": [ + "fit_result_fvector = minuit2.optimize(estimator_fvector, initial_parameters_fvector)\n", + "assert fit_result_fvector.minimum_valid\n", + "fit_result_fvector" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "intensity_func_fvector.update_parameters(fit_result_fvector.parameter_values)\n", + "intensity_func_bw.update_parameters(fit_result_bw.parameter_values)\n", + "compare_model(\"m_01\", data, phsp, intensity_func_fvector, intensity_func_bw)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Fit result comparison" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Functions for inspecting fit result" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def compute_aic_bic(fit_result: FitResult) -> tuple[float, float]:\n", + " n_real_par = fit_result.count_number_of_parameters(complex_twice=True)\n", + " n_events = len(next(iter(data.values())))\n", + " log_likelihood = -fit_result.estimator_value\n", + " aic = 2 * n_real_par - 2 * log_likelihood\n", + " bic = n_real_par * np.log(n_events) - 2 * log_likelihood\n", + " return aic, bic\n", + "\n", + "\n", + "def compare_parameters(initial: dict, optimized: dict, expected: dict) -> pd.DataFrame:\n", + " parameters = sorted(set(initial) | set(optimized))\n", + " df = pd.DataFrame(\n", + " {\n", + " f\"${p}$\": (\n", + " f\"{initial[p]:.3g}\",\n", + " f\"{optimized[p]:.3g}\",\n", + " f\"{expected[p]:.3g}\",\n", + " f\"{100 * abs((optimized[p] - expected[p]) / expected[p]):.1f}%\",\n", + " )\n", + " for p in parameters\n", + " },\n", + " ).T\n", + " df.columns = (\"initial\", \"fit result\", \"expected\", \"deviation\")\n", + " return df" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### P vector" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "compute_aic_bic(fit_result_fvector)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "compare_parameters(\n", + " initial=initial_parameters_fvector,\n", + " optimized=fit_result_fvector.parameter_values,\n", + " expected=original_parameters_fvector,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Breit–Wigner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "compute_aic_bic(fit_result_bw)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "compare_parameters(\n", + " initial=initial_parameters_bw,\n", + " optimized=fit_result_bw.parameter_values,\n", + " expected=original_parameters_bw,\n", + ")" + ] + } + ], + "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.10.14" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/report/032.ipynb b/docs/report/032.ipynb new file mode 100644 index 00000000..1bdf2969 --- /dev/null +++ b/docs/report/032.ipynb @@ -0,0 +1,1411 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "```{autolink-concat}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [ + "dynamics", + "K-matrix" + ] + }, + "source": [ + "::::{margin}\n", + ":::{card} Coupled-channel fit with $P$-vector dynamics for one single pole\n", + "TR-032\n", + "^^^\n", + "Illustration of how to formulate an amplitude model for two channels with P-vector dynamics. A combined fit is performed over the sum of the likelihood over both distributions. The example uses a single pole, but can easily be extended to multiple poles.\n", + "+++\n", + "🚧 [compwa.github.io#278](https://github.com/ComPWA/compwa.github.io/pull/278)\n", + ":::\n", + "::::" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "# P-vector fit comparison" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-cell" + ] + }, + "outputs": [], + "source": [ + "%pip install -q 'qrules[viz]==0.10.2' 'tensorwaves[jax,phsp]==0.4.12' ampform==0.15.4 pandas==2.2.2 sympy==1.12" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import logging\n", + "import os\n", + "import re\n", + "from collections import defaultdict\n", + "from functools import lru_cache\n", + "from itertools import product\n", + "from typing import Any, Iterable, Mapping\n", + "\n", + "import ampform\n", + "import attrs\n", + "import graphviz\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pandas as pd\n", + "import qrules\n", + "import sympy as sp\n", + "from ampform.dynamics.builder import TwoBodyKinematicVariableSet\n", + "from ampform.helicity import HelicityModel\n", + "from ampform.io import aslatex, improve_latex_rendering\n", + "from ampform.kinematics.phasespace import Kallen\n", + "from ampform.sympy import perform_cached_doit, unevaluated\n", + "from attrs import define, field, frozen\n", + "from IPython.display import Math, display\n", + "from qrules.particle import Particle, ParticleCollection\n", + "from qrules.transition import ReactionInfo\n", + "from sympy import Abs\n", + "from sympy.matrices.expressions.matexpr import MatrixElement\n", + "from tensorwaves.data import (\n", + " IntensityDistributionGenerator,\n", + " SympyDataTransformer,\n", + " TFPhaseSpaceGenerator,\n", + " TFUniformRealNumberGenerator,\n", + " TFWeightedPhaseSpaceGenerator,\n", + ")\n", + "from tensorwaves.estimator import UnbinnedNLL\n", + "from tensorwaves.function.sympy import create_parametrized_function\n", + "from tensorwaves.interface import (\n", + " DataSample,\n", + " Estimator,\n", + " FitResult,\n", + " Function,\n", + " ParameterValue,\n", + ")\n", + "from tensorwaves.optimizer import Minuit2\n", + "\n", + "improve_latex_rendering()\n", + "logging.getLogger(\"absl\").setLevel(logging.ERROR)\n", + "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", + "plt.rc(\"font\", size=12)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "%config InlineBackend.figure_formats = ['svg']" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Studied decay" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Define N* resonances" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@lru_cache(maxsize=1)\n", + "def create_particle_database() -> ParticleCollection:\n", + " particles = qrules.load_default_particles()\n", + " for nstar in particles.filter(lambda p: p.name.startswith(\"N\")):\n", + " particles.remove(nstar)\n", + " particles += create_nstar(mass=1.82, width=0.6, parity=+1, spin=1.5, idx=1)\n", + " return particles\n", + "\n", + "\n", + "def create_nstar(\n", + " mass: float, width: float, parity: int, spin: float, idx: int\n", + ") -> Particle:\n", + " spin = sp.Rational(spin)\n", + " parity_symbol = \"⁺\" if parity > 0 else \"⁻\"\n", + " unicode_subscripts = list(\"₀₁₂₃₄₅₆₇₈₉\")\n", + " return Particle(\n", + " name=f\"N{unicode_subscripts[idx]}({spin}{parity_symbol})\",\n", + " latex=Rf\"N_{idx}({spin.numerator}/{spin.denominator}^-)\",\n", + " pid=2024_05_00_00 + 100 * bool(parity + 1) + idx,\n", + " mass=mass,\n", + " width=width,\n", + " baryon_number=1,\n", + " charge=+1,\n", + " isospin=(0.5, +0.5),\n", + " parity=parity,\n", + " spin=1.5,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "FINAL_STATES: list[tuple[str, ...]] = [\n", + " [\"K0\", \"Sigma+\", \"p~\"],\n", + " [\"eta\", \"p\", \"p~\"],\n", + "]\n", + "REACTIONS: list[ReactionInfo] = [\n", + " qrules.generate_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=final_state,\n", + " allowed_intermediate_particles=[\"N\"],\n", + " allowed_interaction_types=[\"strong\"],\n", + " formalism=\"helicity\",\n", + " particle_db=create_particle_database(),\n", + " )\n", + " for final_state in FINAL_STATES\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "for reaction in REACTIONS:\n", + " src = qrules.io.asdot(reaction, collapse_graphs=True)\n", + " graph = graphviz.Source(src)\n", + " display(graph)\n", + " del reaction, src, graph" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Amplitude builder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Dynamics builder with X symbols of J^PC channels" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@define\n", + "class DynamicsSymbolBuilder:\n", + " collected_symbols: set[sp.Symbol, tuple[Particle, TwoBodyKinematicVariableSet]] = (\n", + " field(factory=lambda: defaultdict(set))\n", + " )\n", + "\n", + " def __call__(\n", + " self, resonance: Particle, variable_pool: TwoBodyKinematicVariableSet\n", + " ) -> tuple[sp.Expr, dict[sp.Symbol, float]]:\n", + " jp = render_jp(resonance)\n", + " charge = resonance.charge\n", + " if variable_pool.angular_momentum is not None:\n", + " L = sp.Rational(variable_pool.angular_momentum)\n", + " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}^{{l={L}}}\")\n", + " else:\n", + " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}\")\n", + " self.collected_symbols[X].add((resonance, variable_pool))\n", + " parameter_defaults = {}\n", + " return X, parameter_defaults\n", + "\n", + "\n", + "def render_jp(particle: Particle) -> str:\n", + " spin = sp.Rational(particle.spin)\n", + " j = (\n", + " str(spin)\n", + " if spin.denominator == 1\n", + " else Rf\"\\frac{{{spin.numerator}}}{{{spin.denominator}}}\"\n", + " )\n", + " if particle.parity is None:\n", + " return f\"J={j}\"\n", + " p = \"-\" if particle.parity < 0 else \"+\"\n", + " return f\"J^P={{{j}}}^{{{p}}}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "MODELS: list[HelicityModel] = []\n", + "for reaction in REACTIONS:\n", + " builder = ampform.get_builder(reaction)\n", + " builder.adapter.permutate_registered_topologies()\n", + " builder.config.scalar_initial_state_mass = True\n", + " builder.config.stable_final_state_ids = [0, 1, 2]\n", + " create_dynamics_symbol = DynamicsSymbolBuilder()\n", + " for resonance in reaction.get_intermediate_particles():\n", + " builder.set_dynamics(resonance.name, create_dynamics_symbol)\n", + " MODELS.append(builder.formulate())\n", + " del builder, reaction, resonance" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "selected_amplitudes = {\n", + " k: v for i, (k, v) in enumerate(MODELS[0].amplitudes.items()) if i == 0\n", + "}\n", + "Math(aslatex(selected_amplitudes, terms_per_line=1))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "src = R\"\\begin{array}{cll}\" \"\\n\"\n", + "for symbol, resonances in create_dynamics_symbol.collected_symbols.items():\n", + " src += Rf\" {symbol} \\\\\" \"\\n\"\n", + " for p, _ in resonances:\n", + " src += Rf\" {p.latex} & m={p.mass:g}\\text{{ GeV}} & \\Gamma={p.width:g}\\text{{ GeV}} \\\\\"\n", + " src += \"\\n\"\n", + "src += R\"\\end{array}\"\n", + "Math(src)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Dynamics parametrization" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Phasespace factor" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + ":::{seealso}\n", + "**[TR-026](./026.ipynb)** and **[TR-027](./027.ipynb)** on analyticity and Riemann sheets.\n", + ":::" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Expression classes for phase space factors" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@unevaluated(real=False)\n", + "class PhaseSpaceCM(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\n", + " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " s, m1, m2 = self.args\n", + " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", + "\n", + "\n", + "@unevaluated(real=False)\n", + "class ChewMandelstam(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\n", + " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " s, m1, m2 = self.args\n", + " q = BreakupMomentum(s, m1, m2)\n", + " return (\n", + " (2 * q / sp.sqrt(s))\n", + " * sp.log(Abs((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", + "\n", + "\n", + "@unevaluated(real=False)\n", + "class BreakupMomentum(sp.Expr):\n", + " s: Any\n", + " m1: Any\n", + " m2: Any\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(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "s, m1, m2 = sp.symbols(\"s m1 m2\", nonnegative=True)\n", + "exprs = [\n", + " PhaseSpaceCM(s, m1, m2),\n", + " ChewMandelstam(s, m1, m2),\n", + " BreakupMomentum(s, m1, m2),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### $K$-matrix formalism" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "n_channels = len(REACTIONS)\n", + "I = sp.Identity(n_channels)\n", + "K = sp.MatrixSymbol(\"K\", n_channels, n_channels)\n", + "P = sp.MatrixSymbol(\"P\", n_channels, 1)\n", + "F = sp.MatrixSymbol(\"F\", n_channels, 1)\n", + "rho = sp.MatrixSymbol(\"rho\", n_channels, n_channels)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Find decay products per channel" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def get_decay_products(reaction: ReactionInfo) -> DecayProducts:\n", + " some_transition, *_ = reaction.transitions\n", + " decay_product_ids = some_transition.topology.get_edge_ids_outgoing_from_node(1)\n", + " for transition in reaction.transitions:\n", + " if decay_product_ids != transition.topology.get_edge_ids_outgoing_from_node(1):\n", + " msg = \"Reaction contains multiple sub-systems\"\n", + " raise ValueError(msg)\n", + " child1_id, child2_id = sorted(decay_product_ids)\n", + " return DecayProducts(\n", + " child1=reaction.final_state[child1_id],\n", + " child2=reaction.final_state[child2_id],\n", + " )\n", + "\n", + "\n", + "@frozen\n", + "class DecayProducts:\n", + " child1: Particle\n", + " child2: Particle\n", + "\n", + " @property\n", + " def children(self) -> tuple[Particle, Particle]:\n", + " return self.child1, self.child2\n", + "\n", + "\n", + "DECAYS = tuple(get_decay_products(m.reaction_info) for m in MODELS)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [] + }, + "outputs": [], + "source": [ + "PARAMETERS_DEFAULTS = {}\n", + "for model in MODELS:\n", + " PARAMETERS_DEFAULTS.update(model.parameter_defaults)\n", + " del model\n", + "PARAMETERS_DEFAULTS = {\n", + " par: value\n", + " for par, value in PARAMETERS_DEFAULTS.items()\n", + " if not re.match(r\"^m_\\d+$\", par.name)\n", + "}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### $K$-matrix parametrization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def formulate_k_matrix(\n", + " resonances: list[tuple[Particle, int]], n_channels: int\n", + ") -> dict[MatrixElement, sp.Expr]:\n", + " expressions = {}\n", + " for i, j in product(range(n_channels), range(n_channels)):\n", + " resonance_contributions = []\n", + " for res, _ in resonances:\n", + " s = sp.Symbol(\"m_01\", real=True) ** 2\n", + " g_Ri = sp.Symbol(Rf\"g_{{{res.latex},{i}}}\")\n", + " g_Rj = sp.Symbol(Rf\"g_{{{res.latex},{j}}}\")\n", + " m_R = sp.Symbol(Rf\"m_{{{res.latex}}}\")\n", + " parameter_defaults = {\n", + " m_R: res.mass,\n", + " g_Ri: 1,\n", + " g_Rj: 0.1,\n", + " }\n", + " PARAMETERS_DEFAULTS.update(parameter_defaults)\n", + " expr = (g_Ri * g_Rj) / (m_R**2 - s)\n", + " resonance_contributions.append(expr)\n", + " expressions[K[i, j]] = sum(resonance_contributions)\n", + " return expressions\n", + "\n", + "\n", + "K_expressions = formulate_k_matrix(resonances, n_channels=len(REACTIONS))\n", + "K_matrix = K.as_explicit()\n", + "K.as_explicit().xreplace(K_expressions)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### $P$-vector parametrization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def formulate_p_vector(\n", + " resonances: list[tuple[Particle, int]], n_channels: int\n", + ") -> dict[MatrixElement, sp.Expr]:\n", + " expressions = {}\n", + " for i in range(n_channels):\n", + " resonance_contributions = []\n", + " for res, _ in resonances:\n", + " s = sp.Symbol(\"m_01\", real=True) ** 2\n", + " g_Ri = sp.Symbol(Rf\"g_{{{res.latex},{i}}}\")\n", + " beta_R = sp.Symbol(Rf\"\\beta_{{{res.latex}}}\")\n", + " m_R = sp.Symbol(Rf\"m_{{{res.latex}}}\")\n", + " parameter_defaults = {\n", + " m_R: res.mass,\n", + " beta_R: 1 + 0j,\n", + " g_Ri: 1,\n", + " }\n", + " PARAMETERS_DEFAULTS.update(parameter_defaults)\n", + " expr = (beta_R * g_Ri) / (m_R**2 - s)\n", + " resonance_contributions.append(expr)\n", + " expressions[P[i, 0]] = sum(resonance_contributions)\n", + " return expressions\n", + "\n", + "\n", + "P_expressions = formulate_p_vector(resonances, n_channels=len(REACTIONS))\n", + "P_vector = P.as_explicit()\n", + "P.as_explicit().xreplace(P_expressions)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Phase space factor parametrization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def formulate_phsp_factor_matrix(n_channels: int) -> dict[sp.MatrixElement, sp.Expr]:\n", + " expressions = {}\n", + " for i in range(n_channels):\n", + " for j in range(n_channels):\n", + " if i == j:\n", + " m_a_i = sp.Symbol(Rf\"m_{{0,{i}}}\")\n", + " m_b_i = sp.Symbol(Rf\"m_{{1,{i}}}\")\n", + " s = sp.Symbol(\"m_01\", real=True) ** 2\n", + " rho_i = PhaseSpaceCM(s, m_a_i, m_b_i)\n", + " expressions[rho[i, j]] = rho_i\n", + " parameter_defaults = {\n", + " m_a_i: DECAYS[i].child1.mass,\n", + " m_b_i: DECAYS[i].child2.mass,\n", + " }\n", + " PARAMETERS_DEFAULTS.update(parameter_defaults)\n", + " else:\n", + " expressions[rho[i, j]] = 0\n", + " return expressions\n", + "\n", + "\n", + "rho_expressions = formulate_phsp_factor_matrix(n_channels=len(REACTIONS))\n", + "rho.as_explicit().xreplace(rho_expressions)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### $F$-vector construction" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + ":::{note}\n", + "For some reason one has to leave out the multiplication of $\\rho$ by $i$ within the calculation of the $F$ vector\n", + ":::" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "F = (I - sp.I * K * rho).inv() * P\n", + "F" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "full-width" + ] + }, + "outputs": [], + "source": [ + "F_vector = F.as_explicit()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "full-width" + ] + }, + "outputs": [], + "source": [ + "parametrizations = {**K_expressions, **rho_expressions, **P_expressions}\n", + "F_exprs = F_vector.xreplace(parametrizations)\n", + "F_exprs[0].simplify(doit=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Create numerical functions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "F_unfolded_exprs = np.array([perform_cached_doit(expr) for expr in F_exprs])\n", + "DYNAMICS_EXPRESSIONS_FVECTOR = [\n", + " {\n", + " symbol: F_unfolded_exprs[i]\n", + " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", + " }\n", + " for i in range(n_channels)\n", + "]\n", + "MODELS_FVECTOR = [\n", + " attrs.evolve(\n", + " model,\n", + " parameter_defaults=PARAMETERS_DEFAULTS,\n", + " )\n", + " for model in MODELS\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "FULL_EXPRESSIONS_FVECTOR = [\n", + " perform_cached_doit(MODELS_FVECTOR[i].expression).xreplace(\n", + " DYNAMICS_EXPRESSIONS_FVECTOR[i]\n", + " )\n", + " for i in range(n_channels)\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "INTENSITY_FUNCS_FVECTOR = [\n", + " create_parametrized_function(\n", + " expression=perform_cached_doit(FULL_EXPRESSIONS_FVECTOR[i]),\n", + " backend=\"jax\",\n", + " parameters=MODELS_FVECTOR[i].parameter_defaults,\n", + " )\n", + " for i in range(n_channels)\n", + "]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Generate data" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Phase space sample" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "HELICITY_TRANSFORMERS = [\n", + " SympyDataTransformer.from_sympy(model.kinematic_variables, backend=\"jax\")\n", + " for model in MODELS_FVECTOR\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "PHSP = []\n", + "ε = 1e-8\n", + "for i in range(n_channels):\n", + " rng = TFUniformRealNumberGenerator(seed=0)\n", + " phsp_generator = TFPhaseSpaceGenerator(\n", + " initial_state_mass=REACTIONS[i].initial_state[-1].mass,\n", + " final_state_masses={it: p.mass for it, p in REACTIONS[i].final_state.items()},\n", + " )\n", + " phsp_momenta = phsp_generator.generate(100_000, rng)\n", + " phsp = HELICITY_TRANSFORMERS[i](phsp_momenta)\n", + " phsp = {k: v.real for k, v in phsp.items()}\n", + " phsp = {k: v + ε * 1j if re.match(r\"^m_\\d\\d$\", k) else v for k, v in phsp.items()}\n", + " PHSP.append(phsp)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set parameters for toy model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Function for plotting histograms with JAX" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def fast_histogram(\n", + " data: jnp.ndarray,\n", + " weights: jnp.ndarray | None = None,\n", + " bins: int = 100,\n", + " density: bool | None = None,\n", + " fill: bool = True,\n", + " ax=plt,\n", + " **plot_kwargs,\n", + ") -> None:\n", + " bin_values, bin_edges = jnp.histogram(\n", + " data,\n", + " bins=bins,\n", + " density=density,\n", + " weights=weights,\n", + " )\n", + " if fill:\n", + " bin_rights = bin_edges[1:]\n", + " ax.fill_between(bin_rights, bin_values, step=\"pre\", **plot_kwargs)\n", + " else:\n", + " bin_mids = (bin_edges[:-1] + bin_edges[1:]) / 2\n", + " ax.step(bin_mids, bin_values, **plot_kwargs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Functions for indicated resonances and thresholds" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def indicate_masses(ax, intensity_func, set_labels: bool = True):\n", + " mass_pars = {\n", + " k: v for k, v in intensity_func.parameters.items() if k.startswith(\"m_{N\")\n", + " }\n", + " for i, (k, v) in enumerate(mass_pars.items()):\n", + " label = f\"${k}$\" if set_labels else None\n", + " ax.axvline(v, c=f\"C{i + n_channels}\", label=label, ls=\"dashed\")\n", + "\n", + "\n", + "def indicate_thresholds(ax, set_labels: bool = True) -> None:\n", + " for i, decay in enumerate(DECAYS):\n", + " m_thr = sum(p.mass for p in decay.children)\n", + " label = None\n", + " if set_labels:\n", + " label = f\"${'+'.join(f'm_{{{p.latex}}}' for p in decay.children)}$\"\n", + " ax.axvline(m_thr, c=f\"C{i}\", label=label, ls=\"dotted\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "toy_parameters = {\n", + " R\"\\beta_{N_1(3/2^-)}\": 1 + 0j,\n", + " R\"m_{N_1(3/2^-)}\": 1.71,\n", + " R\"g_{N_1(3/2^-),0}\": 3.2,\n", + " R\"g_{N_1(3/2^-),1}\": 1.5,\n", + "}\n", + "for func in INTENSITY_FUNCS_FVECTOR:\n", + " func.update_parameters(toy_parameters)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(9, 4))\n", + "ax.set_title(\"Model rendering from phase space\")\n", + "ax.set_xlabel(R\"$m_{p\\eta/K\\Sigma}$ [GeV]\")\n", + "for i in range(n_channels):\n", + " intensity = np.real(INTENSITY_FUNCS_FVECTOR[i](PHSP[i]))\n", + " fast_histogram(\n", + " np.real(PHSP[i][\"m_01\"]),\n", + " weights=intensity,\n", + " alpha=0.5,\n", + " bins=200,\n", + " density=True,\n", + " label=f\"${' '.join(p.latex for p in DECAYS[i].children)}$\",\n", + " ax=ax,\n", + " )\n", + "indicate_thresholds(ax)\n", + "indicate_masses(ax, INTENSITY_FUNCS_FVECTOR[i])\n", + "ax.legend()\n", + "ax.set_ylim(0, None)\n", + "fig.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Toy data sample" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "DATA = []\n", + "for i in range(n_channels):\n", + " weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(\n", + " initial_state_mass=MODELS[i].reaction_info.initial_state[-1].mass,\n", + " final_state_masses={\n", + " i: p.mass for i, p in MODELS[i].reaction_info.final_state.items()\n", + " },\n", + " )\n", + " data_generator = IntensityDistributionGenerator(\n", + " domain_generator=weighted_phsp_generator,\n", + " function=INTENSITY_FUNCS_FVECTOR[i],\n", + " domain_transformer=HELICITY_TRANSFORMERS[i],\n", + " )\n", + " data_momenta = data_generator.generate(50_000, rng)\n", + " data = HELICITY_TRANSFORMERS[i](data_momenta)\n", + " DATA.append(data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(9, 4))\n", + "ax.set_title(\"Toy data sample\")\n", + "ax.set_xlabel(R\"$m_{p\\eta/K\\Sigma}$ [GeV]\")\n", + "for i in range(n_channels):\n", + " fast_histogram(\n", + " np.real(DATA[i][\"m_01\"]),\n", + " alpha=0.5,\n", + " bins=200,\n", + " density=True,\n", + " label=f\"${' '.join(p.latex for p in DECAYS[i].children)}$\",\n", + " ax=ax,\n", + " )\n", + "indicate_thresholds(ax)\n", + "indicate_masses(ax, INTENSITY_FUNCS_FVECTOR[i])\n", + "ax.legend()\n", + "ax.set_ylim(0, None)\n", + "fig.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Perform fit" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Estimator definition" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "class EstimatorSum(Estimator):\n", + " def __init__(self, estimators: Iterable[Estimator]) -> None:\n", + " self.__estimators = tuple(estimators)\n", + "\n", + " def __call__(self, parameters: Mapping[str, ParameterValue]) -> float:\n", + " return sum(estimator(parameters) for estimator in self.__estimators)\n", + "\n", + " def gradient(\n", + " self, parameters: Mapping[str, ParameterValue]\n", + " ) -> dict[str, ParameterValue]:\n", + " raise NotImplementedError" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "combined_estimators = EstimatorSum(\n", + " UnbinnedNLL(\n", + " INTENSITY_FUNCS_FVECTOR[i],\n", + " data=DATA[i],\n", + " phsp=PHSP[i],\n", + " backend=\"jax\",\n", + " )\n", + " for i in range(n_channels)\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Initial parameters " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Functions for comparing model to data" + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def compare_models(functions: list[Function], title: str, bins: int = 100):\n", + " fig, axes = plt.subplots(figsize=(8.5, 4.5), nrows=2, sharex=True)\n", + " axes[0].set_title(title)\n", + " for ax in axes:\n", + " ax.set_yticks([])\n", + " for i in range(n_channels):\n", + " _plot_comparison(\n", + " axes[i],\n", + " decay_id=i,\n", + " variable_name=\"m_01\",\n", + " function=functions[i],\n", + " bins=bins,\n", + " color=f\"C{i}\",\n", + " legend=(i == 1),\n", + " )\n", + " fig.legend()\n", + " fig.tight_layout()\n", + " fig.show()\n", + "\n", + "\n", + "def _plot_comparison(\n", + " ax,\n", + " decay_id: int,\n", + " variable_name: str,\n", + " function: Function[DataSample, np.ndarray],\n", + " bins: int,\n", + " color: str,\n", + " legend: bool,\n", + "):\n", + " phsp = PHSP[decay_id]\n", + " fast_histogram(\n", + " DATA[decay_id][variable_name].real,\n", + " alpha=0.5,\n", + " bins=bins,\n", + " color=color,\n", + " density=True,\n", + " label=f\"Data ${' '.join(p.latex for p in DECAYS[decay_id].children)}$\",\n", + " ax=ax,\n", + " )\n", + " fast_histogram(\n", + " phsp[variable_name].real,\n", + " weights=function(phsp),\n", + " bins=bins,\n", + " color=\"red\",\n", + " density=True,\n", + " fill=False,\n", + " label=\"Fit model\" if legend else None,\n", + " ax=ax,\n", + " )\n", + " indicate_thresholds(ax, set_labels=legend)\n", + " indicate_masses(ax, function, set_labels=legend)\n", + " ax.set_ylim(0, None)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "initial_parameters = {\n", + " R\"m_{N_1(3/2^-)}\": 1.9,\n", + " R\"g_{N_1(3/2^-),0}\": 2.8,\n", + " R\"g_{N_1(3/2^-),1}\": 1.6,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "ORIGINAL_PARAMETERS_F = []\n", + "for i in range(n_channels):\n", + " ORIGINAL_PARAMETERS_F.append(dict(INTENSITY_FUNCS_FVECTOR[i].parameters))\n", + " INTENSITY_FUNCS_FVECTOR[i].update_parameters(initial_parameters)\n", + "compare_models(INTENSITY_FUNCS_FVECTOR, title=\"Model with starting parameters\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Optimize parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "minuit2 = Minuit2()\n", + "fit_result = minuit2.optimize(combined_estimators, initial_parameters)\n", + "assert fit_result.minimum_valid\n", + "fit_result" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "for i in range(n_channels):\n", + " INTENSITY_FUNCS_FVECTOR[i].update_parameters(fit_result.parameter_values)\n", + "compare_models(INTENSITY_FUNCS_FVECTOR, title=\"Model with optimized parameters\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Fit quality check" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Functions for inspecting fit result" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def compute_aic_bic(fit_result: FitResult) -> tuple[float, float]:\n", + " n_real_par = fit_result.count_number_of_parameters(complex_twice=True)\n", + " n_events = len(next(iter(data.values())))\n", + " log_likelihood = -fit_result.estimator_value\n", + " aic = 2 * n_real_par - 2 * log_likelihood\n", + " bic = n_real_par * np.log(n_events) - 2 * log_likelihood\n", + " return aic, bic\n", + "\n", + "\n", + "def compare_parameters(initial: dict, optimized: dict, expected: dict) -> pd.DataFrame:\n", + " parameters = sorted(set(initial) | set(optimized))\n", + " df = pd.DataFrame(\n", + " {\n", + " f\"${p}$\": (\n", + " f\"{initial[p]:.3g}\",\n", + " f\"{optimized[p]:.3g}\",\n", + " f\"{expected[p]:.3g}\",\n", + " f\"{100 * abs((optimized[p] - expected[p]) / expected[p]):.1f}%\",\n", + " )\n", + " for p in parameters\n", + " },\n", + " ).T\n", + " df.columns = (\"initial\", \"fit result\", \"expected\", \"deviation\")\n", + " return df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "compute_aic_bic(fit_result)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "compare_parameters(\n", + " initial=initial_parameters,\n", + " optimized=fit_result.parameter_values,\n", + " expected={\n", + " **ORIGINAL_PARAMETERS_F[0],\n", + " **ORIGINAL_PARAMETERS_F[1],\n", + " },\n", + ")" + ] + } + ], + "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.10.14" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +}