def load_trees(tree_path, read_name_mapping=False, max_trees=None):
with open(tree_path, 'r') as tree_file:
nexus_str = tree_file.read().lower().strip()
if read_name_mapping:
name_map = read_nexus_name_mapping(nexus_str)
else:
name_map = None
trees = []
for line in nexus_str.split('\n'):
if line.strip().startswith('tree '):
newick_str = line.split(' = ')[-1]
if newick_str.startswith(r'[&r] '):
newick_str = newick_str[len(r'[&r] '):]
newick_str = newick_str.replace(']:[',',')
newick_str = newick_str.replace(']',']:')
newick_str = newick_str.replace(':;', ';')
tree = Tree.from_newick(newick_str, translate=name_map)
trees.append(tree)
if max_trees is not None and len(trees) >= max_trees:
return trees
return trees
def load_tree_from_nexus(tree_path, location_key='location',
use_translation_table=False, name_mapping=None):
if use_translation_table == True:
name_mapping = read_translation_table(tree_path)
print(name_mapping)
with open(tree_path, 'r') as tree_file:
nexus_str = tree_file.read()
newick_str = extract_newick_from_nexus(nexus_str)
tree = Tree.from_newick(newick_str, location_key=location_key,
translate=name_mapping)
return tree
def write_bantu_xml(xml_path, chain_length, root=None, exclude_outgroup=False,
movement_model='rrw', adapt_tree=False, adapt_height=False):
with open(NEWICK_TREE_PATH, 'r') as tree_file:
tree_str = tree_file.read().lower().strip()
tree = Tree.from_newick(tree_str.strip())
tree.load_locations_from_csv(LOCATIONS_PATH, swap_xy=True)
if exclude_outgroup:
tree.remove_nodes_by_name(OUTGROUP_NAMES)
# tree = tree.big_child().big_child().big_child()
tree.write_beast_xml(xml_path, chain_length, root=root,
diffusion_on_a_sphere=True, movement_model=movement_model,
adapt_tree=adapt_tree, adapt_height=adapt_height)
def write_bantu_sample_xml(xml_path, chain_length, root=None, exclude_outgroup=False,
movement_model='rrw', adapt_tree=False, adapt_height=False):
with open(POSTERIOR_PATH, 'r') as tree_file:
nexus_str = tree_file.read().lower().strip()
name_map = read_nexus_name_mapping(nexus_str)
tree_lines = [line.split(' = ')[-1] for line in nexus_str.split('\n') if line.startswith('\t\ttree')]
tree_str = random.choice(tree_lines)
tree = Tree.from_newick(tree_str.strip())
for node in tree.iter_descendants():
if node.name in name_map:
node.name = name_map[node.name]
tree.load_locations_from_csv(LOCATIONS_PATH, swap_xy=True)
leafs_without_locations = [node.name for node in tree.iter_leafs() if node.location is None]
tree.remove_nodes_by_name(leafs_without_locations)
if exclude_outgroup:
tree.remove_nodes_by_name(OUTGROUP_NAMES)
tree.write_beast_xml(xml_path, chain_length, root=root,
diffusion_on_a_sphere=True, movement_model=movement_model,
adapt_tree=adapt_tree, adapt_height=adapt_height)