From 25dd68be4e5020b3faee7df46d1b5ebc981aba4d Mon Sep 17 00:00:00 2001
From: dengzq1234 <dengziqi1234@gmail.com>
Date: Wed, 3 Apr 2024 11:48:51 +0200
Subject: [PATCH] improve autopruning by finding lca

---
 treeprofiler/src/utils.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/treeprofiler/src/utils.py b/treeprofiler/src/utils.py
index 5167e2c..8485fcf 100644
--- a/treeprofiler/src/utils.py
+++ b/treeprofiler/src/utils.py
@@ -204,11 +204,21 @@ def ete4_parse(newick, internal_parser="name"):
 # pruning
 def taxatree_prune(tree, rank_limit='subspecies'):
     for node in tree.traverse("preorder"):
-        if node.props.get('rank') == rank_limit:
+        rank = node.props.get('rank')
+        if rank == rank_limit:
             children = node.children.copy()
             for ch in children:
                 print("prune", ch.name)
                 remove(ch)
+        lca_dict = node.props.get('lca')
+        if lca_dict:
+            lca = lca_dict.get(rank_limit, None)
+            if lca:
+                node.name = lca
+                children = node.children.copy()
+                for ch in children:
+                    print("prune", ch.name)
+                    remove(ch)
     return tree
 
 def conditional_prune(tree, conditions_input, prop2type):