{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amplitude model with LS-couplings"
   ]
  },
  {
   "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 jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import sympy as sp\n",
    "from IPython.display import Latex\n",
    "from sympy.core.symbol import Str\n",
    "from tensorwaves.interface import Function\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from polarimetry.amplitude import (\n",
    "    AmplitudeModel,\n",
    "    get_indexed_base,\n",
    "    simplify_latex_rendering,\n",
    ")\n",
    "from polarimetry.data import (\n",
    "    create_data_transformer,\n",
    "    generate_meshgrid_sample,\n",
    "    generate_phasespace_sample,\n",
    ")\n",
    "from polarimetry.decay import Particle\n",
    "from polarimetry.function import integrate_intensity, sub_intensity\n",
    "from polarimetry.io import (\n",
    "    as_latex,\n",
    "    display_latex,\n",
    "    mute_jax_warnings,\n",
    "    perform_cached_doit,\n",
    "    perform_cached_lambdify,\n",
    ")\n",
    "from polarimetry.lhcb import (\n",
    "    get_conversion_factor_ls,\n",
    "    load_model_builder,\n",
    "    load_model_parameters,\n",
    ")\n",
    "from polarimetry.lhcb.particle import load_particles\n",
    "from polarimetry.plot import use_mpl_latex_fonts\n",
    "\n",
    "mute_jax_warnings()\n",
    "simplify_latex_rendering()\n",
    "MODEL_FILE = \"../../data/model-definitions.yaml\"\n",
    "PARTICLES = load_particles(\"../../data/particle-definitions.yaml\")\n",
    "\n",
    "NO_TQDM = \"EXECUTE_NB\" in os.environ\n",
    "if NO_TQDM:\n",
    "    logging.getLogger().setLevel(logging.ERROR)\n",
    "    logging.getLogger(\"polarimetry.io\").setLevel(logging.ERROR)\n",
    "    logging.getLogger(\"tensorwaves.data\").setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Model inspection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "def formulate_model(title: str) -> AmplitudeModel:\n",
    "    builder = load_model_builder(MODEL_FILE, PARTICLES, title)\n",
    "    imported_parameters = load_model_parameters(\n",
    "        MODEL_FILE, builder.decay, title, PARTICLES\n",
    "    )\n",
    "    model = builder.formulate()\n",
    "    model.parameter_defaults.update(imported_parameters)\n",
    "    return model\n",
    "\n",
    "\n",
    "def simplify_notation(expr: sp.Expr) -> sp.Expr:\n",
    "    def substitute_node(node):\n",
    "        if isinstance(node, sp.Indexed):\n",
    "            if node.indices[2:] == (0, 0):\n",
    "                return sp.Indexed(node.base, *node.indices[:2])\n",
    "        return node\n",
    "\n",
    "    for node in sp.preorder_traversal(expr):\n",
    "        new_node = substitute_node(node)\n",
    "        expr = expr.xreplace({node: new_node})\n",
    "    return expr\n",
    "\n",
    "\n",
    "LS_MODEL = formulate_model(\"Alternative amplitude model obtained using LS couplings\")\n",
    "simplify_notation(LS_MODEL.intensity.args[0].args[0].args[0].cleanup())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "display_latex({simplify_notation(k): v for k, v in LS_MODEL.amplitudes.items()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Show conversion factors for LHCb-PAPER-2022-002"
    },
    "tags": [
     "hide-input",
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "H_prod = get_indexed_base(\"production\", min_ls=False)\n",
    "\n",
    "latex = R\"\"\"\n",
    "\\begin{array}{c|ccc|c}\n",
    "  \\textbf{Decay} & \\textbf{coupling} & & & \\textbf{factor} \\\\\n",
    "  \\hline\n",
    "\"\"\"\n",
    "for chain in LS_MODEL.decay.chains:\n",
    "    R = Str(chain.resonance.name)\n",
    "    L = chain.incoming_ls.L\n",
    "    S = chain.incoming_ls.S\n",
    "    symbol = H_prod[R, L, S]\n",
    "    value = sp.sympify(LS_MODEL.parameter_defaults[symbol])\n",
    "    factor = get_conversion_factor_ls(chain.resonance, L, S)\n",
    "    coupling_value = f\"{as_latex(symbol)} &=& {as_latex(value.n(3))}\"\n",
    "    latex += Rf\"  {as_latex(chain)} & {coupling_value} & {factor:+d} \\\\\" \"\\n\"\n",
    "latex += R\"\\end{array}\"\n",
    "Latex(f\"{latex}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "source": [
    "It is asserted that these amplitude expressions to not evaluate to $0$ once the Clebsch-Gordan coefficients are evaluated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def assert_non_zero_amplitudes(model: AmplitudeModel) -> None:\n",
    "    for amplitude in tqdm(model.amplitudes.values(), disable=NO_TQDM):\n",
    "        assert amplitude.doit() != 0\n",
    "\n",
    "\n",
    "assert_non_zero_amplitudes(LS_MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "See {ref}`amplitude-model:Resonances and LS-scheme` for the allowed $LS$-values.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Convert expressions to numerical functions"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def lambdify(model: AmplitudeModel) -> sp.Expr:\n",
    "    intensity_expr = unfold_intensity(model)\n",
    "    pars = model.parameter_defaults\n",
    "    free_parameters = {s: v for s, v in pars.items() if \"production\" in str(s)}\n",
    "    fixed_parameters = {s: v for s, v in pars.items() if s not in free_parameters}\n",
    "    subs_intensity_expr = intensity_expr.xreplace(fixed_parameters)\n",
    "    return perform_cached_lambdify(subs_intensity_expr, free_parameters)\n",
    "\n",
    "\n",
    "def unfold_intensity(model: AmplitudeModel) -> sp.Expr:\n",
    "    unfolded_intensity = perform_cached_doit(model.intensity)\n",
    "    return perform_cached_doit(unfolded_intensity.xreplace(model.amplitudes))\n",
    "\n",
    "\n",
    "NOMINAL_MODEL = formulate_model(\"Default amplitude model\")\n",
    "NOMINAL_INTENSITY_FUNC = lambdify(NOMINAL_MODEL)\n",
    "LS_INTENSITY_FUNC = lambdify(LS_MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Create phase space grid"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "GRID = generate_meshgrid_sample(NOMINAL_MODEL.decay, resolution=300)\n",
    "TRANSFORMER = create_data_transformer(NOMINAL_MODEL)\n",
    "GRID.update(TRANSFORMER(GRID))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "def compare_2d_distributions() -> None:\n",
    "    NOMINAL_INTENSITIES = compute_normalized_intensity(NOMINAL_INTENSITY_FUNC)\n",
    "    LS_INTENSITIES = compute_normalized_intensity(LS_INTENSITY_FUNC)\n",
    "    max_intensity = max(\n",
    "        jnp.nanmax(NOMINAL_INTENSITIES),\n",
    "        jnp.nanmax(LS_INTENSITIES),\n",
    "    )\n",
    "    use_mpl_latex_fonts()\n",
    "    fig, axes = plt.subplots(\n",
    "        dpi=200,\n",
    "        figsize=(12, 5),\n",
    "        ncols=2,\n",
    "    )\n",
    "    for ax in axes:\n",
    "        ax.set_box_aspect(1)\n",
    "    ax1, ax2 = axes\n",
    "    ax1.set_title(\"Nominal model\")\n",
    "    ax2.set_title(\"LS-model\")\n",
    "    ax1.pcolormesh(\n",
    "        GRID[\"sigma1\"],\n",
    "        GRID[\"sigma2\"],\n",
    "        NOMINAL_INTENSITIES,\n",
    "        vmax=max_intensity,\n",
    "    )\n",
    "    ax2.pcolormesh(\n",
    "        GRID[\"sigma1\"],\n",
    "        GRID[\"sigma2\"],\n",
    "        LS_INTENSITIES,\n",
    "        vmax=max_intensity,\n",
    "    )\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def compute_normalized_intensity(func: Function) -> jnp.ndarray:\n",
    "    intensities = func(GRID)\n",
    "    integral = jnp.nansum(intensities)\n",
    "    return intensities / integral\n",
    "\n",
    "\n",
    "compare_2d_distributions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decay rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "def to_regex(text: str) -> str:\n",
    "    text = text.replace(\"(\", r\"\\(\")\n",
    "    text = text.replace(\")\", r\"\\)\")\n",
    "    return text\n",
    "\n",
    "\n",
    "def compute_decay_rates() -> dict[Particle, tuple[float, float]]:\n",
    "    decay_rates = {}\n",
    "    nominal_I_tot = integrate_intensity(NOMINAL_INTENSITY_FUNC(PHSP))\n",
    "    LS_I_tot = integrate_intensity(LS_INTENSITY_FUNC(PHSP))\n",
    "    for chain in tqdm(NOMINAL_MODEL.decay.chains, disable=NO_TQDM):\n",
    "        filter_ = [to_regex(chain.resonance.name)]\n",
    "        LS_I_sub = sub_intensity(LS_INTENSITY_FUNC, PHSP, filter_)\n",
    "        nominal_I_sub = sub_intensity(NOMINAL_INTENSITY_FUNC, PHSP, filter_)\n",
    "        decay_rates[chain.resonance] = (\n",
    "            float(nominal_I_sub / nominal_I_tot),\n",
    "            float(LS_I_sub / LS_I_tot),\n",
    "        )\n",
    "    return decay_rates\n",
    "\n",
    "\n",
    "PHSP = generate_phasespace_sample(NOMINAL_MODEL.decay, n_events=100_000, seed=0)\n",
    "PHSP = TRANSFORMER(PHSP)\n",
    "DECAY_RATES = compute_decay_rates()\n",
    "src = R\"\"\"\n",
    "\\begin{array}{l|rr|r}\n",
    "  \\textbf{Resonance} & \\textbf{Nominal} & \\textbf{LS-model} & \\textbf{Difference}\\\\\n",
    "  \\hline\n",
    "\"\"\"\n",
    "for res, (nominal_rate, ls_rate) in DECAY_RATES.items():\n",
    "    nominal_rate *= 100\n",
    "    ls_rate *= 100\n",
    "    src += (\n",
    "        Rf\"  {res.latex} & {nominal_rate:.2f} & {ls_rate:.2f} &\"\n",
    "        rf\" {ls_rate - nominal_rate:+.2f} \\\\\"\n",
    "        \"\\n\"\n",
    "    )\n",
    "    del res, nominal_rate, ls_rate\n",
    "src += R\"\\end{array}\"\n",
    "Latex(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip}\n",
    "Compare with the values with uncertainties as reported in {ref}`uncertainties:Decay rates`.\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
}