 codex/gaussian-square-gradient-noise.py

"""Original analytic variance comparison for J(mu)=E[(mu+2 epsilon)^2].

Write PNG into cwd. Python3.14.4, NumPy2.3.5, Matplotlib3.10.7.
Curves are exact formulas, not empirical measurements or copied paper data.
"""
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

mu, sigma = 1.0, 2.0
h = np.geomspace(0.02, 1.0, 240)
independent = (sigma**4 + 2*mu**2*sigma**2)/h**2 + 2*sigma**2
shared = np.full_like(h, 4*sigma**2)
assert np.allclose(independent, 24/h**2 + 8)
assert np.allclose(shared, 16)
assert np.isclose((sigma**4 + 2*mu**2*sigma**2)/0.1**2 + 2*sigma**2, 2408)
epsilon = np.array([-2.5, -1, 0, 1, 3.0])
x = mu + sigma*epsilon
for step in [0.02, 0.1, 1.0]:
    central = ((x+step)**2 - (x-step)**2)/(2*step)
    assert np.allclose(central, 2*x)

plt.rcParams.update({"font.size": 11, "figure.facecolor": "white"})
fig, ax = plt.subplots(figsize=(8, 4.4), dpi=100)
fig.subplots_adjust(left=0.105, right=0.97, top=0.80, bottom=0.25)
fig.suptitle("Freezing the noise changes gradient-estimator variance", fontsize=14, y=0.96)
ax.loglog(h, independent, color="#b0521e", lw=2,
          label="Independent noise: 24/h² + 8")
ax.loglog(h, shared, color="#176b9c", lw=2,
          label="Shared noise / pathwise: 16")
ax.scatter([0.1], [2408], color="#b0521e", s=30, zorder=4)
ax.annotate("h = 0.1: variance 2408", xy=(0.1, 2408), xytext=(0.18, 9000),
            arrowprops={"arrowstyle": "->", "color": "#777777"}, fontsize=10)
ax.set(xlim=(0.02, 1), ylim=(8, 100000), xlabel="Central finite-difference step h",
       ylabel="Estimator variance")
ax.legend(loc="lower left", bbox_to_anchor=(0.0, 0.12), frameon=False, fontsize=10)
ax.grid(which="major", alpha=0.2)
ax.spines[["top", "right"]].set_visible(False)
fig.text(0.5, 0.115, "Chosen quadratic model: mu = 1, sigma = 2; both estimates have mean 2.",
         ha="center", fontsize=10)
fig.text(0.5, 0.055, "Analytic curves. Shared noise is held fixed while the parameter changes.",
         ha="center", fontsize=10)
fig.savefig(Path("gaussian-square-gradient-noise.png"), facecolor="white")
plt.close(fig)