treeflow.traversal.phylo_likelihood module

treeflow.traversal.phylo_likelihood.move_indices_to_outside(x, start, size)

treeflow.traversal.phylo_likelihood.phylogenetic_likelihood(topology: TensorflowTreeTopology, sequences_onehot: Tensor, transition_probs: Tensor, frequencies: Tensor, batch_shape=(), use_matvec: bool = False, unroll: bool | str = 'auto')

Per-site phylogenetic likelihood, on the generic postorder_node_traversal.

Assumes all parameters are broadcastable w.r.t. batch shape.

topology

The tree topology (provides postorder/child indices; no batch dimensions).

sequences_onehot

Tensor with shape […, leaf, state]

transition_probs

Tensor with shape […, node, state, state]; broadcast over sites.

frequencies

Tensor with shape […, state]

use_matvec

See _combine_child_partials(). True is forward-faster but slower for value+gradient; leave False when gradients are needed.

unroll

Forwarded to postorder_node_traversal(): "auto" unrolls when the topology is statically known, True forces it, False keeps the dynamic tf.while_loop.

treeflow.traversal.phylo_likelihood.phylogenetic_log_likelihood_rescaled(topology: TensorflowTreeTopology, sequences_onehot: Tensor, transition_probs: Tensor, frequencies: Tensor, batch_shape=(), use_matvec: bool = False, unroll: bool | str = 'auto')

Numerically stable per-site phylogenetic LOG likelihood.

Identical recursion to phylogenetic_likelihood(), but each internal node’s partials are divided by their per-site maximum and the log of that scale factor is accumulated, preventing underflow on large/deep trees. The scale is treated as a constant (stop_gradient), so the gradient matches the unrescaled likelihood.

The running scale is carried as a second component of the per-node output structure (partials, log_scale) (leaves contribute 0) and summed over nodes afterwards. Parameters match phylogenetic_likelihood().