added type alias for spaces (with states as a special case)

balm/SuccessionDiagram.py

@@ -19,6 +19,7 @@
 from balm.petri_net_translation import network_to_petrinet
 from balm.space_utils import percolate_space, space_unique_key
 from balm.trappist_core import trappist
+from balm.types import space_type
 
 # Enables helpful "progress" messages.
 DEBUG = False
@@ -136,7 +137,7 @@ def __setstate__(
         self.petri_net = cast(nx.DiGraph, state["petri net"])
         self.nfvs = cast(list[str], state["nfvs"])
         self.dag = cast(nx.DiGraph, state["G"])  # type: ignore
-        self.node_indices = cast(dict[str, int], state["node_indices"])  # type: ignore
+        self.node_indices = cast(space_type, state["node_indices"])  # type: ignore
 
     def __len__(self) -> int:
         """
@@ -173,7 +174,7 @@ def from_file(path: str) -> SuccessionDiagram:
         """
         return SuccessionDiagram(BooleanNetwork.from_file(path))
 
-    def expanded_attractor_seeds(self) -> list[list[dict[str, int]]]:
+    def expanded_attractor_seeds(self) -> list[list[space_type]]:
         return [self.node_attractor_seeds(id) for id in self.expanded_ids()]
 
     def summary(self) -> str:
@@ -262,7 +263,7 @@ def minimal_trap_spaces(self) -> list[int]:
         """
         return [i for i in self.expanded_ids() if self.node_is_minimal(i)]
 
-    def find_node(self, node_space: dict[str, int]) -> int | None:
+    def find_node(self, node_space: space_type) -> int | None:
         """
         Return the ID of the node matching the provided `node_space`, or `None`
         if no such node exists in this succession diagram.
@@ -335,14 +336,14 @@ def node_depth(self, node_id: int) -> int:
         """
         return cast(int, self.dag.nodes[node_id]["depth"])
 
-    def node_space(self, node_id: int) -> dict[str, int]:
+    def node_space(self, node_id: int) -> space_type:
         """
         Get the sub-space associated with the provided `node_id`.
 
         Note that this is the space *after* percolation. Hence it can hold that
         `|node_space(child)| < |node_space(parent)| + |stable_motif(parent, child)|`.
         """
-        return cast(dict[str, int], self.dag.nodes[node_id]["space"])
+        return cast(space_type, self.dag.nodes[node_id]["space"])
 
     def node_is_expanded(self, node_id: int) -> bool:
         """
@@ -387,7 +388,7 @@ def node_successors(self, node_id: int, compute: bool = False) -> list[int]:
 
     def node_attractor_seeds(
         self, node_id: int, compute: bool = False
-    ) -> list[dict[str, int]]:
+    ) -> list[space_type]:
         """
         Return the list of attractor seed states corresponding to the given
         `node_id`. Similar to `node_successors`, the method either computes the
@@ -401,7 +402,7 @@ def node_attractor_seeds(
         """
         node = cast(dict[str, Any], self.dag.nodes[node_id])
 
-        attractors = cast(list[dict[str, int]] | None, node["attractors"])
+        attractors = cast(list[space_type] | None, node["attractors"])
 
         if attractors is None and not compute:
             raise KeyError(f"Attractor data not computed for node {node_id}.")
@@ -414,7 +415,7 @@ def node_attractor_seeds(
 
     def edge_stable_motif(
         self, parent_id: int, child_id: int, reduced: bool = False
-    ) -> dict[str, int]:
+    ) -> space_type:
         """
         Return the *stable motif* associated with the specified parent-child
         edge. If `reduced` is set to `False` (default), the unpercolated stable
@@ -428,15 +429,15 @@ def edge_stable_motif(
 
         if reduced:
             return cast(
-                dict[str, int],
+                space_type,
                 {
                     k: v
                     for k, v in self.dag.edges[parent_id, child_id]["motif"].items()  # type: ignore
                     if k not in self.node_space(parent_id)
                 },
             )
         else:
-            return cast(dict[str, int], self.dag.edges[parent_id, child_id]["motif"])
+            return cast(space_type, self.dag.edges[parent_id, child_id]["motif"])
 
     def build(self):
         """
@@ -537,7 +538,7 @@ def expand_attractor_seeds(self, size_limit: int | None = None) -> bool:
         return expand_attractor_seeds(self, size_limit)
 
     def expand_to_target(
-        self, target: dict[str, int], size_limit: int | None = None
+        self, target: space_type, size_limit: int | None = None
     ) -> bool:
         """
         Expands the succession diagram using BFS in such a way that only nodes
@@ -618,7 +619,7 @@ def _expand_one_node(self, node_id: int):
             if DEBUG:
                 print(f"[{node_id}] Created edge into node {child_id}.")
 
-    def _ensure_node(self, parent_id: int | None, stable_motif: dict[str, int]) -> int:
+    def _ensure_node(self, parent_id: int | None, stable_motif: space_type) -> int:
         """
         Internal method that ensures the provided node is present in this
         succession diagram as a child of the given `parent_id`.

balm/_sd_algorithms/compute_attractor_seeds.py

@@ -10,12 +10,13 @@
 from balm.motif_avoidant import detect_motif_avoidant_attractors, make_retained_set
 from balm.terminal_restriction_space import get_terminal_restriction_space
 from balm.trappist_core import compute_fixed_point_reduced_STG
+from balm.types import space_type
 
 
 def compute_attractor_seeds(
     sd: SuccessionDiagram,
     node_id: int,
-) -> list[dict[str, int]]:
+) -> list[space_type]:
     """
     Compute the list of vertices such that each attractor within the subspace of
     the given `node_id` is covered by exactly one vertex.

balm/_sd_algorithms/expand_source_SCCs.py

@@ -13,6 +13,7 @@
 from balm.interaction_graph_utils import infer_signed_interaction_graph
 from balm.petri_net_translation import extract_variable_names, network_to_petrinet
 from balm.space_utils import percolate_network, percolate_space
+from balm.types import space_type
 
 if TYPE_CHECKING:
     expander_function_type = Callable[
@@ -71,7 +72,7 @@ def expand_source_SCCs(
     if len(source_nodes) != 0:
         bin_values_iter = it.product(range(2), repeat=len(source_nodes))
         for bin_values in bin_values_iter:
-            source_comb = dict(zip(source_nodes, bin_values))
+            source_comb = cast(space_type, dict(zip(source_nodes, bin_values)))
 
             sub_space = source_comb
             sub_space.update(perc_space)
@@ -89,7 +90,7 @@ def expand_source_SCCs(
 
         # each level consists of one round of fixing all source SCCs
         for node_id in current_level:
-            sub_space = cast(dict[str, int], sd.dag.nodes[node_id]["space"])
+            sub_space = cast(space_type, sd.dag.nodes[node_id]["space"])
 
             # find source SCCs
             clean_bnet, clean_bn = perc_and_remove_constants_from_bn(perc_bn, sub_space)
@@ -183,7 +184,7 @@ def find_source_nodes(network: BooleanNetwork | DiGraph) -> list[str]:
 
 
 def perc_and_remove_constants_from_bn(
-    bn: BooleanNetwork, space: dict[str, int]
+    bn: BooleanNetwork, space: space_type
 ) -> tuple[str, BooleanNetwork]:
     """
     Take a BooleanNetwork and percolate given space.
@@ -317,11 +318,11 @@ def find_scc_sd(
     # delete the implicit parameters from the node subspaces and the edge motifs
     for node_id in scc_sd.node_ids():
         for implicit in implicit_parameters:
-            cast(dict[str, int], scc_sd.dag.nodes[node_id]["space"]).pop(implicit, None)
+            cast(space_type, scc_sd.dag.nodes[node_id]["space"]).pop(implicit, None)
 
     for x, y in cast(Iterable[tuple[int, int]], scc_sd.dag.edges):
         for implicit in implicit_parameters:
-            cast(dict[str, int], scc_sd.dag.edges[x, y]["motif"]).pop(implicit, None)
+            cast(space_type, scc_sd.dag.edges[x, y]["motif"]).pop(implicit, None)
 
     return scc_sd, exist_maa
 
@@ -363,7 +364,7 @@ def attach_scc_sd(
             parent_id = size_before_attach + scc_parent_id - 1
 
         motif = scc_sd.edge_stable_motif(scc_parent_id, scc_node_id)
-        motif.update(cast(dict[str, int], sd.dag.nodes[branch]["space"]))
+        motif.update(cast(space_type, sd.dag.nodes[branch]["space"]))
 
         child_id = sd._ensure_node(parent_id, motif)  # type: ignore
         if check_maa:
@@ -384,7 +385,7 @@ def attach_scc_sd(
         scc_child_ids = cast(list[int], list(scc_sd.dag.successors(scc_node_id)))  # type: ignore
         for scc_child_id in scc_child_ids:
             motif = scc_sd.edge_stable_motif(scc_node_id, scc_child_id)
-            motif.update(cast(dict[str, int], sd.dag.nodes[branch]["space"]))
+            motif.update(cast(space_type, sd.dag.nodes[branch]["space"]))
 
             child_id = sd._ensure_node(parent_id, motif)  # type: ignore
             assert child_id == size_before_attach + scc_child_id - 1

balm/_sd_algorithms/expand_to_target.py

@@ -6,10 +6,11 @@
     from balm.SuccessionDiagram import SuccessionDiagram
 
 from balm.space_utils import intersect, is_subspace
+from balm.types import space_type
 
 
 def expand_to_target(
-    sd: SuccessionDiagram, target: dict[str, int], size_limit: int | None = None
+    sd: SuccessionDiagram, target: space_type, size_limit: int | None = None
 ):
     """
     See `SuccessionDiagram.exapnd_to_target` for documentation.

balm/control.py

@@ -1,16 +1,14 @@
 from __future__ import annotations
 
 from itertools import combinations, product
-from typing import cast
+from typing import Literal, cast
 
 import networkx as nx  # type: ignore
 from biodivine_aeon import BooleanNetwork
 
 from balm.space_utils import is_subspace, percolate_space
 from balm.SuccessionDiagram import SuccessionDiagram
-
-SuccessionType = list[dict[str, int]]  # sequence of stable motifs
-ControlType = list[dict[str, int]]  # ways of locking in an individual stable motif
+from balm.types import ControlType, SuccessionType, space_type
 
 
 def controls_are_equal(a: ControlType, b: ControlType) -> bool:
@@ -107,7 +105,7 @@ def __str__(self):
 
 def succession_control(
     bn: BooleanNetwork,
-    target: dict[str, int],
+    target: space_type,
     strategy: str = "internal",
     succession_diagram: SuccessionDiagram | None = None,
     max_drivers_per_succession_node: int | None = None,
@@ -120,7 +118,7 @@ def succession_control(
     ----------
     bn : BooleanNetwork
         The network to analyze, which contains the Boolean update functions.
-    target : dict[str, int]
+    target : space_type
         The target subspace.
     strategy : str, optional
         The searching strategy to use to look for driver nodes. Options are
@@ -174,7 +172,7 @@ def succession_control(
 
 def successions_to_target(
     succession_diagram: SuccessionDiagram,
-    target: dict[str, int],
+    target: space_type,
     expand_diagram: bool = True,
 ) -> list[SuccessionType]:
     """Find lists of nested trap spaces (successions) that lead to the
@@ -184,7 +182,7 @@ def successions_to_target(
     ----------
     succession_diagram : SuccessionDiagram
         The succession diagram from which successions will be extracted.
-    target : dict[str, int]
+    target : space_type
         The target subspace.
     expand_diagram: bool
         Whether to ensure that the succession diagram is expanded enough to
@@ -228,7 +226,7 @@ def successions_to_target(
 
 def drivers_of_succession(
     bn: BooleanNetwork,
-    succession: list[dict[str, int]],
+    succession: list[space_type],
     strategy: str = "internal",
     max_drivers_per_succession_node: int | None = None,
     forbidden_drivers: set[str] | None = None,
@@ -239,7 +237,7 @@ def drivers_of_succession(
     ----------
     bn : BooleanNetwork
         The network to analyze, which contains the Boolean update functions.
-    succession : list[dict[str, int]]
+    succession : list[space_type]
         A list of sequentially nested trap spaces that specify the target.
     strategy: str
         The searching strategy to use to look for driver nodes. Options are
@@ -260,7 +258,7 @@ def drivers_of_succession(
         of drivers for the corresponding trap space in the succession.
     """
     control_strategies: list[ControlType] = []
-    assume_fixed: dict[str, int] = {}
+    assume_fixed: space_type = {}
     for ts in succession:
         control_strategies.append(
             find_drivers(
@@ -280,9 +278,9 @@ def drivers_of_succession(
 
 def find_drivers(
     bn: BooleanNetwork,
-    target_trap_space: dict[str, int],
+    target_trap_space: space_type,
     strategy: str = "internal",
-    assume_fixed: dict[str, int] | None = None,
+    assume_fixed: space_type | None = None,
     max_drivers_per_succession_node: int | None = None,
     forbidden_drivers: set[str] | None = None,
 ) -> ControlType:
@@ -292,7 +290,7 @@ def find_drivers(
     ----------
     bn : BooleanNetwork
         The network to analyze, which contains the Boolean update functions.
-    target_trap_space : dict[str, int]
+    target_trap_space : space_type
         The trap space we want to find drivers for.
     strategy: str
         The searching strategy to use to look for driver nodes. Options are
@@ -343,7 +341,10 @@ def find_drivers(
                 continue
 
             if strategy == "internal":
-                driver_dict = {k: target_trap_space_inner[k] for k in driver_set}
+                driver_dict: space_type = {
+                    k: cast(Literal[0, 1], target_trap_space_inner[k])
+                    for k in driver_set
+                }
                 ldoi = percolate_space(
                     bn, driver_dict | assume_fixed, strict_percolation=False
                 )
@@ -352,7 +353,8 @@ def find_drivers(
             elif strategy == "all":
                 for vals in product([0, 1], repeat=driver_set_size):
                     driver_dict = {
-                        driver: value for driver, value in zip(driver_set, vals)
+                        driver: cast(Literal[0, 1], value)
+                        for driver, value in zip(driver_set, vals)
                     }
                     ldoi = percolate_space(
                         bn, driver_dict | assume_fixed, strict_percolation=False

balm/drivers.py

@@ -1,16 +1,19 @@
 from __future__ import annotations
 
+from typing import Literal, cast
+
 from biodivine_aeon import BooleanNetwork
 
 from balm.space_utils import percolate_space
+from balm.types import space_type
 
 
-def find_single_node_LDOIs(bn: BooleanNetwork) -> dict[tuple[str, int], dict[str, int]]:
+def find_single_node_LDOIs(bn: BooleanNetwork) -> dict[tuple[str, int], space_type]:
     """
     finds LDOIs of every single node state
     TODO: take an initial set of LDOIs (e.g., of the original system) as an argument for speed-up
     """
-    LDOIs: dict[tuple[str, int], dict[str, int]] = {}
+    LDOIs: dict[tuple[str, int], space_type] = {}
     for var in bn.variables():
         name = bn.get_variable_name(var)
         function = bn.get_update_function(var)
@@ -19,16 +22,16 @@ def find_single_node_LDOIs(bn: BooleanNetwork) -> dict[tuple[str, int], dict[str
             continue
         for i in range(2):
             fix = (name, i)
-            space = {name: i}
+            space: space_type = {name: cast(Literal[0, 1], i)}
             LDOIs[fix] = percolate_space(bn, space)
 
     return LDOIs
 
 
 def find_single_drivers(
-    target_subspace: dict[str, int],
+    target_subspace: space_type,
     bn: BooleanNetwork,
-    LDOIs: dict[tuple[str, int], dict[str, int]] | None = None,
+    LDOIs: dict[tuple[str, int], space_type] | None = None,
 ) -> set[tuple[str, int]]:
     """
     find all the single node drivers for a given target_subspace,