{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell",
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "# WARNING: advised to install a specific version, e.g. qrules==0.1.2\n",
    "%pip install -q qrules[doc,viz] IPython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conservation rules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import attrs\n",
    "import graphviz\n",
    "from IPython.display import display\n",
    "\n",
    "import qrules\n",
    "from qrules.conservation_rules import (\n",
    "    SpinEdgeInput,\n",
    "    SpinNodeInput,\n",
    "    parity_conservation,\n",
    "    spin_conservation,\n",
    "    spin_magnitude_conservation,\n",
    ")\n",
    "from qrules.quantum_numbers import Parity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "QRules generates {class}`.MutableTransition`s, populates them with quantum numbers (edge properties representing states and nodes properties representing interactions), then checks whether the generated {class}`.MutableTransition`s comply with the rules formulated in the {mod}`.conservation_rules` module.\n",
    "\n",
    "The {mod}`.conservation_rules` module can also be used separately. In this notebook, we will illustrate this by checking spin and parity conservation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parity conservation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See {func}`.parity_conservation`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parity_conservation(\n",
    "    ingoing_edge_qns=[Parity(-1)],\n",
    "    outgoing_edge_qns=[Parity(+1), Parity(+1)],\n",
    "    l_magnitude=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Spin conservation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "{func}`.spin_conservation`, [`tests/unit/conservation_rules/test_spin.py`](https://github.com/ComPWA/qrules/blob/ffa91f5308f59bd729b25d1584827ac61a56d2de/tests/unit/conservation_rules/test_spin.py), {pdg-review}`Quark Model`, and [these lecture notes](http://www.curtismeyer.com/material/lecture.pdf) by Curtis Meyer.\n",
    "\n",
    ":::\n",
    "\n",
    "{func}`.spin_conservation` checks whether spin magnitude and spin projections are conserved. In addition, it checks whether the Clebsch-Gordan coefficients are non-zero, meaning that the coupled spins on the interaction nodes are valid as well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### No spin and angular momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spin_conservation(\n",
    "    ingoing_spins=[\n",
    "        SpinEdgeInput(0, 0),\n",
    "    ],\n",
    "    outgoing_spins=[\n",
    "        SpinEdgeInput(0, 0),\n",
    "        SpinEdgeInput(0, 0),\n",
    "    ],\n",
    "    interaction_qns=SpinNodeInput(\n",
    "        l_magnitude=0,  # false if 1\n",
    "        l_projection=0,\n",
    "        s_magnitude=0,\n",
    "        s_projection=0,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Non-zero example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spin_conservation(\n",
    "    ingoing_spins=[\n",
    "        SpinEdgeInput(1, 0),\n",
    "    ],\n",
    "    outgoing_spins=[\n",
    "        SpinEdgeInput(1, +1),\n",
    "        SpinEdgeInput(1, -1),\n",
    "    ],\n",
    "    interaction_qns=SpinNodeInput(\n",
    "        l_magnitude=1,\n",
    "        l_projection=0,\n",
    "        s_magnitude=2,\n",
    "        s_projection=0,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example with a {obj}`.StateTransition`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, generate some {obj}`.StateTransition`s with {func}`.generate_transitions`, then select one of them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=\"J/psi(1S)\",\n",
    "    final_state=[\"K0\", \"Sigma+\", \"p~\"],\n",
    "    allowed_interaction_types=\"strong\",\n",
    "    formalism=\"canonical\",\n",
    ")\n",
    "transition = reaction.transitions[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, have a look at the edge and node properties, and use the underlying {class}`.Topology` to extract one of the node {class}`.InteractionProperties` with the surrounding states (these are {obj}`tuple`s of a {class}`.Particle` and a {obj}`float` spin projection)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(transition, render_node=True)\n",
    "display(graphviz.Source(dot))\n",
    "\n",
    "dot = qrules.io.asdot(\n",
    "    transition.topology,\n",
    "    render_node=True,\n",
    "    render_resonance_id=True,\n",
    "    render_initial_state_id=True,\n",
    ")\n",
    "display(graphviz.Source(dot))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We select node $(0)$, which has incoming state ID $-1$ and outgoing state IDs $0$ and $3$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topology = transition.topology\n",
    "node_id = 0\n",
    "in_id, *_ = topology.get_edge_ids_ingoing_to_node(node_id)\n",
    "out_id1, out_id2, *_ = topology.get_edge_ids_outgoing_from_node(node_id)\n",
    "\n",
    "incoming_state = transition.states[in_id]\n",
    "outgoing_state1 = transition.states[out_id1]\n",
    "outgoing_state2 = transition.states[out_id2]\n",
    "interaction = transition.interactions[node_id]\n",
    "\n",
    "spin_magnitude_conservation(\n",
    "    ingoing_spin_magnitudes=[incoming_state.particle.spin],\n",
    "    outgoing_spin_magnitudes=[\n",
    "        outgoing_state1.particle.spin,\n",
    "        outgoing_state2.particle.spin,\n",
    "    ],\n",
    "    interaction_qns=interaction,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Contrary to expectations, this transition does not conserve spin **projection** and therefore {func}`.spin_conservation` returns {obj}`False`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spin_conservation(\n",
    "    ingoing_spins=[\n",
    "        SpinEdgeInput(\n",
    "            spin_magnitude=incoming_state.particle.spin,\n",
    "            spin_projection=incoming_state.spin_projection,\n",
    "        )\n",
    "    ],\n",
    "    outgoing_spins=[\n",
    "        SpinEdgeInput(\n",
    "            spin_magnitude=outgoing_state1.particle.spin,\n",
    "            spin_projection=outgoing_state1.spin_projection,\n",
    "        ),\n",
    "        SpinEdgeInput(\n",
    "            spin_magnitude=outgoing_state2.particle.spin,\n",
    "            spin_projection=outgoing_state2.spin_projection,\n",
    "        ),\n",
    "    ],\n",
    "    interaction_qns=interaction,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reason is that AmpForm formulates the {class}`~ampform.helicity.HelicityModel` with the helicity formalism first and then uses a transformation to get the model in the canonical basis (see {func}`~ampform.helicity.formulate_isobar_cg_coefficients`). The canonical basis does not conserve helicity (taken to be {attr}`.State.spin_projection`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modifying {obj}`.StateTransition`s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When checking conservation rules, you may want to modify the properties on the {obj}`.StateTransition`s. However, a {obj}`.StateTransition` is a {class}`.FrozenTransition`, so it is not possible to modify its {attr}`~.FrozenTransition.interactions` and {attr}`~.FrozenTransition.states`. The only way around this is to create a new instance with {func}`attrs.evolve`.\n",
    "\n",
    "First, we get the instance (in this case one of the {class}`.InteractionProperties`) and substitute its {attr}`.InteractionProperties.l_magnitude`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_interaction = attrs.evolve(transition.interactions[node_id], l_magnitude=2)\n",
    "new_interaction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then again use {func}`attrs.evolve` to substitute the {attr}`.Transition.interactions` of the original {obj}`.StateTransition`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_interaction_dict = dict(transition.interactions)  # make mutable\n",
    "new_interaction_dict[node_id] = new_interaction\n",
    "new_transition = attrs.evolve(transition, interactions=new_interaction_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(new_transition, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}