def numpy_tree_from_ratio_test_data(
ratio_test_data: RatioTestData) -> NumpyRootedTree:
return NumpyRootedTree(
sampling_times=ratio_test_data.sampling_times,
node_heights=ratio_test_data.heights,
topology=numpy_topology_from_ratio_test_data(ratio_test_data),
)
def parse_newick(newick_file: str,
remove_zero_edges: bool = True,
epsilon: float = _EPSILON) -> NumpyRootedTree:
"""Return leaves followed by nodes (postorder)"""
t = ete3.Tree(newick_file)
ordered_nodes = sorted(t.traverse("postorder"),
key=lambda n: not n.is_leaf())
indices = {n: i for i, n in enumerate(ordered_nodes)}
parent_indices = np.array([indices[n.up] for n in ordered_nodes[:-1]])
root_distances = np.array([t.get_distance(n) for n in ordered_nodes],
dtype=DEFAULT_FLOAT_DTYPE_NP) # TODO: Optimise
root_height = max(root_distances)
heights = root_height - root_distances
taxon_count = (len(ordered_nodes) + 1) // 2
taxon_set = DictTaxonSet([x.name for x in ordered_nodes[:taxon_count]])
tree = NumpyRootedTree(
heights=heights,
parent_indices=parent_indices,
taxon_set=taxon_set,
)
if remove_zero_edges:
tree = _remove_zero_edges_func(tree, epsilon=epsilon)
return tree
def data_to_tensor_tree(tree_test_data: TreeTestData) -> TensorflowRootedTree:
numpy_tree = NumpyRootedTree(
node_heights=tree_test_data.node_heights,
sampling_times=tree_test_data.sampling_times,
parent_indices=tree_test_data.parent_indices,
)
tf_tree = convert_tree_to_tensor(numpy_tree)
return tf_tree
def remove_zero_edges(tree: NumpyRootedTree,
epsilon: float = _EPSILON) -> NumpyRootedTree:
heights = np.array(tree.heights)
child_indices = tree.topology.child_indices
for node_index in tree.topology.postorder_node_indices:
child_heights = heights[child_indices[node_index]]
heights[node_index] = np.max(
(heights[node_index], np.max(child_heights) + epsilon))
return NumpyRootedTree(heights=heights, topology=tree.topology)
def test_numpy_tree_get_branch_lengths(tree_test_data: TreeTestData):
expected_branch_lengths = tree_test_data.branch_lengths
tree = NumpyRootedTree(
sampling_times=tree_test_data.sampling_times,
node_heights=tree_test_data.node_heights,
parent_indices=tree_test_data.parent_indices,
)
branch_lengths = tree.branch_lengths
assert_allclose(branch_lengths, expected_branch_lengths)
def test_remove_zero_edges():
node_height = 0.5
epsilon = 0.01
parent_indices = np.array([4, 4, 5, 6, 5, 6])
heights = np.array(
[0.0, 0.1, 0.4, 0.2, node_height, node_height, node_height])
expected_heights = np.concatenate([
heights[:4],
[node_height, node_height + epsilon, node_height + 2 * epsilon]
])
tree = NumpyRootedTree(heights, topology=NumpyTreeTopology(parent_indices))
res = remove_zero_edges(tree, epsilon=epsilon)
assert_allclose(res.heights, expected_heights)
def get_tree_info(inst) -> tp.Tuple[NumpyRootedTree, np.ndarray]:
bito_tree = inst.tree_collection.trees[0]
parent_indices = np.array(bito_tree.parent_id_vector())
node_heights = np.array(bito_tree.node_heights)
tree = NumpyRootedTree(
heights=node_heights,
topology=NumpyTreeTopology(parent_indices=parent_indices),
)
node_bounds = np.array(bito_tree.node_bounds)[tree.taxon_count:]
return (
tree,
node_bounds,
)
def transition_prob_tree(flat_tree_test_data):
tree = convert_tree_to_tensor(
NumpyRootedTree(
heights=flat_tree_test_data.heights,
parent_indices=flat_tree_test_data.parent_indices,
)
).get_unrooted_tree()
state_count = 5
transition_probs = tf.fill(
tree.branch_lengths.shape + (state_count, state_count), 1.0 / state_count
)
return TensorflowUnrootedTree(
branch_lengths=transition_probs,
topology=numpy_topology_to_tensor(tree.topology),
)
def tree_from_arrays(
node_heights: tp.Union[np.ndarray, tf.Tensor],
parent_indices: tp.Union[np.ndarray, tf.Tensor],
sampling_times: tp.Optional[tp.Union[np.ndarray, tf.Tensor]],
taxon_set: tp.Optional[TaxonSet] = None,
height_dtype: tf.DType = DEFAULT_FLOAT_DTYPE_TF,
) -> TensorflowRootedTree:
node_heights_np = (node_heights.numpy() if isinstance(
node_heights, tf.Tensor) else node_heights)
sampling_times_np = (sampling_times.numpy() if isinstance(
sampling_times, tf.Tensor) else sampling_times)
parent_indices_np = (parent_indices.numpy() if isinstance(
parent_indices, tf.Tensor) else parent_indices)
numpy_tree = NumpyRootedTree(
sampling_times=sampling_times_np,
node_heights=node_heights_np,
parent_indices=parent_indices_np,
taxon_set=taxon_set,
)
return convert_tree_to_tensor(numpy_tree, height_dtype=height_dtype)
def numpy(self) -> NumpyRootedTree:
return NumpyRootedTree(
node_heights=self.node_heights.numpy(),
sampling_times=self.sampling_times.numpy(),
topology=self.topology.numpy(),
)