failed_by_peak = stats[:,0] < 10
failed_by_time = stats[:,1] < 10
print(failed_by_peak.mean(), failed_by_time.mean())

stats = sir_sim_stats(1000, 100, [99,1,0], 1, .25, rng)
fig, axes = plt.subplots(1, 2, figsize=(9, 4), layout="constrained")
axes[0].hist(stats[:,0], bins=np.arange(.5,65.6,2))
axes[1].hist(stats[:,1], bins=np.arange(0,61,2))
axes[0].set(xlabel="Peak infected population", ylabel="Number of outbreaks")
axes[1].set(xlabel="Final recorded time (days)", ylabel="Number of outbreaks")
plt.show()

def sir_sim_stats(n_sim,t_max,X0,beta,gamma,rng=None):
    result=np.empty((n_sim,2))
    for run in range(n_sim):
        t,s=sir_sim(t_max,X0,beta,gamma,rng)
        result[run]=s[:,1].max(),t[-1]
    return result

import numpy as np
import matplotlib.pyplot as plt

EPIDEMIC_RNG = np.random.default_rng(2345)

def sir_sim(t_max,X0,beta,gamma,rng=None):
    if rng is None:rng=EPIDEMIC_RNG
    state=np.array(X0,dtype=int);population=state.sum();time=0.
    times=[time];states=[state.copy()]
    while time<t_max and state[1]>0:
        infection=beta*state[0]*state[1]/population;recovery=gamma*state[1]
        rate=infection+recovery
        if rate==0:
            times.append(float(t_max));states.append(state.copy());break
        next_time=time+rng.exponential(1/rate)
        if next_time>t_max:
            times.append(float(t_max));states.append(state.copy());break
        if rng.random()<infection/rate:state+=np.array([-1,1,0])
        else:state+=np.array([0,-1,1])
        time=next_time;times.append(time);states.append(state.copy())
    return np.array(times),np.array(states)

rng = np.random.default_rng(2345)
trajectories = [sir_sim(100, [99,1,0], 1, .25, rng) for _ in range(10)]
for times, states in trajectories:
    plt.step(times, states[:,1], where="post", alpha=.8)
plt.xlabel("Time (days)")
plt.ylabel("Infected individuals")
plt.show()

prob = initialise_prob(leaves)
felsenstein(prob, mu_max, len(leaves), left, right, left_time, right_time)
root_posterior = prob[-1] / prob[-1].sum()
print(dict(zip("ACGT", root_posterior)))

def evaluate(mu):
    p = initialise_prob(leaves)
    return felsenstein(p, mu, len(leaves), left, right, left_time, right_time)

coarse = np.linspace(.1, 1., 901)
centre = coarse[np.argmax([evaluate(mu) for mu in coarse])]
fine = np.linspace(max(.1, centre-.001), min(1., centre+.001), 2001)
mu_max = float(fine[np.argmax([evaluate(mu) for mu in fine])])
print(f"mu = {mu_max:.3g}; log-likelihood = {evaluate(mu_max):.8f}")

def felsenstein(prob, mu, L, left, right, left_time, right_time):
    # All four entries are overwritten at every internal node.
    for node in range(L, 2*L-1):
        Tl = np.array([[jukes(a,b,left_time[node],mu)
                        for b in range(4)] for a in range(4)])
        Tr = np.array([[jukes(a,b,right_time[node],mu)
                        for b in range(4)] for a in range(4)])
        prob[node] = (Tl @ prob[left[node]]) * (Tr @ prob[right[node]])
    return float(np.log(prob[-1].mean()))

print(felsenstein(prob, .5, len(leaves), left, right, left_time, right_time))

def initialise_prob(leafList):
    alphabet = {base: j for j, base in enumerate("ACGT")}
    result = np.zeros((2*len(leafList)-1, 4), dtype=float)
    for i, base in enumerate(leafList):
        result[i, alphabet[base]] = 1.
    return result

leaves = ["G", "A", "T", "G", "G", "G", "C", "G"]
prob = initialise_prob(leaves)
print(prob[:len(leaves)])

0 0 1 0
1 0 0 0
0 0 0 1
0 0 1 0
0 0 1 0
0 0 1 0
0 1 0 0
0 0 1 0

def jukes(a, b, t, mu):
    decay = np.exp(-4*mu*t/3)
    return .25 + .75*decay if a == b else .25*(1-decay)

print(f"{jukes(2,2,1,.5):.3f} {jukes(2,3,1,.5):.3f}")

insert_columns = np.flatnonzero(~match_mask)
runs = np.split(insert_columns, np.flatnonzero(np.diff(insert_columns) > 1) + 1)
insert_emissions = []
for run in runs:
    if len(run):
        counts, n = emission_counts(A[:, run].ravel())
        insert_emissions.append({aa: count/n for aa, count in counts.items()})
        print(run+1, n, insert_emissions[-1])

information = []
for j in range(A.shape[1]):
    counts, n = emission_counts(A[:, j])
    p = np.array(list(counts.values()), dtype=float) / n
    H = -np.sum(p * np.log2(p))
    information.append(float(np.log2(20) - H))
print(np.round(information, 4))

01–10: 1.9025, 1.5964, 1.9025, 3.8187, 1.9025, 1.5964, 1.9025, 3.5577, 1.9025, 1.9025
11–20: 2.4420, 3.3355, 1.5964, 4.3219, 3.8187, 1.6803, 3.0975, 3.3355, 3.0975, 2.4420
21–30: 2.5692, 2.8219, 3.5106, 3.3219, 4.3219, 4.3219, 2.8000, 2.0703, 2.6574, 1.8003
31–40: 2.1938, 1.9163, 3.7784, 3.2606, 1.6803, 1.6803, 2.6642, 2.5692, 2.8753, 2.2631
41–50: 2.8753, 3.3355, 1.6803, 2.9703, 3.3355, 2.7914, 2.2086, 2.5692, 4.3219, 1.3742
51–60: 2.2631, 4.3219, 4.3219, 3.8187, 4.3219, 3.5577, 2.8753, 1.9025, 2.7914, 2.0409
61–70: 3.3355, 2.7370, 2.4309, 2.4309, 2.5692, 2.0409, 2.5692, 3.3355, 2.9298, 3.5577
71–79: 1.9025, 1.8187, 4.3219, 2.2631, 3.8187, 2.4309, 1.9025, 1.9025, 1.9025

def emission_counts(column):
    counts = Counter(str(x) for x in column if x != ".")
    n = sum(counts.values())
    if n == 0:
        raise ValueError("an all-gap column has no empirical emission distribution")
    return counts, n

match_emissions = {}
for state, j in enumerate(match_columns, 1):
    counts, n = emission_counts(A[:, j])
    match_emissions[state] = {aa: count/n for aa, count in counts.items()}
    print(state, j+1, match_emissions[state])

from collections import Counter
A = np.asarray(msa_matrix)
gap_fraction = np.mean(A == ".", axis=0)
match_mask = gap_fraction < .30
match_columns = np.flatnonzero(match_mask)
print(match_columns + 1)  # human-readable, one-based alignment columns

for name, a in alignments.items():
    print(f"{name}: query length={a['query_length']}, "
          f"reference length={a['reference_length']}, "
          f"alignment columns={a['alignment_length']}, score={a['score']}")
    print(f"Query {a['query_start']} {a['query']} {a['query_end']}")
    print(f"Ref   {a['reference_start']} {a['reference']} {a['reference_end']}")
    print(f"Identity {a['identity_percent']:.2f}%; "
          f"similarity {a['similarity_percent']:.2f}%; "
          f"gaps {a['gap_percent']:.2f}%")

for name, a in alignments.items():
    print(name, a["alignment_length"], a["identities"],
          a["similarities"], a["gaps"])
    print(a["identity_percent"], a["similarity_percent"], a["gap_percent"])