{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%config Completer.use_jedi = False\n",
    "%config InlineBackend.figure_formats = ['svg']\n",
    "\n",
    "# Install on Google Colab\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "from IPython import get_ipython\n",
    "\n",
    "install_packages = \"google.colab\" in str(get_ipython())\n",
    "if install_packages:\n",
    "    for package in [\"expertsystem\", \"graphviz\"]:\n",
    "        subprocess.check_call(\n",
    "            [sys.executable, \"-m\", \"pip\", \"install\", package]\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize decay topologies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{warning}\n",
    "The {doc}`PWA Expert System <index>` has been split up into {doc}`QRules <qrules:index>` and {doc}`AmpForm <ampform:index>`. Please use these packages instead!\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "The formulas and figures on this page have been generated with the lineshapes in the {mod}`.lineshape` module, so as to [glue](https://myst-nb.readthedocs.io/en/latest/use/glue.html) them back in to the API of that module.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {mod}`~expertsystem.io` module allows you to convert {class}`.StateTransitionGraph` and {class}`.Topology` instances 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}`.Result` object with a {class}`.list` of {class}`.StateTransitionGraph` instances (see {doc}`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": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import graphviz\n",
    "\n",
    "from expertsystem import io\n",
    "from expertsystem.reaction.topology import (\n",
    "    create_isobar_topologies,\n",
    "    create_n_body_topology,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topology = create_n_body_topology(2, 4)\n",
    "graphviz.Source(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(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(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 = 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": [
    "## {class}`.StateTransitionGraph`s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, 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 expertsystem as es\n",
    "\n",
    "result = es.reaction.generate(\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}`~.Result.transitions` contain all spin projection combinations (which is necessary for the {mod}`~expertsystem.amplitude` module). 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}`~.Result.transitions`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dot = es.io.asdot(result.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}`.StateTransitionGraph` instances can then be visualized with a third-party library, for instance, with {class}`graphviz.Source`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{margin}\n",
    "```{warning}\n",
    "[graphviz](graphviz.Source) requires your system to have DOT installed, see {doc}`Installation <graphviz:index>`.\n",
    "```\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import graphviz\n",
    "\n",
    "dot = es.io.asdot(\n",
    "    result.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 [io.write](expertsystem.io.write). The file extension for a DOT file is `.gv`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "es.io.write(result, \"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}`~.Result.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 = es.io.asdot(result.transitions[:3], strip_spin=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "By default, `.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 = es.io.asdot(result, collapse_graphs=True, render_node=False)\n",
    "graphviz.Source(dot)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}