{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nominal amplitude model"
   ]
  },
  {
   "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": [
    "import logging\n",
    "import os\n",
    "from textwrap import dedent\n",
    "\n",
    "import sympy as sp\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "from polarimetry.amplitude import simplify_latex_rendering\n",
    "from polarimetry.io import (\n",
    "    as_latex,\n",
    "    as_markdown_table,\n",
    "    display_latex,\n",
    "    perform_cached_doit,\n",
    ")\n",
    "from polarimetry.lhcb import load_model_builder, load_model_parameters\n",
    "from polarimetry.lhcb.particle import K, Λc, Σ, load_particles, p, π\n",
    "\n",
    "simplify_latex_rendering()\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resonances and LS-scheme"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Particle definitions for $\\Lambda_c^+$ and $p, \\pi^+, K^-$ in the sequential order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "decay_particles = [Λc, p, π, K, Σ]\n",
    "Markdown(as_markdown_table(decay_particles))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Particle definitions as defined in {download}`particle-definitions.yaml<../data/particle-definitions.yaml>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "particles = load_particles(\"../data/particle-definitions.yaml\")\n",
    "resonances = [p for p in particles.values() if p not in set(decay_particles)]\n",
    "src = as_markdown_table(resonances)\n",
    "Markdown(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "{doc}`appendix/ls-model`\n",
    ":::\n",
    "\n",
    "Most models work take the **minimal $L$-value** in each $LS$-coupling (only model 17 works in the full $LS$-basis. The generated $LS$-couplings look as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def sort_chains(chains):\n",
    "    return sorted(\n",
    "        chains,\n",
    "        key=lambda c: (c.resonance.name[0], c.resonance.mass),\n",
    "    )\n",
    "\n",
    "\n",
    "def render_ls_table(with_jp: bool) -> None:\n",
    "    all_ls_chains = load_model_builder(\n",
    "        model_file=\"../data/model-definitions.yaml\",\n",
    "        particle_definitions=particles,\n",
    "        model_id=17,\n",
    "    ).decay.chains\n",
    "    min_ls_chains = load_model_builder(\n",
    "        model_file=\"../data/model-definitions.yaml\",\n",
    "        particle_definitions=particles,\n",
    "        model_id=0,\n",
    "    ).decay.chains\n",
    "\n",
    "    all_ls_chains = sort_chains(all_ls_chains)\n",
    "    min_ls_chains = sort_chains(min_ls_chains)\n",
    "\n",
    "    n_all_ls = len(all_ls_chains)\n",
    "    n_min_ls = len(min_ls_chains)\n",
    "\n",
    "    src = Rf\"\"\"\n",
    "    | Only minimum $LS$ ({n_min_ls}) | All $LS$-couplings ({n_all_ls}) |\n",
    "    |:------------------------------:|:-------------------------------:|\n",
    "    \"\"\"\n",
    "    src = dedent(src).strip() + \"\\n\"\n",
    "    min_ls_chain_iter = iter(min_ls_chains)\n",
    "    min_ls_chain = None\n",
    "    for all_ls_chain in all_ls_chains:\n",
    "        min_ls_chain_latex = \"\"\n",
    "        if (\n",
    "            min_ls_chain is None\n",
    "            or min_ls_chain.resonance.name != all_ls_chain.resonance.name\n",
    "        ):\n",
    "            try:\n",
    "                min_ls_chain = next(min_ls_chain_iter)\n",
    "                min_ls_chain_latex = f\"${as_latex(min_ls_chain, with_jp=with_jp)}$\"\n",
    "            except StopIteration:\n",
    "                pass\n",
    "        all_ls_chain_latex = f\"${as_latex(all_ls_chain, with_jp=with_jp)}$\"\n",
    "        src += f\"| {min_ls_chain_latex} | {all_ls_chain_latex} |\\n\"\n",
    "    display(Markdown(src))\n",
    "\n",
    "\n",
    "render_ls_table(with_jp=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or with $J^P$-values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "render_ls_table(with_jp=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Amplitude"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spin-alignment amplitude"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The full intensity of the amplitude model is obtained by summing the following aligned amplitude over all helicity values $\\lambda_i$ in the initial state $0$ and final states $1, 2, 3$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_choice = 0\n",
    "amplitude_builder = load_model_builder(\n",
    "    model_file=\"../data/model-definitions.yaml\",\n",
    "    particle_definitions=particles,\n",
    "    model_id=model_choice,\n",
    ")\n",
    "model = amplitude_builder.formulate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "def simplify_notation(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",
    "display(simplify_notation(model.intensity.args[0].args[0].args[0].cleanup()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Note that we simplified notation here: the amplitude indices for the spinless states are not rendered and their corresponding Wigner-$d$ alignment functions are simply $1$.\n",
    "\n",
    "The relevant $\\zeta^i_{j(k)}$ angles are {doc}`defined as</appendix/angles>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "display_latex({k: v for k, v in model.variables.items() if \"zeta\" in str(k)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sub-system amplitudes"
   ]
  },
  {
   "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 model.amplitudes.items()})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "The $\\theta_{ij}$ angles are {doc}`defined as</appendix/angles>`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "display_latex({k: v for k, v in model.variables.items() if \"theta\" in str(k)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Definitions for the $\\phi_{ij}$ angles can be found under {doc}`/appendix/angles`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter definitions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameter values are provided in {download}`model-definitions.yaml<../data/model-definitions.yaml>`, but the **keys** of the helicity couplings have to remapped to the helicity **symbols** that are used in this amplitude model. The function {func}`.parameter_key_to_symbol` implements this remapping, following the [supplementary material](https://cds.cern.ch/record/2824328/files) of {cite}`LHCb-PAPER-2022-002`. It is asserted below that:\n",
    "1. the keys are mapped to symbols that exist in the nominal amplitude model\n",
    "2. all parameter symbols in the nominal amplitude model have a value assigned to them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Collect all symbols in the amplitude moddel"
    },
    "tags": [
     "hide-cell",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "imported_parameter_values = load_model_parameters(\n",
    "    \"../data/model-definitions.yaml\",\n",
    "    amplitude_builder.decay,\n",
    "    model_choice,\n",
    "    particle_definitions=particles,\n",
    ")\n",
    "unfolded_intensity_expr = perform_cached_doit(model.full_expression)\n",
    "model_symbols = unfolded_intensity_expr.free_symbols\n",
    "\n",
    "non_existent = set(imported_parameter_values) - set(model_symbols)\n",
    "error_message = \"Imported symbols that don't exist in model:\\n  \"\n",
    "error_message += \"\\n  \".join(map(str, sorted(non_existent, key=str)))\n",
    "assert non_existent == set(), error_message\n",
    "\n",
    "undefined = (\n",
    "    set(model_symbols)\n",
    "    - set(imported_parameter_values)\n",
    "    - set(model.parameter_defaults)\n",
    "    - set(model.variables)\n",
    "    - set(sp.symbols(\"sigma1:4\", nonnegative=True))\n",
    ")\n",
    "undefined = {\n",
    "    s\n",
    "    for s in undefined\n",
    "    if not str(s).endswith(\"{decay}\")\n",
    "    if not str(s).endswith(\"production}\")\n",
    "}\n",
    "error_message = \"Symbols in model that don't have a definition:\\n  \"\n",
    "error_message += \"\\n  \".join(map(str, sorted(undefined, key=str)))\n",
    "assert undefined == set(), error_message\n",
    "model.parameter_defaults.update(imported_parameter_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Helicity coupling values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Production couplings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "production_couplings = {\n",
    "    key: value\n",
    "    for key, value in model.parameter_defaults.items()\n",
    "    if isinstance(key, sp.Indexed)\n",
    "    if \"production\" in str(key.base)\n",
    "    if str(value) != \"1\"\n",
    "}\n",
    "display_latex(production_couplings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Decay couplings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "decay_couplings = {\n",
    "    key: value\n",
    "    for key, value in model.parameter_defaults.items()\n",
    "    if isinstance(key, sp.Indexed)\n",
    "    if \"decay\" in str(key.base)\n",
    "}\n",
    "display_latex(decay_couplings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Non-coupling parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "couplings = set(production_couplings) | set(decay_couplings)\n",
    "non_coupling_parameters = {\n",
    "    symbol: model.parameter_defaults[symbol]\n",
    "    for symbol in sorted(model.parameter_defaults, key=str)\n",
    "    if not isinstance(symbol, sp.Indexed)\n",
    "}\n",
    "display_latex(non_coupling_parameters)"
   ]
  }
 ],
 "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
}