{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Polarimeter vector field" ] }, { "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 math\n", "import os\n", "import shutil\n", "from functools import reduce\n", "from warnings import filterwarnings\n", "\n", "import jax\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import svgutils.compose as sc\n", "import sympy as sp\n", "from IPython.display import SVG, display\n", "from matplotlib import cm\n", "from matplotlib.axes import Axes\n", "from matplotlib.collections import LineCollection\n", "from matplotlib.colors import LogNorm\n", "from matplotlib.patches import Patch\n", "from tensorwaves.interface import DataSample\n", "from tqdm.auto import tqdm\n", "\n", "from polarimetry import formulate_polarimetry\n", "from polarimetry.data import create_data_transformer, generate_meshgrid_sample\n", "from polarimetry.function import compute_sub_function\n", "from polarimetry.io import (\n", " mute_jax_warnings,\n", " perform_cached_doit,\n", " perform_cached_lambdify,\n", ")\n", "from polarimetry.lhcb import (\n", " flip_production_coupling_signs,\n", " load_model_builder,\n", " load_model_parameters,\n", ")\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", "filterwarnings(\"ignore\")\n", "logging.getLogger(\"polarimetry.function\").setLevel(logging.INFO)\n", "mute_jax_warnings()\n", "\n", "NO_TQDM = \"EXECUTE_NB\" in os.environ\n", "if NO_TQDM:\n", " logging.getLogger().setLevel(logging.ERROR)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Final state IDs:\n", "1. $p$\n", "2. $\\pi^+$\n", "3. $K^-$\n", "\n", "Sub-system definitions:\n", "1. $K^{**} \\to \\pi^+ K^-$\n", "2. $\\Lambda^{**} \\to p K^-$\n", "3. $\\Delta^{**} \\to p \\pi^+$" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Formulate amplitude models" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "model_choice = 0\n", "model_file = \"../data/model-definitions.yaml\"\n", "particles = load_particles(\"../data/particle-definitions.yaml\")\n", "amplitude_builder = load_model_builder(model_file, particles, model_choice)\n", "imported_parameter_values = load_model_parameters(\n", " model_file, amplitude_builder.decay, model_choice, particles\n", ")\n", "models = {}\n", "for reference_subsystem in [1, 2, 3]:\n", " models[reference_subsystem] = amplitude_builder.formulate(\n", " reference_subsystem, cleanup_summations=True\n", " )\n", " models[reference_subsystem].parameter_defaults.update(imported_parameter_values)\n", "del reference_subsystem\n", "\n", "models[2] = flip_production_coupling_signs(models[2], subsystem_names=[\"K\", \"L\"])\n", "models[3] = flip_production_coupling_signs(models[3], subsystem_names=[\"K\", \"D\"])\n", "\n", "DECAY = models[1].decay\n", "FINAL_STATE = {\n", " 1: \"p\",\n", " 2: R\"\\pi^+\",\n", " 3: \"K^-\",\n", "}" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Unfold symbolic expressions" }, "tags": [ "hide-input", "remove-output" ] }, "outputs": [], "source": [ "unfolded_polarimetry_exprs = {}\n", "unfolded_intensity_expr = {}\n", "for i, model in tqdm(models.items(), \"Unfolding expressions\", disable=NO_TQDM):\n", " reference_subsystem = i\n", " polarimetry_exprs = formulate_polarimetry(amplitude_builder, reference_subsystem)\n", " unfolded_polarimetry_exprs[i] = [\n", " perform_cached_doit(expr.doit().xreplace(model.amplitudes))\n", " for expr in tqdm(polarimetry_exprs, disable=NO_TQDM, leave=False)\n", " ]\n", " unfolded_intensity_expr[i] = perform_cached_doit(model.full_expression)\n", "del i, polarimetry_exprs, reference_subsystem" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Convert to numerical functions" }, "tags": [ "scroll-input", "hide-input", "remove-output" ] }, "outputs": [], "source": [ "polarimetry_funcs = {}\n", "intensity_func = {}\n", "for i, model in tqdm(models.items(), \"Lambdifying to JAX\", disable=NO_TQDM):\n", " production_couplings = {\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 production_couplings\n", " }\n", " polarimetry_funcs[i] = [\n", " perform_cached_lambdify(\n", " expr.xreplace(fixed_parameters),\n", " parameters=production_couplings,\n", " backend=\"jax\",\n", " )\n", " for expr in tqdm(unfolded_polarimetry_exprs[i], disable=NO_TQDM, leave=False)\n", " ]\n", " intensity_func[i] = perform_cached_lambdify(\n", " unfolded_intensity_expr[i].xreplace(fixed_parameters),\n", " parameters=production_couplings,\n", " backend=\"jax\",\n", " )\n", "\n", "del fixed_parameters, model, production_couplings" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Generate grid phase space sample" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "data_sample = generate_meshgrid_sample(DECAY, resolution=400)\n", "X = data_sample[\"sigma1\"]\n", "Y = data_sample[\"sigma2\"]\n", "for model in models.values():\n", " transformer = create_data_transformer(model)\n", " data_sample.update(transformer(data_sample))\n", "del model, transformer" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Dominant contributions" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Code for indicating sub-regions" }, "tags": [ "hide-cell" ] }, "outputs": [], "source": [ "def create_dominant_region_contours(\n", " decay, data_sample: DataSample, threshold: float\n", ") -> dict[str, jax.Array]:\n", " I_tot = intensity_func[1](data_sample)\n", " resonance_names = [chain.resonance.name for chain in decay.chains]\n", " region_filters = {}\n", " progress_bar = tqdm(\n", " desc=\"Computing dominant region contours\",\n", " disable=NO_TQDM,\n", " total=len(resonance_names),\n", " )\n", " for resonance_name in resonance_names:\n", " progress_bar.postfix = resonance_name\n", " regex_filter = resonance_name.replace(\"(\", r\"\\(\").replace(\")\", r\"\\)\")\n", " I_sub = compute_sub_function(intensity_func[1], data_sample, [regex_filter])\n", " ratio = I_sub / I_tot\n", " selection = jnp.select(\n", " [jnp.isnan(ratio), ratio < threshold, True],\n", " [0, 0, 1],\n", " )\n", " progress_bar.update()\n", " if jnp.all(selection == 0):\n", " continue\n", " region_filters[resonance_name] = selection\n", " contour_arrays = {}\n", " for contour_level, subsystem in enumerate([\"K\", \"L\", \"D\"], 1):\n", " contour_array = reduce(\n", " jnp.bitwise_or,\n", " (a for k, a in region_filters.items() if k.startswith(subsystem)),\n", " )\n", " contour_array *= contour_level\n", " contour_arrays[subsystem] = contour_array\n", " return contour_arrays\n", "\n", "\n", "def indicate_dominant_regions(\n", " contour_arrays, ax: Axes, selected_subsystems=None\n", ") -> dict[str, LineCollection]:\n", " if selected_subsystems is None:\n", " selected_subsystems = {\"K\", \"L\", \"D\"}\n", " selected_subsystems = set(selected_subsystems)\n", " colors = dict(K=\"red\", L=\"blue\", D=\"green\")\n", " labels = dict(K=\"K^{**}\", L=R\"\\Lambda^{**}\", D=R\"\\Delta^{**}\")\n", " legend_elements = {}\n", " for subsystem, Z in contour_arrays.items():\n", " if subsystem not in selected_subsystems:\n", " continue\n", " contour_set = ax.contour(X, Y, Z, colors=\"none\")\n", " stylize_contour(\n", " contour_set,\n", " edgecolor=colors[subsystem],\n", " linewidth=0.5,\n", " label=f\"${labels[subsystem]}$\",\n", " )\n", " line_collection = get_contour_line(contour_set)\n", " legend_elements[subsystem] = line_collection\n", " return legend_elements" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "scroll-input", "full-width", "hide-input" ] }, "outputs": [], "source": [ "%%time\n", "%config InlineBackend.figure_formats = ['png']\n", "subsystem_identifiers = [\"K\", \"L\", \"D\"]\n", "subsystem_labels = [\"K^{**}\", R\"\\Lambda^{**}\", R\"\\Delta^{**}\"]\n", "nrows = 4\n", "ncols = 5\n", "scale = 3.0\n", "aspect_ratio = 1.05\n", "plt.rcdefaults()\n", "use_mpl_latex_fonts()\n", "plt.rc(\"font\", size=15)\n", "fig, axes = plt.subplots(\n", " dpi=200,\n", " figsize=scale * np.array([ncols, aspect_ratio * nrows]),\n", " gridspec_kw={\"width_ratios\": (ncols - 1) * [1] + [1.24]},\n", " ncols=ncols,\n", " nrows=nrows,\n", " sharex=True,\n", " sharey=True,\n", ")\n", "plt.subplots_adjust(wspace=0.05)\n", "\n", "s1_label = R\"$m^2\\left(K^-\\pi^+\\right)$ [GeV$^2$]\"\n", "s2_label = R\"$m^2\\left(pK^-\\right)$ [GeV$^2$]\"\n", "for subsystem in range(nrows):\n", " for i in range(ncols):\n", " ax = axes[subsystem, i]\n", " if i == 0:\n", " alpha_str = R\"I_\\mathrm{tot}\"\n", " elif i == 1:\n", " alpha_str = R\"|\\alpha|\"\n", " else:\n", " xyz = i - 2\n", " alpha_str = Rf\"\\alpha_{'xyz'[xyz]}\"\n", " title = alpha_str\n", " if subsystem > 0:\n", " label = subsystem_labels[subsystem - 1]\n", " title = Rf\"{title}\\left({label}\\right)\"\n", " ax.set_title(f\"${title}$\")\n", " if ax is axes[-1, i]:\n", " ax.set_xlabel(s1_label)\n", " if i == 0:\n", " ax.set_ylabel(s2_label)\n", "\n", "intensity_arrays = []\n", "polarimetry_arrays = []\n", "for subsystem in range(nrows):\n", " # alpha_xyz distributions\n", " alpha_xyz_arrays = []\n", " for i in range(2, ncols):\n", " xyz = i - 2\n", " if subsystem == 0:\n", " z_values = polarimetry_funcs[1][xyz](data_sample)\n", " polarimetry_arrays.append(z_values)\n", " else:\n", " identifier = subsystem_identifiers[subsystem - 1]\n", " z_values = compute_sub_function(\n", " polarimetry_funcs[1][xyz], data_sample, identifier\n", " )\n", " z_values = np.real(z_values)\n", " alpha_xyz_arrays.append(z_values)\n", " mesh = axes[subsystem, i].pcolormesh(X, Y, z_values, cmap=cm.coolwarm)\n", " mesh.set_clim(vmin=-1, vmax=+1)\n", " if xyz == 2:\n", " c_bar = fig.colorbar(mesh, ax=axes[subsystem, i])\n", " c_bar.set_ticks([-1, 0, +1])\n", " c_bar.set_ticklabels([\"-1\", \"0\", \"+1\"])\n", " # absolute value of alpha_xyz vector\n", " alpha_abs = np.sqrt(np.sum(np.array(alpha_xyz_arrays) ** 2, axis=0))\n", " mesh = axes[subsystem, 1].pcolormesh(X, Y, alpha_abs, cmap=cm.coolwarm)\n", " mesh.set_clim(vmin=-1, vmax=+1)\n", " # total intensity\n", " if subsystem == 0:\n", " z_values = intensity_func[1](data_sample)\n", " else:\n", " identifier = subsystem_identifiers[subsystem - 1]\n", " z_values = compute_sub_function(intensity_func[1], data_sample, identifier)\n", " intensity_arrays.append(z_values)\n", " axes[subsystem, 0].pcolormesh(X, Y, z_values, norm=LogNorm())\n", "\n", "threshold = 0.7\n", "contour_arrays = create_dominant_region_contours(DECAY, data_sample, threshold)\n", "\n", "for ax in axes[0]:\n", " legend_elements = indicate_dominant_regions(contour_arrays, ax)\n", " if ax is axes[0, -1]:\n", " leg = ax.legend(\n", " handles=legend_elements.values(),\n", " title=Rf\"$>{100*threshold:.0f}\\%$\",\n", " bbox_to_anchor=(0.9, 0.88, 1.0, 0.1),\n", " framealpha=1,\n", " )\n", "\n", "for subsystem, ax_row in zip([\"K\", \"L\", \"D\"], axes[1:]):\n", " for ax in ax_row:\n", " indicate_dominant_regions(\n", " contour_arrays, ax, selected_subsystems=[subsystem]\n", " )\n", "\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "full-width", "hide-input" ] }, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['png']\n", "plt.rcdefaults()\n", "use_mpl_latex_fonts()\n", "plt.rc(\"font\", size=16)\n", "fig, axes = plt.subplots(\n", " dpi=200,\n", " figsize=(13, 5),\n", " gridspec_kw={\"width_ratios\": [1, 1, 1.2]},\n", " ncols=3,\n", " sharey=True,\n", " tight_layout=True,\n", ")\n", "axes[0].set_ylabel(s2_label)\n", "I_times_alpha = jnp.array(\n", " [array * intensity_arrays[0] for array in polarimetry_arrays]\n", ")\n", "global_min_max = float(jnp.nanmax(jnp.abs(I_times_alpha)))\n", "for ax, z_values, xyz in zip(axes, I_times_alpha, \"xyz\"):\n", " ax.set_title(Rf\"$\\alpha_{xyz} \\cdot I$\")\n", " ax.set_xlabel(s1_label)\n", " mesh = ax.pcolormesh(X, Y, np.real(z_values), cmap=cm.RdYlGn_r)\n", " mesh.set_clim(vmin=-global_min_max, vmax=global_min_max)\n", " if ax is axes[-1]:\n", " fig.colorbar(mesh, ax=ax, pad=0.02)\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Total polarimetry vector field" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input", "scroll-input" ] }, "outputs": [], "source": [ "def plot_field(\n", " reference_subsystem: int,\n", " contour_arrays: dict[str, jnp.array] | None = None,\n", " threshold: float | None = None,\n", " add_title: bool = False,\n", " watermark: bool = False,\n", " show: bool = False,\n", ") -> None:\n", " plt.ioff()\n", " plt.rcdefaults()\n", " use_mpl_latex_fonts()\n", " plt.rc(\"font\", size=18)\n", " fig, ax = plt.subplots(\n", " figsize=(8, 6.8),\n", " tight_layout=True,\n", " )\n", " if add_title:\n", " ax.set_title(f\"Reference subsystem {reference_subsystem}\", y=1.02)\n", " ax.set_box_aspect(1)\n", " ax.set_xlabel(X_LABEL_ALPHA)\n", " ax.set_ylabel(Y_LABEL_ALPHA)\n", "\n", " polarimetry_arrays = [\n", " func(data_sample) for func in polarimetry_funcs[reference_subsystem]\n", " ]\n", " polarimetry_arrays = jnp.array(polarimetry_arrays).real\n", " mesh = plot_polarimetry_field(polarimetry_arrays, ax, strides=14)\n", " color_bar = fig.colorbar(mesh, ax=ax, pad=0.01)\n", " color_bar.set_label(R\"$\\left|\\vec{\\alpha}\\right|$\")\n", " if contour_arrays is not None:\n", " color_bar.ax.set_zorder(-10)\n", " xlim = ax.get_xlim()\n", " ylim = ax.get_ylim()\n", " _add_contours(ax, contour_arrays, threshold)\n", " ax.set_xlim(*xlim)\n", " ax.set_ylim(*ylim)\n", "\n", " if watermark:\n", " x_pos = 0.05 if contour_arrays is None else 0.2\n", " add_watermark(ax, x_pos, 0.04, fontsize=18)\n", "\n", " subsystem_id_to_name = {1: \"K\", 2: \"L\", 3: \"D\"}\n", " subsystem_name = subsystem_id_to_name[reference_subsystem]\n", " suffixes = [\n", " \"-contours\" if contour_arrays else \"\",\n", " \"-title\" if add_title else \"\",\n", " \"-watermark\" if watermark else \"\",\n", " ]\n", " suffix = \"\".join(suffixes)\n", " base_file = f\"_static/images/polarimetry-field-{subsystem_name}{suffix}.svg\"\n", " fig.savefig(base_file)\n", " plt.close(fig)\n", " plt.ion()\n", "\n", " overlay_file = f\"_images/orientation-{subsystem_name}.svg\"\n", " output_file = base_file.replace(\".svg\", \"-inset.svg\")\n", " y_pos = 0.08 if add_title else 0.058\n", " svg = overlay_inset(\n", " base_file,\n", " overlay_file,\n", " output_file,\n", " position=(0.353, y_pos),\n", " )\n", " if show:\n", " display(svg)\n", "\n", "\n", "def _add_contours(\n", " ax,\n", " contour_arrays: dict[str, jnp.array],\n", " threshold: float,\n", ") -> None:\n", " colors = dict(K=\"red\", L=\"blue\", D=\"green\")\n", " labels = dict(K=\"K^{**}\", L=R\"\\Lambda^{**}\", D=R\"\\Delta^{**}\")\n", " patch_transparency = 0.1\n", " for subsystem, Z in contour_arrays.items():\n", " contour_set = ax.contour(X, Y, Z, colors=\"none\", zorder=-5)\n", " stylize_contour(\n", " contour_set,\n", " label=f\"${labels[subsystem]}$\",\n", " linewidth=0,\n", " )\n", " contour_line = contour_set.collections[0]\n", " contour_line.set_alpha(patch_transparency)\n", " contour_line.set_facecolor(colors[subsystem])\n", " legend_elements = [\n", " Patch(\n", " alpha=patch_transparency,\n", " facecolor=color,\n", " label=f\"${labels[subsystem]}$\",\n", " )\n", " for subsystem, color in colors.items()\n", " ]\n", " ax.legend(\n", " bbox_to_anchor=(0.20, 0.25),\n", " framealpha=1,\n", " handles=legend_elements,\n", " loc=\"upper right\",\n", " prop={\"size\": 19},\n", " title=Rf\"$>{100*threshold:.0f}\\%$\",\n", " )\n", "\n", "\n", "def plot_polarimetry_field(polarimetry_arrays, ax, strides=12, cmap=cm.viridis_r):\n", " alpha_abs = jnp.sqrt(jnp.sum(polarimetry_arrays**2, axis=0))\n", " mesh = ax.quiver(\n", " X[::strides, ::strides],\n", " Y[::strides, ::strides],\n", " np.real(polarimetry_arrays[2][::strides, ::strides]),\n", " np.real(polarimetry_arrays[0][::strides, ::strides]),\n", " np.real(alpha_abs[::strides, ::strides]),\n", " cmap=cmap,\n", " )\n", " mesh.set_clim(vmin=0, vmax=+1)\n", " return mesh\n", "\n", "\n", "def overlay_inset(\n", " base_file: str,\n", " overlay_file: str,\n", " output_file: str | None = None,\n", " position: tuple[float, float] = (0.355, 0.08),\n", " scale: float = 1 / 240,\n", " show: bool = False,\n", ") -> SVG:\n", " if output_file is None:\n", " output_file = base_file\n", " if \"_static/images/\" not in base_file:\n", " base_file = f\"_static/images/{base_file}\"\n", " if \"_images/\" not in overlay_file:\n", " overlay_file = f\"_images/{overlay_file}\"\n", " if \"_static/images/\" not in output_file:\n", " output_file = f\"_static/images/{output_file}\"\n", " base_figure = sc.SVG(base_file)\n", " overlay_figure = sc.SVG(overlay_file)\n", " factor = 1.1\n", " w = factor * base_figure._width.value\n", " h = factor * base_figure._height.value\n", " new_x = position[0] * w\n", " new_y = position[1] * h\n", " figure = sc.Figure(\n", " w,\n", " h,\n", " sc.Panel(base_figure),\n", " sc.Panel(overlay_figure).scale(scale * w).move(new_x, new_y),\n", " ).scale(1.4)\n", " figure.save(output_file)\n", " plt.close(fig)\n", " svg = SVG(output_file)\n", " if show:\n", " display(svg)\n", " return svg\n", "\n", "\n", "%config InlineBackend.figure_formats = ['svg']\n", "X_LABEL_ALPHA = s1_label + R\",$\\quad \\alpha_z$\"\n", "Y_LABEL_ALPHA = s2_label + R\",$\\quad \\alpha_x$\"\n", "threshold = 0.7\n", "contour_arrays = create_dominant_region_contours(DECAY, data_sample, threshold)\n", "for ref in tqdm([1, 2, 3], leave=False):\n", " args = (ref, contour_arrays, threshold)\n", " plot_field(*args, add_title=True, watermark=False, show=True)\n", " plot_field(*args, add_title=True, watermark=True)\n", " plot_field(*args, add_title=False, watermark=False)\n", " plot_field(*args, add_title=False, watermark=True)\n", " plot_field(ref, add_title=True, watermark=False)\n", " plot_field(ref, add_title=True, watermark=True)\n", " plot_field(ref, add_title=False, watermark=False)\n", " plot_field(ref, add_title=False, watermark=True)\n", " del args, ref" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ ":::{only} latex\n", "{{ FIG_POLARIMETER_TOTAL }}\n", ":::" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Figure for the paper" }, "tags": [ "remove-cell" ] }, "outputs": [], "source": [ "_ = shutil.copyfile(\n", " \"_static/images/polarimetry-field-K.svg\",\n", " \"_static/images/total-polarimetry-field-no-inset.svg\",\n", ")\n", "_ = shutil.copyfile(\n", " \"_static/images/polarimetry-field-K-inset.svg\",\n", " \"_static/images/total-polarimetry-field.svg\",\n", ")\n", "_ = shutil.copyfile(\n", " \"_static/images/polarimetry-field-K-watermark.svg\",\n", " \"_static/images/total-polarimetry-field-watermark-no-inset.svg\",\n", ")\n", "_ = shutil.copyfile(\n", " \"_static/images/polarimetry-field-K-watermark-inset.svg\",\n", " \"_static/images/total-polarimetry-field-watermark.svg\",\n", ")\n", "SVG(\"_static/images/total-polarimetry-field-watermark.svg\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Aligned vector fields per chain" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input", "full-width", "scroll-input" ] }, "outputs": [], "source": [ "def to_regex(text: str) -> str:\n", " regex = text\n", " regex = regex.replace(\"(\", R\"\\(\")\n", " return regex.replace(\")\", R\"\\)\")\n", "\n", "\n", "def plot_field_per_resonance(reference_subsystem: int, watermark: bool) -> None:\n", " spectator = FINAL_STATE[reference_subsystem]\n", " subsystem_name = subsystem_identifiers[reference_subsystem - 1]\n", " subsystem_resonances = [\n", " chain.resonance\n", " for chain in DECAY.chains\n", " if chain.resonance.name.startswith(subsystem_name)\n", " ]\n", " ncols = 3\n", " nrows = math.ceil(len(subsystem_resonances) / ncols)\n", " fig, axes = plt.subplots(\n", " figsize={1: (13, 5), 2: (13, 9.0)}[nrows],\n", " gridspec_kw={\"width_ratios\": [1, 1, 1.06]},\n", " ncols=3,\n", " nrows=nrows,\n", " sharex=True,\n", " sharey=True,\n", " tight_layout=True,\n", " )\n", " fig.suptitle(\n", " f\"Polarimetry field, aligned to ${spectator}$\",\n", " y={1: 0.95, 2: 0.97}[nrows],\n", " )\n", " for i, (ax, resonance) in enumerate(zip(axes.flatten(), subsystem_resonances)):\n", " ax.set_box_aspect(1)\n", " non_zero_couplings = [to_regex(resonance.name)]\n", " polarimetry_field = [\n", " compute_sub_function(func, data_sample, non_zero_couplings)\n", " for func in polarimetry_funcs[reference_subsystem]\n", " ]\n", " polarimetry_field = jnp.array(polarimetry_field).real\n", " abs_alpha = jnp.sqrt(jnp.sum(polarimetry_field**2, axis=0))\n", " mesh = plot_polarimetry_field(\n", " polarimetry_field,\n", " ax=ax,\n", " strides=22,\n", " )\n", " mean = jnp.nanmean(abs_alpha)\n", " std = jnp.nanstd(abs_alpha)\n", "\n", " text = Rf\"$\\overline{{\\left|\\vec\\alpha\\right|}} = {mean:.3f}$\"\n", " if round(std, 3) != 0:\n", " text = text.replace(\"=\", R\"\\approx\")\n", " ax.text(\n", " x=1.80,\n", " y=4.44,\n", " s=text,\n", " fontsize=16,\n", " horizontalalignment=\"right\",\n", " )\n", " ax.set_title(f\"${resonance.latex}$\")\n", " if i // 3 == nrows - 1:\n", " ax.set_xlabel(X_LABEL_ALPHA)\n", " if i % 3 == 0:\n", " ax.set_ylabel(Y_LABEL_ALPHA)\n", " if i % 3 == 2:\n", " color_bar = fig.colorbar(mesh, ax=ax, fraction=0.0472, pad=0.01)\n", " color_bar.set_label(R\"$\\left|\\vec{\\alpha}\\right|$\")\n", " if watermark:\n", " add_watermark(ax, fontsize=14)\n", " output_file = f\"polarimetry-{subsystem_name}-chains\"\n", " if watermark:\n", " output_file += \"-watermark\"\n", " fig.savefig(f\"_static/images/{output_file}.svg\", bbox_inches=\"tight\")\n", " if watermark:\n", " plt.show()\n", " plt.close(fig)\n", " plt.ion()\n", "\n", "\n", "%config InlineBackend.figure_formats = ['svg']\n", "for reference_subsystem in tqdm([1, 2, 3], disable=NO_TQDM):\n", " plot_field_per_resonance(reference_subsystem, watermark=False)\n", " plot_field_per_resonance(reference_subsystem, watermark=True)\n", " del reference_subsystem" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ":::{only} latex\n", "{{ FIG_POLARIMETER_CHAIN }}\n", ":::" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "tags": [ "hide-input", "full-width", "scroll-input" ] }, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "fig, axes = plt.subplots(\n", " figsize=(13, 4.5),\n", " gridspec_kw={\"width_ratios\": [1, 1, 1.14]},\n", " ncols=3,\n", " sharey=True,\n", " tight_layout=True,\n", ")\n", "fig.suptitle(\"Polarimetry field per sub-system\", y=0.95)\n", "items = zip(axes, [1, 2, 3], subsystem_identifiers, subsystem_labels)\n", "for ax, reference_subsystem, subsystem_name, subsystem_label in items:\n", " ax.set_box_aspect(1)\n", " non_zero_couplings = [subsystem_name]\n", " polarimetry_field = [\n", " compute_sub_function(func, data_sample, non_zero_couplings)\n", " for func in polarimetry_funcs[reference_subsystem]\n", " ]\n", " polarimetry_field = jnp.array(polarimetry_field).real\n", " abs_alpha = jnp.sqrt(jnp.sum(polarimetry_field**2, axis=0))\n", " mesh = plot_polarimetry_field(\n", " polarimetry_field,\n", " ax=ax,\n", " strides=18,\n", " )\n", " mean = jnp.nanmean(abs_alpha)\n", " std = jnp.nanstd(abs_alpha)\n", "\n", " ax.text(\n", " x=1.8,\n", " y=4.4,\n", " s=Rf\"$\\overline{{\\left|\\vec\\alpha\\right|}} = {mean:.3f} \\pm {std:.3f}$\",\n", " fontsize=12,\n", " horizontalalignment=\"right\",\n", " )\n", " spectator = FINAL_STATE[reference_subsystem]\n", " ax.set_title(f\"${subsystem_label}$ (aligned to ${spectator}$)\")\n", " if ax is axes[-1]:\n", " color_bar = fig.colorbar(mesh, ax=ax, pad=0.01)\n", " color_bar.set_label(R\"$\\left|\\vec{\\alpha}\\right|$\")\n", "\n", "fig.savefig(\"_static/images/polarimetry-per-subsystem.svg\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ ":::{only} latex\n", "{{ FIG_POLARIMETER_SUBSYSTEM }}\n", ":::" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Export figures for the paper" }, "tags": [ "scroll-input", "hide-input", "remove-output" ] }, "outputs": [], "source": [ "def plot_figure2(watermark: bool) -> None:\n", " reference_subsystem = 1\n", " fig, ax = plt.subplots(\n", " figsize=(8, 6.8),\n", " tight_layout=True,\n", " )\n", " ax.set_box_aspect(1)\n", " ax.set_xlabel(X_LABEL_ALPHA)\n", " ax.set_ylabel(Y_LABEL_ALPHA)\n", " resonance = next(\n", " c.resonance for c in DECAY.chains if c.resonance.name == \"K(892)\"\n", " )\n", " non_zero_couplings = [to_regex(resonance.name)]\n", " polarimetry_field = [\n", " compute_sub_function(func, data_sample, non_zero_couplings)\n", " for func in polarimetry_funcs[reference_subsystem]\n", " ]\n", " polarimetry_field = jnp.array(polarimetry_field).real\n", " mesh = plot_polarimetry_field(polarimetry_field, ax=ax, strides=14)\n", " color_bar = fig.colorbar(mesh, ax=ax, pad=0.01)\n", " color_bar.set_label(R\"$\\left|\\vec{\\alpha}\\right|$\")\n", "\n", " output_filename = \"polarimetry-field-K892\"\n", " if watermark:\n", " output_filename += \"-watermark\"\n", " add_watermark(ax, fontsize=24)\n", " output_filename += \"-no-inset.svg\"\n", " fig.savefig(f\"_static/images/{output_filename}\", transparent=True)\n", " overlay_inset(\n", " output_filename,\n", " \"orientation-K.svg\",\n", " output_filename.replace(\"-no-inset\", \"\"),\n", " position=(0.34, 0.05),\n", " scale=4.4e-3,\n", " show=watermark,\n", " )\n", " plt.close(fig)\n", "\n", "\n", "def plot_figure3(watermark: bool, reference_subsystem: int) -> None:\n", " fig, ax = plt.subplots(figsize=(6, 5), tight_layout=True)\n", " ax.set_box_aspect(1)\n", " ax.set_xlabel(X_LABEL_ALPHA)\n", " ax.set_ylabel(Y_LABEL_ALPHA)\n", " resonances = [c.resonance for c in DECAY.chains if c.resonance.name == \"L(1520)\"]\n", " resonance = resonances[0]\n", " non_zero_couplings = [to_regex(resonance.name)]\n", " polarimetry_field = [\n", " compute_sub_function(func, data_sample, non_zero_couplings)\n", " for func in polarimetry_funcs[reference_subsystem]\n", " ]\n", " polarimetry_field = jnp.array(polarimetry_field).real\n", " mesh = plot_polarimetry_field(polarimetry_field, ax=ax, strides=22)\n", " color_bar = fig.colorbar(mesh, ax=ax, pad=0.01)\n", " color_bar.set_label(R\"$\\left|\\vec{\\alpha}\\right|$\")\n", "\n", " output_filename = \"polarimetry-field-L1520\"\n", " if reference_subsystem == 2:\n", " output_filename += \"-aligned\"\n", " else:\n", " output_filename += \"-unaligned\"\n", " if watermark:\n", " output_filename += \"-watermark\"\n", " add_watermark(ax, 0.033, 0.04, fontsize=18)\n", " output_filename += \"-no-inset.svg\"\n", " fig.savefig(f\"_static/images/{output_filename}\", transparent=True)\n", " subsystem_id = {1: \"K\", 2: \"L\", 3: \"D\"}[reference_subsystem]\n", " overlay_inset(\n", " output_filename,\n", " f\"orientation-{subsystem_id}.svg\",\n", " output_filename.replace(\"-no-inset\", \"\"),\n", " position=(0.34, 0.065),\n", " scale=4.1e-3,\n", " show=watermark,\n", " )\n", " plt.close(fig)\n", "\n", "\n", "%config InlineBackend.figure_formats = ['svg']\n", "plt.ioff()\n", "for use_watermark in [False, True]:\n", " plot_figure2(use_watermark)\n", " plot_figure3(use_watermark, reference_subsystem=1)\n", " plot_figure3(use_watermark, reference_subsystem=2)\n", " del use_watermark\n", "_ = plt.ion()" ] } ], "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.16" }, "myst": { "all_links_external": true } }, "nbformat": 4, "nbformat_minor": 4 }