Collapse rank-2 vertices after pruning.

src/nni/startle.py

@@ -485,16 +485,19 @@ def infer_mode(args):
     random.seed(73)
 
     seed_tree = from_newick_get_nx_tree(args.seed_tree)
-    seed_tree = arbitrarily_resolve_polytomies(seed_tree)
+    logger.info(f"Unpruned tree size: {len(seed_tree.nodes)}")
 
     character_matrix = pd.read_csv(args.character_matrix, index_col=[0], dtype=str)
     character_matrix = character_matrix.replace('-', '-1')
     character_matrix.index = character_matrix.index.map(str)
 
-    logger.info(f"Unpruned tree size: {len(seed_tree.nodes)}")
+    # Prune tree and character matrix to only include distinct cells
     eq_class_dict, character_matrix = compute_equivalence_classes(character_matrix)
-    seed_tree = prune_tree(seed_tree, set(eq_class_dict.keys()))
     character_matrix = character_matrix.astype(str)
+
+    seed_tree = prune_tree(seed_tree, set(eq_class_dict.keys()))
+    seed_tree = arbitrarily_resolve_polytomies(seed_tree)
+
     logger.info(f"Pruned tree size: {len(seed_tree.nodes)}")
 
     mutation_prior_dict = load_mutation_prior_dict(args)

src/nni/utilities.py

@@ -294,6 +294,16 @@ def prune_tree(tree, eq_class_leaders):
         for leaf in leaves_to_prune:
             T.remove_node(leaf)
 
+    # remove all unifurcations
+    while True:
+        unifurcations = [v for v in T.nodes if T.out_degree(v) == 1 and T.in_degree(v) == 1]
+        if not unifurcations: break
+        for unifurcation in unifurcations:
+            parent = list(T.predecessors(unifurcation))[0]
+            child = list(T.successors(unifurcation))[0]
+            T.add_edge(parent, child)
+            T.remove_node(unifurcation)
+
     return T