refactor!: implement StateTransition classes (#75)

* feat: define ReactionInfo, State, StateTransition, and StateTransitionCollection classes * feat: implement as/fromdict for ReactionInfo etc * feat: provide pretty printer for FrozenDict * feat: provide pretty printer for Topology * fix: small improvement to Spin error message * refactor!: remove Result class * refactor: extract _dot.__strip_spin * style: clean up ignore statements on top * style: update notebook language version
ComPWA · Jun 18, 2021 · fe764e1 · fe764e1
1 parent 9fa8183
commit fe764e1
Show file tree

Hide file tree

Showing 24 changed files with 856 additions and 262 deletions.
diff --git a/docs/abbreviate_signature.py b/docs/abbreviate_signature.py
@@ -1,3 +1,6 @@
+# cspell:ignore docutils
+# pylint: disable=import-error
+# pyright: reportMissingImports=false
 """Abbreviated the annotations generated by sphinx-autodoc.
 
 It's not necessary to generate the full path of type hints, because they are
@@ -6,9 +9,6 @@
 See also https://github.com/sphinx-doc/sphinx/issues/5868.
 """
 
-# cspell:ignore docutils
-# pylint: disable=import-error
-# pyright: reportMissingImports=false
 import sphinx.domains.python
 from docutils import nodes
 from sphinx import addnodes

diff --git a/docs/usage.ipynb b/docs/usage.ipynb
@@ -88,7 +88,7 @@
    "source": [
     "import qrules\n",
     "\n",
-    "result = qrules.generate_transitions(\n",
+    "reaction = qrules.generate_transitions(\n",
     "    initial_state=\"J/psi(1S)\",\n",
     "    final_state=[\"K0\", \"Sigma+\", \"p~\"],\n",
     "    allowed_interaction_types=\"strong\",\n",
@@ -104,7 +104,7 @@
    "source": [
     "import graphviz\n",
     "\n",
-    "dot = qrules.io.asdot(result, collapse_graphs=True)\n",
+    "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n",
     "graphviz.Source(dot)"
    ]
   },
@@ -271,7 +271,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

diff --git a/docs/usage/particle.ipynb b/docs/usage/particle.ipynb
@@ -467,7 +467,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

diff --git a/docs/usage/reaction.ipynb b/docs/usage/reaction.ipynb
@@ -215,14 +215,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "result = stm.find_solutions(problem_sets)"
+    "reaction = stm.find_solutions(problem_sets)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The {meth}`~.StateTransitionManager.find_solutions` method returns a {class}`.Result` object from which you can extract the {attr}`~.Result.transitions`. Now, you can use {meth}`~.Result.get_intermediate_particles` to print the names of the intermediate states that the {class}`.StateTransitionManager` found:"
+    "The {meth}`~.StateTransitionManager.find_solutions` method returns a {class}`.ReactionInfo` object from which you can extract the {attr}`~.ReactionInfo.transitions`. Now, you can use {meth}`~.ReactionInfo.get_intermediate_particles` to print the names of the intermediate states that the {class}`.StateTransitionManager` found:"
    ]
   },
   {
@@ -231,8 +231,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(\"found\", len(result.transitions), \"solutions!\")\n",
-    "result.get_intermediate_particles().names"
+    "print(\"found\", len(reaction.transitions), \"solutions!\")\n",
+    "reaction.get_intermediate_particles().names"
    ]
   },
   {
@@ -244,9 +244,9 @@
     "class: dropdown\n",
     "----\n",
     "\n",
-    "The \"number of {attr}`~.Result.transitions`\" is the total number of allowed {obj}`.StateTransitionGraph` instances that the {class}`.StateTransitionManager` has found. This also includes all allowed **spin projection combinations**. In this channel, we for example consider a $J/\\psi$ with spin projection $\\pm1$ that decays into a $\\gamma$ with spin projection $\\pm1$, which already gives us four possibilities.\n",
+    "The \"number of {attr}`~.ReactionInfo.transitions`\" is the total number of allowed {obj}`.StateTransitionGraph` instances that the {class}`.StateTransitionManager` has found. This also includes all allowed **spin projection combinations**. In this channel, we for example consider a $J/\\psi$ with spin projection $\\pm1$ that decays into a $\\gamma$ with spin projection $\\pm1$, which already gives us four possibilities.\n",
     "\n",
-    "On the other hand, the intermediate state names that was extracted with {meth}`.Result.get_intermediate_particles`, is just a {obj}`set` of the state names on the intermediate edges of the list of {attr}`~.Result.transitions`, regardless of spin projection.\n",
+    "On the other hand, the intermediate state names that was extracted with {meth}`.ReactionInfo.get_intermediate_particles`, is just a {obj}`set` of the state names on the intermediate edges of the list of {attr}`~.ReactionInfo.transitions`, regardless of spin projection.\n",
     "````"
    ]
   },
@@ -272,10 +272,10 @@
    "source": [
     "stm.set_allowed_interaction_types([InteractionType.STRONG])\n",
     "problem_sets = stm.create_problem_sets()\n",
-    "result = stm.find_solutions(problem_sets)\n",
+    "reaction = stm.find_solutions(problem_sets)\n",
     "\n",
-    "print(\"found\", len(result.transitions), \"solutions!\")\n",
-    "result.get_intermediate_particles().names"
+    "print(\"found\", len(reaction.transitions), \"solutions!\")\n",
+    "reaction.get_intermediate_particles().names"
    ]
   },
   {
@@ -297,10 +297,10 @@
    "source": [
     "stm.set_allowed_interaction_types([InteractionType.STRONG, InteractionType.EM])\n",
     "problem_sets = stm.create_problem_sets()\n",
-    "result = stm.find_solutions(problem_sets)\n",
+    "reaction = stm.find_solutions(problem_sets)\n",
     "\n",
-    "print(\"found\", len(result.transitions), \"solutions!\")\n",
-    "result.get_intermediate_particles().names"
+    "print(\"found\", len(reaction.transitions), \"solutions!\")\n",
+    "reaction.get_intermediate_particles().names"
    ]
   },
   {
@@ -339,10 +339,10 @@
     "# i.e. f2 will find all f2's and f all f's independent of their spin\n",
     "stm.set_allowed_intermediate_particles([\"f(0)\", \"f(2)\"])\n",
     "\n",
-    "result = stm.find_solutions(problem_sets)\n",
+    "reaction = stm.find_solutions(problem_sets)\n",
     "\n",
-    "print(\"found\", len(result.transitions), \"solutions!\")\n",
-    "result.get_intermediate_particles().names"
+    "print(\"found\", len(reaction.transitions), \"solutions!\")\n",
+    "reaction.get_intermediate_particles().names"
    ]
   },
   {
@@ -366,7 +366,7 @@
     "\n",
     "from qrules import io\n",
     "\n",
-    "dot = io.asdot(result, collapse_graphs=True, render_node=False)\n",
+    "dot = io.asdot(reaction, collapse_graphs=True, render_node=False)\n",
     "graphviz.Source(dot)"
    ]
   },
@@ -392,7 +392,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The {class}`.Result`, {class}`.StateTransitionGraph`, and {class}`.Topology` can be serialized to and from a {obj}`dict` with {func}`.io.asdict` and {func}`.io.fromdict`:"
+    "The {class}`.ReactionInfo`, {class}`.StateTransitionGraph`, and {class}`.Topology` can be serialized to and from a {obj}`dict` with {func}`.io.asdict` and {func}`.io.fromdict`:"
    ]
   },
   {
@@ -403,8 +403,7 @@
    "source": [
     "from qrules import io\n",
     "\n",
-    "graph = result.transitions[0]\n",
-    "io.asdict(graph.topology)"
+    "io.asdict(reaction.transition_groups[0].topology)"
    ]
   },
   {
@@ -415,7 +414,7 @@
     "YAML is more human-readable than JSON, but reading and writing JSON is faster.\n",
     "```\n",
     "\n",
-    "This also means that the {obj}`.Result` can be written to JSON or YAML format with {func}`.io.write` and loaded again with {func}`.io.load`:"
+    "This also means that the {obj}`.ReactionInfo` can be written to JSON or YAML format with {func}`.io.write` and loaded again with {func}`.io.load`:"
    ]
   },
   {
@@ -424,9 +423,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "io.write(result, \"transitions.json\")\n",
-    "imported_result = io.load(\"transitions.json\")\n",
-    "assert imported_result == result"
+    "io.write(reaction, \"transitions.json\")\n",
+    "imported_reaction = io.load(\"transitions.json\")\n",
+    "assert imported_reaction == reaction"
    ]
   },
   {
@@ -436,15 +435,6 @@
     "Handy if it takes a lot of computation time to re-generate the transitions!"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "```{warning}\n",
-    "It's not possible to {mod}`pickle` a {class}`.Result`, because {class}`.StateTransitionGraph` makes use of {class}`~typing.Generic`.\n",
-    "```"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -472,7 +462,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

diff --git a/docs/usage/visualize.ipynb b/docs/usage/visualize.ipynb
@@ -51,7 +51,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The {mod}`~qrules.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}`/usage/reaction`)."
+    "The {mod}`~qrules.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}`.ReactionInfo` object with a {class}`.list` of {class}`.StateTransitionGraph` instances (see {doc}`/usage/reaction`)."
    ]
   },
   {
@@ -173,7 +173,7 @@
    "source": [
     "import qrules\n",
     "\n",
-    "result = qrules.generate_transitions(\n",
+    "reaction = qrules.generate_transitions(\n",
     "    initial_state=\"psi(2S)\",\n",
     "    final_state=[\"gamma\", \"eta\", \"eta\"],\n",
     "    allowed_interaction_types=\"EM\",\n",
@@ -184,7 +184,7 @@
    "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}`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}`~.Result.transitions`:"
+    "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`:"
    ]
   },
   {
@@ -193,7 +193,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dot = qrules.io.asdot(result.transitions[::50][:3])  # just some selection"
+    "dot = qrules.io.asdot(reaction.transitions[::50][:3])  # just some selection"
    ]
   },
   {
@@ -223,7 +223,7 @@
     "import graphviz\n",
     "\n",
     "dot = qrules.io.asdot(\n",
-    "    result.transitions[::50][:3], render_node=False\n",
+    "    reaction.transitions[::50][:3], render_node=False\n",
     ")  # just some selection\n",
     "graphviz.Source(dot)"
    ]
@@ -241,7 +241,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "qrules.io.write(result, \"decay_topologies_with_spin.gv\")"
+    "qrules.io.write(reaction, \"decay_topologies_with_spin.gv\")"
    ]
   },
   {
@@ -255,7 +255,7 @@
    "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):"
+    "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):"
    ]
   },
   {
@@ -264,7 +264,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dot = qrules.io.asdot(result.transitions[:3], strip_spin=True)\n",
+    "dot = qrules.io.asdot(reaction.transitions[:3], strip_spin=True)\n",
     "graphviz.Source(dot)"
    ]
   },
@@ -290,7 +290,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dot = qrules.io.asdot(result, collapse_graphs=True, render_node=False)\n",
+    "dot = qrules.io.asdot(reaction, collapse_graphs=True, render_node=False)\n",
     "graphviz.Source(dot)"
    ]
   }
@@ -311,7 +311,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

diff --git a/src/qrules/__init__.py b/src/qrules/__init__.py
@@ -1,5 +1,4 @@
 # pylint: disable=too-many-lines
-
 """A rule based system that facilitates particle reaction analysis.
 
 QRules generates allowed particle transitions from a set of conservation rules
@@ -74,7 +73,7 @@
 from .transition import (
     EdgeSettings,
     ProblemSet,
-    Result,
+    ReactionInfo,
     StateTransitionManager,
 )
 
@@ -299,7 +298,7 @@ def generate_transitions(  # pylint: disable=too-many-arguments
     max_spin_magnitude: float = 2.0,
     topology_building: str = "isobar",
     number_of_threads: Optional[int] = None,
-) -> Result:
+) -> ReactionInfo:
     """Generate allowed transitions between an initial and final state.
 
     Serves as a facade to the `.StateTransitionManager` (see
@@ -357,7 +356,7 @@ def generate_transitions(  # pylint: disable=too-many-arguments
     arguments) would be:
 
     >>> import qrules
-    >>> result = qrules.generate_transitions(
+    >>> reaction = qrules.generate_transitions(
     ...     initial_state="D0",
     ...     final_state=["K~0", "K+", "K-"],
     ...     allowed_intermediate_particles=["a(0)(980)", "a(2)(1320)-"],
@@ -366,7 +365,9 @@ def generate_transitions(  # pylint: disable=too-many-arguments
     ...     particle_db=qrules.load_pdg(),
     ...     topology_building="isobar",
     ... )
-    >>> len(result.transitions)
+    >>> len(reaction.transition_groups)
+    3
+    >>> len(reaction.transitions)
     4
     """
     if isinstance(initial_state, str) or (