{ "cells": [ { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "# Intensity distribution" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```{autolink-concat}\n", "```" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Import Python libraries" }, "tags": [ "hide-cell", "scroll-input" ] }, "outputs": [], "source": [ "from __future__ import annotations\n", "\n", "import logging\n", "import os\n", "from io import BytesIO\n", "from itertools import product\n", "from urllib.request import urlopen\n", "from zipfile import ZipFile\n", "\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.helicity.naming import natural_sorting\n", "from IPython.display import Markdown\n", "from matplotlib.image import imread\n", "from matplotlib.patches import Rectangle\n", "from tensorwaves.interface import DataSample\n", "from tqdm.auto import tqdm\n", "\n", "from polarimetry.data import (\n", " create_data_transformer,\n", " generate_meshgrid_sample,\n", " generate_phasespace_sample,\n", " generate_sub_meshgrid_sample,\n", ")\n", "from polarimetry.decay import Particle\n", "from polarimetry.function import (\n", " compute_sub_function,\n", " integrate_intensity,\n", " interference_intensity,\n", " sub_intensity,\n", ")\n", "from polarimetry.io import (\n", " mute_jax_warnings,\n", " perform_cached_doit,\n", " perform_cached_lambdify,\n", ")\n", "from polarimetry.lhcb import load_model\n", "from polarimetry.lhcb.particle import load_particles\n", "from polarimetry.plot import (\n", " add_watermark,\n", " get_contour_line,\n", " stylize_contour,\n", " use_mpl_latex_fonts,\n", ")\n", "\n", "mute_jax_warnings()\n", "particles = load_particles(\"../data/particle-definitions.yaml\")\n", "model = load_model(\"../data/model-definitions.yaml\", particles, model_id=0)\n", "\n", "NO_TQDM = \"EXECUTE_NB\" in os.environ\n", "if NO_TQDM:\n", " logging.getLogger().setLevel(logging.ERROR)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [] }, "outputs": [], "source": [ "unfolded_intensity_expr = perform_cached_doit(model.full_expression)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "remove-input" ] }, "outputs": [], "source": [ "def assert_all_symbols_defined(expr: sp.Expr) -> None:\n", " sigmas = sp.symbols(\"sigma1:4\", nonnegative=True)\n", " remaining_symbols = expr.xreplace(model.parameter_defaults).free_symbols\n", " remaining_symbols -= set(model.variables)\n", " remaining_symbols -= set(sigmas)\n", " assert not remaining_symbols, remaining_symbols\n", "\n", "\n", "assert_all_symbols_defined(unfolded_intensity_expr)\n", "Markdown(\n", " \"The complete intensity expression contains\"\n", " f\" **{sp.count_ops(unfolded_intensity_expr):,} mathematical operations**.\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Definition of free parameters" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [] }, "outputs": [], "source": [ "free_parameters = {\n", " symbol: value\n", " for symbol, value in model.parameter_defaults.items()\n", " if isinstance(symbol, sp.Indexed)\n", " if \"production\" in str(symbol)\n", "}\n", "fixed_parameters = {\n", " symbol: value\n", " for symbol, value in model.parameter_defaults.items()\n", " if symbol not in free_parameters\n", "}\n", "subs_intensity_expr = unfolded_intensity_expr.xreplace(fixed_parameters)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "remove-input" ] }, "outputs": [], "source": [ "Markdown(\n", " \"After substituting the parameters that are not production couplings, the total\"\n", " \" intensity expression contains\"\n", " f\" **{sp.count_ops(subs_intensity_expr):,} operations**.\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Distribution" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "intensity_func = perform_cached_lambdify(\n", " subs_intensity_expr,\n", " parameters=free_parameters,\n", " backend=\"jax\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "transformer = create_data_transformer(model)\n", "grid_sample = generate_meshgrid_sample(model.decay, resolution=1_000)\n", "grid_sample = transformer(grid_sample)\n", "X = grid_sample[\"sigma1\"]\n", "Y = grid_sample[\"sigma2\"]" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['png']\n", "s1_label = R\"$\\sigma_1=m^2\\left(K^-\\pi^+\\right)$ [GeV$^2$]\"\n", "s2_label = R\"$\\sigma_2=m^2\\left(pK^-\\right)$ [GeV$^2$]\"\n", "s3_label = R\"$\\sigma_3=m^2\\left(p\\pi^+\\right)$ [GeV$^2$]\"\n", "\n", "plt.rcdefaults()\n", "use_mpl_latex_fonts()\n", "plt.rc(\"font\", size=21)\n", "fig, ax = plt.subplots(dpi=200, figsize=(8.22, 7), tight_layout=True)\n", "ax.set_xlabel(s1_label)\n", "ax.set_ylabel(s2_label)\n", "\n", "INTENSITIES = intensity_func(grid_sample)\n", "INTENSITY_INTEGRAL = jnp.nansum(INTENSITIES)\n", "NORMALIZED_INTENSITIES = INTENSITIES / INTENSITY_INTEGRAL\n", "np.testing.assert_almost_equal(jnp.nansum(NORMALIZED_INTENSITIES), 1.0)\n", "mesh = ax.pcolormesh(X, Y, NORMALIZED_INTENSITIES)\n", "c_bar = fig.colorbar(mesh, ax=ax, pad=0.02)\n", "c_bar.ax.set_ylabel(\"Normalized intensity\")\n", "add_watermark(ax, 0.7, 0.82, fontsize=24)\n", "fig.savefig(\"_static/images/intensity-distribution.png\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "{{ DOWNLOAD_INTENSITY_DISTRIBUTION }}" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Comparison with [Figure 2](https://arxiv.org/pdf/2208.03262.pdf#page=9) from the original LHCb study {cite}`LHCb-PAPER-2022-002`:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "full-width", "hide-input" ] }, "outputs": [], "source": [ "def download_lhcb_intensity() -> None:\n", " figure_path = \"LHCb-PAPER-2022-002-figures/Fig2.png\"\n", " if os.path.exists(figure_path):\n", " return\n", " url = \"https://cds.cern.ch/record/2824328/files/LHCb-PAPER-2022-002-figures.zip?version=1\"\n", " http_response = urlopen(url) # noqa: S310\n", " zipfile = ZipFile(BytesIO(http_response.read()))\n", " zipfile.extract(figure_path)\n", "\n", "\n", "def plot_horizontal_intensity(ax) -> None:\n", " ax.set_xlabel(\"$\" + s2_label[10:])\n", " ax.set_ylabel(\"$\" + s1_label[10:])\n", " ax.set_xlim(1.79, 4.95)\n", " ax.set_ylim(0.18, 2.05)\n", " add_watermark(ax, 0.7, 0.78, fontsize=18)\n", " mesh = ax.pcolormesh(\n", " grid_sample[\"sigma2\"],\n", " grid_sample[\"sigma1\"],\n", " NORMALIZED_INTENSITIES,\n", " )\n", " c_bar = fig.colorbar(mesh, ax=ax, pad=0.02)\n", " c_bar.ax.set_ylabel(\"Normalized intensity\")\n", "\n", "\n", "%config InlineBackend.figure_formats = ['png']\n", "download_lhcb_intensity()\n", "plt.rcdefaults()\n", "use_mpl_latex_fonts()\n", "plt.rc(\"font\", size=16)\n", "fig = plt.figure(dpi=200, figsize=(11.5, 4))\n", "ax1 = fig.add_axes((0.1, 0.17, 0.40, 0.78))\n", "ax2 = fig.add_axes((0.5, 0.0, 0.5, 1.0))\n", "ax2.axis(\"off\")\n", "plot_horizontal_intensity(ax1)\n", "image = imread(\"LHCb-PAPER-2022-002-figures/Fig2.png\")\n", "ax2.imshow(image)\n", "fig.savefig(\"_static/images/intensity-distribution-comparison.png\")\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Export left part of the intensity comparison" }, "tags": [ "hide-cell" ] }, "outputs": [], "source": [ "plt.ioff()\n", "fig, ax = plt.subplots(dpi=200, figsize=(5.8, 4))\n", "plot_horizontal_intensity(ax)\n", "fig.savefig(\n", " \"_static/images/intensity-distribution-comparison-left.png\",\n", " bbox_inches=\"tight\",\n", ")\n", "plt.ion()\n", "plt.close(fig)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input", "full-width" ] }, "outputs": [], "source": [ "def set_ylim_to_zero(ax):\n", " _, y_max = ax.get_ylim()\n", " ax.set_ylim(0, y_max)\n", "\n", "\n", "%config InlineBackend.figure_formats = ['svg']\n", "plt.rcdefaults()\n", "use_mpl_latex_fonts()\n", "plt.rc(\"font\", size=18)\n", "fig, (ax1, ax2) = plt.subplots(\n", " ncols=2,\n", " figsize=(12, 5),\n", " tight_layout=True,\n", " sharey=True,\n", ")\n", "ax1.set_xlabel(s1_label)\n", "ax2.set_xlabel(s2_label)\n", "ax1.set_ylabel(\"Normalized intensity\")\n", "\n", "subsystem_identifiers = [\"K\", \"L\", \"D\"]\n", "subsystem_labels = [\"K^{**}\", R\"\\Lambda^{**}\", R\"\\Delta^{**}\"]\n", "x, y = X[0], Y[:, 0]\n", "ax1.fill(x, jnp.nansum(NORMALIZED_INTENSITIES, axis=0), alpha=0.3)\n", "ax2.fill(y, jnp.nansum(NORMALIZED_INTENSITIES, axis=1), alpha=0.3)\n", "\n", "original_parameters = dict(intensity_func.parameters)\n", "for label, identifier in zip(subsystem_labels, subsystem_identifiers):\n", " label = f\"${label}$\"\n", " sub_intensities = compute_sub_function(intensity_func, grid_sample, [identifier])\n", " sub_intensities /= INTENSITY_INTEGRAL\n", " ax1.plot(x, jnp.nansum(sub_intensities, axis=0), label=label)\n", " ax2.plot(y, jnp.nansum(sub_intensities, axis=1), label=label)\n", " del sub_intensities\n", " intensity_func.update_parameters(original_parameters)\n", "set_ylim_to_zero(ax1)\n", "set_ylim_to_zero(ax2)\n", "ax2.legend()\n", "plt.savefig(\"_images/intensity-distributions-1D.svg\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ":::{only} latex\n", "{{ FIG_INTENSITY }}\n", ":::" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Decay rates" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "integration_sample = generate_phasespace_sample(\n", " model.decay, n_events=100_000, seed=0\n", ")\n", "integration_sample = transformer(integration_sample)\n", "I_tot = integrate_intensity(intensity_func(integration_sample))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "remove-cell" ] }, "outputs": [], "source": [ "np.testing.assert_allclose(\n", " I_tot,\n", " sub_intensity(intensity_func, integration_sample, [\"K\", \"L\", \"D\"]),\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [] }, "outputs": [], "source": [ "I_K = sub_intensity(intensity_func, integration_sample, non_zero_couplings=[\"K\"])\n", "I_Λ = sub_intensity(intensity_func, integration_sample, non_zero_couplings=[\"L\"])\n", "I_Δ = sub_intensity(intensity_func, integration_sample, non_zero_couplings=[\"D\"])\n", "I_ΛΔ = interference_intensity(intensity_func, integration_sample, [\"L\"], [\"D\"])\n", "I_KΔ = interference_intensity(intensity_func, integration_sample, [\"K\"], [\"D\"])\n", "I_KΛ = interference_intensity(intensity_func, integration_sample, [\"K\"], [\"L\"])\n", "np.testing.assert_allclose(I_tot, I_K + I_Λ + I_Δ + I_ΛΔ + I_KΔ + I_KΛ)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Functions for computing the decay rate" }, "tags": [ "hide-cell" ] }, "outputs": [], "source": [ "def compute_decay_rates(func, integration_sample: DataSample):\n", " decay_rates = np.zeros(shape=(n_resonances, n_resonances))\n", " combinations = list(product(enumerate(resonances), enumerate(resonances)))\n", " progress_bar = tqdm(\n", " desc=\"Calculating rate matrix\",\n", " disable=NO_TQDM,\n", " total=(len(combinations) + n_resonances) // 2,\n", " )\n", " I_tot = integrate_intensity(intensity_func(integration_sample))\n", " for (i, resonance1), (j, resonance2) in combinations:\n", " if j < i:\n", " continue\n", " progress_bar.postfix = f\"{resonance1.name} × {resonance2.name}\"\n", " res1 = to_regex(resonance1.name)\n", " res2 = to_regex(resonance2.name)\n", " if res1 == res2:\n", " I_sub = sub_intensity(\n", " func, integration_sample, non_zero_couplings=[res1]\n", " )\n", " else:\n", " I_sub = interference_intensity(func, integration_sample, [res1], [res2])\n", " decay_rates[i, j] = I_sub / I_tot\n", " if i != j:\n", " decay_rates[j, i] = decay_rates[i, j]\n", " progress_bar.update()\n", " progress_bar.close()\n", " return decay_rates\n", "\n", "\n", "def to_regex(text: str) -> str:\n", " text = text.replace(\"(\", r\"\\(\")\n", " return text.replace(\")\", r\"\\)\")\n", "\n", "\n", "def sort_resonances(resonance: Particle):\n", " KDL = {\"L\": 1, \"D\": 2, \"K\": 3}\n", " return KDL[resonance.name[0]], natural_sorting(resonance.name)\n", "\n", "\n", "resonances = sorted(\n", " (chain.resonance for chain in model.decay.chains),\n", " key=sort_resonances,\n", " reverse=True,\n", ")\n", "n_resonances = len(resonances)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "def visualize_decay_rates(decay_rates, title=R\"Rate matrix for isobars (\\%)\"):\n", " vmax = jnp.max(jnp.abs(decay_rates))\n", " plt.rcdefaults()\n", " use_mpl_latex_fonts()\n", " plt.rc(\"font\", size=14)\n", " plt.rc(\"axes\", titlesize=24)\n", " plt.rc(\"xtick\", labelsize=16)\n", " plt.rc(\"ytick\", labelsize=16)\n", " fig, ax = plt.subplots(figsize=(9, 9))\n", " ax.set_title(title)\n", " ax.matshow(\n", " jnp.rot90(decay_rates).T, cmap=plt.cm.coolwarm, vmin=-vmax, vmax=+vmax\n", " )\n", "\n", " resonance_latex = [f\"${p.latex}$\" for p in resonances]\n", " ax.set_xticks(range(n_resonances))\n", " ax.set_xticklabels(reversed(resonance_latex), rotation=90)\n", " ax.set_yticks(range(n_resonances))\n", " ax.set_yticklabels(resonance_latex)\n", " for i in range(n_resonances):\n", " for j in range(n_resonances):\n", " if j < i:\n", " continue\n", " rate = decay_rates[i, j]\n", " ax.text(\n", " n_resonances - j - 1,\n", " i,\n", " f\"{100 * rate:.2f}\",\n", " va=\"center\",\n", " ha=\"center\",\n", " )\n", " fig.tight_layout()\n", " return fig\n", "\n", "\n", "%config InlineBackend.figure_formats = ['svg']\n", "decay_rates = compute_decay_rates(intensity_func, integration_sample)\n", "fig = visualize_decay_rates(decay_rates)\n", "fig.savefig(\"_images/rate-matrix.svg\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ":::{only} latex\n", "{{ FIG_RATE_MATRIX }}\n", ":::" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Function for computing the total over all rates" }, "tags": [ "hide-cell" ] }, "outputs": [], "source": [ "def compute_sum_over_decay_rates(decay_rate_matrix) -> float:\n", " decay_rate_sum = 0.0\n", " for i in range(len(resonances)):\n", " for j in range(len(resonances)):\n", " if j < i:\n", " continue\n", " decay_rate_sum += decay_rate_matrix[i, j]\n", " return decay_rate_sum" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [] }, "outputs": [], "source": [ "np.testing.assert_almost_equal(compute_sum_over_decay_rates(decay_rates), 1.0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Dominant decays" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "threshold = 0.5\n", "percentage = int(100 * threshold)\n", "I_tot = intensity_func(grid_sample)\n", "\n", "plt.rcdefaults()\n", "use_mpl_latex_fonts()\n", "plt.rc(\"font\", size=18)\n", "fig, ax = plt.subplots(figsize=(9.1, 7), sharey=True, tight_layout=True)\n", "ax.set_ylabel(s2_label)\n", "ax.set_xlabel(s1_label)\n", "fig.suptitle(\n", " Rf\"Regions where the resonance has a decay ratio of $\\geq {percentage}$\\%\",\n", " y=0.95,\n", ")\n", "\n", "phsp_region = jnp.select(\n", " [I_tot > 0, True],\n", " (1, 0),\n", ")\n", "contour_set = ax.contour(X, Y, phsp_region, colors=\"none\")\n", "stylize_contour(contour_set, edgecolor=\"black\", linewidth=0.2)\n", "\n", "resonances = [c.resonance for c in model.decay.chains]\n", "contour_levels = [i for i, _ in enumerate(resonances, 1)]\n", "colors = [plt.cm.rainbow(x) for x in np.linspace(0, 1, len(resonances))]\n", "linestyles = {\n", " \"K\": \"dotted\",\n", " \"L\": \"dashed\",\n", " \"D\": \"solid\",\n", "}\n", "items = list(zip(contour_levels, resonances, colors)) # tqdm requires len\n", "progress_bar = tqdm(\n", " desc=\"Computing dominant region contours\",\n", " disable=NO_TQDM,\n", " total=len(items),\n", ")\n", "legend_elements = []\n", "for res_id, resonance, color in items:\n", " progress_bar.postfix = resonance.name\n", " regex_filter = resonance.name.replace(\"(\", r\"\\(\").replace(\")\", r\"\\)\")\n", " I_sub = compute_sub_function(intensity_func, grid_sample, [regex_filter])\n", " ratio = I_sub / I_tot\n", " selection = jnp.select(\n", " [jnp.isnan(ratio), ratio < threshold, True],\n", " [0, 0, res_id],\n", " )\n", " progress_bar.update()\n", " if jnp.all(selection == 0):\n", " continue\n", " contour_set = ax.contour(X, Y, selection, colors=\"none\")\n", " contour_set.set_clim(vmin=1, vmax=len(model.decay.chains))\n", " stylize_contour(\n", " contour_set,\n", " label=f\"${resonance.latex}$\",\n", " edgecolor=color,\n", " linestyle=linestyles[resonance.name[0]],\n", " )\n", " line_collection = get_contour_line(contour_set)\n", " legend_elements.append(line_collection)\n", "progress_bar.close()\n", "\n", "\n", "sub_region_x_range = (0.68, 0.72)\n", "sub_region_y_range = (2.52, 2.60)\n", "region_indicator = Rectangle(\n", " xy=(\n", " sub_region_x_range[0],\n", " sub_region_y_range[0],\n", " ),\n", " width=sub_region_x_range[1] - sub_region_x_range[0],\n", " height=sub_region_y_range[1] - sub_region_y_range[0],\n", " edgecolor=\"black\",\n", " facecolor=\"none\",\n", " label=\"Sub-region\",\n", " linewidth=0.5,\n", ")\n", "ax.add_patch(region_indicator)\n", "legend_elements.append(region_indicator)\n", "\n", "leg = plt.legend(\n", " handles=legend_elements,\n", " bbox_to_anchor=(1.38, 1),\n", " framealpha=1,\n", " loc=\"upper right\",\n", ")\n", "fig.savefig(\"_images/sub-regions.svg\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ":::{only} latex\n", "{{ FIG_SUB_REGIONS }}\n", ":::" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "sub_sample = generate_sub_meshgrid_sample(\n", " model.decay,\n", " resolution=50,\n", " x_range=sub_region_x_range,\n", " y_range=sub_region_y_range,\n", ")\n", "sub_sample = transformer(sub_sample)\n", "sub_decay_rates = compute_decay_rates(intensity_func, sub_sample)\n", "fig = visualize_decay_rates(sub_decay_rates, title=\"Rate matrix over sub-region\")\n", "fig.savefig(\"_images/rate-matrix-sub-region.svg\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ":::{only} latex\n", "{{ FIG_RATE_MATRIX_SUB }}\n", ":::" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "np.testing.assert_almost_equal(compute_sum_over_decay_rates(sub_decay_rates), 1.0)" ] } ], "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.8.15" }, "myst": { "all_links_external": true } }, "nbformat": 4, "nbformat_minor": 4 }