FEAT: inspect quantum number solutions (#168)

* DX: add `tox -e docnb-force` job

* DX: do not format notebooks with Prettier

* DX: install `ipywidgets` with `jupyter` (needed for `tqdm`)

* FEAT: add DOT repr for dictionaries
  This is especially useful to represent objects returned by
  `StateTransitionManager._solve()`

* FEAT: extract find_quantum_number_transitions

* MAINT: add key-word arguments to clarify call

* MAINT: define variable for CSP solutions result

* MAINT: make allowed_intermediate states immutable
  Also fixes some of the naming: "particles" becomes "quantum_numbers"

* MAINT: simplify logics with `defaultdict`

.prettierignore

@@ -1 +1,2 @@
+*.ipynb
 LICENSE

docs/conf.py

@@ -56,6 +56,15 @@ def get_branch_name() -> str:
     return branch_name
 
 
+def get_execution_mode() -> str:
+    if "FORCE_EXECUTE_NB" in os.environ:
+        print("\033[93;1mWill run ALL Jupyter notebooks!\033[0m")
+        return "force"
+    if "EXECUTE_NB" in os.environ:
+        return "cache"
+    return "off"
+
+
 BRANCH = get_branch_name()
 
 try:
@@ -363,15 +372,11 @@ def get_version(package_name: str) -> str:
 ]
 
 # Settings for myst_nb
+nb_execution_mode = get_execution_mode()
 nb_execution_show_tb = True
 nb_execution_timeout = -1
 nb_output_stderr = "remove"
 
-nb_execution_mode = "off"
-if "EXECUTE_NB" in os.environ:
-    print("\033[93;1mWill run Jupyter notebooks!\033[0m")
-    nb_execution_mode = "cache"
-
 # Settings for myst-parser
 myst_enable_extensions = [
     "amsmath",

docs/usage/reaction.ipynb

@@ -137,9 +137,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "```{eval-rst}\n",
     ".. admonition:: `.StateTransitionManager`\n",
@@ -266,13 +264,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "If you are happy with the default settings generated by the {class}`.StateTransitionManager`, just start with solving directly!"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
+    "If you are happy with the default settings generated by the {class}`.StateTransitionManager`, just start with solving directly!\n",
+    "\n",
     "```{toggle}\n",
     "This step takes about 23 sec on an Intel(R) Core(TM) i7-6820HQ CPU of 2.70GHz running, multi-threaded.\n",
     "```"
@@ -291,6 +284,10 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    ":::{tip}\n",
+    "See {ref}`usage/visualize:Quantum number solutions` for a visualization of the intermediate steps.\n",
+    ":::\n",
+    "\n",
     "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:"
    ]
   },
@@ -399,9 +396,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# particles are found by name comparison,\n",
@@ -513,6 +508,18 @@
    "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.16"
   }
  },
  "nbformat": 4,

docs/usage/visualize.ipynb

@@ -63,7 +63,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Visualize decay topologies"
+    "# Visualize solutions"
    ]
   },
   {
@@ -187,6 +187,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "(problem-sets)=\n",
     "## {class}`.ProblemSet`s"
    ]
   },
@@ -203,11 +204,14 @@
    "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()"
    ]
   },
@@ -227,35 +231,101 @@
     "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": [
-    "len(problem_sets[60.0])"
+    "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"
+     "full-width",
+     "hide-input"
     ]
    },
    "outputs": [],
    "source": [
-    "problem_set = problem_sets[60.0][0]\n",
-    "dot = qrules.io.asdot(problem_set, render_node=True)\n",
+    "dot = qrules.io.asdot(qn_problem_set, render_node=True)\n",
     "graphviz.Source(dot)"
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
-    "tags": []
+    "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": [
     "## {obj}`.StateTransition`s"
    ]
@@ -264,7 +334,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Here, we'll visualize the allowed transitions for the decay $\\psi' \\to \\gamma\\eta\\eta$ as an example."
+    "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."
    ]
   },
   {
@@ -369,9 +439,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "dot = qrules.io.asdot(reaction.transitions[:3], strip_spin=True, render_node=True)\n",
@@ -533,8 +601,16 @@
    "name": "python3"
   },
   "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
    "name": "python",
-   "version": "3.8.12"
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.16"
   }
  },
  "nbformat": 4,

setup.cfg

@@ -119,6 +119,7 @@ sty =
     pre-commit >=1.4.0
 jupyter =
     aquirdturtle-collapsible-headings
+    ipywidgets
     jupyterlab
     jupyterlab-code-formatter
     jupyterlab-myst; python_version >="3.7.0"

src/qrules/io/_dot.py

@@ -16,7 +16,7 @@
 
 from qrules.particle import Particle, ParticleWithSpin, Spin
 from qrules.quantum_numbers import InteractionProperties, _to_fraction
-from qrules.solving import EdgeSettings, NodeSettings
+from qrules.solving import EdgeSettings, NodeSettings, QNProblemSet, QNResult
 from qrules.topology import FrozenTransition, MutableTransition, Topology, Transition
 from qrules.transition import ProblemSet, ReactionInfo, State
 
@@ -77,13 +77,15 @@ def _create_preface(self) -> List[str]:
         ]
 
     def _render(self, obj: Any) -> List[str]:
+        if isinstance(obj, QNResult):
+            obj = obj.solutions
         if isinstance(obj, ReactionInfo):
             obj = obj.transitions
         if isinstance(obj, abc.Iterable):
             return self._render_multiple_transitions(obj)
-        if isinstance(obj, (ProblemSet, Topology, Transition)):
+        if isinstance(obj, (ProblemSet, QNProblemSet, Topology, Transition)):
             return self._render_transition(obj)
-        raise NotImplementedError
+        raise NotImplementedError(f"No DOT rendering for type {type(obj).__name__}")
 
     def _render_multiple_transitions(self, obj: Iterable) -> List[str]:
         if self.collapse_graphs:
@@ -102,15 +104,17 @@ def _render_multiple_transitions(self, obj: Iterable) -> List[str]:
 
     def _render_transition(
         self,
-        obj: Union[ProblemSet, Topology, Transition],
+        obj: Union[ProblemSet, QNProblemSet, Topology, Transition],
         prefix: str = "",
     ) -> List[str]:
         # pylint: disable=too-many-branches,too-many-locals,too-many-statements
         lines: List[str] = []
         if isinstance(obj, tuple) and len(obj) == 2:
             topology: Topology = obj[0]
-            rendered_graph: Union[ProblemSet, Topology, Transition] = obj[1]
-        elif isinstance(obj, (ProblemSet, Transition)):
+            rendered_graph: Union[ProblemSet, QNProblemSet, Topology, Transition] = obj[
+                1
+            ]
+        elif isinstance(obj, (ProblemSet, QNProblemSet, Transition)):
             rendered_graph = obj
             topology = obj.topology
         elif isinstance(obj, Topology):
@@ -137,7 +141,7 @@ def _render_transition(
             else:
                 label = _create_edge_label(rendered_graph, i, self.render_resonance_id)
                 lines += [self._create_graphviz_edge(from_node, to_node, label)]
-        if isinstance(obj, ProblemSet):
+        if isinstance(obj, (ProblemSet, QNProblemSet)):
             node_settings = obj.solving_settings.interactions
             for node_id, settings in node_settings.items():
                 label = ""
@@ -238,22 +242,22 @@ def _create_same_rank_line(node_edge_ids: Iterable[int], prefix: str = "") -> st
 
 
 def _create_edge_label(
-    graph: Union[ProblemSet, Topology, Transition],
+    graph: Union[ProblemSet, QNProblemSet, Topology, Transition],
     edge_id: int,
     render_edge_id: bool,
 ) -> str:
     if isinstance(graph, Topology):
         if render_edge_id:
             return str(edge_id)
         return ""
-    if isinstance(graph, ProblemSet):
+    if isinstance(graph, (ProblemSet, QNProblemSet)):
         edge_setting = graph.solving_settings.states.get(edge_id)
         initial_fact = graph.initial_facts.states.get(edge_id)
         edge_property: Optional[Union[EdgeSettings, ParticleWithSpin]] = None
         if edge_setting:
             edge_property = edge_setting
         if initial_fact:
-            edge_property = initial_fact
+            edge_property = initial_fact  # type: ignore[assignment]
         return __render_edge_with_id(edge_id, edge_property, render_edge_id)
     edge_prop = graph.states.get(edge_id)
     return __render_edge_with_id(edge_id, edge_prop, render_edge_id)
@@ -291,6 +295,18 @@ def as_string(obj: Any) -> str:
 as_string.register(str, lambda _: _)  # avoid warning for str type
 
 
+@as_string.register(dict)
+def _(obj: dict) -> str:
+    lines = []
+    for key, value in obj.items():
+        if isinstance(key, type) or callable(key):
+            key_repr = key.__name__
+        else:
+            key_repr = key
+        lines.append(f"{key_repr} = {value}")
+    return "\n".join(lines)
+
+
 @as_string.register(InteractionProperties)
 def _(obj: InteractionProperties) -> str:
     lines = []