#!/usr/bin/env python3 """Gain sweep for the nonlinear mean-estimation objective of experiment 1. Estimator: mhat_0=0, mhat_t=(1-a/t)*mhat_{t-1}+(a/t)*X_t. SGD: theta_{t+1}=theta_t-0.5*t^(-0.65)*(theta_t+tanh(theta_t)-mhat_t). All gains share the same sample stream within each independent replication. """ 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 from pr_histograms import optimum, summarize def plot(report, args): plt.rcParams.update({'font.family': 'serif', 'font.size': 11, 'axes.spines.top': False, 'axes.spines.right': False, 'pdf.fonttype': 42}) fig, axes = plt.subplots(1, 2, figsize=(9.5, 3.8)) curve = np.geomspace(min(args.gains), max(args.gains), 700) h = report['theory']['H'] axes[0].plot(curve, curve**2/(2*curve-1), color='#333333', linewidth=1.8, label='Asymptotic theory') axes[1].plot(curve, (2*curve/(2*curve-1))/h**2, color='#333333', linewidth=1.8, label='Asymptotic theory') colors = ['#7198ba', '#c46b35'] horizons = args.horizons if len(args.horizons) <= 2 else [args.horizons[0], args.horizons[-1]] for t, color in zip(horizons, colors): rows = [r for r in report['results'] if r['T'] == t] axes[0].errorbar([r['gain'] for r in rows], [r['single_step_scaled_mse'] for r in rows], yerr=[1.96*r['single_step_scaled_mse_mcse'] for r in rows], fmt='o--', color=color, linewidth=1, markersize=4, capsize=2, label=f'Empirical, T={t:,}') axes[1].errorbar([r['gain'] for r in rows], [r['pr']['variance'] for r in rows], yerr=[1.96*r['pr']['variance_mcse'] for r in rows], fmt='o--', color=color, linewidth=1, markersize=4, capsize=2, label=f'Empirical, T={t:,}') axes[0].axvline(1.0, color='#777777', linestyle=':', linewidth=1) axes[1].axhline(1/h**2, color='#777777', linestyle=':', linewidth=1, label='Asymptotic lower bound') axes[0].set_ylabel(r'$T\,\mathbb{E}(\widehat m_T-m)^2$') axes[1].set_ylabel(r'$\operatorname{Var}[\sqrt{T}(\bar\theta_T-\theta^*)]$') axes[0].set_title('Single-step estimation error') axes[1].set_title('Downstream PR variance') for ax in axes: ax.set_xscale('log') ax.set_xlabel('Estimator gain a') ax.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter()) ticks = [.6, 1, 2, 4, 8] ax.set_xticks([v for v in ticks if min(args.gains) <= v <= max(args.gains)]) ax.set_xticklabels([str(v) for v in ticks if min(args.gains) <= v <= max(args.gains)]) ax.legend(frameon=False, fontsize=8) ax.grid(axis='y', alpha=.15) fig.suptitle(f'Same stochastic objective | {args.reps:,} paired replications per gain', fontsize=12) fig.tight_layout() for ext in ['pdf', 'png']: fig.savefig(args.output/f'gain_sweep.{ext}', dpi=220, bbox_inches='tight') plt.close(fig) def main(): parser=argparse.ArgumentParser(description=__doc__) parser.add_argument('--gains', nargs='+', type=float, default=[.6,.75,1,1.25,1.5,2,3,5,8]) parser.add_argument('--horizons', nargs='+', type=int, default=[50000,200000]) parser.add_argument('--reps', type=int, default=5000) parser.add_argument('--seed', type=int, default=20261017) 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('--output', type=Path, default=Path(__file__).resolve().parent/'results'/'experiment_2') args=parser.parse_args() args.gains=sorted(set(args.gains)) args.horizons=sorted(set(args.horizons)) if min(args.gains)<=.5 or min(args.horizons)<1 or args.reps<2 or not .5 0.5, positive horizons/step-scale, reps >= 2, and 0.5 < alpha < 1.') args.output.mkdir(parents=True,exist_ok=True) gains=np.array(args.gains)[:,None] theta=np.full((len(args.gains),args.reps),args.initial) mhat=np.zeros_like(theta) theta_sum=np.zeros_like(theta) rng=np.random.default_rng(args.seed) star=optimum(args.mean) h=float(1+1/np.cosh(star)**2) rows=[] draws={} started=time.perf_counter() progress_stride=10000 for t in range(1,max(args.horizons)+1): theta_sum += theta x=rng.normal(args.mean,1,args.reps) mhat += (gains/t)*(x[None,:]-mhat) if t in args.horizons: z=np.sqrt(t)*(theta_sum/t-star) errors=np.sqrt(t)*(mhat-args.mean) draws[f'scaled_pr_T{t}']=z.copy() draws[f'scaled_estimator_T{t}']=errors.copy() for i,gain in enumerate(args.gains): losses=errors[i]**2 rows.append({'T':t,'gain':gain,'pr':summarize(z[i]), 'pr_variance_theory':float((2*gain/(2*gain-1))/h**2), 'single_step_scaled_mse':float(losses.mean()), 'single_step_scaled_mse_mcse':float(losses.std(ddof=1)/np.sqrt(args.reps)), 'single_step_scaled_mse_theory':float(gain**2/(2*gain-1))}) print(f'Checkpoint T={t:,}: saved all {len(args.gains)} gains',flush=True) eta=args.step_scale*t**(-args.alpha) theta -= eta*(theta+np.tanh(theta)-mhat) if t % progress_stride == 0: print(f'T={t:,}/{max(args.horizons):,}; elapsed={time.perf_counter()-started:.1f}s',flush=True) config={k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()} report={'config':config,'theory':{'theta_star':star,'H':h,'lower_bound':h**-2}, 'results':rows,'elapsed_seconds':time.perf_counter()-started, 'numpy_version':np.__version__,'matplotlib_version':matplotlib.__version__} np.savez_compressed(args.output/'gain_draws.npz',**draws,gains=np.array(args.gains)) (args.output/'summary.json').write_text(json.dumps(report,indent=2)+'\n') flat=[] for r in rows: flat.append({k:v for k,v in r.items() if k!='pr'}|{f'pr_{k}':v for k,v in r['pr'].items()}) with (args.output/'summary.csv').open('w',newline='') as handle: writer=csv.DictWriter(handle,fieldnames=list(flat[0])) writer.writeheader();writer.writerows(flat) plot(report,args) print('Final-horizon results:',flush=True) for r in rows: if r['T']==max(args.horizons): print(f'a={r["gain"]:4g} single MSE*T={r["single_step_scaled_mse"]:.4f} ' f'(theory {r["single_step_scaled_mse_theory"]:.4f}); ' f'PR variance={r["pr"]["variance"]:.4f} ' f'(theory {r["pr_variance_theory"]:.4f})',flush=True) print(f'Elapsed {report["elapsed_seconds"]:.1f}s; output: {args.output}',flush=True) if __name__=='__main__': main()