#!/usr/bin/env python3 """Paired Monte Carlo for fresh-sample vs historical-mean gradients. Population loss: E[0.5*(theta-X)**2 + log(cosh(theta))], X ~ N(m, 1). Each replication uses one independent sample stream; within a replication, algorithms share that stream. PR averages include theta_1,...,theta_T, BEFORE each update, exactly as defined in the manuscript. """ from __future__ import annotations import argparse import csv import json import time from pathlib import Path import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt def optimum(mean: float) -> float: lo, hi = mean - 2.0, mean + 2.0 for _ in range(100): mid = (lo + hi) / 2 if mid + np.tanh(mid) < mean: lo = mid else: hi = mid return (lo + hi) / 2 def summarize(z: np.ndarray) -> dict: n = z.size center = z - z.mean() variance = float(z.var(ddof=1)) variance_mcse = float(np.sqrt(np.mean((center**2 - np.mean(center**2))**2) / n)) return {'mean': float(z.mean()), 'mean_mcse': float(np.sqrt(variance / n)), 'variance': variance, 'variance_mcse': variance_mcse, 'scaled_mse': float(np.mean(z**2))} def simulate(args): rng = np.random.default_rng(args.seed) theta_star = optimum(args.mean) curvature = 1 + 1 / np.cosh(theta_star)**2 theta_fresh = np.full(args.reps, args.initial, dtype=float) theta_history = theta_fresh.copy() running_mean = np.zeros(args.reps) sum_fresh = np.zeros(args.reps) sum_history = np.zeros(args.reps) samples = {} rows = [] started = time.perf_counter() for t in range(1, max(args.horizons) + 1): # This sum is the requested PR average of pre-update iterates. sum_fresh += theta_fresh sum_history += theta_history x = rng.normal(args.mean, 1.0, args.reps) running_mean += (x - running_mean) / t if t in args.horizons: zf = np.sqrt(t) * (sum_fresh / t - theta_star) zh = np.sqrt(t) * (sum_history / t - theta_star) samples[f'fresh_T{t}'] = zf samples[f'history_T{t}'] = zh row = {'T': t, 'fresh': summarize(zf), 'history': summarize(zh), 'fresh_gradient_mse': float(np.mean((x - args.mean)**2)), 'history_gradient_mse': float(np.mean((running_mean - args.mean)**2)), 'fresh_gradient_mse_theory': 1.0, 'history_gradient_mse_theory': 1.0 / t} rows.append(row) print(f'T={t:>6}: fresh var={row["fresh"]["variance"]:.6f}, ' f'history var={row["history"]["variance"]:.6f}, ' f'ratio={row["history"]["variance"]/row["fresh"]["variance"]:.3f}', flush=True) eta = args.step_scale * t**(-args.alpha) # Derivatives of the sample loss, using fresh or estimated E[X]. theta_fresh -= eta * (theta_fresh + np.tanh(theta_fresh) - x) theta_history -= eta * (theta_history + np.tanh(theta_history) - running_mean) return samples, rows, {'theta_star': theta_star, 'H': float(curvature), 'fresh_variance': float(curvature**-2), 'history_variance': float(2 * curvature**-2), 'elapsed_seconds': time.perf_counter() - started} def plot(samples, rows, theory, args): plt.rcParams.update({'font.family': 'serif', 'font.size': 11, 'axes.spines.top': False, 'axes.spines.right': False, 'pdf.fonttype': 42, 'ps.fonttype': 42}) t = max(args.horizons) draws = [samples[f'fresh_T{t}'], samples[f'history_T{t}']] colors = ['#3569a8', '#c46b35'] variances = [theory['fresh_variance'], theory['history_variance']] titles = ['Fresh-sample gradient', 'Historical-mean gradient'] extent = max(max(float(np.max(np.abs(z))) for z in draws), 4.3 * np.sqrt(max(variances))) * 1.03 edges = np.linspace(-extent, extent, args.bins + 1) grid = np.linspace(-extent, extent, 1000) fig, axes = plt.subplots(1, 2, figsize=(9.4, 3.7), sharex=True, sharey=True) for ax, z, color, variance, title in zip(axes, draws, colors, variances, titles): ax.hist(z, bins=edges, density=True, color=color, alpha=.58, edgecolor='white', linewidth=.35, label='Monte Carlo') density = np.exp(-grid**2 / (2 * variance)) / np.sqrt(2 * np.pi * variance) ax.plot(grid, density, color=color, linewidth=2, label='Asymptotic normal') stat = summarize(z) ax.text(.97, .96, f'Empirical variance: {stat["variance"]:.4f}\n' f'Asymptotic variance: {variance:.4f}\n' f'Empirical mean: {stat["mean"]:.4f}', transform=ax.transAxes, ha='right', va='top', fontsize=9, bbox={'facecolor': 'white', 'alpha': .85, 'edgecolor': 'none'}) ax.set_title(title) ax.set_xlabel(r'$\sqrt{T}(\bar\theta_T-\theta^*)$') ax.set_xlim(-extent, extent) ax.legend(loc='upper left', fontsize=8, frameon=False) axes[0].set_ylabel('Density') fig.suptitle(f'Nonlinear stochastic optimization | T = {t:,} | {args.reps:,} replications', fontsize=12) fig.tight_layout() for suffix in ['png', 'pdf']: fig.savefig(args.output / f'pr_histograms.{suffix}', dpi=220, bbox_inches='tight') plt.close(fig) fig, ax = plt.subplots(figsize=(6.5, 3.7)) for key, color, variance, title in zip(['fresh', 'history'], colors, variances, titles): ax.errorbar([r['T'] for r in rows], [r[key]['variance'] for r in rows], yerr=[1.96 * r[key]['variance_mcse'] for r in rows], fmt='o-', color=color, capsize=3, label=title) ax.axhline(variance, color=color, linestyle='--', linewidth=1.3) ax.set_xscale('log') ax.set_xlabel('Horizon T') ax.set_ylabel(r'Variance of $\sqrt{T}(\bar\theta_T-\theta^*)$') ax.set_title('Variance convergence (bars: 1.96 Monte Carlo SE)') ax.legend(frameon=False, fontsize=9) fig.tight_layout() for suffix in ['png', 'pdf']: fig.savefig(args.output / f'variance_convergence.{suffix}', dpi=220, bbox_inches='tight') plt.close(fig) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--reps', type=int, default=10000) parser.add_argument('--horizons', type=int, nargs='+', default=[1000, 5000, 20000, 50000, 100000, 200000]) parser.add_argument('--seed', type=int, default=20261016) parser.add_argument('--alpha', type=float, default=.65) parser.add_argument('--step-scale', type=float, default=.5) parser.add_argument('--initial', type=float, default=2.0) parser.add_argument('--mean', type=float, default=0.0) parser.add_argument('--bins', type=int, default=65) parser.add_argument('--output', type=Path, default=Path(__file__).resolve().parent / 'results' / 'experiment_1') args = parser.parse_args() if args.reps < 2 or min(args.horizons) < 1 or not .5 < args.alpha < 1 or args.step_scale <= 0: parser.error('Require reps >= 2, positive horizons/step-scale, and 0.5 < alpha < 1.') args.horizons = sorted(set(args.horizons)) args.output.mkdir(parents=True, exist_ok=True) samples, rows, theory = simulate(args) np.savez_compressed(args.output / 'scaled_pr_draws.npz', **samples) config = {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()} report = {'config': config, 'theory': theory, 'numpy_version': np.__version__, 'results': rows, 'interpretation': 'Total PR variances of two gradient strategies for the same objective; ' 'their difference is not an additive perturbation comparison.'} (args.output / 'summary.json').write_text(json.dumps(report, indent=2) + '\n') with (args.output / 'summary.csv').open('w', newline='') as handle: fields = ['T', 'algorithm', 'mean', 'mean_mcse', 'variance', 'variance_mcse', 'scaled_mse', 'theory_variance', 'gradient_mse', 'theory_gradient_mse'] writer = csv.DictWriter(handle, fieldnames=fields) writer.writeheader() for r in rows: for key in ['fresh', 'history']: writer.writerow({'T': r['T'], 'algorithm': key, **r[key], 'theory_variance': theory[f'{key}_variance'], 'gradient_mse': r[f'{key}_gradient_mse'], 'theory_gradient_mse': r[f'{key}_gradient_mse_theory']}) plot(samples, rows, theory, args) print(f'Elapsed: {theory["elapsed_seconds"]:.1f}s; output: {args.output}', flush=True) if __name__ == '__main__': main()