diff --git a/docs/lambda-k-pi/index.md b/docs/lambda-k-pi/index.md index 78a8251..e0ee1b3 100644 --- a/docs/lambda-k-pi/index.md +++ b/docs/lambda-k-pi/index.md @@ -3,4 +3,5 @@ ```{toctree} ampform ampform-dpd +manual-symbolic ``` diff --git a/docs/lambda-k-pi/manual-symbolic.ipynb b/docs/lambda-k-pi/manual-symbolic.ipynb new file mode 100644 index 0000000..3f35443 --- /dev/null +++ b/docs/lambda-k-pi/manual-symbolic.ipynb @@ -0,0 +1,910 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Manually modelling amplitudes using `sympy`" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In this notebook, we formulate the amplitude model for the $\\gamma p \\to K^+ \\pi^0 \\Lambda $ symbolically by adapting the model originally for the $\\gamma p \\to \\eta\\pi^0 p$ channel example as described in [Reaction and Models](../eta-pi-p/reaction-model.md).\n", + "\n", + "The model we want to implement is" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$$\n", + "\\begin{array}{rcl}\n", + "I &=& \\left|A^{12} + A^{23} + A^{31}\\right|^2 \\\\\n", + "A^{12} &=& \\frac{\\sum a_m Y_2^m (\\Omega_1)}{s_{12}-m^2_{K^{*+}_2}+im_{K^{*+}_2} \\Gamma_{K^{*+}_2}} \\\\\n", + "A^{23} &=& \\frac{\\sum b_m Y_1^m (\\Omega_2)}{s_{23}-m^2_{\\Sigma^*}+im_{\\Sigma^*} \\Gamma_{\\Sigma^*}} \\\\\n", + "A^{31} &=& \\frac{c_0}{s_{31}-m^2_{N^{*+}}+im_{N^{*+}} \\Gamma_{N^{*+}}} \\,,\n", + "\\end{array}\n", + "$$\n", + "\n", + "where $1=K^+$, $2=\\pi^0$, and $3=\\Lambda$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-cell" + ] + }, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import logging\n", + "import os\n", + "import warnings\n", + "from collections import defaultdict\n", + "\n", + "import ipywidgets as w\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import sympy as sp\n", + "from ampform.io import aslatex\n", + "from ampform.kinematics.angles import Phi, Theta\n", + "from ampform.kinematics.lorentz import (\n", + " ArrayMultiplication,\n", + " ArraySize,\n", + " BoostZMatrix,\n", + " Energy,\n", + " EuclideanNorm,\n", + " FourMomentumSymbol,\n", + " RotationYMatrix,\n", + " RotationZMatrix,\n", + " ThreeMomentum,\n", + " three_momentum_norm,\n", + ")\n", + "from ampform.sympy import unevaluated\n", + "from ampform.sympy._array_expressions import ArraySum\n", + "from IPython.display import SVG, Image, Latex, display\n", + "from tensorwaves.data import (\n", + " SympyDataTransformer,\n", + " TFPhaseSpaceGenerator,\n", + " TFUniformRealNumberGenerator,\n", + ")\n", + "from tensorwaves.function.sympy import create_parametrized_function\n", + "\n", + "STATIC_PAGE = \"EXECUTE_NB\" in os.environ\n", + "\n", + "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", + "logging.disable(logging.WARNING)\n", + "warnings.filterwarnings(\"ignore\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Model implementation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "l_max = 2" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### $A^{12}$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "s12, m_Kstar2, Gamma_Kstar2, l12 = sp.symbols(r\"s_{12} m_{K^*_2} \\Gamma_{K^*_2} l_{12}\")\n", + "theta1, phi1 = sp.symbols(\"theta_1 phi_1\")\n", + "a = sp.IndexedBase(\"a\")\n", + "m = sp.symbols(\"m\", cls=sp.Idx)\n", + "A12 = sp.Sum(a[m] * sp.Ynm(l12, m, theta1, phi1), (m, -l12, l12)) / (\n", + " s12 - m_Kstar2**2 + sp.I * m_Kstar2 * Gamma_Kstar2\n", + ")\n", + "A12" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "A12_funcs = [\n", + " sp.lambdify(\n", + " [\n", + " s12,\n", + " *(a[j] for j in range(-l_max, l_max + 1)),\n", + " m_Kstar2,\n", + " Gamma_Kstar2,\n", + " theta1,\n", + " phi1,\n", + " ],\n", + " expr=A12.subs(l12, i).doit().expand(func=True),\n", + " )\n", + " for i in range(l_max + 1)\n", + "]\n", + "A12_funcs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### $A^{23}$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "s23, m_Sigma, Gamma_Sigma, l23 = sp.symbols(\n", + " r\"s_{23} m_{\\Sigma^{*+}} \\Gamma_{\\Sigma^{*+}} l_{23}\"\n", + ")\n", + "b = sp.IndexedBase(\"b\")\n", + "m = sp.symbols(\"m\", cls=sp.Idx)\n", + "theta2, phi2 = sp.symbols(\"theta_2 phi_2\")\n", + "A23 = sp.Sum(b[m] * sp.Ynm(l23, m, theta2, phi2), (m, -l23, l23)) / (\n", + " s23 - m_Sigma**2 + sp.I * m_Sigma * Gamma_Sigma\n", + ")\n", + "A23" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "A23_funcs = [\n", + " sp.lambdify(\n", + " [\n", + " s23,\n", + " *(b[j] for j in range(-l_max, l_max + 1)),\n", + " m_Sigma,\n", + " Gamma_Sigma,\n", + " theta2,\n", + " phi2,\n", + " ],\n", + " A23.subs(l23, i).doit().expand(func=True),\n", + " )\n", + " for i in range(l_max + 1)\n", + "]\n", + "A23_funcs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### $A^{31}$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "s31, m_Nstar, Gamma_Nstar = sp.symbols(r\"s_{31} m_{N^*} \\Gamma_{N^*}\")\n", + "c = sp.IndexedBase(\"c\")\n", + "theta3, phi3, l31 = sp.symbols(\"theta_3 phi_3 l_{31}\")\n", + "A31 = sp.Sum(c[m] * sp.Ynm(l31, m, theta3, phi3), (m, -l31, l31)) / (\n", + " s31 - m_Nstar**2 + sp.I * m_Nstar * Gamma_Nstar\n", + ")\n", + "A31" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "A31_funcs = [\n", + " sp.lambdify(\n", + " [\n", + " s31,\n", + " *(c[j] for j in range(-l_max, l_max + 1)),\n", + " m_Nstar,\n", + " Gamma_Nstar,\n", + " theta3,\n", + " phi3,\n", + " ],\n", + " A31.subs(l31, i).doit().expand(func=True),\n", + " )\n", + " for i in range(l_max + 1)\n", + "]\n", + "A31_funcs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### $I = |A|^2 = |A^{12}+A^{23}+A^{31}|^2$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "intensity_expr = sp.Abs(A12 + A23 + A31) ** 2\n", + "intensity_expr" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Phase Space Generation" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Mass for $p\\gamma$ system" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "E_lab_gamma = 8.5\n", + "m_proton = 0.938\n", + "m_0 = np.sqrt(2 * E_lab_gamma * m_proton + m_proton**2)\n", + "m_lambda = 1.12\n", + "m_k = 0.494\n", + "m_pi = 0.135\n", + "m_0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rng = TFUniformRealNumberGenerator(seed=0)\n", + "phsp_generator = TFPhaseSpaceGenerator(\n", + " initial_state_mass=m_0,\n", + " final_state_masses={1: m_k, 2: m_pi, 3: m_lambda},\n", + ")\n", + "phsp_momenta = phsp_generator.generate(500_000, rng)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Kinematic variables" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "@unevaluated\n", + "class SquaredInvariantMass(sp.Expr):\n", + " momentum: sp.Basic\n", + " _latex_repr_ = \"m_{{{momentum}}}^2\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " p = self.momentum\n", + " p_xyz = ThreeMomentum(p)\n", + " return Energy(p) ** 2 - EuclideanNorm(p_xyz) ** 2\n", + "\n", + "\n", + "def formulate_helicity_angles(\n", + " pi: FourMomentumSymbol, pj: FourMomentumSymbol\n", + ") -> tuple[Theta, Phi]:\n", + " pij = ArraySum(pi, pj)\n", + " beta = three_momentum_norm(pij) / Energy(pij)\n", + " Rz = RotationZMatrix(-Phi(pij), n_events=ArraySize(Phi(pij)))\n", + " Ry = RotationYMatrix(-Theta(pij), n_events=ArraySize(Theta(pij)))\n", + " Bz = BoostZMatrix(beta, n_events=ArraySize(beta))\n", + " pi_boosted = ArrayMultiplication(Bz, Ry, Rz, pi)\n", + " return Theta(pi_boosted), Phi(pi_boosted)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "p1 = FourMomentumSymbol(\"p1\", shape=[])\n", + "p2 = FourMomentumSymbol(\"p2\", shape=[])\n", + "p3 = FourMomentumSymbol(\"p3\", shape=[])\n", + "p12 = ArraySum(p1, p2)\n", + "p23 = ArraySum(p2, p3)\n", + "p31 = ArraySum(p3, p1)\n", + "\n", + "theta1_expr, phi1_expr = formulate_helicity_angles(p1, p2)\n", + "theta2_expr, phi2_expr = formulate_helicity_angles(p2, p3)\n", + "theta3_expr, phi3_expr = formulate_helicity_angles(p3, p1)\n", + "\n", + "s12_expr = SquaredInvariantMass(p12)\n", + "s23_expr = SquaredInvariantMass(p23)\n", + "s31_expr = SquaredInvariantMass(p31)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "kinematic_variables = {\n", + " theta1: theta1_expr,\n", + " theta2: theta2_expr,\n", + " theta3: theta3_expr,\n", + " phi1: phi1_expr,\n", + " phi2: phi2_expr,\n", + " phi3: phi3_expr,\n", + " s12: s12_expr,\n", + " s23: s23_expr,\n", + " s31: s31_expr,\n", + "}\n", + "\n", + "Latex(aslatex(kinematic_variables))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "helicity_transformer = SympyDataTransformer.from_sympy(\n", + " kinematic_variables, backend=\"jax\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "phsp = helicity_transformer(phsp_momenta)\n", + "list(phsp)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "a_vals = [0, 1.0, 3.0, 3.5, 2.0] # Slight adjustment to emphasize higher waves\n", + "b_vals = [0, -1.0, 3.5, 0.5, 0] # Adjust for new final state coupling\n", + "c_vals = [0, 0, 3.0, 0, 0] # Adjust if more s-wave or p-wave is expected\n", + "\n", + "m_Kstar2_val = 1.43 # K*2(1430)\n", + "m_Sigma_val = 1.66 # Sigma(1385) Sigma(1660), (1670)\n", + "m_Nstar_val = 1.90 # N(1710) N(1720) N(1900)\n", + "\n", + "Gamma_Kstar2_val = 0.1\n", + "Gamma_Sigma_val = 0.1\n", + "Gamma_Nstar_val = 0.1\n", + "\n", + "l12_val = 2 # I still use 2 assuming K^*_2\n", + "l23_val = 1\n", + "l31_val = 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "parameters_default = {\n", + " m_Kstar2: m_Kstar2_val,\n", + " m_Sigma: m_Sigma_val,\n", + " m_Nstar: m_Nstar_val,\n", + " Gamma_Kstar2: Gamma_Kstar2_val,\n", + " Gamma_Sigma: Gamma_Sigma_val,\n", + " Gamma_Nstar: Gamma_Nstar_val,\n", + " l12: l12_val,\n", + " l23: l23_val,\n", + " l31: l31_val,\n", + "}\n", + "\n", + "a_dict = {a[i]: a_vals[i + l_max] for i in range(-l_max, l_max + 1)}\n", + "b_dict = {b[i]: b_vals[i + l_max] for i in range(-l_max, l_max + 1)}\n", + "c_dict = {c[i]: c_vals[i + l_max] for i in range(-l_max, l_max + 1)}\n", + "parameters_default.update(a_dict)\n", + "parameters_default.update(b_dict)\n", + "parameters_default.update(c_dict)\n", + "\n", + "Latex(aslatex(parameters_default))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Visulization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "sliders = {}\n", + "categorized_sliders_m = defaultdict(list)\n", + "categorized_sliders_gamma = defaultdict(list)\n", + "categorized_cphi_pair = defaultdict(list)\n", + "categorized_sliders_l = defaultdict(list)\n", + "\n", + "for symbol, value in parameters_default.items():\n", + " if symbol.name.startswith(R\"\\Gamma_{\"):\n", + " slider = w.FloatSlider(\n", + " description=Rf\"\\({sp.latex(symbol)}\\)\",\n", + " min=0.0,\n", + " max=1.0,\n", + " step=0.01,\n", + " value=value,\n", + " continuous_update=False,\n", + " )\n", + " sliders[symbol.name] = slider\n", + " if symbol.name.startswith(R\"\\Gamma_{N\"):\n", + " categorized_sliders_gamma[0].append(slider)\n", + " elif symbol.name.startswith(R\"\\Gamma_{\\S\"):\n", + " categorized_sliders_gamma[1].append(slider)\n", + " elif symbol.name.startswith(R\"\\Gamma_{K\"):\n", + " categorized_sliders_gamma[2].append(slider)\n", + "\n", + " elif symbol.name.startswith(\"m_{\"):\n", + " slider = w.FloatSlider(\n", + " description=Rf\"\\({sp.latex(symbol)}\\)\",\n", + " min=0.63,\n", + " max=4,\n", + " step=0.01,\n", + " value=value,\n", + " continuous_update=False,\n", + " )\n", + " sliders[symbol.name] = slider\n", + " if symbol.name.startswith(\"m_{N\"):\n", + " categorized_sliders_m[0].append(slider)\n", + " elif symbol.name.startswith(R\"m_{\\S\"):\n", + " categorized_sliders_m[1].append(slider)\n", + " elif symbol.name.startswith(\"m_{K\"):\n", + " categorized_sliders_m[2].append(slider)\n", + "\n", + " elif isinstance(symbol, sp.Indexed):\n", + " c_latex = sp.latex(symbol)\n", + " phi_latex = Rf\"\\phi_{{{c_latex}}}\"\n", + " slider_c = w.FloatSlider(\n", + " description=Rf\"\\({c_latex}\\)\",\n", + " min=0,\n", + " max=10,\n", + " step=0.01,\n", + " value=abs(value),\n", + " continuous_update=False,\n", + " )\n", + " slider_phi = w.FloatSlider(\n", + " description=Rf\"\\({phi_latex}\\)\",\n", + " min=-np.pi,\n", + " max=+np.pi,\n", + " step=0.01,\n", + " value=np.angle(value),\n", + " continuous_update=False,\n", + " )\n", + "\n", + " sliders[symbol.name] = slider_c\n", + " sliders[f\"phi_{symbol.name}\"] = slider_phi\n", + " cphi_hbox = w.HBox([slider_c, slider_phi])\n", + " if symbol.base is a:\n", + " categorized_cphi_pair[2].append(cphi_hbox)\n", + " elif symbol.base is b:\n", + " categorized_cphi_pair[1].append(cphi_hbox)\n", + " elif symbol.base is c:\n", + " categorized_cphi_pair[0].append(cphi_hbox)\n", + " else:\n", + " raise NotImplementedError(symbol.name)\n", + "\n", + " elif symbol.name.startswith(\"l_{12}\"):\n", + " slider = w.IntSlider(\n", + " value=2,\n", + " min=0,\n", + " max=2,\n", + " step=1,\n", + " description=\"ℓ₁₂\",\n", + " disabled=False,\n", + " continuous_update=False,\n", + " orientation=\"horizontal\",\n", + " readout=True,\n", + " readout_format=\"d\",\n", + " )\n", + " sliders[symbol.name] = slider\n", + " categorized_sliders_l[2].append(slider)\n", + " elif symbol.name.startswith(\"l_{23}\"):\n", + " slider = w.IntSlider(\n", + " value=1,\n", + " min=0,\n", + " max=2,\n", + " step=1,\n", + " description=\"ℓ₂₃\",\n", + " disabled=False,\n", + " continuous_update=False,\n", + " orientation=\"horizontal\",\n", + " readout=True,\n", + " readout_format=\"d\",\n", + " )\n", + " sliders[symbol.name] = slider\n", + " categorized_sliders_l[1].append(slider)\n", + " elif symbol.name.startswith(\"l_{31}\"):\n", + " slider = w.IntSlider(\n", + " value=0,\n", + " min=0,\n", + " max=2,\n", + " step=1,\n", + " description=\"ℓ₃₁\",\n", + " disabled=False,\n", + " continuous_update=False,\n", + " orientation=\"horizontal\",\n", + " readout=True,\n", + " readout_format=\"d\",\n", + " )\n", + " sliders[symbol.name] = slider\n", + " categorized_sliders_l[0].append(slider)\n", + " else:\n", + " raise NotImplementedError(symbol.name)\n", + "tab_contents = []\n", + "resonances_name = [\"N* [ΛK⁺ (31)]\", \"Σ* [π⁰Λ(23)]\", \"K* [K⁺π⁰(12)]\"]\n", + "for i in range(len(resonances_name)):\n", + " tab_content = w.VBox([\n", + " w.HBox(categorized_sliders_m[i] + categorized_sliders_gamma[i]),\n", + " w.VBox(categorized_cphi_pair[i]),\n", + " ])\n", + " tab_contents.append(tab_content)\n", + "UI_lsliders = w.HBox([\n", + " *categorized_sliders_l[0],\n", + " *categorized_sliders_l[1],\n", + " *categorized_sliders_l[2],\n", + "])\n", + "UI_1 = w.Tab(tab_contents, titles=resonances_name)\n", + "UI = w.VBox([UI_lsliders, UI_1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "intensity_funcs = np.array([\n", + " [\n", + " [\n", + " create_parametrized_function(\n", + " expression=intensity_expr.xreplace({l12: i, l23: j, l31: k})\n", + " .doit()\n", + " .expand(func=True),\n", + " parameters=parameters_default,\n", + " backend=\"jax\",\n", + " )\n", + " for i in range(l_max + 1)\n", + " ]\n", + " for j in range(l_max + 1)\n", + " ]\n", + " for k in range(l_max + 1)\n", + "])\n", + "intensity_funcs.shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def insert_phi(parameters: dict) -> dict:\n", + " updated_parameters = {}\n", + " for key, value in parameters.items():\n", + " if key.startswith(\"phi_\"):\n", + " continue\n", + " if key.startswith((\"a\", \"b\", \"c\")):\n", + " phi_key = f\"phi_{key}\"\n", + " if phi_key in parameters:\n", + " phi = parameters[phi_key]\n", + " value *= np.exp(1j * phi) # noqa:PLW2901\n", + " updated_parameters[key] = value\n", + "\n", + " return updated_parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "%matplotlib widget\n", + "%config InlineBackend.figure_formats = ['png']\n", + "fig_2d, ax_2d = plt.subplots(dpi=150)\n", + "ax_2d.set_title(\"Model-weighted Phase space Dalitz Plot\")\n", + "ax_2d.set_xlabel(R\"$m^2(\\Lambda K^+)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n", + "ax_2d.set_ylabel(R\"$m^2(K^+ \\pi^0)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n", + "\n", + "mesh = None\n", + "\n", + "\n", + "def update_histogram(**parameters):\n", + " global mesh\n", + "\n", + " l12para = parameters[\"l_{12}\"]\n", + " l23para = parameters[\"l_{23}\"]\n", + " l31para = parameters[\"l_{31}\"]\n", + " intensity_func = intensity_funcs[l12para, l23para, l31para]\n", + " parameters = insert_phi(parameters)\n", + "\n", + " intensity_func.update_parameters(parameters)\n", + " intensity_weights = intensity_func(phsp)\n", + " bin_values, xedges, yedges = jnp.histogram2d(\n", + " phsp[\"s_{31}\"],\n", + " phsp[\"s_{12}\"],\n", + " bins=200,\n", + " weights=intensity_weights,\n", + " density=True,\n", + " )\n", + " bin_values = jnp.where(bin_values < 1e-10, jnp.nan, bin_values)\n", + " x, y = jnp.meshgrid(xedges[:-1], yedges[:-1])\n", + " if mesh is None:\n", + " mesh = ax_2d.pcolormesh(x, y, bin_values.T, cmap=\"jet\", vmax=0.15)\n", + " else:\n", + " mesh.set_array(bin_values.T)\n", + " fig_2d.canvas.draw_idle()\n", + "\n", + "\n", + "interactive_plot = w.interactive_output(\n", + " update_histogram,\n", + " {**sliders},\n", + ")\n", + "fig_2d.tight_layout()\n", + "fig_2d.colorbar(mesh, ax=ax_2d)\n", + "\n", + "if STATIC_PAGE:\n", + " filename = \"dalitz-plot-man.png\"\n", + " fig_2d.savefig(filename)\n", + " plt.close(fig_2d)\n", + " display(UI, Image(filename))\n", + "else:\n", + " display(UI, interactive_plot)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "%matplotlib widget\n", + "%config InlineBackend.figure_formats = ['svg']\n", + "\n", + "theta_subtitles = [\n", + " R\"$\\cos (\\theta_{1}^{{12}}) \\equiv \\cos (\\theta_{K^+}^{{K^+ \\pi^0}})$\",\n", + " R\"$\\cos (\\theta_{2}^{{23}}) \\equiv \\cos (\\theta_{\\pi^0}^{{\\pi^0 \\Lambda}})$\",\n", + " R\"$\\cos (\\theta_{3}^{{31}}) \\equiv \\cos (\\theta_{\\Lambda}^{{\\Lambda K^+}})$\",\n", + "]\n", + "phi_subtitles = [\n", + " R\"$\\phi_1^{12} \\equiv \\phi_{K^+}^{{K^+ \\pi^0}}$\",\n", + " R\"$\\phi_2^{23} \\equiv \\phi_{\\pi^0}^{{\\pi^0 \\Lambda}}$\",\n", + " R\"$\\phi_3^{31} \\equiv \\phi_{\\Lambda}^{{\\Lambda K^+}}$\",\n", + "]\n", + "mass_subtitles = [\n", + " R\"$m_{12} \\equiv m_{K^+ \\pi^0}$\",\n", + " R\"$m_{23} \\equiv m_{\\pi^0 \\Lambda}$\",\n", + " R\"$m_{31} \\equiv m_{\\Lambda K^+}$\",\n", + "]\n", + "\n", + "fig, (theta_axes, phi_axes, mass_axes) = plt.subplots(figsize=(12, 8), ncols=3, nrows=3)\n", + "fig.canvas.toolbar_visible = False\n", + "fig.canvas.header_visible = False\n", + "fig.canvas.footer_visible = False\n", + "\n", + "for i, ax in enumerate(theta_axes):\n", + " ax.set_title(theta_subtitles[i])\n", + " ax.set_xticks([-1, 0, 1])\n", + "\n", + "for i, ax in enumerate(phi_axes):\n", + " ax.set_title(phi_subtitles[i])\n", + " ax.set_xticks([-np.pi, 0, np.pi])\n", + " ax.set_xticklabels([R\"-$\\pi$\", 0, R\"$\\pi$\"])\n", + "\n", + "for i, ax in enumerate(mass_axes):\n", + " ax.set_title(mass_subtitles[i])\n", + " ax.set_xlabel(\"Mass [GeV]\")\n", + "\n", + "LINES = 3 * [None]\n", + "THETA_LINES = 3 * [None]\n", + "PHI_LINES = 3 * [None]\n", + "RESONANCE_LINES = defaultdict(dict)\n", + "RESONANCES_MASS_NAME = [m_Kstar2, m_Sigma, m_Nstar]\n", + "\n", + "\n", + "def update_plot(**parameters): # noqa: C901, PLR0912, PLR0914\n", + " l12para = parameters[\"l_{12}\"]\n", + " l23para = parameters[\"l_{23}\"]\n", + " l31para = parameters[\"l_{31}\"]\n", + " intensity_func = intensity_funcs[l12para, l23para, l31para]\n", + " parameters = insert_phi(parameters)\n", + " intensity_func.update_parameters(parameters)\n", + " intensities = intensity_func(phsp)\n", + " max_value_theta = 0.0\n", + " max_value_phi = 0.0\n", + " max_value_mass = 0.0\n", + " theta_keys = [\"theta_1\", \"theta_2\", \"theta_3\"]\n", + " phi_keys = [\"phi_1\", \"phi_2\", \"phi_3\"]\n", + " m2_keys = [\"s_{12}\", \"s_{23}\", \"s_{31}\"]\n", + " line_labels = [R\"$m_{K^{*+}_2}$\", R\"$m_{\\Sigma^*}$\", R\"$m_{N^*}$\"]\n", + " line_colors = [\"r\", \"g\", \"b\"]\n", + " plot_style = dict(bins=120, weights=intensities, density=True)\n", + "\n", + " for i, ax in enumerate(mass_axes):\n", + " bin_values, bin_edges = jnp.histogram(np.sqrt(phsp[m2_keys[i]]), **plot_style)\n", + " max_value_mass = max(max_value_mass, bin_values.max())\n", + "\n", + " if LINES[i] is None:\n", + " LINES[i] = ax.step(bin_edges[:-1], bin_values, alpha=0.5)[0]\n", + " else:\n", + " LINES[i].set_ydata(bin_values)\n", + "\n", + " symbol_key = sp.latex([m_Kstar2, m_Sigma, m_Nstar][i])\n", + " val = parameters[symbol_key]\n", + " resonance_line = RESONANCE_LINES[i].get(symbol_key)\n", + " if resonance_line is None:\n", + " RESONANCE_LINES[i][symbol_key] = ax.axvline(\n", + " val, color=line_colors[i], linestyle=\"--\", label=line_labels[i]\n", + " )\n", + " else:\n", + " resonance_line.set_xdata([val, val])\n", + "\n", + " for i, ax in enumerate(theta_axes):\n", + " bin_values, bin_edges = jnp.histogram(np.cos(phsp[theta_keys[i]]), **plot_style)\n", + " max_value_theta = max(max_value_theta, bin_values.max())\n", + " if THETA_LINES[i] is None:\n", + " THETA_LINES[i] = ax.step(bin_edges[:-1], bin_values, alpha=0.5)[0]\n", + " else:\n", + " THETA_LINES[i].set_ydata(bin_values)\n", + "\n", + " for i, ax in enumerate(phi_axes):\n", + " bin_values, bin_edges = jnp.histogram(phsp[phi_keys[i]], **plot_style)\n", + " max_value_phi = max(max_value_phi, bin_values.max())\n", + " if PHI_LINES[i] is None:\n", + " PHI_LINES[i] = ax.step(bin_edges[:-1], bin_values, alpha=0.5)[0]\n", + " else:\n", + " PHI_LINES[i].set_ydata(bin_values)\n", + "\n", + " for ax in mass_axes:\n", + " ax.set_ylim(0, max_value_mass * 1.1)\n", + " ax.legend()\n", + "\n", + " for ax in theta_axes:\n", + " ax.set_ylim(0, max_value_theta * 1.1)\n", + "\n", + " for ax in phi_axes:\n", + " ax.set_ylim(0, max_value_phi * 1.1)\n", + "\n", + "\n", + "interactive_plot = w.interactive_output(update_plot, sliders)\n", + "fig.tight_layout()\n", + "\n", + "if STATIC_PAGE:\n", + " filename = \"1d-histograms-man.svg\"\n", + " fig.savefig(filename)\n", + " plt.close(fig)\n", + " display(UI, SVG(filename))\n", + "else:\n", + " display(UI, interactive_plot)" + ] + } + ], + "metadata": { + "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.12.5" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +}