treeflow.tree.topology.numpy_tree_topology module

class treeflow.tree.topology.numpy_tree_topology.NumpyTreeTopology(parent_indices: ndarray, taxon_set: TaxonSet | None = None)

Bases: NumpyTreeTopologyAttrs

Class representing a bifurcating tree topology as a composition of integer NumPy arrays.

For a phylogenetic tree with n taxa at the leaves, the representation maintains a labelling of the 1n-1 nodes with integer indices. The labelling convention is that the leaves are the first n indices and the root is at the last index (2n-2).

property taxon_count: int

Number of leaf taxa the tree is based on (n)

property node_count: int
property postorder_node_indices: ndarray
property child_indices: ndarray

Array of length 2n-2 representing the parent-child structure of a tree topology on n taxa.

The i th element of this array is the index of the parent of the i th indexed node in the tree, for every node except the root.

property preorder_indices: ndarray

Array of indices of length 2n-1 that form a pre-order traversal of the tree.

A pre-order traversal is an ordering of the nodes where every node is visited before its children, starting at the root.

property taxon_set: TaxonSet | None
parent_indices: ndarray
class treeflow.tree.topology.numpy_tree_topology.StaticNumpyTreeTopology(parent_indices: ndarray, taxon_set: TaxonSet | None = None)

Bases: object

A NumpyTreeTopology variant with every derived index array pre-computed at construction.

NumpyTreeTopology is an attrs class, which tf.nest expands into its single parent_indices leaf. Passing one into a tf.function therefore replaces parent_indices with a symbolic placeholder, and the lazily-computed NumPy properties (child_indices etc.) then fail on the symbolic tensor.

This class is a plain Python object instead, so tf.nest treats it as an atom: it threads through traced code as a captured constant, and its index arrays stay concrete NumPy arrays inside the trace. That lets tf.get_static_value fold the topology regardless of tree size, so the node-height ratio-transform traversal unrolls rather than falling back to a tf.while_loop.

It is meant to be used as a fixed-topology pin. It is not a Tensor, so it is not placed directly into a JointDistribution’s tree value (which TFP coerces to tensors); FixedTopologyRootedTreeBijector rebuilds it as an in-graph-constant TensorflowTreeTopology for that purpose (see its _forward).

classmethod from_numpy_topology(topology: NumpyTreeTopology) → StaticNumpyTreeTopology
property parent_indices: ndarray
property child_indices: ndarray
property preorder_indices: ndarray
property taxon_count: int
property node_count: int
property postorder_node_indices: ndarray
property node_child_indices: ndarray
property preorder_node_indices: ndarray
property taxon_set: TaxonSet | None
has_batch_dimensions() → bool
to_constant_tensor_topology()

Return a TensorflowTreeTopology whose index arrays are ``tf.constant``s.

IMPORTANT: call this inside a tf.function trace (e.g. a VI loss step, the fixed-topology bijector’s forward, or a profiler timing function). The arrays then become in-graph Const ops, which gives both properties we need at once:

  • a normal tensor topology that a JointDistribution accepts in its tree value (it coerces the value to tensors), and

  • a topology that still folds via tf.get_static_value at any tree size, so the downstream traversals unroll.

Calling it eagerly and capturing the result into a later trace loses the second property: a captured constant can be re-materialised as a Placeholder for large trees (the >64-taxon non-fold), which is exactly what building the Const inside the trace avoids.

numpy() → NumpyTreeTopology