Hidden Markov models and the forward-backward algorithm
faithful = sns.load_dataset("geyser")
fig, ax = plt.subplots(figsize=(13, 4))
ax.plot(faithful["duration"].to_numpy(), color="0.7", lw=1, zorder=1)
sns.scatterplot(x=np.arange(len(faithful)), y=faithful["duration"],
hue=faithful["kind"], s=32, ax=ax, zorder=2)
ax.set_xlabel("eruption number"); ax.set_ylabel("duration (min)")
ax.set_xlim(0, 80)
plt.show()
Last week we fit a two component mixture to these durations. To do that we treated the 272 eruptions as exchangeable: we could have shuffled them without changing the model.
def plot_rolls(ax, y, z=None, t0=0, t1=None):
t1 = len(y) if t1 is None else t1
tt = np.arange(t0, t1)
if z is not None:
ax.fill_between(tt, 0.5, 6.5, where=z[t0:t1] == 1, color=M2, alpha=0.15,
step="mid", label="loaded die in use")
six = y[t0:t1] == 6
ax.scatter(tt[~six], y[t0:t1][~six], s=14, color="0.45")
ax.scatter(tt[six], y[t0:t1][six], s=30, color=M2, zorder=3, label="a six")
ax.set_yticks(range(1, 7)); ax.set_ylim(0.5, 6.5); ax.set_ylabel("roll")
fig, axes = plt.subplots(2, 1, figsize=(13, 5.4), sharex=True)
plot_rolls(axes[0], y_cas)
axes[0].set_title("what we see", loc="left")
plot_rolls(axes[1], y_cas, x_cas)
axes[1].set_title("what we do not see", loc="left"); axes[1].set_xlabel("$t$")
axes[1].legend(loc="lower right", fontsize=10, ncol=2)
plt.tight_layout(); plt.show()
Even knowing where the loaded die was used, the rolls give it away only weakly: a six is evidence, but fair dice roll sixes too. The stretches are what carry the information.
Copy number changes in a tumour. The hidden state is the number of copies in each bin, it changes rarely, and each bin reports it through a noisy read-depth ratio. The “time” index is position along the genome.
Shaded nodes are observed. A single chain of hidden states, each emitting one observation. Compare with the mixture model’s graph: the plate is gone, and in its place is an arrow from each \(z_t\) to the next.
def simulate_hmm(T, pi, A, mu, sd, rng):
K = len(pi)
z = np.zeros(T, dtype=int)
z[0] = rng.choice(K, p=pi)
for t in range(1, T):
z[t] = rng.choice(K, p=A[z[t-1]])
return z, rng.normal(mu[z], sd[z])
pi0 = np.array([0.5, 0.5])
mu, sd = np.array([0.0, 4.0]), np.array([0.8, 0.8])
settings = {
"sticky: A = [[.95,.05],[.05,.95]]": np.array([[0.95, 0.05], [0.05, 0.95]]),
"alternating: A = [[.10,.90],[.90,.10]]": np.array([[0.10, 0.90], [0.90, 0.10]]),
}
fig, axes = plt.subplots(2, 1, figsize=(13, 6), sharex=True, sharey=True)
for ax, (name, A) in zip(axes, settings.items()):
z, y = simulate_hmm(120, pi0, A, mu, sd, np.random.default_rng(1))
ax.plot(y, color="0.75", lw=1, zorder=1)
ax.scatter(np.arange(len(y)), y, c=z, cmap="coolwarm", s=26, zorder=2)
ax.set_title(name, fontsize=13, loc="left")
axes[1].set_xlabel("$t$")
plt.tight_layout(); plt.show()
Same emissions, same stationary distribution \((0.5, 0.5)\). Only \(A\) differs.
The likelihood requires marginalizing every hidden state:
\[ p(y_{1:T} \mid \theta) = \sum_{z_1=1}^{K} \sum_{z_2=1}^{K} \cdots \sum_{z_T=1}^{K} \; p(z_1) \prod_{t=2}^{T} p(z_t \mid z_{t-1}) \prod_{t=1}^{T} p(y_t \mid z_t). \]
logB_cas = np.log(B_cas[:, y_cas - 1].T) # T x K
alpha_cas, c_cas = forward(logB_cas, pi_cas, A_cas)
fig, axes = plt.subplots(2, 1, figsize=(13, 5.6), sharex=True,
gridspec_kw={"height_ratios": [1, 1.3]})
plot_rolls(axes[0], y_cas, x_cas)
axes[1].plot(alpha_cas[:, 1], color=M1, lw=1.8, label=r"filtered $p(z_t{=}\mathrm{loaded} \mid y_{1:t})$")
axes[1].plot(x_cas, color="0.55", lw=1, ls="--", label="true state")
axes[1].set_ylim(-0.05, 1.05); axes[1].set_xlabel("$t$"); axes[1].legend(loc="upper right", fontsize=10)
plt.tight_layout(); plt.show()
print(f"log p(y_1:300) = {np.log(c_cas).sum():.2f}")
print(f"filtering accuracy (argmax vs true state): {(alpha_cas.argmax(1) == x_cas).mean():.1%}")
log p(y_1:300) = -530.56
filtering accuracy (argmax vs true state): 73.0%The filter reacts to sixes and decays back towards fair between them. It is always a little late: at a switch it has to wait for evidence.
Back to the chromosome. Around a copy number change, the filter has a structural problem: it can only use the bins it has already seen.
So we need a second quantity that summarizes the future the way \(\alpha_t\) summarizes the past.
breakpoint at bin 297: filter settles 4 bins late, smoother 2 bin(s) off
whole chromosome, 600 bins: filtering accuracy 97.3%, smoothing accuracy 99.5%The same breakpoint as before. The bins after the change have now told the bins at the change what they were. The smoother is still off by a bin or two, and honestly so: look at the data at bins 297 and 298, where one point sits at two copies, and the smoothed posterior there is a genuine maybe rather than a confident no. Over the whole chromosome the gain in accuracy is a few percent; at the breakpoints, which is where the answer lives, it is most of the answer.
gamma_cas, xi_cas, ll_cas, _ = forward_backward(logB_cas, pi_cas, A_cas)
fig, axes = plt.subplots(2, 1, figsize=(13, 5.6), sharex=True,
gridspec_kw={"height_ratios": [1, 1.3]})
plot_rolls(axes[0], y_cas, x_cas)
axes[1].plot(alpha_cas[:, 1], color=M1, lw=1.4, alpha=0.7, label=r"filtered $p(z_t{=}\mathrm{loaded} \mid y_{1:t})$")
axes[1].plot(gamma_cas[:, 1], color=M2, lw=2.2, label=r"smoothed $p(z_t{=}\mathrm{loaded} \mid y_{1:T})$")
axes[1].plot(x_cas, color="0.55", lw=1, ls="--", label="true state")
axes[1].set_ylim(-0.05, 1.05); axes[1].set_xlabel("$t$"); axes[1].legend(loc="upper right", fontsize=10)
plt.tight_layout(); plt.show()
print(f"filtering accuracy : {(alpha_cas.argmax(1) == x_cas).mean():.1%}")
print(f"smoothing accuracy : {(gamma_cas.argmax(1) == x_cas).mean():.1%}")
filtering accuracy : 73.0%
smoothing accuracy : 80.3%Smoothing is sharper and its switches are better placed: it has already seen the sixes that the filter is still waiting for. Neither is perfect. A loaded die that happens to roll no sixes for a while is indistinguishable from a fair one.