def __init__(
self,
heights: tp.Optional[np.ndarray] = None,
node_heights: tp.Optional[np.ndarray] = None,
sampling_times: tp.Optional[np.ndarray] = None,
topology: tp.Optional[NumpyTreeTopology] = None,
parent_indices: tp.Optional[np.ndarray] = None,
taxon_set: tp.Optional[TaxonSet] = None,
):
if topology is None:
assert (
parent_indices is not None
), "Either `topology` or `parent_indices` must be specified"
topology = NumpyTreeTopology(
parent_indices=parent_indices, taxon_set=taxon_set
)
taxon_count = topology.taxon_count
if node_heights is not None:
if sampling_times is None:
final_sampling_times = np.zeros(
node_heights.shape[:-1] + (taxon_count,),
dtype=node_heights.dtype,
)
else:
final_sampling_times = sampling_times
else:
assert heights is not None
final_sampling_times = heights[..., :taxon_count]
node_heights = heights[..., taxon_count:]
super().__init__(
node_heights=node_heights,
sampling_times=final_sampling_times,
topology=topology,
)
def test_TensorflowTreeTopology_function_arg(flat_tree_test_data):
numpy_topology = NumpyTreeTopology(
parent_indices=flat_tree_test_data.parent_indices)
tf_topology = numpy_topology_to_tensor(numpy_topology)
res = function_tf(tf_topology)
expected = function_np(numpy_topology)
assert_allclose(res.numpy(), expected)
def test_TensorflowTreeTopology_get_prefer_static_rank(flat_tree_test_data):
numpy_topology = NumpyTreeTopology(
parent_indices=flat_tree_test_data.parent_indices)
tf_topology = numpy_topology_to_tensor(numpy_topology)
rank = tf_topology.get_prefer_static_rank()
assert isinstance(rank.parent_indices, np.ndarray)
assert isinstance(rank.preorder_indices, np.ndarray)
assert isinstance(rank.child_indices, np.ndarray)
assert rank.parent_indices == 1
assert rank.preorder_indices == 1
assert rank.child_indices == 2
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 numpy(self) -> NumpyTreeTopology:
return NumpyTreeTopology(parent_indices=self.parent_indices,
taxon_set=self.taxon_set)
def numpy_topology_from_ratio_test_data(
ratio_test_data: RatioTestData, ) -> NumpyTreeTopology:
return NumpyTreeTopology(parent_indices=ratio_test_data.parent_indices)
def test_TensorflowTreeTopology_nest_map(flat_tree_test_data):
numpy_topology = NumpyTreeTopology(
parent_indices=flat_tree_test_data.parent_indices)
tf_topology = numpy_topology_to_tensor(numpy_topology)
res = tf.nest.map_structure(tf.shape, tf_topology)