codexbranching-generation-distributions.py
branching-generation-distributions.py
"""Exact Galton-Watson generation PMFs for p=(1/4, 1/4, 1/2).
Tested with Python 3.14.4, NumPy 2.3.5 and Matplotlib 3.10.7.
Build: make _media/branching-generation-distributions.png
Writes an opaque white same-basename PNG to the working directory.
"""
from fractions import Fraction
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
def generation_distributions(last=6):
"""Return exact rational coefficients of G_n=F composed n times.
Independently check normalization, extinction iteration and mean.
Coefficient k is P(Z_n=k); no Monte Carlo or truncated support.
"""
quarter, half = Fraction(1, 4), Fraction(1, 2)
pmf = [Fraction(0), Fraction(1)]
extinction = Fraction(0)
distributions = []
for n in range(last + 1):
assert len(pmf) == 2**n + 1
assert all(p >= 0 for p in pmf)
assert sum(pmf) == 1
assert pmf[0] == extinction
mean = sum(k * p for k, p in enumerate(pmf))
assert mean == Fraction(5, 4)**n
if n == 2:
assert pmf == [Fraction(11, 32), Fraction(1, 8),
Fraction(9, 32), Fraction(1, 8), Fraction(1, 8)]
distributions.append(pmf)
if n == last:
break
# G_{n+1} = 1/4 + G_n/4 + G_n^2/2, computed over rationals.
square = [Fraction(0)] * (2 * len(pmf) - 1)
for i, a in enumerate(pmf):
for j, b in enumerate(pmf):
square[i + j] += a * b
next_pmf = [half * coefficient for coefficient in square]
for k, p in enumerate(pmf):
next_pmf[k] += quarter * p
next_pmf[0] += quarter
pmf = next_pmf
extinction = quarter + quarter * extinction + half * extinction**2
return distributions
def main():
distributions = generation_distributions()
for n, pmf in enumerate(distributions):
mean = sum(k * p for k, p in enumerate(pmf))
print(f"n={n}: full support 0..{len(pmf)-1}, exact mass=1, "
f"q_n={float(pmf[0]):.12f}, mean={float(mean):.12f}")
plt.rcParams.update({"font.size": 11, "axes.titlesize": 13})
fig, axes = plt.subplots(2, 2, figsize=(10, 6.8), dpi=120,
sharey=True, facecolor="white")
for ax, n in zip(axes.flat, (0, 2, 4, 6)):
pmf = distributions[n]
sizes = np.arange(len(pmf))
masses = np.array([float(p) for p in pmf])
colors = ["#c55c28"] + ["#237a9c"] * (len(pmf) - 1)
ax.bar(sizes, masses, color=colors, width=0.8, zorder=3)
mean = float(Fraction(5, 4)**n)
ax.axvline(mean, color="#324254", ls="--", lw=1.4, zorder=4)
# One ancestor is certain at n=0; every later PMF has all sizes 0..2^n.
ax.set(title=f"Generation n = {n}", xlabel=r"Population size $k$",
ylabel=r"Probability $P(Z_n=k)$", ylim=(0, 1.08),
xlim=(-0.8, max(2.8, 2**n + 0.8)))
if n == 0:
ax.set_xticks([0, 1, 2])
elif n == 2:
ax.set_xticks(range(5))
else:
ax.set_xticks(np.arange(0, 2**n + 1, 2**n // 4))
ax.text(0.97, 0.94,
f"Empty: {masses[0]:.6f}\nMean: {mean:.6f}",
transform=ax.transAxes, ha="right", va="top", fontsize=11,
bbox={"facecolor": "white", "edgecolor": "none"})
ax.grid(axis="y", color="#e6ebf0", zorder=0)
ax.set_axisbelow(True)
ax.spines[["top", "right"]].set_visible(False)
fig.suptitle("Exact generation distributions: growth can coexist with extinction",
fontsize=15, y=0.97)
fig.text(0.5, 0.91, r"Offspring: $p_0=1/4$, $p_1=1/4$, $p_2=1/2$; one initial ancestor",
ha="center", fontsize=12)
fig.text(0.5, 0.025,
"Orange: empty families. Blue: nonempty families. Dashed line: mean.\n"
"Each panel shows its full support, on a different horizontal scale; eventual extinction = 1/2.",
ha="center", fontsize=10)
fig.subplots_adjust(left=0.08, right=0.98, top=0.83, bottom=0.14,
hspace=0.44, wspace=0.22)
fig.savefig(Path.cwd() / Path(__file__).with_suffix(".png").name,
facecolor="white", transparent=False)
plt.close(fig)
if __name__ == "__main__":
main()