iss272: faulty filters for QNProblemSet

tests/unit/test_solving.py

@@ -8,17 +8,18 @@
 
 import qrules.particle
 import qrules.quantum_numbers
-import qrules.solving
 import qrules.system_control
 import qrules.transition
 from qrules.conservation_rules import (
     GraphElementRule,
-    isospin_validity,
+    c_parity_conservation,
     parity_conservation,
     spin_magnitude_conservation,
     spin_validity,
 )
 from qrules.quantum_numbers import EdgeQuantumNumbers, NodeQuantumNumbers
+from qrules.solving import CSPSolver, EdgeSettings, NodeSettings, QNProblemSet
+from qrules.topology import MutableTransition
 
 if TYPE_CHECKING:
     from qrules.argument_handling import (
@@ -28,48 +29,138 @@
     )
 
 
-def test_solve_with_new_quantum_number_problem_set(
-    particle_database: qrules.particle.ParticleCollection,
-    qn_problem_set: qrules.solving.QNProblemSet,
+def test_solve(
+    all_particles: qrules.particle.ParticleCollection,
+    quantum_number_problem_set: QNProblemSet,
 ) -> None:
-    allowed_intermediate_states = [
-        qrules.system_control.create_edge_properties(part) for part in particle_database
-    ]
-    solver = qrules.solving.CSPSolver(allowed_intermediate_states)
-    new_qn_problem_set = quantum_number_problem_set_with_new_properties(
-        qn_problem_set,
-        graph_edge_property_map={
-            EdgeQuantumNumbers.spin_magnitude: 1,
-            EdgeQuantumNumbers.parity: -1,
-            EdgeQuantumNumbers.c_parity: 1,
-        },
-        graph_node_property_map={
-            NodeQuantumNumbers.s_magnitude: 1,
-            NodeQuantumNumbers.l_magnitude: 0,
+    solver = CSPSolver(all_particles)
+    result = solver.find_solutions(quantum_number_problem_set)
+    assert len(result.solutions) != 0
+
+
+def test_solve_with_filtered_quantum_number_problem_set(
+    all_particles: qrules.particle.ParticleCollection,
+    quantum_number_problem_set: QNProblemSet,
+) -> None:
+    solver = CSPSolver(all_particles)
+    new_quantum_number_problem_set = filter_quantum_number_problem_set_settings(
+        quantum_number_problem_set,
+        edge_rules={spin_validity},
+        node_rules={
+            spin_magnitude_conservation,
+            parity_conservation,
+            c_parity_conservation,
         },
+        edge_domains=(
+            EdgeQuantumNumbers.spin_magnitude,
+            EdgeQuantumNumbers.parity,
+            EdgeQuantumNumbers.c_parity,
+        ),
+        node_domains=(NodeQuantumNumbers.l_magnitude, NodeQuantumNumbers.s_magnitude),
+    )
+    new_quantum_number_problem_set = filter_quantum_number_problem_set_properties(
+        new_quantum_number_problem_set,
+        edge_properties=(
+            EdgeQuantumNumbers.spin_magnitude,
+            EdgeQuantumNumbers.parity,
+            EdgeQuantumNumbers.c_parity,
+        ),
+        node_properties=(
+            NodeQuantumNumbers.l_magnitude,
+            NodeQuantumNumbers.s_magnitude,
+        ),
+    )
+    result = solver.find_solutions(new_quantum_number_problem_set)
+
+    assert len(result.solutions) != 0
+
+
+def remove_quantum_number_problem_set_settings(
+    quantum_number_problem_set: QNProblemSet,
+    edge_rules_to_be_removed: set[GraphElementRule],
+    node_rules_to_be_removed: set[Rule],
+    edge_domains_to_be_removed: tuple[Any, ...],
+    node_domains_to_be_removed: tuple[Any, ...],
+) -> QNProblemSet:
+    old_edge_settings = quantum_number_problem_set.solving_settings.states
+    old_node_settings = quantum_number_problem_set.solving_settings.interactions
+    new_edge_settings = {
+        edge_id: EdgeSettings(
+            conservation_rules=edge_setting.conservation_rules
+            - edge_rules_to_be_removed,
+            rule_priorities=edge_setting.rule_priorities,
+            qn_domains={
+                key: val
+                for key, val in edge_setting.qn_domains.items()
+                if key not in edge_domains_to_be_removed
+            },
+        )
+        for edge_id, edge_setting in old_edge_settings.items()
+    }
+    new_node_settings = {
+        node_id: NodeSettings(
+            conservation_rules=node_setting.conservation_rules
+            - node_rules_to_be_removed,
+            rule_priorities=node_setting.rule_priorities,
+            qn_domains={
+                key: val
+                for key, val in node_setting.qn_domains.items()
+                if key not in node_domains_to_be_removed
+            },
+        )
+        for node_id, node_setting in old_node_settings.items()
+    }
+    new_mutable_transition = MutableTransition(
+        topology=quantum_number_problem_set.solving_settings.topology,
+        states=new_edge_settings,
+        interactions=new_node_settings,
     )
-    new_qn_problem_set = quantum_number_problem_set_with_new_settings(
-        new_qn_problem_set,
-        edge_rules={spin_validity, isospin_validity},
-        node_rules={spin_magnitude_conservation, parity_conservation},
-        edge_domains={"spin_magnitude": [1], "parity": [-1, 1], "c_parity": [-1, 1]},
-        node_domains={"l_projection": [0]},
+    return attrs.evolve(
+        quantum_number_problem_set, solving_settings=new_mutable_transition
     )
-    solutions = solver.find_solutions(new_qn_problem_set)
 
-    assert isinstance(solutions, qrules.solving.QNResult)
+
+def remove_quantum_number_problem_set_properties(
+    quantum_number_problem_set: QNProblemSet,
+    edge_properties_to_be_removed: tuple[EdgeQuantumNumbers],
+    node_properties_to_be_removed: tuple[NodeQuantumNumbers],
+) -> QNProblemSet:
+    old_edge_properties = quantum_number_problem_set.initial_facts.states
+    old_node_properties = quantum_number_problem_set.initial_facts.interactions
+    new_edge_properties = {
+        edge_id: {
+            edge_quantum_number: scalar
+            for edge_quantum_number, scalar in graph_edge_property_map.items()
+            if edge_quantum_number not in edge_properties_to_be_removed
+        }
+        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 not in node_properties_to_be_removed
+        }
+        for node_id, graph_node_property_map in old_node_properties.items()
+    }
+    new_mutable_transition = MutableTransition(
+        topology=quantum_number_problem_set.initial_facts.topology,
+        states=new_edge_properties,
+        interactions=new_node_properties,
+    )
+    return attrs.evolve(
+        quantum_number_problem_set, initial_facts=new_mutable_transition
+    )
 
 
 def test_inner_dicts_unchanged(
-    qn_problem_set: qrules.solving.QNProblemSet,
-    graph_edge_property_map: GraphEdgePropertyMap,
-    graph_node_property_map: GraphNodePropertyMap,
+    quantum_number_problem_set: QNProblemSet,
 ) -> None:
     old_inner_graph_edge_property_map = copy.deepcopy(
-        qn_problem_set.initial_facts.states
+        quantum_number_problem_set.initial_facts.states
     )
     old_inner_graph_node_property_map = copy.deepcopy(
-        qn_problem_set.initial_facts.interactions
+        quantum_number_problem_set.initial_facts.interactions
     )
     graph_edge_property_map = {
         EdgeQuantumNumbers.spin_magnitude: 1,
@@ -81,29 +172,118 @@ def test_inner_dicts_unchanged(
         NodeQuantumNumbers.l_magnitude: 0,
     }
     quantum_number_problem_set_with_new_properties(
-        qn_problem_set, graph_edge_property_map, graph_node_property_map
+        quantum_number_problem_set, graph_edge_property_map, graph_node_property_map
+    )
+    assert (
+        old_inner_graph_edge_property_map
+        == quantum_number_problem_set.initial_facts.states
     )
-    assert old_inner_graph_edge_property_map == qn_problem_set.initial_facts.states
     assert (
-        old_inner_graph_node_property_map == qn_problem_set.initial_facts.interactions
+        old_inner_graph_node_property_map
+        == quantum_number_problem_set.initial_facts.interactions
+    )
+
+
+def filter_quantum_number_problem_set_settings(
+    quantum_number_problem_set: QNProblemSet,
+    edge_rules: set[GraphElementRule],
+    node_rules: set[Rule],
+    edge_domains: tuple[Any, ...],
+    node_domains: tuple[Any, ...],
+    keep_domains: bool = True,
+) -> QNProblemSet:
+    old_edge_settings = quantum_number_problem_set.solving_settings.states
+    old_node_settings = quantum_number_problem_set.solving_settings.interactions
+    new_edge_settings = {
+        edge_id: EdgeSettings(
+            conservation_rules=edge_rules,
+            rule_priorities=edge_setting.rule_priorities,
+            qn_domains=(
+                edge_setting.qn_domains
+                if keep_domains
+                else {
+                    key: val
+                    for key, val in edge_setting.qn_domains.items()
+                    if key in edge_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=(
+                node_setting.qn_domains
+                if keep_domains
+                else {
+                    key: val
+                    for key, val in node_setting.qn_domains.items()
+                    if key in node_domains
+                }
+            ),
+        )
+        for node_id, node_setting in old_node_settings.items()
+    }
+    new_mutable_transition = MutableTransition(
+        topology=quantum_number_problem_set.solving_settings.topology,
+        states=new_edge_settings,
+        interactions=new_node_settings,
+    )
+    return attrs.evolve(
+        quantum_number_problem_set, solving_settings=new_mutable_transition
+    )
+
+
+def filter_quantum_number_problem_set_properties(
+    quantum_number_problem_set: QNProblemSet,
+    edge_properties: tuple[EdgeQuantumNumbers],
+    node_properties: tuple[NodeQuantumNumbers],
+) -> QNProblemSet:
+    old_edge_properties = quantum_number_problem_set.initial_facts.states
+    old_node_properties = quantum_number_problem_set.initial_facts.interactions
+    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
+        }
+        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
+        }
+        for node_id, graph_node_property_map in old_node_properties.items()
+    }
+    new_mutable_transition = MutableTransition(
+        topology=quantum_number_problem_set.initial_facts.topology,
+        states=new_edge_properties,
+        interactions=new_node_properties,
+    )
+    return attrs.evolve(
+        quantum_number_problem_set, initial_facts=new_mutable_transition
     )
 
 
 def quantum_number_problem_set_with_new_settings(
-    qn_problem_set: qrules.solving.QNProblemSet,
+    quantum_number_problem_set: QNProblemSet,
     edge_rules: set[GraphElementRule],
     node_rules: set[Rule],
     edge_domains: dict[Any, list],
     node_domains: dict[Any, list],
-) -> qrules.solving.QNProblemSet:
+) -> QNProblemSet:
     def qnp_with_new_rules(
-        qn_problem_set: qrules.solving.QNProblemSet,
+        quantum_number_problem_set: QNProblemSet,
         edge_rules: set[GraphElementRule],
         node_rules: set[Rule],
-    ) -> qrules.solving.QNProblemSet:
-        old_settings = qn_problem_set.solving_settings
+    ) -> QNProblemSet:
+        old_settings = quantum_number_problem_set.solving_settings
         new_settings = attrs.evolve(
-            qn_problem_set.solving_settings,
+            quantum_number_problem_set.solving_settings,
             states={
                 edge_id: attrs.evolve(
                     setting, conservation_rules=setting.conservation_rules & edge_rules
@@ -117,14 +297,14 @@ def qnp_with_new_rules(
                 for node_id, setting in old_settings.interactions.items()
             },
         )
-        return attrs.evolve(qn_problem_set, solving_settings=new_settings)
+        return attrs.evolve(quantum_number_problem_set, solving_settings=new_settings)
 
     def qnp_with_new_domains(
-        qn_problem_set: qrules.solving.QNProblemSet,
+        quantum_number_problem_set: QNProblemSet,
         edge_domains: dict[Any, list],
         node_domains: dict[Any, list],
-    ) -> qrules.solving.QNProblemSet:
-        old_settings = qn_problem_set.solving_settings
+    ) -> QNProblemSet:
+        old_settings = quantum_number_problem_set.solving_settings
         new_settings = attrs.evolve(
             old_settings,
             states={
@@ -136,21 +316,21 @@ def qnp_with_new_domains(
                 for node_id, setting in old_settings.interactions.items()
             },
         )
-        return attrs.evolve(qn_problem_set, solving_settings=new_settings)
+        return attrs.evolve(quantum_number_problem_set, solving_settings=new_settings)
 
     return qnp_with_new_rules(
-        qnp_with_new_domains(qn_problem_set, edge_domains, node_domains),
+        qnp_with_new_domains(quantum_number_problem_set, edge_domains, node_domains),
         edge_rules,
         node_rules,
     )
 
 
 def quantum_number_problem_set_with_new_properties(
-    qn_problem_set: qrules.solving.QNProblemSet,
+    quantum_number_problem_set: QNProblemSet,
     graph_edge_property_map: GraphEdgePropertyMap,
     graph_node_property_map: GraphNodePropertyMap,
-) -> qrules.solving.QNProblemSet:
-    old_facts = qn_problem_set.initial_facts
+) -> QNProblemSet:
+    old_facts = quantum_number_problem_set.initial_facts
     new_facts = attrs.evolve(
         old_facts,
         states={
@@ -170,11 +350,19 @@ def quantum_number_problem_set_with_new_properties(
             for node_id, prop_map in old_facts.interactions.items()
         },
     )
-    return attrs.evolve(qn_problem_set, initial_facts=new_facts)
+    return attrs.evolve(quantum_number_problem_set, initial_facts=new_facts)
+
+
+@pytest.fixture(scope="session")
+def all_particles():
+    return [
+        qrules.system_control.create_edge_properties(part)
+        for part in qrules.particle.load_pdg()
+    ]
 
 
 @pytest.fixture(scope="session")
-def qn_problem_set() -> qrules.solving.QNProblemSet:
+def quantum_number_problem_set(request) -> QNProblemSet:
     stm = qrules.StateTransitionManager(
         initial_state=["psi(2S)"],
         final_state=["gamma", "eta", "eta"],