{
 "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": [
    ":::{warning}\n",
    "Currently the main user-interface is the ```StateTransitionManager```. There is work in progress to remove it and split its functionality into several functions/classes to separate concerns\n",
    "and to facilitate the modification of intermediate results like the filtering of ```QNProblemSet```s, setting allowed interaction types, etc. (see below)\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{margin}\n",
    "```{warning}\n",
    "{class}`graphviz.Source` requires your system to have DOT installed, see {doc}`Installation <graphviz:index>`.\n",
    "```\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {mod}`~qrules.io` module allows you to convert {class}`.MutableTransition`, {class}`.Topology` instances, and {class}`.ProblemSet`s to [DOT language](https://graphviz.org/doc/info/lang.html) with {func}`.asdot`. You can visualize its output with third-party libraries, such as [Graphviz](https://graphviz.org). This is particularly useful after running {meth}`~.StateTransitionManager.find_solutions`, which produces a {class}`.ReactionInfo` object with a {class}`.list` of {class}`.MutableTransition` instances (see {doc}`/usage/reaction`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Topologies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all, here are is an example of how to visualize a group of {class}`.Topology` instances. We use {func}`.create_isobar_topologies` and {func}`.create_n_body_topology` to create a few standard topologies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import graphviz\n",
    "from IPython.display import display\n",
    "\n",
    "import qrules\n",
    "from qrules.conservation_rules import (\n",
    "    parity_conservation,\n",
    "    spin_magnitude_conservation,\n",
    "    spin_validity,\n",
    ")\n",
    "from qrules.particle import Spin\n",
    "from qrules.quantum_numbers import EdgeQuantumNumbers, NodeQuantumNumbers\n",
    "from qrules.solving import (\n",
    "    CSPSolver,\n",
    "    dict_set_intersection,\n",
    "    filter_quantum_number_problem_set,\n",
    ")\n",
    "from qrules.topology import create_isobar_topologies, create_n_body_topology\n",
    "from qrules.transition import State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topology = create_n_body_topology(2, 4)\n",
    "graphviz.Source(qrules.io.asdot(topology, render_initial_state_id=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the IDs of the {attr}`~.Topology.nodes` is also rendered if there is more than node:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topologies = create_isobar_topologies(4)\n",
    "graphviz.Source(qrules.io.asdot(topologies))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can be turned on or off with the arguments of {func}`.asdot`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topologies = create_isobar_topologies(3)\n",
    "graphviz.Source(qrules.io.asdot(topologies, render_node=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`.asdot` provides other options as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topologies = create_isobar_topologies(5)\n",
    "dot = qrules.io.asdot(\n",
    "    topologies[0],\n",
    "    render_final_state_id=False,\n",
    "    render_resonance_id=True,\n",
    "    render_node=False,\n",
    ")\n",
    "display(graphviz.Source(dot))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(problem-sets)=\n",
    "## {class}`.ProblemSet`s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As noted in {doc}`reaction`, the {class}`.StateTransitionManager` provides more control than the façade function {func}`.generate_transitions`. One advantages, is that the {class}`.StateTransitionManager` first generates a set of {class}`.ProblemSet`s with {meth}`.create_problem_sets` that you can further configure if you wish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qrules.settings import InteractionType\n",
    "\n",
    "stm = qrules.StateTransitionManager(\n",
    "    initial_state=[\"J/psi(1S)\"],\n",
    "    final_state=[\"K0\", \"Sigma+\", \"p~\"],\n",
    "    formalism=\"canonical-helicity\",\n",
    ")\n",
    "stm.set_allowed_interaction_types([InteractionType.STRONG, InteractionType.EM])\n",
    "problem_sets = stm.create_problem_sets()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the output of {meth}`.create_problem_sets` is a {obj}`dict` with {obj}`float` values as keys (representing the interaction strength) and {obj}`list`s of {obj}`.ProblemSet`s as values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted(problem_sets, reverse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "problem_set = problem_sets[60.0][0]\n",
    "dot = qrules.io.asdot(problem_set, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantum number solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As noted in {ref}`usage/reaction:3. Find solutions`, a {obj}`.ProblemSet` can be fed to {meth}`.StateTransitionManager.find_solutions` directly to get a {obj}`.ReactionInfo` object. {obj}`.ReactionInfo` is a final result that consists of {obj}`.Particle`s, but in the intermediate steps, QRules works with sets of quantum numbers. One can inspect these intermediate generated quantum numbers by using {meth}`.find_quantum_number_transitions` and inspecting is output. Note that the resulting object is again a {obj}`dict` with strengths as keys and a list of solution as values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qn_solutions = stm.find_quantum_number_transitions(problem_sets)\n",
    "{strength: len(values) for strength, values in qn_solutions.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The list of solutions consist of a {obj}`tuple` of a {obj}`.QNProblemSet` (compare {ref}`problem-sets`) and a {obj}`.QNResult`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strong_qn_solutions = qn_solutions[3600.0]\n",
    "qn_problem_set, qn_result = strong_qn_solutions[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(qn_problem_set, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(qn_result, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filtering quantum number problem sets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes, only a certain subset of quantum numbers and conservation rules are relevant, or the number of solutions the {class}`.StateTransitionManager` gives by default is too large for the follow-up analysis.\n",
    "The {func}`.filter_quantum_number_problem_set` function can be used to produce a {class}`.QNProblemSet` where only the desired quantum numbers and conservation rules are considered when fed back to the solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "desired_edge_properties = {EdgeQuantumNumbers.spin_magnitude, EdgeQuantumNumbers.parity}\n",
    "desired_node_properties = {\n",
    "    NodeQuantumNumbers.l_magnitude,\n",
    "    NodeQuantumNumbers.s_magnitude,\n",
    "}  # has to be reused in the CSPSolver-constructor\n",
    "filtered_qn_problem_set = filter_quantum_number_problem_set(\n",
    "    qn_problem_set,\n",
    "    edge_rules={spin_validity},\n",
    "    node_rules={spin_magnitude_conservation, parity_conservation},\n",
    "    edge_properties=desired_edge_properties,\n",
    "    node_properties=desired_node_properties,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(filtered_qn_problem_set, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{warning}\n",
    "The next cell will use some (currently) internal functionality. As stated at the top, a workflow similar to this will be used in future versions of {mod}`qrules`, see e.g. [ComPWA/qrules#305](https://github.com/ComPWA/qrules/issues/305). Manual setup of the {obj}`.CSPSolver` like in here will then also not be necessary.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = CSPSolver([\n",
    "    dict_set_intersection(\n",
    "        qrules.system_control.create_edge_properties(part),\n",
    "        desired_edge_properties,\n",
    "    )\n",
    "    for part in qrules.particle.load_pdg()\n",
    "])\n",
    "\n",
    "filtered_qn_solutions = solver.find_solutions(filtered_qn_problem_set)\n",
    "filtered_qn_result = filtered_qn_solutions.solutions[6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(filtered_qn_result, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## {obj}`.StateTransition`s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After finding the {ref}`usage/visualize:Quantum number solutions`, QRules finds {obj}`.Particle` definitions that match these quantum numbers. All these steps are hidden in the convenience functions {meth}`.StateTransitionManager.find_solutions` and {func}`.generate_transitions`. In the following, we'll visualize the allowed transitions for the decay $\\psi' \\to \\gamma\\eta\\eta$ as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qrules\n",
    "\n",
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=\"psi(2S)\",\n",
    "    final_state=[\"gamma\", \"eta\", \"eta\"],\n",
    "    allowed_interaction_types=\"EM\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As noted in {ref}`usage/reaction:3. Find solutions`, the {attr}`~.ReactionInfo.transitions` contain all spin projection combinations (which is necessary for the {mod}`ampform` package). It is possible to convert all these solutions to DOT language with {func}`~.asdot`. To avoid visualizing all solutions, we just take a subset of the {attr}`~.ReactionInfo.transitions`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(reaction.transitions[::50][:3])  # just some selection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This {class}`str` of [DOT language](https://graphviz.org/doc/info/lang.html) for the list of {class}`.MutableTransition` instances can then be visualized with a third-party library, for instance, with {class}`graphviz.Source`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import graphviz\n",
    "\n",
    "dot = qrules.io.asdot(\n",
    "    reaction.transitions[::50][:3], render_node=False\n",
    ")  # just some selection\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also serialize the DOT string to file with {func}`.io.write`. The file extension for a DOT file is `.gv`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qrules.io.write(reaction, \"decay_topologies_with_spin.gv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Collapse graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since this list of all possible spin projections {attr}`~.ReactionInfo.transitions` is rather long, it is often useful to use `strip_spin=True` or `collapse_graphs=True` to bundle comparable graphs. First, {code}`strip_spin=True` allows one collapse (ignore) the spin projections (we again show a selection only):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(reaction.transitions[:3], strip_spin=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or, with stripped node properties:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(reaction.transitions[:3], strip_spin=True, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "By default, {func}`.asdot` renders edge IDs, because they represent the (final) state IDs as well. In the example above, we switched this off.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If that list is still too much, there is {code}`collapse_graphs=True`, which bundles all graphs with the same final state groupings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(reaction, collapse_graphs=True, render_node=False)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Other state renderings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {meth}`~.FrozenTransition.convert` method makes it possible to convert the types of its {attr}`~.FrozenTransition.states`. This for instance allows us to only render the spin states on in a {class}`.Transition`:\n",
    "\n",
    "::::{margin}\n",
    "\n",
    ":::{tip}\n",
    "\n",
    "We use the fact that a {obj}`.StateTransition` is frozen (and therefore hashable) to remove any duplicate transitions.\n",
    "\n",
    ":::\n",
    "\n",
    "::::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spin_transitions = sorted({\n",
    "    t.convert(lambda s: Spin(s.particle.spin, s.spin_projection))\n",
    "    for t in reaction.transitions\n",
    "})\n",
    "some_selection = spin_transitions[::67][:3]\n",
    "dot = qrules.io.asdot(some_selection, render_node=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or any other properties of a {class}`.State`, such as masses or $J^{PC}(I^G)$ numbers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "scroll-output"
    ]
   },
   "outputs": [],
   "source": [
    "def render_mass(state: State, digits: int = 3) -> str:\n",
    "    mass = round(state.particle.mass, digits)\n",
    "    width = round(state.particle.width, digits)\n",
    "    if width == 0:\n",
    "        return str(mass)\n",
    "    return f\"{mass}±{width}\"\n",
    "\n",
    "\n",
    "mass_transitions = sorted({\n",
    "    t.convert(\n",
    "        state_converter=render_mass,\n",
    "        interaction_converter=lambda _: None,\n",
    "    )\n",
    "    for t in reaction.transitions\n",
    "})\n",
    "dot = qrules.io.asdot(mass_transitions[::10])\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "scroll-input",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from fractions import Fraction\n",
    "\n",
    "\n",
    "def render_jpc_ig(state: State) -> str:\n",
    "    particle = state.particle\n",
    "    text = render_fraction(particle.spin)\n",
    "    if particle.parity is not None:\n",
    "        text += render_sign(particle.parity)\n",
    "    if particle.c_parity is not None:\n",
    "        text += render_sign(particle.c_parity)\n",
    "    if particle.isospin is not None and particle.g_parity is not None:\n",
    "        text += \"(\"\n",
    "        text += f\"{render_fraction(particle.isospin.magnitude)}\"\n",
    "        text += f\"{render_sign(particle.g_parity)}\"\n",
    "        text += \")\"\n",
    "    return text\n",
    "\n",
    "\n",
    "def render_fraction(value: float) -> str:\n",
    "    fraction = Fraction(value)\n",
    "    if fraction.denominator == 1:\n",
    "        return str(fraction.numerator)\n",
    "    return f\"{fraction.numerator}/{fraction.denominator}\"\n",
    "\n",
    "\n",
    "def render_sign(parity: int) -> str:\n",
    "    if parity == -1:\n",
    "        return \"⁻\"\n",
    "    if parity == +1:\n",
    "        return \"⁺\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "\n",
    "jpc_ig_transitions = sorted({\n",
    "    t.convert(\n",
    "        state_converter=render_jpc_ig,\n",
    "        interaction_converter=lambda _: None,\n",
    "    )\n",
    "    for t in reaction.transitions\n",
    "})\n",
    "dot = qrules.io.asdot(jpc_ig_transitions, collapse_graphs=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{tip}\n",
    "Note that collapsing the graphs also works for custom edge properties.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Styling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- cspell:ignore darkgreen fontcolor fontsize penwidth -->\n",
    "The {func}`.asdot` function also takes [Graphviz attributes](https://graphviz.org/doc/info/attrs.html). These can be used to modify the layout of the whole figure. Examples are the [`size`](https://graphviz.org/docs/attrs/size), [`color`](https://graphviz.org/docs/attrs/color), and [`fontcolor`](https://graphviz.org/docs/attrs/fontcolor). Edges and nodes can be styled with `edge_style` and `node_style` respectively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(\n",
    "    reaction.transitions[0],\n",
    "    render_node=True,\n",
    "    size=12,\n",
    "    bgcolor=\"white\",\n",
    "    edge_style={\n",
    "        \"color\": \"red\",\n",
    "        \"arrowhead\": \"open\",\n",
    "        \"fontcolor\": \"blue\",\n",
    "        \"fontsize\": 25,\n",
    "    },\n",
    "    node_style={\n",
    "        \"color\": \"gray\",\n",
    "        \"penwidth\": 2,\n",
    "        \"shape\": \"ellipse\",\n",
    "        \"style\": \"dashed\",\n",
    "    },\n",
    ")\n",
    "display(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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}