def test_treemap(self):
def _lower(x):
if x.label:
x.label = x.label.lower()
return x
self.assertEqual(
alg.treemap(sp.p_tree.parse("(B,(A,C,E),D);").tree, _lower),
sp.Node(kids=[
sp.Node(label="b"),
sp.Node(kids=[
sp.Node(label="a"),
sp.Node(label="c"),
sp.Node(label="e"),
]),
sp.Node(label="d"),
]),
)
def rm_color(newick, tree):
"""
Remove all color annotations from a tree
"""
import smot.algorithm as alg
tree = read_tree(tree)
tree.colmap = dict()
def _fun(d):
if d.form and "!color" in d.form:
del d.form["!color"]
return d
tree.tree = alg.treemap(tree.tree, _fun)
if newick:
print(sf.newick(tree))
else:
print(sf.nexus(tree))
def tipsed(pattern, replacement, newick, tree):
"""
Search and replace patterns in tip labels.
"""
import smot.algorithm as alg
import re
pat = re.compile(pattern)
def fun_(nodeData):
if nodeData.label:
nodeData.label = re.sub(pat, replacement, nodeData.label)
return nodeData
tree = read_tree(tree)
tree.tree = alg.treemap(tree.tree, fun_)
if newick:
print(sf.newick(tree))
else:
print(sf.nexus(tree))
def filter_cmd(
# conditions
all_match,
some_match,
none_match,
larger_than,
smaller_than,
# actions
remove,
color,
sample,
replace,
# factor methods
factor_by_capture,
factor_by_field,
factor_by_table,
default,
# phylogenetic options
patristic,
seed,
# boilerplate
newick,
tree,
):
"""
An advanced tool for performaing actions (remove, color, sample, or
replace) on monophyletic groups that meet specified conditions (all-match,
some-match, etc.
"""
import smot.algorithm as alg
import re
tree = read_tree(tree)
tree.tree = factorTree(
node=tree.tree,
factor_by_capture=factor_by_capture,
factor_by_field=factor_by_field,
factor_by_table=factor_by_table,
default=default,
patristic=patristic,
)
def condition(node):
tips = alg.tips(node)
return ((not larger_than or len(tips) > larger_than)
and (not smaller_than or len(tips) < smaller_than)
and (not all_match or all([
all([re.search(pat, tip) for tip in tips])
for pat in all_match
])) and (not some_match or all([
any([re.search(pat, tip) for tip in tips])
for pat in some_match
])) and (not none_match or all([
all([not re.search(pat, tip) for tip in tips])
for pat in none_match
])))
if remove:
action = lambda x: None
elif color:
action = lambda x: alg.colorTree(x, color)
elif sample:
action = lambda x: alg.sampleProportional(x,
proportion=sample,
scale=None,
minTips=3,
keep_regex="",
seed=seed)
elif replace:
def _fun(d):
d.label = re.sub(replace[0], replace[1], d.label)
return d
action = lambda x: alg.treemap(x, _fun)
tree.tree = alg.filterMono(tree.tree, condition=condition, action=action)
tree.tree = alg.clean(tree.tree)
if newick:
print(sf.newick(tree))
else:
print(sf.nexus(tree))
def factor(
method,
factor_by_capture,
factor_by_field,
factor_by_table,
default,
impute,
patristic,
newick,
tree,
):
"""
Impute, annotate with, and/or tabulate factors. The --impute option will
fill in missing factors in monophyletic branches. This is useful, for
example, for inferring clades given a few references in a tree. There are
three modes: 'table' prints a TAB-delimited table of tip names and factors,
'prepend' adds the factor to the beginning of the tiplabel (delimited with
'|'), 'append' adds it to the end.
"""
import smot.algorithm as alg
tree = read_tree(tree)
tree.tree = factorTree(
node=tree.tree,
factor_by_capture=factor_by_capture,
factor_by_field=factor_by_field,
factor_by_table=factor_by_table,
default=default,
impute=impute,
patristic=patristic,
)
# create TAB-delimited, table with columns for the tip labels and the
# (possibly imputed) factor
if method.lower() == "table":
def _fun(b, x):
if x.isLeaf:
if x.factor is None:
factor = default
else:
factor = x.factor
b.append(f"{x.label}\t{factor}")
return b
for row in alg.treefold(tree.tree, _fun, []):
print(row)
# prepend or append the factor to the tip labels and print the resulting tree
else:
def _fun(x):
if x.isLeaf:
if x.factor is None:
x.factor = default
if method.lower() == "prepend":
x.label = f"{x.factor}|{x.label}"
else:
x.label = f"{x.label}|{x.factor}"
return x
tree.tree = alg.treemap(tree.tree, _fun)
if newick:
print(sf.newick(tree))
else:
print(sf.nexus(tree))