Assignment 3: Bayesian phylogenetics with Markov chain Monte Carlo

BST 570-1: Probabilistic Machine Learning, Fall 2026

Submission cutoff: Nov 18th Noon. Submit two files on a branch of your course repository and open a pull request into main: a3.py, which holds all of your code, and a3.qmd, which holds your derivations, explanations, tables and figures. The samplers take a long time to run, so do not run them inside a3.qmd. Run them once from a3.py, save what they produce (parameter traces, acceptance counts, the visited topologies) to files with np.savez, commit those files, and have a3.qmd import from a3.py and load the saved results to make its tables and figures. Problem 1 is quick and may run inside a3.qmd. Check that quarto render a3.qmd succeeds and takes well under a minute; failure to render will result in a 10% grade deduction.

Saving and loading, with the visited topologies turned into strings so that the file needs no pickling:

# in a3.py, after a run
np.savez("results/problem2_sigma_0.5.npz", trace=trace, accepted=accepted)
np.savez("results/problem3_chain1.npz", loglik=lls,
         visited=np.array(["|".join("+".join(p) for p in key) for key in visited]))

# in a3.qmd
run = np.load("results/problem3_chain1.npz")
counts = Counter(run["visited"][burn_in:])        # keys like "Chimpanzee+Human|Gibbon+Orangutan"

How this assignment is organized. Each part has some background and then a numbered list of deliverables, which is everything you have to hand in for that part. Number your answers the same way. The code behind every number, table and plot must be in a3.py, and the saved results must be the ones the document loads; derivations and explanations are a few lines of your own algebra or prose.

Ground rules on AI are in the syllabus: discuss freely, but do not submit generated code verbatim, and write the explanations yourself. You will be asked to explain your solution in class.

The data

Are humans more closely related to chimpanzees or to gorillas? The three lineages split within a short span of time, so the answer was contested for decades. Download the data file from the course website and place it under data/a3_hominoid.fasta in your course repository. The file holds a classic data set on the question: an 895-site stretch of mitochondrial DNA from a human, a chimpanzee, a gorilla, an orangutan and a gibbon, aligned, with no gaps (Brown, Prager, Wang and Wilson, Journal of Molecular Evolution 18, 1982; distributed with Ziheng Yang’s PAML software).

import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from scipy import stats
from scipy.linalg import expm

def read_fasta(path):
    seqs, name = {}, None
    for line in open(path).read().split():
        if line.startswith(">"):
            name = line[1:]; seqs[name] = ""
        else:
            seqs[name] += line
    return {k: np.array(["ACGT".index(c) for c in s]) for k, s in seqs.items()}

aln = read_fasta("data/a3_hominoid.fasta")
pi = np.bincount(np.concatenate(list(aln.values())), minlength=4) / (5 * len(aln["Human"]))
print(f"{len(aln['Human'])} sites; taxa {list(aln)}")
print(f"base frequencies A C G T = {np.round(pi, 4)}")
print("".join("ACGT"[b] for b in aln["Human"][:60]), "...")

895 sites; taxa ['Human', 'Chimpanzee', 'Gorilla', 'Orangutan', 'Gibbon']
base frequencies A C G T = [0.312  0.3289 0.1059 0.2532]
AAGCTTCACCGGCGCAGTCATTCTCATAATCGCCCACGGACTTACATCCTCATTACTATT ...

Each sequence is an integer array with \(0, 1, 2, 3\) for \(A, C, G, T\).

The model. Sites evolve independently under the HKY85 model. The bases are far from equally frequent, and transitions, substitutions within the purines (\(A \leftrightarrow G\)) or within the pyrimidines (\(C \leftrightarrow T\)), are much more common than transversions between the two groups. HKY85 has a parameter for each. The rate of substitution from base \(a\) to base \(b \neq a\) is \[ Q_{ab} \propto \begin{cases} \kappa\, \pi_b & \text{if } a \to b \text{ is a transition,} \\ \pi_b & \text{if it is a transversion,} \end{cases} \] with the diagonal chosen so that each row sums to zero, and \(Q\) scaled so that a branch of length \(t\) carries \(t\) expected substitutions per site. The transition matrix is \(P(t) = \exp(Qt)\) and the root distribution is \(\pi\). With \(\kappa = 1\) and \(\pi\) uniform, HKY85 is the Jukes-Cantor model. Throughout, fix \(\pi\) at the base frequencies printed above; \(\kappa\) is unknown.

def hky_Q(kappa, pi):
    """HKY85 rate matrix, scaled to one expected substitution per unit of branch length."""
    Q = np.tile(pi, (4, 1)).astype(float)                # rate a -> b is pi_b ...
    for a, b in [(0, 2), (2, 0), (1, 3), (3, 1)]:         # ... times kappa for A<->G and C<->T
        Q[a, b] *= kappa
    np.fill_diagonal(Q, 0.0)
    np.fill_diagonal(Q, -Q.sum(axis=1))
    return Q / -(pi @ np.diag(Q))

Use your Node class and traversals from Lab 4, and apply_nni, undo and sample_nni_move from Lab 5. Problem 1 asks for bottom_up and log_likelihood from Lab 5, and adds the top-down pass.

Trees. Every unrooted binary tree on five taxa consists of two pairs and a middle taxon, and is stored as a Node rooted at the internal node next to the middle taxon, with three children. The seven branches are the five pendant branches, named by their taxon, and the two internal branches, named by the pair they separate off. The two functions below build such a tree and name its topology by its two pairs. They assume your Node(name, branch_length) and preorder from Lab 4.

def tree_from_pairs(pair1, middle, pair2, lengths, aln):
    """((pair1), middle, (pair2)) rooted at u. lengths: one entry per taxon, and 'x' and 'y'
    for the internal branches above pair1 and pair2. Attaches the sequences in aln to the leaves."""
    u = Node("u")
    x = u.add_child(Node("x", lengths["x"]))
    u.add_child(Node(middle, lengths[middle]))
    y = u.add_child(Node("y", lengths["y"]))
    for a in pair1:
        x.add_child(Node(a, lengths[a]))
    for a in pair2:
        y.add_child(Node(a, lengths[a]))
    for v in preorder(u):
        if v.is_leaf():
            v.seq = aln[v.name]
    return u

def split_key(root):
    """The topology as its two pairs, sorted, e.g. (('Chimpanzee', 'Human'), ('Gibbon', 'Orangutan'))."""
    taxa = {v.name for v in preorder(root) if v.is_leaf()}
    key = []
    for v in preorder(root):
        if not v.is_leaf() and not v.is_root():
            below = {w.name for w in preorder(v) if w.is_leaf()}
            side = below if len(below) == 2 else taxa - below
            key.append(tuple(sorted(side)))
    return tuple(sorted(key))

As a check, on \(((\text{Human}, \text{Chimpanzee}), \text{Gorilla}, (\text{Orangutan}, \text{Gibbon}))\) with every branch length \(0.1\), the log likelihood is \(-2933.5771\) with \(\kappa = 1\) and \(-2744.2495\) with \(\kappa = 10\).

The prior. Every branch length has an independent Exponential prior with mean \(0.1\), \(\kappa\) has an Exponential prior with mean \(10\), and in Problem 3 the topology has a uniform prior over the \(15\) unrooted topologies.

Problem 1: Two passes on a tree (30 points)

Felsenstein’s pruning algorithm from Lab 5 is the bottom-up pass of belief propagation on the tree. Its mirror image, the top-down pass, gives the posterior of the base at every ancestral node. This problem asks for both, on the tree \(((\text{Human}, \text{Chimpanzee}), \text{Gorilla}, (\text{Orangutan}, \text{Gibbon}))\) with every branch length \(0.1\), and it is where the likelihood used by the samplers in Problems 2 and 3 comes from.

Bottom-up. The likelihood table of node \(v\) is the \(4 \times n\) matrix \(L_v\) with entries \[ L_v(a, s) = p\big(\mathbf{y}^{(s)}_{\lfloor v \rfloor} \mid z_v^{(s)} = a\big), \] the bases at the leaves below \(v\) at site \(s\), given that \(v\) carries base \(a\). At a leaf it is the indicator of the observed base. At an internal node it is the elementwise product over the children \(c\) of \(P(t_c)\, L_c\). The site likelihoods are the entries of \(\pi^\top L_{\text{root}}\), and the log likelihood of the alignment is the sum of their logs. This is Lab 5 with a general rate matrix in place of Jukes-Cantor.

Top-down. The above table of node \(v\) is the \(4 \times n\) matrix \(A_v\) with entries \[ A_v(a, s) = p\big(\mathbf{y}^{(s)}_{\lceil v \rceil},\; z_v^{(s)} = a\big), \] the bases at the leaves above \(v\), meaning every leaf not in its subtree, at site \(s\), jointly with the base at \(v\). Column \(s\) is the top-down message \(m_{u \to v}\) from the parent \(u\) of \(v\) in the belief propagation lecture, for that site. At the root nothing is above, so \(A_{\text{root}}(a, s) = \pi_a\) for every \(s\). The top-down message rule gives \(A_v\) from the above table of its parent and the likelihood tables of its siblings. Once both tables are in hand, the posterior of the base at \(v\) is \[ p\big(z_v^{(s)} = a \mid \mathbf{y}^{(s)}\big) = \frac{A_v(a, s)\, L_v(a, s)}{p(\mathbf{y}^{(s)})}. \]

Deliverables.

(10 points) Implement bottom_up(root, Q), which sets v.likelihood_table for every node and returns the root’s table, and log_likelihood(table, pi), for a rate matrix Q and a root with any number of children. Reproduce the two check values above.
(5 points) Write \(A_v\) as a matrix expression in \(A_u\), \(P(t_v)\), and \(P(t_{s}) L_{s}\) for the siblings \(s\) of \(v\) under the parent \(u\), using \(\odot\) for the elementwise product. Say which traversal order computes the above tables and why, and show that \(\sum_a A_v(a, s)\, L_v(a, s) = p(\mathbf{y}^{(s)})\) for every node \(v\).
(10 points) Implement top_down(root, Q, pi), which sets v.above_table for every node. With \(\kappa = 10\), verify the identity of the previous item at every node and every site, to machine precision.
(5 points) With \(\kappa = 10\), compute the posterior of the base at the ancestor of Human and Chimpanzee at every site. At the sites where Human and Chimpanzee differ, the evidence below the ancestor is split between two bases. Report how many such sites there are and at how many of them the posterior mode is Gorilla’s base, and explain what settles the base at those sites and how that information reaches the ancestor.

Problem 2: The posterior over branch lengths and \(\kappa\) (30 points)

This problem keeps the topology fixed at \(((\text{Human}, \text{Chimpanzee}), \text{Gorilla}, (\text{Orangutan}, \text{Gibbon}))\). The posterior \(p(\mathbf{t}, \kappa \mid \mathbf{y})\) over the seven branch lengths and \(\kappa\) has no closed form, so we sample from it.

The sampler. Update one parameter at a time with a Metropolis-Hastings step: propose a new value for that parameter alone, accept or reject, and move to the next. One sweep updates each of the seven branch lengths once and then \(\kappa\) once. The proposal is yours to design. Your choice of proposal, the justification you give for it, and how well the resulting chain performs are all part of the grade for this problem. A branch length is positive, so the proposal must keep \(t' > 0\), and it must have a step size you can tune; \(\kappa\) is positive too, so the same kind of update applies to it. Each proposal changes one parameter, so recomputing the whole likelihood is simplest and fast enough here.

Effective sample size. Successive states of a Markov chain are correlated, so \(S\) draws carry less information than \(S\) independent ones. The effective sample size is \(S / \tau\), where \(\tau = 1 + 2 \sum_{k \ge 1} \rho_k\) and \(\rho_k\) is the autocorrelation at lag \(k\). The sum is truncated with Geyer’s initial positive sequence: add the autocorrelations in pairs \(\rho_{2m} + \rho_{2m+1}\) and stop at the first pair that is not positive. Use this function:

def ess(x):
    """Effective sample size of a chain, by Geyer's initial positive sequence."""
    x = np.asarray(x, float)
    n = len(x)
    x = x - x.mean()
    f = np.fft.rfft(x, 2 * n)
    acov = np.fft.irfft(f * np.conj(f))[:n] / n
    if acov[0] == 0:
        return float("nan")                  # a constant chain: no variance to estimate
    rho = acov / acov[0]
    s, m = 0.0, 0
    while 2 * m + 1 < n:
        pair = rho[2 * m] + rho[2 * m + 1]
        if pair <= 0:
            break
        s += pair
        m += 1
    return n / (-1 + 2 * s)                  # tau = -1 + 2 * (sum of the positive pairs)

print(f"ESS of 5000 independent draws: {ess(np.random.default_rng(0).normal(size=5000)):.0f}")

ESS of 5000 independent draws: 4925

For every run, start from all branch lengths \(0.1\) and \(\kappa = 1\), and discard the first \(1000\) sweeps.

Deliverables.

(10 points) Write down \(\log p(\mathbf{t}, \kappa \mid \mathbf{y})\) up to an additive constant. State your proposal \(q(t' \mid t)\) for a single branch length, say how it keeps \(t' > 0\) and why you chose it over the alternatives, and derive its acceptance probability, stating whether the proposal is symmetric in \(t\) and including the Hastings correction if it is not. Then discuss the update of \(\kappa\): whether the same proposal is appropriate given that \(\kappa\) lives on a different scale from the branch lengths, what changes in the acceptance probability, and what has to be recomputed when \(\kappa\) changes.
(10 points) Implement the sampler. Choose four step sizes that span from clearly too small to clearly too large, use each one for every parameter, run \(5000\) sweeps per step size, and tabulate the acceptance rate and the effective sample size of each of the eight parameters. Explain why the step size that gives the largest effective sample size for the branch above \((\text{Human}, \text{Chimpanzee})\) is not the one that does for the gibbon’s branch.
(10 points) Choose a step size, run \(20000\) sweeps, and tabulate the posterior mean and the \(95\%\) central interval of each of the eight parameters. Plot a histogram of the posterior draws of \(\kappa\) with its prior density on top, and say what the posterior of \(\kappa\) tells you about transitions and transversions in these sequences, and whether the data support the Jukes-Cantor assumption \(\kappa = 1\).

Problem 3: The posterior over topologies (40 points)

Now the topology \(\tau\) is unknown as well, and the target is the joint posterior \(p(\tau, \mathbf{t}, \kappa \mid \mathbf{y})\). Its marginal \(p(\tau \mid \mathbf{y})\) answers “which tree?”, and the posterior probability of a pair answers “are humans and chimpanzees sisters?” with every topology that contains that pair counted.

The sampler. Each iteration does two things.

A topology move. Draw a move with sample_nni_move, apply it with apply_nni, and accept or reject; on rejection, undo it. The branch lengths travel with their subtrees, the internal branch keeps its length, and \(\kappa\) does not change.
A sweep. One sweep of Problem 2 on the current tree, with your step size from Problem 2.

Record split_key(root) after every iteration, save the list as in the submission instructions, and count topologies with a Counter:

visited = []
for it in range(n_iter):
    ...                                        # the topology move, then a sweep
    visited.append(split_key(root))
counts = Counter(visited[burn_in:])            # counts[key] / (n_iter - burn_in) is the posterior probability of key

Run two chains. Chain 1 starts from \(((\text{Human}, \text{Gibbon}), \text{Chimpanzee}, (\text{Gorilla}, \text{Orangutan}))\) and chain 2 from \(((\text{Chimpanzee}, \text{Orangutan}), \text{Human}, (\text{Gorilla}, \text{Gibbon}))\); neither shares a pair with the topology of Problem 2. Both start from all branch lengths \(0.1\) and \(\kappa = 1\), run for \(8000\) iterations, and discard the first \(2000\).

Deliverables.

(10 points) Write the Metropolis-Hastings ratio for the topology move and show that the acceptance probability is the minimum of \(1\) and the likelihood ratio. Name each term of the Metropolis-Hastings ratio that cancels and the reason it does; for the proposal, compute the probability that sample_nni_move proposes a given neighbouring topology, and the probability of the reverse move.
(10 points) Run both chains. Tabulate, side by side, each chain’s posterior probability of every topology it visits more than \(1\%\) of the time, plot both chains’ log likelihood traces, and give each chain’s effective sample size for the indicator that Human and Chimpanzee are a pair.
(10 points) Run both chains again under the Jukes-Cantor model: \(\pi\) uniform and \(\kappa\) fixed at \(1\), with no \(\kappa\) update. Tabulate the posterior probability that Human and Chimpanzee, Chimpanzee and Gorilla, or Human and Gorilla are a pair, under each model, pooling the two chains.
(10 points) Which species is the closest relative of humans, according to each model? Explain, using your posterior of \(\kappa\) from Problem 2, why the two models can disagree on the same data, and what this says about how to read a posterior probability of a tree.