1#!/usr/bin/env python3
  2"""Exact toy verification of Fisher-geodesic finite-step annealing.
  3
  4Path: p_theta = Normal(0, q(theta)), q(theta)=q0*exp(log_ratio*theta).
  5For this path, I(theta)=0.5*log_ratio**2 and equal Fisher arc length is
  6therefore uniform theta.  To make the comparison nontrivial, the baseline
  7uses a linear variance path q(theta)=q0+(q1-q0)*theta, whose Fisher speed
  8varies strongly; the proposed schedule is the inverse-arc schedule on that
  9same path (equivalently, geometric spacing of q).
 10"""
 11import csv
 12import json
 13import math
 14from pathlib import Path
 15import numpy as np
 16
 17
 18def q_linear(theta, q0, q1):
 19    return q0 + (q1-q0)*np.asarray(theta)
 20
 21
 22def fisher_speed(theta, q0, q1):
 23    q = q_linear(theta, q0, q1)
 24    dq = q1-q0
 25    return abs(dq)/(math.sqrt(2.0)*q)
 26
 27
 28def fisher_length(q0, q1):
 29    return abs(math.log(q1/q0))/math.sqrt(2.0)
 30
 31
 32def geodesic_schedule(n, q0, q1):
 33    # Equal arc length: log(q_k/q0)=k*log(q1/q0)/n.
 34    k = np.arange(n+1, dtype=float)
 35    q = q0 * (q1/q0)**(k/n)
 36    return (q-q0)/(q1-q0)
 37
 38
 39def linear_schedule(n):
 40    return np.linspace(0.0, 1.0, n+1)
 41
 42
 43def kl_zero_gaussians(q_from, q_to):
 44    # KL(N(0,q_from) || N(0,q_to)).
 45    r = q_from/q_to
 46    return 0.5*(r - 1.0 - math.log(r))
 47
 48
 49def local_kls(schedule, q0, q1):
 50    q = q_linear(schedule, q0, q1)
 51    return np.array([kl_zero_gaussians(q[i], q[i+1]) for i in range(len(q)-1)])
 52
 53
 54def arc_steps(schedule, q0, q1):
 55    q = q_linear(schedule, q0, q1)
 56    return np.abs(np.log(q[1:]/q[:-1]))/math.sqrt(2.0)
 57
 58
 59def summarize(schedule, q0, q1):
 60    kls = local_kls(schedule, q0, q1)
 61    return {
 62        "cumulative_kl": float(kls.sum()),
 63        "local_kl_cv": float(kls.std()/kls.mean()),
 64        "arc_step_cv": float(arc_steps(schedule,q0,q1).std()/arc_steps(schedule,q0,q1).mean()),
 65        "local_kls": kls.tolist(),
 66    }
 67
 68
 69def fit_slope(ns, vals):
 70    return float(np.polyfit(np.log(ns), np.log(vals), 1)[0])
 71
 72
 73def run():
 74    # Exact Fisher identity check against Monte Carlo score estimates.
 75    rng = np.random.default_rng(2733)
 76    q0, q1 = 0.05, 20.0
 77    theta = 0.37
 78    q = q_linear(theta,q0,q1)
 79    dq = q1-q0
 80    x = rng.normal(0.0, math.sqrt(q), size=1_000_000)
 81    # d/dtheta log N(x;0,q(theta)) = q'*(x^2-q)/(2q^2)
 82    score = dq*(x*x-q)/(2*q*q)
 83    mc_fisher = float(np.mean(score*score))
 84    exact_fisher = fisher_speed(theta,q0,q1)**2
 85    fisher_relerr = abs(mc_fisher-exact_fisher)/exact_fisher
 86
 87    ns = [4,8,16,32,64,128]
 88    ranges = [(1.0,2.0),(1.0,10.0),(0.05,20.0)]
 89    rows=[]
 90    for a,b in ranges:
 91        L=fisher_length(a,b)
 92        for n in ns:
 93            base=summarize(linear_schedule(n),a,b)
 94            geo=summarize(geodesic_schedule(n,a,b),a,b)
 95            rows.append({"q0":a,"q1":b,"N":n,"L":L,
 96                         "baseline_cumulative_kl":base["cumulative_kl"],
 97                         "geodesic_cumulative_kl":geo["cumulative_kl"],
 98                         "baseline_local_kl_cv":base["local_kl_cv"],
 99                         "geodesic_local_kl_cv":geo["local_kl_cv"],
100                         "baseline_arc_cv":base["arc_step_cv"],
101                         "geodesic_arc_cv":geo["arc_step_cv"],
102                         "predicted_leading_kl":L*L/(2*n)})
103
104    # Fit the asymptotic exponent on the finer four grids for the wide range.
105    wide=[r for r in rows if r["q0"]==0.05 and r["N"]>=16]
106    slope_geo=fit_slope(np.array([r["N"] for r in wide]),np.array([r["geodesic_cumulative_kl"] for r in wide]))
107    slope_base=fit_slope(np.array([r["N"] for r in wide]),np.array([r["baseline_cumulative_kl"] for r in wide]))
108
109    # Mechanism summaries at N=32, including the predicted leading coefficient.
110    mech=[]
111    for a,b in ranges:
112        r=[z for z in rows if z["q0"]==a and z["q1"]==b and z["N"]==32][0]
113        mech.append({"range":f"{a:g}->{b:g}","fisher_length":r["L"],
114                     "predicted_equal_arc_cv":0.0,
115                     "observed_geodesic_arc_cv":r["geodesic_arc_cv"],
116                     "observed_geodesic_local_kl_cv":r["geodesic_local_kl_cv"],
117                     "observed_linear_local_kl_cv":r["baseline_local_kl_cv"],
118                     "predicted_kl":r["predicted_leading_kl"],
119                     "observed_kl":r["geodesic_cumulative_kl"],
120                     "ratio_linear_over_geodesic":r["baseline_cumulative_kl"]/r["geodesic_cumulative_kl"]})
121
122    out={"fisher_check":{"mc_fisher":mc_fisher,"exact_fisher":exact_fisher,"relative_error":fisher_relerr},
123         "predictions":{"equal_arc_arc_cv":mech,
124                         "wide_range_loglog_slope_geodesic":slope_geo,
125                         "wide_range_loglog_slope_linear_baseline":slope_base,
126                         "asymptotic_target_slope":-1.0},
127         "rows":rows}
128    Path("toy_results.json").write_text(json.dumps(out,indent=2))
129    with open("toy_results.csv","w",newline="") as f:
130        w=csv.DictWriter(f,fieldnames=rows[0].keys()); w.writeheader(); w.writerows(rows)
131    print(json.dumps({"fisher_check":out["fisher_check"],"predictions":out["predictions"]},indent=2))
132
133if __name__ == "__main__": run()