{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Alignment consistency"
   ]
  },
  {
   "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",
    "\n",
    "import cairosvg\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import sympy as sp\n",
    "from numpy.testing import assert_almost_equal\n",
    "from tensorwaves.data import SympyDataTransformer\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from polarimetry.amplitude import AmplitudeModel, simplify_latex_rendering\n",
    "from polarimetry.data import create_data_transformer, generate_meshgrid_sample\n",
    "from polarimetry.io import (\n",
    "    display_latex,\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 add_watermark, use_mpl_latex_fonts\n",
    "\n",
    "mute_jax_warnings()\n",
    "simplify_latex_rendering()\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": [
    "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",
    "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\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "display_latex(m.intensity.cleanup() for m in models.values())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See {doc}`/appendix/angles` for the definition of each $\\zeta^i_{j(k)}$.\n",
    "\n",
    "Note that a change in reference sub-system requires the production couplings for certain sub-systems to flip sign:\n",
    "- **Sub-system 2** as reference system: flip signs of $\\mathcal{H}^\\mathrm{production}_{K^{**}}$ and  $\\mathcal{H}^\\mathrm{production}_{L^{**}}$\n",
    "- **Sub-system 3** as reference system: flip signs of $\\mathcal{H}^\\mathrm{production}_{K^{**}}$ and  $\\mathcal{H}^\\mathrm{production}_{D^{**}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "coupling = [\n",
    "    symbol\n",
    "    for symbol in models[1].parameter_defaults\n",
    "    if str(symbol) == R\"\\mathcal{H}^\\mathrm{production}[K(892), -1, -1/2]\"\n",
    "][0]\n",
    "assert (\n",
    "    models[2].parameter_defaults[coupling] == -models[1].parameter_defaults[coupling]\n",
    ")\n",
    "assert (\n",
    "    models[3].parameter_defaults[coupling] == -models[1].parameter_defaults[coupling]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "unfolded_intensity_exprs = {\n",
    "    reference_subsystem: perform_cached_doit(model.full_expression)\n",
    "    for reference_subsystem, model in tqdm(models.items(), disable=NO_TQDM)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "def assert_all_symbols_defined(expr: sp.Expr, model: AmplitudeModel) -> 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",
    "for reference_subsystem in unfolded_intensity_exprs:\n",
    "    assert_all_symbols_defined(\n",
    "        expr=unfolded_intensity_exprs[reference_subsystem],\n",
    "        model=models[reference_subsystem],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "subs_intensity_exprs = {\n",
    "    reference_subsystem: expr.xreplace(\n",
    "        models[reference_subsystem].parameter_defaults\n",
    "    )\n",
    "    for reference_subsystem, expr in unfolded_intensity_exprs.items()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "intensity_funcs = {\n",
    "    reference_subsystem: perform_cached_lambdify(expr, backend=\"jax\")\n",
    "    for reference_subsystem, expr in tqdm(\n",
    "        subs_intensity_exprs.items(), disable=NO_TQDM\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "transformer = {}\n",
    "for reference_subsystem in tqdm([1, 2, 3], disable=NO_TQDM):\n",
    "    model = models[reference_subsystem]\n",
    "    transformer.update(create_data_transformer(model).functions)\n",
    "transformer = SympyDataTransformer(transformer)\n",
    "grid_sample = generate_meshgrid_sample(model.decay, resolution=400)\n",
    "grid_sample = transformer(grid_sample)\n",
    "intensity_grids = {i: func(grid_sample) for i, func in intensity_funcs.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "{i: jnp.nansum(grid) for i, grid in intensity_grids.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert_almost_equal(\n",
    "    jnp.nansum(intensity_grids[2] - intensity_grids[1]), 0, decimal=6\n",
    ")\n",
    "assert_almost_equal(\n",
    "    jnp.nansum(intensity_grids[2] - intensity_grids[1]), 0, decimal=6\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "def convert_svg_to_png(input_file: str, dpi: int) -> None:\n",
    "    output_file = input_file.replace(\".svg\", \".png\").replace(\".SVG\", \".png\")\n",
    "    with open(input_file) as f:\n",
    "        src = f.read()\n",
    "    cairosvg.svg2png(bytestring=src, write_to=output_file, dpi=dpi)\n",
    "\n",
    "\n",
    "def overlay_inset(\n",
    "    png_file: str, ax, position: tuple[float, float], width: float\n",
    ") -> None:\n",
    "    image = plt.imread(png_file)\n",
    "    res_x, res_y, _ = image.shape\n",
    "    x_min, x_max = ax.get_xlim()\n",
    "    y_min, y_max = ax.get_ylim()\n",
    "    aspect_ratio = res_x / res_y\n",
    "    aspect_ratio /= (x_max - x_min) / (y_max - y_min)\n",
    "    extent = [\n",
    "        position[0],\n",
    "        position[0] + width,\n",
    "        position[1],\n",
    "        position[1] + width / aspect_ratio,\n",
    "    ]\n",
    "    ax.imshow(image, aspect=\"auto\", extent=extent, zorder=2)\n",
    "    ax.set_xlim(x_min, x_max)\n",
    "    ax.set_ylim(y_min, y_max)\n",
    "\n",
    "\n",
    "for subsystem in [\"K\", \"D\", \"L\"]:\n",
    "    convert_svg_to_png(f\"../_images/orientation-{subsystem}.svg\", dpi=200)\n",
    "    del subsystem\n",
    "\n",
    "\n",
    "def plot_comparison(use_watermark: bool) -> None:\n",
    "    plt.ioff()\n",
    "    x_label = R\"$m^2\\left(K^-\\pi^+\\right)$ [GeV$^2$]\"\n",
    "    y_label = R\"$m^2\\left(pK^-\\right)$ [GeV$^2$]\"\n",
    "\n",
    "    plt.rcdefaults()\n",
    "    use_mpl_latex_fonts()\n",
    "    fig, axes = plt.subplots(\n",
    "        dpi=200,\n",
    "        figsize=(20, 6),\n",
    "        ncols=3,\n",
    "        sharey=True,\n",
    "        gridspec_kw={\"width_ratios\": [1, 1, 1.21]},\n",
    "    )\n",
    "    normalized_intensities = {\n",
    "        i: I / jnp.nansum(I) for i, I in intensity_grids.items()\n",
    "    }\n",
    "    global_max = max(map(jnp.nanmax, normalized_intensities.values()))\n",
    "    axes[0].set_ylabel(y_label)\n",
    "    subsystem_names = [\"K\", \"L\", \"D\"]\n",
    "    for i, (ax, name) in enumerate(zip(axes, subsystem_names), 1):\n",
    "        ax.set_xlabel(x_label)\n",
    "        ax.set_box_aspect(1)\n",
    "        mesh = ax.pcolormesh(\n",
    "            grid_sample[\"sigma1\"],\n",
    "            grid_sample[\"sigma2\"],\n",
    "            normalized_intensities[i],\n",
    "        )\n",
    "        mesh.set_clim(vmax=global_max)\n",
    "        if ax is axes[-1]:\n",
    "            c_bar = fig.colorbar(mesh, ax=ax)\n",
    "            c_bar.ax.set_ylabel(\"Normalized intensity\")\n",
    "        if use_watermark:\n",
    "            add_watermark(ax)\n",
    "        overlay_inset(\n",
    "            f\"../_images/orientation-{name}.png\",\n",
    "            ax=ax,\n",
    "            position=(1.05, 3.85),\n",
    "            width=0.75,\n",
    "        )\n",
    "    fig.subplots_adjust(wspace=0)\n",
    "    output_filename = \"intensity-alignment-consistency\"\n",
    "    if use_watermark:\n",
    "        output_filename += \"-watermark\"\n",
    "    fig.savefig(f\"../_static/images/{output_filename}.png\", bbox_inches=\"tight\")\n",
    "    if use_watermark:\n",
    "        plt.show()\n",
    "    plt.close(fig)\n",
    "    plt.ion()\n",
    "\n",
    "\n",
    "plt.rc(\"font\", size=18)\n",
    "plot_comparison(use_watermark=False)\n",
    "plot_comparison(use_watermark=True)"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}