Merge pull request #129 from dengzq1234/main

update

Author: dengzq1234

treeprofiler/layouts/staple_layouts.py

@@ -287,17 +287,20 @@ def __init__(self, name=None, column=0, width=70, height=None,
 
     def set_tree_style(self, tree, tree_style):
         super().set_tree_style(tree, tree_style)
-
+        first_key = min(self.color_range.keys())
+        last_key = max(self.color_range.keys())
+        middle_key = sorted(self.color_range.keys())[len(self.color_range) // 2]
+        
         text = TextFace(self.prop,  padding_x=self.padding_x, width=self.width, rotation=315)
         tree_style.aligned_panel_header.add_face(text, column=self.column)
         if self.legend:
             tree_style.add_legend(title=self.prop,
                                     variable='continuous',
                                     value_range=self.value_range ,
                                     color_range=[
-                                        self.color_range.get(20),
-                                        self.color_range.get(10),
-                                        self.color_range.get(1),
+                                        self.color_range.get(last_key),
+                                        self.color_range.get(middle_key),
+                                        self.color_range.get(first_key),
                                     ]
                                     )
     def set_node_style(self, node):
@@ -397,8 +400,8 @@ def set_tree_style(self, tree, tree_style):
                                     variable='continuous',
                                     value_range=[self.minval, self.maxval],
                                     color_range=[
-                                        self.color_dict[20], 
-                                        self.color_dict[10], 
+                                        self.color_dict[-1], 
+                                        self.color_dict[len(self.color_dict)//2], 
                                         self.color_dict[1]
                                     ]
                                     )

treeprofiler/src/b64pickle.py treeprofiler/src/ete_format.pytreeprofiler/src/b64pickle.py renamed to treeprofiler/src/ete_format.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from treeprofiler.src import b64pickle
+from treeprofiler.src import ete_format
 from ete4.parser.newick import NewickError
 from ete4.core.operations import remove
 from ete4 import Tree, PhyloTree
@@ -160,7 +160,7 @@ def get_consensus_seq(matrix_string: Path | str, threshold=0.7) -> SeqRecord:
     consensus = summary.dumb_consensus(threshold, "-")
     return consensus
 
-def counter2ratio(node, prop, minimum=0.05):
+def counter2ratio(node, prop, minimum=0.01):
     counter_separator = '||'
     items_separator = '--'
     count_missing = True
@@ -188,11 +188,11 @@ def counter2ratio(node, prop, minimum=0.05):
         ratio = 0
 
     if ratio < minimum and ratio != 0: # show minimum color for too low
-        ratio = 0.05
+        ratio = minimum
 
     return ratio
 
-def categorical2ratio(node, prop, all_values, minimum=0.05):
+def categorical2ratio(node, prop, all_values, minimum=0.01):
     counter_separator = '||'
     items_separator = '--'
     count_missing = True
@@ -210,7 +210,7 @@ def categorical2ratio(node, prop, all_values, minimum=0.05):
             positive = 0
         ratio = positive / total
         if ratio < minimum and ratio != 0: # show minimum color for too low
-            ratio = 0.05
+            ratio = minimum
         ratios.append(ratio)
 
     return ratios
@@ -240,7 +240,7 @@ def validate_tree(tree_path, input_type, internal_parser=None):
         try:
             with open(tree_path, 'r') as f:
                 file_content = f.read()
-            tree = b64pickle.loads(file_content, encoder='pickle', unpack=False)
+            tree = ete_format.loads(file_content, encoder='pickle', unpack=False)
             eteformat_flag = True
         except Exception as e:
             if input_type == 'ete':

treeprofiler/src/utils.py

@@ -24,7 +24,7 @@
 from treeprofiler.src import utils
 from treeprofiler.src.phylosignal import run_acr_discrete, run_acr_continuous, run_delta
 from treeprofiler.src.ls import run_ls
-from treeprofiler.src import b64pickle
+from treeprofiler.src import ete_format
 
 from multiprocessing import Pool
 
@@ -112,7 +112,7 @@ def populate_annotate_args(parser):
     #     help=("<kingdom|phylum|class|order|family|genus|species|subspecies> "
     #           "reference tree from taxonomic database"))
     add('--taxadb', type=str.upper,
-        choices=['NCBI', 'GTDB', 'customdb'],
+        choices=['NCBI', 'GTDB', 'MOTUS', 'customdb'],
         help="<NCBI|GTDB> for taxonomic annotation or fetch taxatree")
     add('--gtdb-version', type=int,
         choices=[95, 202, 207, 214, 220],
@@ -648,6 +648,17 @@ def run_tree_annotate(tree, input_annotated_tree=False,
                 else:
                     logger.info("No specific version or dump file provided; using latest GTDB data...")
                     GTDBTaxa().update_taxonomy_database()
+            elif taxadb == 'MOTUS':
+                if gtdb_version and taxa_dump:
+                    logger.error('Please specify either GTDB version or taxa dump file, not both.')
+                    sys.exit(1)
+                if taxa_dump:
+                    logger.info(f"Loading GTDB database dump file {taxa_dump}...")
+                    GTDBTaxa().update_taxonomy_database(taxa_dump)
+                else:
+                    logger.info("No specific version or dump file provided; using latest GTDB data...")
+                    motus_dump = download_motus_dump()
+                    GTDBTaxa().update_taxonomy_database(motus_dump)
             elif taxadb == 'NCBI':
                 if taxa_dump:
                     logger.info(f"Loading NCBI database dump file {taxa_dump}...")
@@ -924,7 +935,7 @@ def run(args):
 
         ### out ete
         with open(os.path.join(args.outdir, base+'_annotated.ete'), 'w') as f:
-            f.write(b64pickle.dumps(annotated_tree, encoder='pickle', pack=False))
+            f.write(ete_format.dumps(annotated_tree, encoder='pickle', pack=False))
 
         ### out tsv
         prop_keys = list(prop2type.keys())
@@ -1780,7 +1791,7 @@ def merge_dictionaries(dict_ranks, dict_names):
         return merged_dict
 
 
-    if db == "GTDB":
+    if db == "GTDB" or "MOTUS":
         gtdb = GTDBTaxa()
         tree.set_species_naming_function(return_spcode_gtdb)
         gtdb.annotate_tree(tree,  taxid_attr="species", ignore_unclassified=ignore_unclassified)
@@ -1840,28 +1851,27 @@ def merge_dictionaries(dict_ranks, dict_names):
 
     return tree, rank2values
 
-# def annotate_evol_events(tree, sp_delimiter='.', sp_field=0):
-#     def return_spcode(leaf):
-#         try:
-#             return leaf.name.split(sp_delimiter)[sp_field]
-#         except (IndexError, ValueError):
-#             return leaf.name
-
-#     tree.set_species_naming_function(return_spcode)
-
-#     node2species = tree.get_cached_content('species')
-#     for n in tree.traverse():
-#         n.props['species'] = node2species[n]
-#         if len(n.children) == 2:
-#             dup_sp = node2species[n.children[0]] & node2species[n.children[1]]
-#             if dup_sp:
-#                 n.props['evoltype'] = 'D'
-#                 n.props['dup_sp'] = ','.join(dup_sp)
-#                 n.props['dup_percent'] = round(len(dup_sp)/len(node2species[n]), 3) * 100
-#             else:
-#                 n.props['evoltype'] = 'S'
-#         n.del_prop('_speciesFunction')
-#     return tree
+def download_motus_dump():
+    from hashlib import md5
+    import requests
+
+    url = "https://github.com/dengzq1234/ete-data/raw/refs/heads/main/motus_taxonomy/motus_latest_dump.tar.gz"
+    fname = './motus_latest_dump.tar.gz'
+    if not os.path.exists(fname):
+        print(f'Downloading {fname} from {url} ...')
+        with open(fname, 'wb') as f:
+            f.write(requests.get(url).content)
+    else:
+        md5_local = md5(open(fname, 'rb').read()).hexdigest()
+        md5_remote = requests.get(url + '.md5').text.split()[0]
+
+        if md5_local != md5_remote:
+            print(f'Updating {fname} from {url} ...')
+            with open(fname, 'wb') as f:
+                f.write(requests.get(url).content)
+        else:
+            print(f'File {fname} is already up-to-date with {url} .')
+    return fname
 
 def annotate_evol_events(tree, sos_thr=0.0, sp_delimiter='.', sp_field=0):
     def return_spcode(leaf):