From 0cc7ddb4c19f7d4262e501de92742fedb54ea6a6 Mon Sep 17 00:00:00 2001
From: kyuhyongpark <kjp5774@psu.edu>
Date: Wed, 20 Dec 2023 11:00:23 -0500
Subject: [PATCH] allow expand min and att to start from non-root nodes

---
 balm/SuccessionDiagram.py                     | 16 +++++++++----
 balm/_sd_algorithms/expand_attractor_seeds.py | 17 ++++++++++----
 balm/_sd_algorithms/expand_minimal_spaces.py  | 23 ++++++++++++++-----
 3 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/balm/SuccessionDiagram.py b/balm/SuccessionDiagram.py
index 8e6c793d..6f3977b0 100644
--- a/balm/SuccessionDiagram.py
+++ b/balm/SuccessionDiagram.py
@@ -504,7 +504,11 @@ def expand_dfs(
         """
         return expand_dfs(self, node_id, dfs_stack_limit, size_limit)
 
-    def expand_minimal_spaces(self, size_limit: int | None = None) -> bool:
+    def expand_minimal_spaces(
+        self,
+        node_id: int | None = None,
+        size_limit: int | None = None
+    ) -> bool:
         """
         Expands the succession diagram in a way that guarantees every minimal
         trap space to be reachable from the root node, but otherwise (greedily)
@@ -520,9 +524,13 @@ def expand_minimal_spaces(self, size_limit: int | None = None) -> bool:
         procedure still tries to avoid expanding unnecessary nodes, which means
         existing expanded nodes can be prioritised over the "canonical" ones.
         """
-        return expand_minimal_spaces(self, size_limit)
+        return expand_minimal_spaces(self, node_id, size_limit)
 
-    def expand_attractor_seeds(self, size_limit: int | None = None) -> bool:
+    def expand_attractor_seeds(
+        self,
+        node_id: int | None = None,
+        size_limit: int | None = None
+    ) -> bool:
         """
         Expands the succession diagram such that for every asynchronous
         attractor, there is at least one expanded trap space which is the
@@ -535,7 +543,7 @@ def expand_attractor_seeds(self, size_limit: int | None = None) -> bool:
         identification. It is possible that some nodes are expanded spuriously
         and the succession diagram is thus larger than necessary.
         """
-        return expand_attractor_seeds(self, size_limit)
+        return expand_attractor_seeds(self, node_id, size_limit)
 
     def expand_to_target(
         self, target: BooleanSpace, size_limit: int | None = None
diff --git a/balm/_sd_algorithms/expand_attractor_seeds.py b/balm/_sd_algorithms/expand_attractor_seeds.py
index ab360811..aaa425fb 100644
--- a/balm/_sd_algorithms/expand_attractor_seeds.py
+++ b/balm/_sd_algorithms/expand_attractor_seeds.py
@@ -12,7 +12,9 @@
 from balm.trappist_core import compute_fixed_point_reduced_STG
 
 
-def expand_attractor_seeds(sd: SuccessionDiagram, size_limit: int | None = None):
+def expand_attractor_seeds(sd: SuccessionDiagram,
+                           node_id: int | None = None,
+                           size_limit: int | None = None):
     """
     See `SuccessionDiagram.expand_attractor_seeds` for documentation.
     """
@@ -22,16 +24,21 @@ def expand_attractor_seeds(sd: SuccessionDiagram, size_limit: int | None = None)
     # because for every attractor in a minimal trap space, we already have the
     # closest trap space, now we just need to do the same for (potential)
     # motif-avoidant attractors.
-    sd.expand_minimal_spaces(size_limit)
+
+    if node_id is None:
+        node_id = sd.root()
+
+    sd.expand_minimal_spaces(node_id, size_limit)
 
     if balm.SuccessionDiagram.DEBUG:
         print(
             "Minimal trap space expansion finished. Proceeding to attractor expansion."
         )
 
-    root = sd.root()
-    seen = set([root])
-    stack: list[tuple[int, list[int] | None]] = [(root, None)]
+    seen: set[int] = set()
+    seen.add(node_id)
+
+    stack: list[tuple[int, list[int] | None]] = [(node_id, None)]
 
     while len(stack) > 0:
         (node, successors) = stack.pop()
diff --git a/balm/_sd_algorithms/expand_minimal_spaces.py b/balm/_sd_algorithms/expand_minimal_spaces.py
index 772ab8df..28234928 100644
--- a/balm/_sd_algorithms/expand_minimal_spaces.py
+++ b/balm/_sd_algorithms/expand_minimal_spaces.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, cast
 
 if TYPE_CHECKING:
     from balm.SuccessionDiagram import SuccessionDiagram
@@ -9,18 +9,29 @@
 from balm.trappist_core import trappist
 
 
-def expand_minimal_spaces(sd: SuccessionDiagram, size_limit: int | None = None) -> bool:
+def expand_minimal_spaces(
+    sd: SuccessionDiagram,
+    node_id: int | None = None,
+    size_limit: int | None = None
+) -> bool:
     """
     See `SuccessionDiagram.expand_minimal_spaces` for documentation.
     """
 
-    minimal_traps = trappist(sd.petri_net, problem="min")
+    if node_id is None:
+        node_id = sd.root()
 
-    root = sd.root()
+    node = cast(dict[str, Any], sd.dag.nodes[node_id])
+    
+    current_space = node["space"]
 
-    seen = set([root])
+    minimal_traps = trappist(sd.petri_net, problem="min",
+                             ensure_subspace=current_space)
 
-    stack: list[tuple[int, list[int] | None]] = [(root, None)]
+    seen: set[int] = set()
+    seen.add(node_id)
+
+    stack: list[tuple[int, list[int] | None]] = [(node_id, None)]
 
     while len(stack) > 0:
         (node, successors) = stack.pop()