Skip to content

Commit

Permalink
rearranged files and adjusted unit tests to match
Browse files Browse the repository at this point in the history
  • Loading branch information
@KatyBrown
KatyBrown committed Nov 28, 2024
1 parent a47c8a6 commit 6f57d0f
Show file tree
Hide file tree
Showing 6 changed files with 117 additions and 100 deletions.
0 plot_phylo/post_draw.py → plot_phylo/amend_tree.py
View file
File renamed without changes.
63 changes: 32 additions & 31 deletions plot_phylo/draw_tree.py
100644 → 100755
View file
Original file line number Diff line number Diff line change
@@ -1,37 +1,7 @@
#!/usr/bin/env python3
import numpy as np


def collapse_nodes(tree, collapse_list, collapse_names):
cD = dict(zip(collapse_list, collapse_names))
collapseD = dict()
for string in collapse_list:
keeps = set()
collapsed = set()
done = set()
ddD = dict()
for node in tree.traverse():
x = 0
L = list(node.get_leaves())
dd = []
for leaf in L:
if leaf.name.endswith(string) and leaf not in done:
dd.append(leaf.dist)
x += 1
if x == len(L) or (len(L) == 1 and leaf not in done):
keeps.add(L[0].name)
done = done | set(L)
if x > 1:
collapsed.add(L[0])
ddD[L[0]] = np.mean(dd)
tree.prune(keeps)
for leaf in tree.get_leaves():
if leaf in collapsed:
leaf.dist = ddD[leaf]
leaf.name = 'COLLAPSE|%s' % (leaf.name)
collapseD[leaf.name] = cD[string]
return (tree, collapseD)


def draw_tree(tree, ax,
x=0,
y=0,
Expand Down Expand Up @@ -306,3 +276,34 @@ def draw_tree(tree, ax,
va='center', fontsize=appearance['font_size']-2)

return (y, ym, ps)


def collapse_nodes(tree, collapse_list, collapse_names):
cD = dict(zip(collapse_list, collapse_names))
collapseD = dict()
for string in collapse_list:
keeps = set()
collapsed = set()
done = set()
ddD = dict()
for node in tree.traverse():
x = 0
L = list(node.get_leaves())
dd = []
for leaf in L:
if leaf.name.endswith(string) and leaf not in done:
dd.append(leaf.dist)
x += 1
if x == len(L) or (len(L) == 1 and leaf not in done):
keeps.add(L[0].name)
done = done | set(L)
if x > 1:
collapsed.add(L[0])
ddD[L[0]] = np.mean(dd)
tree.prune(keeps)
for leaf in tree.get_leaves():
if leaf in collapsed:
leaf.dist = ddD[leaf]
leaf.name = 'COLLAPSE|%s' % (leaf.name)
collapseD[leaf.name] = cD[string]
return (tree, collapseD)
58 changes: 29 additions & 29 deletions plot_phylo/plot_phylo.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#!/usr/bin/env python3
import ete3
import draw_tree
import post_draw
from . import draw_tree
from . import amend_tree


def plot_phylo(tree, ax,
Expand Down Expand Up @@ -163,42 +163,42 @@ def plot_phylo(tree, ax,
# internally when the function is called recursively but
# are not needed by the user

_, _, ps = draw_tree(T, ax,
x=xpos,
y=-ypos-height,
x0=xpos,
ps=[],
height=height,
width=width,
depth=maxdist,
align_tips=align_tips,
rev_align_tips=rev_align_tips,
branch_lengths=branch_lengths,
reverse=reverse,
appearance=appearance,
collapse=collapse,
collapseD=collapseD)
_, _, ps = draw_tree.draw_tree(T, ax,
x=xpos,
y=-ypos-height,
x0=xpos,
ps=[],
height=height,
width=width,
depth=maxdist,
align_tips=align_tips,
rev_align_tips=rev_align_tips,
branch_lengths=branch_lengths,
reverse=reverse,
appearance=appearance,
collapse=collapse,
collapseD=collapseD)

if rev_align_tips:
ps = post_draw.reverse_align(ax, ps, reverse)
ps = amend_tree.reverse_align(ax, ps, reverse)
# Hide axis
if not show_axis:
ax.set_axis_off()
if scale_bar and branch_lengths:
if not reverse:
post_draw.draw_scale_bar(ax, width, height, maxdist, xpos, ypos,
scale_bar_width=scale_bar_width,
appearance=appearance)
amend_tree.draw_scale_bar(ax, width, height, maxdist, xpos, ypos,
scale_bar_width=scale_bar_width,
appearance=appearance)
else:
post_draw.draw_scale_bar(ax, width, height, maxdist, -xpos, ypos,
scale_bar_width=scale_bar_width,
appearance=appearance)
amend_tree.draw_scale_bar(ax, width, height, maxdist, -xpos, ypos,
scale_bar_width=scale_bar_width,
appearance=appearance)
textobj = [p[1] for p in ps]

if auto_ax:
textobj, ax = post_draw.auto_axis(ax, textobj,
xpos, ypos,
width, height, maxdist,
scale_bar, branch_lengths)
boxes = post_draw.get_boxes(ax, textobj)
textobj, ax = amend_tree.auto_axis(ax, textobj,
xpos, ypos,
width, height, maxdist,
scale_bar, branch_lengths)
boxes = amend_tree.get_boxes(ax, textobj)
return (boxes)
2 changes: 0 additions & 2 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,7 @@
long_description_content_type="text/markdown",
url="https://github.com/KatyBrown/plot_phylo",
packages=setuptools.find_packages(),
package_dir={'plot_phylo':'plot_phylo'},
install_requires=['matplotlib', 'ete3'],
scripts=['plot_phylo/plot_phylo.py'],
classifiers=[
"Programming Language :: Python :: 3",
"License :: OSI Approved :: MIT License",
Expand Down
85 changes: 50 additions & 35 deletions tests/test_plot_phylo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
import matplotlib.pyplot as plt
import matplotlib
import plot_phylo
from plot_phylo import draw_tree
from plot_phylo import amend_tree
import pytest
from test_plot_phylo_data import (test_plot_phylo_vars,
test_plot_phylo_list,
Expand All @@ -20,7 +22,6 @@
import os
import shutil
import numpy as np
matplotlib.use('Agg')


def compare_images(f1, f2, tol):
Expand Down Expand Up @@ -63,7 +64,10 @@ def test_plot_phylo_params(xpos,
line_width,
bold,
expected_figure,
ID, tree, ylim):
ID, tree, ylim,
collapse,
collapse_names,
auto_ax):

tree_stem = tree.split("/")[-1].split(".")[0]

Expand All @@ -90,7 +94,10 @@ def test_plot_phylo_params(xpos,
font_size=font_size,
line_col=line_col,
line_width=line_width,
bold=bold)
bold=bold,
collapse=collapse,
collapse_names=collapse_names,
auto_ax=auto_ax)
try:
os.mkdir("test_temp")
except FileExistsError:
Expand Down Expand Up @@ -128,7 +135,9 @@ def test_draw_tree_params(x,
expected,
ID,
tree,
ylim):
ylim,
collapse,
collapseD):

try:
T = ete3.Tree(tree)
Expand All @@ -150,19 +159,19 @@ def test_draw_tree_params(x,
if nam not in appearance['col_dict']:
appearance['col_dict'][nam] = 'black'

test_obj = plot_phylo.draw_tree(tree=T, ax=a,
x=x,
y=y,
x0=x0,
ps=ps,
height=ylim-1,
width=width,
depth=depth,
align_tips=align_tips,
rev_align_tips=rev_align_tips,
branch_lengths=branch_lengths,
reverse=reverse,
appearance=appearance)
test_obj = draw_tree.draw_tree(tree=T, ax=a,
x=x,
y=y,
x0=x0,
ps=ps,
height=ylim-1,
width=width,
depth=depth,
align_tips=align_tips,
rev_align_tips=rev_align_tips,
branch_lengths=branch_lengths,
reverse=reverse,
appearance=appearance)
plt.close()
ytest = round(test_obj[0], 2)
y2test = round(test_obj[1], 2)
Expand Down Expand Up @@ -203,7 +212,9 @@ def test_reverse_align_params(x,
expected,
ID,
tree,
ylim):
ylim,
collapse,
collapseD):
try:
T = ete3.Tree(tree)
except ete3.parser.newick.NewickError:
Expand All @@ -224,22 +235,22 @@ def test_reverse_align_params(x,
if nam not in appearance['col_dict']:
appearance['col_dict'][nam] = 'black'

_, _, ps = plot_phylo.draw_tree(tree=T, ax=a,
x=x,
y=y,
x0=x0,
ps=[],
height=ylim-1,
width=width,
depth=depth,
align_tips=True,
rev_align_tips=True,
branch_lengths=branch_lengths,
reverse=reverse,
appearance=appearance)
plt.close()
reverse = plot_phylo.reverse_align(a, ps, True)
_, _, ps = draw_tree.draw_tree(tree=T, ax=a,
x=x,
y=y,
x0=x0,
ps=[],
height=ylim-1,
width=width,
depth=depth,
align_tips=True,
rev_align_tips=True,
branch_lengths=branch_lengths,
reverse=reverse,
appearance=appearance)

reverse = amend_tree.reverse_align(a, ps, True)
plt.close()
e0 = expected.replace(".pickle", "_%s.pickle" % tree_stem)

test_dat = []
Expand Down Expand Up @@ -271,7 +282,7 @@ def test_reverse_align_params(x,
test_get_boxes_texts,
test_get_boxes_results)))
def test_get_boxes(ax, texts, expected_result):
boxes = plot_phylo.get_boxes(ax, texts)
boxes = amend_tree.get_boxes(ax, texts)
bclean = dict()
for b, vals in boxes.items():
bclean[b] = dict()
Expand Down Expand Up @@ -303,7 +314,10 @@ def test_bad_tree(xpos,
line_width,
bold,
expected_figure,
ID, tree, ylim):
ID, tree, ylim,
collapse,
collapse_names,
auto_ax):
f = plt.figure(figsize=(10, 20))
a = f.add_subplot(111)
a.set_xlim(-10, 20)
Expand All @@ -330,3 +344,4 @@ def test_bad_tree(xpos,
line_col=line_col,
line_width=line_width,
bold=bold)
plt.close()
9 changes: 6 additions & 3 deletions tests/test_plot_phylo_data.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#!/usr/bin/env python3
import plot_phylo
from plot_phylo import draw_tree
import copy

# Default parameters values
Expand All @@ -16,9 +17,9 @@
varis.remove('height')
variD.pop('height')

defaults_draw_tree = plot_phylo.draw_tree.__defaults__
ac_draw_tree = plot_phylo.draw_tree.__code__.co_argcount
varis_draw_tree = list(plot_phylo.draw_tree.__code__.co_varnames[
defaults_draw_tree = draw_tree.draw_tree.__defaults__
ac_draw_tree = draw_tree.draw_tree.__code__.co_argcount
varis_draw_tree = list(draw_tree.draw_tree.__code__.co_varnames[
ac_draw_tree-len(defaults_draw_tree):ac_draw_tree])
varis_draw_tree.remove('height')

Expand Down Expand Up @@ -94,6 +95,7 @@
'line_col', 'line_width', 'show_support', 'bold']:
curr_dict['appearance'][var] = curr_dict[var]
curr_dict['depth'] = [5, 5, 5]
curr_dict['collapseD'] = dict()
curr_dict.update(test)
pass_vals = [curr_dict[v] for v in varis_draw_tree]
pass_vals.append('tests/test_objects/%s.pickle' % testnam)
Expand All @@ -116,6 +118,7 @@
for var in ['col_dict', 'label_dict', 'font_size',
'line_col', 'line_width', 'show_support', 'bold']:
curr_dict['appearance'][var] = curr_dict[var]
curr_dict['collapseD'] = dict()
curr_dict.update(test)
curr_dict['depth'] = [5, 5, 5]
curr_dict.pop('rev_align_tips')
Expand Down

0 comments on commit 6f57d0f

Please sign in to comment.