{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Serialization\n",
    "\n",
    "```{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 json\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from IPython.display import Markdown\n",
    "from scipy.interpolate import RegularGridInterpolator, griddata\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from polarimetry import formulate_polarimetry\n",
    "from polarimetry.data import (\n",
    "    create_data_transformer,\n",
    "    generate_meshgrid_sample,\n",
    "    generate_phasespace_sample,\n",
    ")\n",
    "from polarimetry.io import (\n",
    "    export_polarimetry_field,\n",
    "    import_polarimetry_field,\n",
    "    mute_jax_warnings,\n",
    "    perform_cached_doit,\n",
    "    perform_cached_lambdify,\n",
    ")\n",
    "from polarimetry.lhcb import load_model_builder, load_model_parameters\n",
    "from polarimetry.lhcb.particle import load_particles\n",
    "from polarimetry.plot import use_mpl_latex_fonts\n",
    "\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "mute_jax_warnings()\n",
    "\n",
    "model_choice = \"Default amplitude model\"\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",
    "reference_subsystem = 1\n",
    "model = amplitude_builder.formulate(reference_subsystem)\n",
    "model.parameter_defaults.update(imported_parameter_values)\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": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Formulate expressions and lambdify"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "polarimetry_exprs = formulate_polarimetry(amplitude_builder, reference_subsystem)\n",
    "unfolded_exprs = [\n",
    "    perform_cached_doit(expr.doit().xreplace(model.amplitudes))\n",
    "    for expr in tqdm(\n",
    "        [model.full_expression, *polarimetry_exprs], disable=NO_TQDM, leave=False\n",
    "    )\n",
    "]\n",
    "actual_funcs = [\n",
    "    perform_cached_lambdify(expr.xreplace(model.parameter_defaults), backend=\"jax\")\n",
    "    for expr in tqdm(\n",
    "        unfolded_exprs, leave=False, desc=\"Lambdifying\", disable=NO_TQDM\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "resolution = 100\n",
    "transformer = create_data_transformer(model)\n",
    "grid_sample = generate_meshgrid_sample(model.decay, resolution)\n",
    "grid_sample = transformer(grid_sample)\n",
    "X = grid_sample[\"sigma1\"]\n",
    "Y = grid_sample[\"sigma2\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## File size checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "mystnb": {
     "code_prompt_show": "Export do different file formats"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "alpha_x_func = actual_funcs[1]\n",
    "alpha_x = alpha_x_func(grid_sample).real\n",
    "df = pd.DataFrame(alpha_x, index=X[0], columns=Y[:, 0])\n",
    "os.makedirs(\"export\", exist_ok=True)\n",
    "df.to_json(\"export/alpha-x-pandas.json\")\n",
    "df.to_json(\"export/alpha-x-pandas-json.zip\", compression={\"method\": \"zip\"})\n",
    "df.to_csv(\"export/alpha-x-pandas.csv\")\n",
    "\n",
    "df_dict = df.to_dict()\n",
    "filtered_df_dict = {\n",
    "    x: {y: v for y, v in row.items() if not math.isnan(v)}\n",
    "    for x, row in df_dict.items()\n",
    "}\n",
    "with open(\"export/alpha-x-python.json\", \"w\") as f:\n",
    "    json.dump(filtered_df_dict, f)\n",
    "\n",
    "json_dict = dict(\n",
    "    x=X[0].tolist(),\n",
    "    y=Y[:, 0].tolist(),\n",
    "    z=alpha_x.tolist(),\n",
    ")\n",
    "with open(\"export/alpha-x-arrays.json\", \"w\") as f:\n",
    "    json.dump(json_dict, f, separators=(\",\", \":\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "def render_kilobytes(path, markdown: bool = False) -> str:\n",
    "    byt = os.path.getsize(path)\n",
    "    kb = f\"{1e-3*byt:.0f}\"\n",
    "    return f\"| {{download}}`{path}` | **{kb} kB** |\\n\"\n",
    "\n",
    "\n",
    "src = f\"File sizes for **{len(X[0])}x{len(Y[:, 0])} grid**:\"\n",
    "src += \"\"\"\n",
    "| File type | Size |\n",
    "|:----------|-----:|\n",
    "\"\"\"\n",
    "src += render_kilobytes(\"export/alpha-x-arrays.json\")\n",
    "src += render_kilobytes(\"export/alpha-x-pandas.json\")\n",
    "src += render_kilobytes(\"export/alpha-x-python.json\")\n",
    "src += render_kilobytes(\"export/alpha-x-pandas-json.zip\")\n",
    "src += render_kilobytes(\"export/alpha-x-pandas.csv\")\n",
    "Markdown(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export polarimetry grids\n",
    "\n",
    "Decided to use the `alpha-x-arrays.json` format. It can be exported with {func}`.export_polarimetry_field`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(\"export\", exist_ok=True)\n",
    "filename = \"export/polarimetry-model-0.json\"\n",
    "export_polarimetry_field(\n",
    "    sigma1=X[0],\n",
    "    sigma2=Y[:, 0],\n",
    "    intensity=actual_funcs[0](grid_sample).real,\n",
    "    alpha_x=actual_funcs[1](grid_sample).real,\n",
    "    alpha_y=actual_funcs[2](grid_sample).real,\n",
    "    alpha_z=actual_funcs[3](grid_sample).real,\n",
    "    filename=filename,\n",
    "    metadata={\"model description\": model_choice},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "byt = os.path.getsize(filename)\n",
    "kb = f\"{1e-3*byt:.0f}\"\n",
    "src = f\"Polarimetry grid can be downloaded here: {{download}}`{filename}` ({kb} kB).\"\n",
    "Markdown(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import and interpolate\n",
    "\n",
    "The arrays in the {ref}`exported JSON files <exported-distributions>` can be used to create a {class}`~scipy.interpolate.RegularGridInterpolator` for the intensity and for each components of $\\vec\\alpha$.\n",
    "\n",
    ":::{margin}\n",
    "\n",
    "{func}`.import_polarimetry_field` returns JAX arrays, which are read-only. {class}`~scipy.interpolate.RegularGridInterpolator` requires modifiable arrays, so we convert them to NumPy.\n",
    "\n",
    "Also note that the `values` array needs to be **transposed**!\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "field_definition = import_polarimetry_field(\"export/polarimetry-model-0.json\")\n",
    "imported_sigma1 = field_definition[\"m^2_Kpi\"]\n",
    "imported_sigma2 = field_definition[\"m^2_pK\"]\n",
    "imported_arrays = (\n",
    "    field_definition[\"intensity\"],\n",
    "    field_definition[\"alpha_x\"],\n",
    "    field_definition[\"alpha_y\"],\n",
    "    field_definition[\"alpha_z\"],\n",
    ")\n",
    "interpolated_funcs = [\n",
    "    RegularGridInterpolator(\n",
    "        points=(\n",
    "            np.array(imported_sigma1),\n",
    "            np.array(imported_sigma2),\n",
    "        ),\n",
    "        values=np.array(z).transpose(),\n",
    "        method=\"linear\",\n",
    "    )\n",
    "    for z in imported_arrays\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a function that can compute an interpolated value of each of these observables for a random point on the Dalitz plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpolated_funcs[1]([0.8, 3.6])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As opposed to SciPy's deprecated {obj}`~scipy.interpolate.interp2d`, {class}`~scipy.interpolate.RegularGridInterpolator` is already in vectorized form, so there is no need to {obj}`~numpy.vectorize` it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_points = 100_000\n",
    "mini_sample = generate_phasespace_sample(model.decay, n_points, seed=0)\n",
    "mini_sample = transformer(mini_sample)\n",
    "x = mini_sample[\"sigma1\"]\n",
    "y = mini_sample[\"sigma2\"]\n",
    "z_interpolated = [func((x, y)) for func in tqdm(interpolated_funcs, disable=NO_TQDM)]\n",
    "z_interpolated[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "use_mpl_latex_fonts()\n",
    "plt.rc(\"font\", size=18)\n",
    "fig, axes = plt.subplots(\n",
    "    dpi=200,\n",
    "    figsize=(15, 11.5),\n",
    "    gridspec_kw={\"width_ratios\": [1, 1, 1, 1.2]},\n",
    "    ncols=4,\n",
    "    nrows=3,\n",
    "    sharex=True,\n",
    "    sharey=True,\n",
    ")\n",
    "plt.subplots_adjust(hspace=0.1, wspace=0.03)\n",
    "fig.suptitle(\"Comparison interpolated and actual values\", y=0.94)\n",
    "\n",
    "points = np.transpose([x, y])\n",
    "for i in tqdm(range(4), disable=NO_TQDM, leave=False):\n",
    "    if i == 0:\n",
    "        title = \"$I$\"\n",
    "        cmap = plt.cm.viridis\n",
    "        clim = None\n",
    "    else:\n",
    "        title = Rf\"$\\alpha_{'xyz'[i-1]}$\"\n",
    "        cmap = plt.cm.coolwarm\n",
    "        clim = (-1, +1)\n",
    "    axes[0, i].set_title(title, y=1.03)\n",
    "\n",
    "    z_actual = actual_funcs[i](mini_sample)\n",
    "    z_diff = 100 * ((z_interpolated[i] - z_actual) / z_actual).real\n",
    "    Z_interpolated = griddata(points, z_interpolated[i], (X, Y))\n",
    "    Z_diff = griddata(points, z_diff, (X, Y))\n",
    "\n",
    "    mesh = (\n",
    "        axes[0, i].pcolormesh(X, Y, actual_funcs[i](grid_sample).real, cmap=cmap),\n",
    "        axes[1, i].pcolormesh(X, Y, Z_interpolated, cmap=cmap),\n",
    "        axes[2, i].pcolormesh(X, Y, Z_diff, clim=(-100, +100), cmap=plt.cm.coolwarm),\n",
    "    )\n",
    "    if i != 0:\n",
    "        mesh[0].set_clim(-1, +1)\n",
    "        mesh[1].set_clim(-1, +1)\n",
    "    if i == 3:\n",
    "        c_bar = [fig.colorbar(mesh[j], ax=axes[j, i], pad=0.015) for j in range(3)]\n",
    "        c_bar[0].ax.set_ylabel(\"Original distribution\", labelpad=3)\n",
    "        c_bar[1].ax.set_ylabel(\"Interpolated distribution\", labelpad=3)\n",
    "        c_bar[2].ax.set_ylabel(\"Difference\", labelpad=-20)\n",
    "        for c in c_bar[:-1]:\n",
    "            c.ax.set_yticks([-1, 0, +1])\n",
    "            c.ax.set_yticklabels([\"$-1$\", \"$0$\", \"$+1$\"])\n",
    "        c_bar[-1].ax.set_yticks([-100, 0, +100])\n",
    "        c_bar[-1].ax.set_yticklabels([R\"$-100\\%$\", R\"$0\\%$\", R\"$+100\\%$\"])\n",
    "        axes[0, i].text(\n",
    "            x=0.96,\n",
    "            y=0.83,\n",
    "            s=f\"grid size:\\n{resolution}x{resolution}\",\n",
    "            fontsize=16,\n",
    "            horizontalalignment=\"right\",\n",
    "            transform=axes[0, i].transAxes,\n",
    "        )\n",
    "        axes[1, i].text(\n",
    "            x=0.96,\n",
    "            y=0.83,\n",
    "            s=f\"phsp size:\\n{n_points:,}\",\n",
    "            fontsize=16,\n",
    "            horizontalalignment=\"right\",\n",
    "            transform=axes[1, i].transAxes,\n",
    "        )\n",
    "plt.show()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "The interpolated values over this phase space sample have been visualized by interpolating again over a {obj}`~numpy.meshgrid` with {obj}`scipy.interpolate.griddata`.\n",
    ":::\n",
    "\n",
    ":::{tip}\n",
    "{doc}`/zz.polarization-fit` shows how this interpolation method can be used to determine the polarization $\\vec{P}$ from a given intensity distribution.\n",
    ":::"
   ]
  }
 ],
 "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
}