repro-bundle-v0/experiments/exp2_pushpull_rate.py

"""Claim 3 (Theorem 3.6): single-level Push-Pull over time-varying directed
graphs, min_x (1/n) sum_i f_i(x) (Eq. 9). Theorem 3.6 uses a FIXED step size
eta_x = O(1) (independent of the horizon K) and claims
min_{0<=k<K-1} ||grad F(x_bar^k)||^2 = O((ab)^{-n} K^{-1}).

We run ONE long trajectory with a fixed step size and read off the
running-min statistic at increasing horizons K -- exactly what the theorem
bounds -- rather than re-tuning per K (which would be the FAB/Theorem 3.4
protocol, not this one).

As in exp1, the noise-free run converges geometrically (over-satisfies the
bound, uninformative about the exponent) once eta is small enough to be
stable; we additionally run a stochastic-gradient variant (persistent i.i.d.
noise) whose fitted slope we compare directly against the claimed -1
exponent. Consensus under *persistent* noise and a *fixed* step size cannot
vanish at all (standard decentralized-SGD bias-floor behavior) -- that half
of the claim is instead confirmed in the noise-free run below, together with
the step-size-threshold ("properly chosen step size") finding.
"""
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from fab import run_pushpull  # noqa: E402
from graphs import TimeVaryingDigraphs  # noqa: E402
from problems import SingleLevelNonconvex  # noqa: E402

RESULTS_DIR = Path(__file__).resolve().parents[1] / "results"
RESULTS_DIR.mkdir(exist_ok=True)

K_MAX = 25600
CHECKPOINTS = [200, 400, 800, 1600, 3200, 6400, 12800, 25600]
SEEDS = [0, 1, 2]
ETA_X = 0.01  # fixed, independent of K -- as Theorem 3.6 prescribes.
# Below a critical threshold ("a properly chosen step size", Thm 3.6) FAB
# vanishes to the exact stationary point; above it, heterogeneous local
# minimizers induce a nonvanishing consensus bias/plateau instead. See
# eta_sensitivity below for the direct confirmation of that qualitative claim.
ETA_SENSITIVITY = [0.05, 0.02, 0.01, 0.005, 0.002]


def main():
    n, d, extra_edges = 8, 10, 2
    grad_runs, cons_runs = [], []
    for seed in SEEDS:
        prob = SingleLevelNonconvex(n=n, d=d, heterogeneity=0.5, seed=seed)
        graphs = TimeVaryingDigraphs(n=n, extra_edges=extra_edges, seed=seed)
        out = run_pushpull(prob, graphs, K=K_MAX, eta_x=ETA_X, seed=seed)
        grad_runs.append(out["grad_norm_running_min"])
        cons_runs.append(out["cons_running_min"])

    grad_mean = np.mean(grad_runs, axis=0)
    cons_mean = np.mean(cons_runs, axis=0)

    rows = []
    for K in CHECKPOINTS:
        rows.append({
            "K": K,
            "grad_norm_sq": float(grad_mean[K - 1]),
            "cons": float(cons_mean[K - 1]),
            "grad_norm_sq_times_K": float(grad_mean[K - 1] * K),  # O(1/K) bound check
            "cons_times_K": float(cons_mean[K - 1] * K),
        })
        print(f"K={K:6d} grad^2={rows[-1]['grad_norm_sq']:.3e} "
              f"(x K = {rows[-1]['grad_norm_sq_times_K']:.3e})  "
              f"cons={rows[-1]['cons']:.3e} (x K = {rows[-1]['cons_times_K']:.3e})")

    Ks = np.array([r["K"] for r in rows])
    slope_grad = float(np.polyfit(np.log(Ks), np.log(np.maximum([r["grad_norm_sq"] for r in rows], 1e-300)), 1)[0])
    slope_cons = float(np.polyfit(np.log(Ks), np.log(np.maximum([r["cons"] for r in rows], 1e-300)), 1)[0])
    print(f"fitted slopes: grad^2={slope_grad:.3f}  cons={slope_cons:.3f}  (theory: both <= -1)")

    # Step-size sensitivity: confirms Theorem 3.6's "properly chosen step size"
    # qualifier is load-bearing -- above a threshold, heterogeneous local minimizers
    # induce a nonvanishing consensus bias; below it, consensus vanishes exactly.
    sens_rows = []
    prob = SingleLevelNonconvex(n=n, d=d, heterogeneity=0.5, seed=0)
    for eta in ETA_SENSITIVITY:
        graphs = TimeVaryingDigraphs(n=n, extra_edges=extra_edges, seed=0)
        out = run_pushpull(prob, graphs, K=K_MAX, eta_x=eta, seed=0)
        cons = out["cons_running_min"]
        sens_rows.append({
            "eta_x": eta,
            "cons_at_1600": float(cons[1599]),
            "cons_at_6400": float(cons[6399]),
            "cons_at_25600": float(cons[-1]),
        })
        print(f"eta={eta}: cons@1600={cons[1599]:.3e} cons@6400={cons[6399]:.3e} cons@25600={cons[-1]:.3e}")

    # Stochastic variant: persistent gradient noise keeps the problem hard across
    # the whole horizon, so the running-min gradient-norm rate actually exercises
    # the theorem's O(K^-1) exponent instead of geometric over-satisfaction.
    stoch_cps = [100, 200, 400, 800, 1600, 3200, 6400, 12800, 25600]
    stoch_eta, stoch_noise = 0.005, 0.3
    stoch_runs = []
    for seed in SEEDS:
        prob = SingleLevelNonconvex(n=n, d=d, heterogeneity=0.5, seed=seed)
        graphs = TimeVaryingDigraphs(n=n, extra_edges=extra_edges, seed=seed)
        out = run_pushpull(prob, graphs, K=max(stoch_cps), eta_x=stoch_eta, seed=seed, noise_std=stoch_noise)
        stoch_runs.append([out["grad_norm_running_min"][k - 1] for k in stoch_cps])
    stoch_mean = np.mean(stoch_runs, axis=0)
    stoch_slope = float(np.polyfit(np.log(stoch_cps), np.log(np.maximum(stoch_mean, 1e-300)), 1)[0])
    print(f"stochastic (eta={stoch_eta}, noise={stoch_noise}) grad^2 slope={stoch_slope:.3f}  (theory: -1)")
    for kk, vv in zip(stoch_cps, stoch_mean):
        print(f"  K={kk:6d} grad^2={vv:.3e}")

    report = {
        "eta_x": ETA_X, "rows": rows,
        "fitted_slopes": {"grad_norm_sq": slope_grad, "cons": slope_cons},
        "eta_sensitivity": sens_rows,
        "stochastic": {
            "eta_x": stoch_eta, "noise_std": stoch_noise,
            "checkpoints": stoch_cps, "grad_norm_sq": stoch_mean.tolist(),
            "fitted_slope": stoch_slope,
        },
    }
    with open(RESULTS_DIR / "exp2_pushpull_rate.json", "w") as f:
        json.dump(report, f, indent=2)
    print("Saved:", RESULTS_DIR / "exp2_pushpull_rate.json")


if __name__ == "__main__":
    main()