From 947db6f3c6c13babbafe67aaa8c27069566b488e Mon Sep 17 00:00:00 2001
From: Tongzhou Xu <tongzhou.xu@hotmail.com>
Date: Wed, 27 Mar 2024 11:55:55 -0400
Subject: [PATCH] bio.phylo

---
 mashpit/__init__.py |  2 +-
 mashpit/query.py    | 23 +++++++++++++++++------
 setup.py            |  3 +--
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/mashpit/__init__.py b/mashpit/__init__.py
index 0a42f48..2bc921c 100755
--- a/mashpit/__init__.py
+++ b/mashpit/__init__.py
@@ -1,3 +1,3 @@
 __author__ = 'Tongzhou Xu'
 __email__ = 'tongzhou.xu@uga.edu'
-__version__ = '0.9.4'
+__version__ = '0.9.41'
diff --git a/mashpit/query.py b/mashpit/query.py
index 027a1af..a1b459f 100755
--- a/mashpit/query.py
+++ b/mashpit/query.py
@@ -7,15 +7,15 @@
 import heapq
 import time
 import pandas as pd
+import matplotlib.pyplot as plt
 
-from phytreeviz import TreeViz
 from skbio import DistanceMatrix
 from skbio.tree import nj
+from Bio import Phylo
 from mashpit.build import create_connection
 from operator import itemgetter
 from sourmash import SourmashSignature, save_signatures, load_one_signature, load_file_as_signatures
 
-
 def get_query_sig(query_path, query_name,hash_number,kmer_size):
     genome = query_path
     mh = sourmash.MinHash(n=hash_number, ksize=kmer_size)
@@ -105,10 +105,21 @@ def generate_mashtree(output_df,min_similarity,query_name,sig_path,added_annotat
         newick_str = nj(dm, result_constructor=str)
         with open(f'{query_name}_annotated_tree.newick','w') as f: 
             f.write(newick_str)
-    
-    tv = TreeViz(f'{query_name}_tree.newick')
-    tv.set_node_label_props(query_name, color="red")
-    tv.savefig(f'{query_name}_tree.png', dpi=300)
+    tree = Phylo.read(f'{query_name}_annotated_tree.newick', "newick")
+    n = len(tree.get_terminals())
+    fig = plt.figure(figsize=(10, n*0.35), dpi=300)
+    axes = fig.add_subplot(1, 1, 1)
+    # disable the axes and borders
+    axes.set_frame_on(False)
+    # remove ticks and labels
+    axes.set_xticks([])
+    axes.set_yticks([])
+    axes.set_xticklabels([])
+    axes.set_yticklabels([])
+    # add white background
+    fig.patch.set_facecolor('white')
+    Phylo.draw(tree, axes=axes,do_show=False)
+    plt.savefig(f'{query_name}_tree.png')
 
 def query(args):
     t = time.localtime()
diff --git a/setup.py b/setup.py
index c90a603..51c9f02 100755
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name='mashpit',
-    version='0.9.4',
+    version='0.9.41',
     url='https://github.com/tongzhouxu/mashpit',
     author='Tongzhou Xu',
     author_email='tongzhou.xu@uga.edu',
@@ -23,7 +23,6 @@
         'pandas',
         'biopython',
         'scikit-bio',
-        'python-dotenv',
         'tqdm',
         'flask',
         'dask[dataframe]',