Implement VBPI subsplit Bayesian network topology inference - #115
Open
christiaanjs wants to merge 4 commits into
Open
Implement VBPI subsplit Bayesian network topology inference#115christiaanjs wants to merge 4 commits into
christiaanjs wants to merge 4 commits into
Conversation
Add a `treeflow.vbpi` subpackage implementing the topology side of variational Bayesian phylogenetic inference (VBPI): a subsplit Bayesian network / conditional clade distribution as the variational family over rooted topologies, an amortised per-split branch-length approximation, and the VIMCO multi-sample gradient estimator that trains the discrete topology parameters. Components: - clade.py: bitset clade/subsplit primitives and decomposition of a rooted topology into (parent clade, child subsplit) factors. - support.py: SubsplitSupport, a compressed-sparse-row (CSR) pointer-array data structure of clades and candidate child subsplits that the SBN probabilities live on; plus complete-enumeration of small topologies. - sbn.py: SubsplitBayesianNetwork with a differentiable segmented-softmax log_prob and ancestral sampling (NumPy reference or native C++). - branch_model.py: SplitLognormalBranchModel, reparameterised per-split diagonal log-normal branch lengths keyed by clade id. - vimco.py: vimco_surrogate, the VIMCO estimator combining the pathwise (reparameterised) and score (leave-one-out control variate) gradients. Native op: - cc/sbn_op.cc + acceleration/native/sbn.py: a compiled TensorFlow custom op (`SbnSample`) that runs the sequential ancestral-sampling traversal of the pointer-array support in C++, sharded across samples. No gradient (discrete topologies); the differentiable log_prob stays in TensorFlow. Wired into build.sh / build.py alongside the existing traversal ops. Tests (test/vbpi + test/acceleration/native/test_native_sbn.py): - SBN log_prob normalises exactly over all enumerated topologies (4-6 taxa). - NumPy and native samplers match the analytic distribution and each other. - VIMCO gradient is unbiased against the exact enumerated K-sample bound. - Branch model shapes, reparameterised gradients, parameter sharing. - End-to-end VBPI loop (sampling + log_prob + branches + VIMCO) increases the bound towards the true log marginal on a toy target. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ka72WR4fMzneXd2wwUAZpX
Add a Metropolis-Hastings MCMC over rooted tree topologies and a rooted Jukes-Cantor likelihood, so VBPI can be compared against a sampling-based posterior. Everything stays in rooted-tree space and reuses existing treeflow components. treeflow/vbpi/likelihood.py: - rooted_jc_log_likelihood: differentiable JC Felsenstein log-likelihood on a rooted tree, using the JC substitution model's eigendecomposition for the transition matrices and the native C++ phylogenetic-likelihood op (falling back to the pure-TensorFlow reference). Root contraction against the equilibrium frequencies gives the exact likelihood for the reversible model with no unrooted conversion. - make_jc_log_likelihood_fn: a tf.function-compiled (topology, branch) closure (traced once) for fast repeated scoring in the MCMC loop. treeflow/vbpi/mcmc.py: - Rooted NNI moves with structural canonicalisation of node labels; every rooted binary tree has 2(n-2) NNI neighbours, so the proposal is symmetric. - TopologyMetropolisHastings: a tfp.mcmc.TransitionKernel over rooted topologies (same abstraction as the HMC code), targeting an arbitrary target_log_prob_fn(parent_indices). - sample_phylogenetic_topologies: Metropolis-within-Gibbs over the joint (topology, branch length) JC posterior, alternating NNI topology moves with a random-walk branch update. Tests: - Likelihood matches an independent NumPy Felsenstein computation (native and reference paths), batched and differentiable; function mode matches eager. - NNI neighbour count/validity/symmetry, canonical-form uniqueness, clade preservation; the TransitionKernel reproduces an enumerable SBN distribution; the joint sampler recovers the uniform topology + exponential branch prior under a constant likelihood and the true topology from simulated data. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ka72WR4fMzneXd2wwUAZpX
examples/vbpi_vs_mcmc_comparison.ipynb simulates a small rooted 5-taxon tree under Jukes-Cantor, then infers the topology posterior two ways over the same joint p(T, b | D): - VBPI with the enumerated subsplit Bayesian network (complete support over all 105 rooted topologies) + per-split log-normal branches, trained with VIMCO; - the rooted-NNI Metropolis-Hastings sampler. It reads off q(T) for every topology, compares it to the MCMC visiting frequencies (per-rooted scatter, top-topology bars, and an aggregation to unrooted topologies that removes the JC rooting non-identifiability), and marks the true tree. In the executed run both methods place ~0.13 on the true rooted topology and >0.86 on the true unrooted tree, spreading mass similarly across the exchangeable rootings. Everything runs in rooted-tree space on treeflow's JC model, native likelihood op, and TFP TransitionKernel MCMC. Also move the shared VBPI test helpers into test/vbpi/vbpi_test_helpers.py so the test files import them under a unique name rather than via `conftest`, which collided with the native suite's conftest when both were collected together. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ka72WR4fMzneXd2wwUAZpX
Rework the comparison to operate on rooted, ultrametric time trees (node
heights, not free branch lengths), reusing TreeFlow's existing time-tree
machinery, and drop the redundant likelihood test.
New treeflow/vbpi/timetree.py:
- build_time_tree / time_tree_jc_log_likelihood: assemble a TensorflowRootedTree
from a topology + node heights and score JC on branch lengths
clock_rate * (parent_height - child_height) via the native pruning op.
- coalescent_prior / yule_prior: thin wrappers over ConstantCoalescent and Yule.
- NodeHeightRatioModel: the VBPI variational q(heights | T). Node heights are
unconstrained through NodeHeightRatioChainBijector (the NodeHeightRatio
transform); a diagonal Gaussian on the ratio latent is keyed by clade id
(shared across topologies), sampled with the reparameterisation trick, and
log q is returned in height space (Gaussian density minus the transform
log-det-Jacobian).
treeflow/vbpi/mcmc.py:
- propose_nni_internal carries the per-internal-node height latent through an NNI
relabelling (the ratio transform keeps every latent a valid tree).
- sample_phylogenetic_time_trees: Metropolis-within-Gibbs over the joint
(topology, node-height) coalescent/Yule posterior, alternating NNI topology
moves with a random-walk height update. The whole target (ratio transform +
Jacobian + native likelihood + tree prior) is one tf.function.
treeflow/vbpi/likelihood.py: drop the JC eigendecomposition cache, which captured
tensors from whichever tf.function first built them and broke reuse from a later
graph ("out of scope").
Notebook (examples/vbpi_vs_mcmc_comparison.ipynb): rebuilt on time trees with a
coalescent prior; VBPI training compiled into a single tf.function. The executed
run shows both methods favour the true unrooted tree (small unrooted TV) while,
at the rooted level, MCMC spreads over the true tree's exchangeable rootings and
VBPI's mode-seeking exclusive-KL objective concentrates on one -- discussed in
the narrative.
Tests (test/vbpi/test_timetree.py): tree building, coalescent/Yule priors and the
time-tree likelihood are differentiable in the heights; the height model produces
valid trees with reparameterised gradients; propose_nni_internal permutes the
latent; the time-tree MCMC reproduces the coalescent prior's ranking-weighted
topology marginal under a constant likelihood and recovers the true topology from
simulated data. Removed test/vbpi/test_likelihood.py (the underlying likelihood
is already well tested in treeflow).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ka72WR4fMzneXd2wwUAZpX
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add a
treeflow.vbpisubpackage implementing the topology side ofvariational Bayesian phylogenetic inference (VBPI): a subsplit Bayesian
network / conditional clade distribution as the variational family over
rooted topologies, an amortised per-split branch-length approximation, and
the VIMCO multi-sample gradient estimator that trains the discrete topology
parameters.
Components:
topology into (parent clade, child subsplit) factors.
data structure of clades and candidate child subsplits that the SBN
probabilities live on; plus complete-enumeration of small topologies.
log_prob and ancestral sampling (NumPy reference or native C++).
diagonal log-normal branch lengths keyed by clade id.
(reparameterised) and score (leave-one-out control variate) gradients.
Native op:
(
SbnSample) that runs the sequential ancestral-sampling traversal of thepointer-array support in C++, sharded across samples. No gradient (discrete
topologies); the differentiable log_prob stays in TensorFlow. Wired into
build.sh / build.py alongside the existing traversal ops.
Tests (test/vbpi + test/acceleration/native/test_native_sbn.py):
bound towards the true log marginal on a toy target.
Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01Ka72WR4fMzneXd2wwUAZpX