small bug fix around node name generator during HybridSolver

YosefLab · Sep 9, 2024 · c568a72 · c568a72
1 parent 78442c4
commit c568a72
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/cassiopeia/solver/HybridSolver.py b/cassiopeia/solver/HybridSolver.py
@@ -200,6 +200,9 @@ def solve(
             for n in subproblem_tree:
                 if n in existing_nodes and n != subproblem_root:
                     mapping[n] = next(node_name_generator)
+                    existing_nodes.append(mapping[n])
+                else:
+                    existing_nodes.append(n)
 
             subproblem_tree = nx.relabel_nodes(subproblem_tree, mapping)
 

diff --git a/cassiopeia/solver/ILPSolver.py b/cassiopeia/solver/ILPSolver.py
@@ -265,10 +265,15 @@ def solve(
         # rename internal nodes such that they are not tuples
         node_name_generator = solver_utilities.node_name_generator()
         internal_node_rename = {}
+        seen_nodes = []
         for i in cassiopeia_tree.internal_nodes:
-            internal_node_rename[i] = next(node_name_generator)
+            new_node_name = next(node_name_generator)
+            while new_node_name in seen_nodes:
+                new_node_name = next(node_name_generator)
+            internal_node_rename[i] = new_node_name
+            seen_nodes.append(internal_node_rename[i])
         cassiopeia_tree.relabel_nodes(internal_node_rename)
-
+        
         cassiopeia_tree.collapse_unifurcations()
 
         # collapse mutationless edges