FEAT: implement quantum problem set filter

docs/conf.py

@@ -47,6 +47,7 @@ def pick_newtype_attrs(some_type: type) -> list:
 add_module_names = False
 api_github_repo = f"{ORGANIZATION}/{REPO_NAME}"
 api_target_substitutions: dict[str, str | tuple[str, str]] = {
+    "EdgeQuantumNumberTypes": ("obj", "qrules.quantum_numbers.EdgeQuantumNumberTypes"),
     "EdgeType": "typing.TypeVar",
     "GraphEdgePropertyMap": ("obj", "qrules.argument_handling.GraphEdgePropertyMap"),
     "GraphElementProperties": ("obj", "qrules.solving.GraphElementProperties"),
@@ -56,11 +57,13 @@ def pick_newtype_attrs(some_type: type) -> list:
     "NewEdgeType": "typing.TypeVar",
     "NewNodeType": "typing.TypeVar",
     "NodeQuantumNumber": ("obj", "qrules.quantum_numbers.NodeQuantumNumber"),
+    "NodeQuantumNumberTypes": ("obj", "qrules.quantum_numbers.NodeQuantumNumberTypes"),
     "NodeType": "typing.TypeVar",
     "ParticleWithSpin": ("obj", "qrules.particle.ParticleWithSpin"),
     "Path": "pathlib.Path",
     "qrules.topology.EdgeType": "typing.TypeVar",
     "qrules.topology.NodeType": "typing.TypeVar",
+    "Rule": ("obj", "qrules.argument_handling.Rule"),
     "SpinFormalism": ("obj", "qrules.transition.SpinFormalism"),
     "StateDefinition": ("obj", "qrules.combinatorics.StateDefinition"),
     "StateTransition": ("obj", "qrules.transition.StateTransition"),

docs/usage/conservation.ipynb

@@ -241,9 +241,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    },
     "tags": [
      "hide-input"
     ]
@@ -417,7 +414,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

docs/usage/reaction.ipynb

@@ -604,7 +604,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.19"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

docs/usage/visualize.ipynb

@@ -50,6 +50,16 @@
     "```"
    ]
   },
+  {
+   "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": {},
@@ -103,7 +113,14 @@
     "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 CSPSolver, filter_quantum_number_problem_set\n",
     "from qrules.topology import create_isobar_topologies, create_n_body_topology\n",
     "from qrules.transition import State"
    ]
@@ -283,9 +300,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    },
     "tags": [
      "full-width",
      "hide-input"
@@ -301,9 +315,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    },
     "tags": [
      "full-width",
      "hide-input"
@@ -315,6 +326,96 @@
     "graphviz.Source(dot)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Filtering Quantum Number Problem Sets\n",
+    "Sometimes only a certain subset of quantum numbers & conservation rules are relevant or the number of solutions the {class}`.StateTransitionManager` gives by default is to 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 & conservation rules are considered when fed back to the {class}`.StateTransitionManager`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "desired_edge_rules = {spin_validity}\n",
+    "desired_node_rules = {spin_magnitude_conservation, parity_conservation}\n",
+    "desired_edge_quantum_numbers = {\n",
+    "    EdgeQuantumNumbers.spin_magnitude,\n",
+    "    EdgeQuantumNumbers.parity,\n",
+    "}\n",
+    "desired_node_quantum_numbers = {\n",
+    "    NodeQuantumNumbers.l_magnitude,\n",
+    "    NodeQuantumNumbers.s_magnitude,\n",
+    "}\n",
+    "\n",
+    "filtered_qn_problem_set = filter_quantum_number_problem_set(\n",
+    "    qn_problem_set,\n",
+    "    desired_edge_rules,\n",
+    "    desired_node_rules,\n",
+    "    desired_edge_quantum_numbers,\n",
+    "    desired_node_quantum_numbers,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-cell"
+    ]
+   },
+   "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 statet at the top, a workflow similar to this will be used in future versions of ```qrules```.\n",
+    ":::"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "solver = CSPSolver([\n",
+    "    qrules.system_control.create_edge_properties(part)\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[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "dot = qrules.io.asdot(filtered_qn_result, render_node=True)\n",
+    "graphviz.Source(dot)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -672,7 +773,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

src/qrules/argument_handling.py

@@ -25,6 +25,7 @@
 Scalar = Union[int, float]
 
 Rule = Union[GraphElementRule, EdgeQNConservationRule, ConservationRule]
+"""Any type of rule"""
 
 _ElementType = TypeVar("_ElementType")
 

src/qrules/quantum_numbers.py

@@ -104,7 +104,6 @@ class EdgeQuantumNumbers:
         edge_qn_type.__module__ = __name__
 
 
-# for static typing
 EdgeQuantumNumber = Union[
     EdgeQuantumNumbers.pid,
     EdgeQuantumNumbers.mass,
@@ -126,8 +125,8 @@ class EdgeQuantumNumbers:
     EdgeQuantumNumbers.c_parity,
     EdgeQuantumNumbers.g_parity,
 ]
+"""Type hint for quantum numbers of edges"""
 
-# for accessing the keys of the dicts in EdgeSettings
 EdgeQuantumNumberTypes = Union[
     type[EdgeQuantumNumbers.pid],
     type[EdgeQuantumNumbers.mass],
@@ -149,6 +148,7 @@ class EdgeQuantumNumbers:
     type[EdgeQuantumNumbers.c_parity],
     type[EdgeQuantumNumbers.g_parity],
 ]
+"""Type-Union for accessing the keys of the dicts in EdgeSettings"""
 
 
 @frozen(init=False)
@@ -186,6 +186,7 @@ class NodeQuantumNumbers:
     type[NodeQuantumNumbers.s_projection],
     type[NodeQuantumNumbers.parity_prefactor],
 ]
+"""Type-Union for accessing the keys of the dicts in NodeSettings"""
 
 
 def _to_optional_float(optional_float: float | None) -> float | None:

src/qrules/solving.py

@@ -101,6 +101,86 @@ def topology(self) -> Topology:
         return self.initial_facts.topology
 
 
+def filter_quantum_number_problem_set(
+    quantum_number_problem_set: QNProblemSet,
+    edge_rules: set[GraphElementRule],
+    node_rules: set[Rule],
+    edge_properties_and_domains: Iterable[EdgeQuantumNumberTypes],
+    node_properties_and_domains: Iterable[NodeQuantumNumberTypes],
+) -> QNProblemSet:
+    """Filter `QNProblemSet` for desired conservation rules, settings and domains.
+
+    Currently it is the responsibility of the user to provide fitting properties
+    and domains for the correspinding conservation rules.
+
+    Args:
+        quantum_number_problem_set: `QNProblemSet` as generated by `CSPSolver`.
+        edge_rules: Conservation rules regarding the edges.
+        node_rules: Conservation rules regarding the nodes.
+        edge_properties_and_domains: Edge settings, properties and domains.
+        node_properties_and_domains: Node settings, properties and domains.
+    """
+    old_edge_settings = quantum_number_problem_set.solving_settings.states
+    old_node_settings = quantum_number_problem_set.solving_settings.interactions
+    old_edge_properties = quantum_number_problem_set.initial_facts.states
+    old_node_properties = quantum_number_problem_set.initial_facts.interactions
+    new_edge_settings = {
+        edge_id: EdgeSettings(
+            conservation_rules=edge_rules,
+            rule_priorities=edge_setting.rule_priorities,
+            qn_domains=({
+                key: val
+                for key, val in edge_setting.qn_domains.items()
+                if key in set(edge_properties_and_domains)
+            }),
+        )
+        for edge_id, edge_setting in old_edge_settings.items()
+    }
+    new_node_settings = {
+        node_id: NodeSettings(
+            conservation_rules=node_rules,
+            rule_priorities=node_setting.rule_priorities,
+            qn_domains=({
+                key: val
+                for key, val in node_setting.qn_domains.items()
+                if key in set(node_properties_and_domains)
+            }),
+        )
+        for node_id, node_setting in old_node_settings.items()
+    }
+    new_combined_settings = MutableTransition(
+        topology=quantum_number_problem_set.solving_settings.topology,
+        states=new_edge_settings,
+        interactions=new_node_settings,
+    )
+    new_edge_properties = {
+        edge_id: {
+            edge_quantum_number: scalar
+            for edge_quantum_number, scalar in graph_edge_property_map.items()
+            if edge_quantum_number in edge_properties_and_domains
+        }
+        for edge_id, graph_edge_property_map in old_edge_properties.items()
+    }
+    new_node_properties = {
+        node_id: {
+            node_quantum_number: scalar
+            for node_quantum_number, scalar in graph_node_property_map.items()
+            if node_quantum_number in node_properties_and_domains
+        }
+        for node_id, graph_node_property_map in old_node_properties.items()
+    }
+    new_combined_properties = MutableTransition(
+        topology=quantum_number_problem_set.initial_facts.topology,
+        states=new_edge_properties,
+        interactions=new_node_properties,
+    )
+    return attrs.evolve(
+        quantum_number_problem_set,
+        solving_settings=new_combined_settings,
+        initial_facts=new_combined_properties,
+    )
+
+
 QuantumNumberSolution = MutableTransition[GraphEdgePropertyMap, GraphNodePropertyMap]