1import json
 2import numpy as np
 3
 4A, Q, R, SIGMA2 = 1.0, 0.80, 1.0, 0.25
 5SIGMA = np.sqrt(SIGMA2)
 6
 7def grad(x): return A*x + Q*x*x/2.0 + R*x*x*x
 8
 9def theoretical_bias(gamma, n, h):
10    hh=(h-1)*(5*h-1)/h
11    return -Q*SIGMA2/(4*A*A)*gamma/n - Q*SIGMA2/(12*A)*hh*gamma*gamma
12
13class BiasEstimator:
14    def __init__(self, delta=0.02, ema=0.08, clip=0.5):
15        self.delta,self.ema,self.clip=delta,ema,clip
16        self.f2,self.f3,self.var=A,0.0,SIGMA2
17    def update(self,x,residuals,gamma,n,h):
18        d=self.delta
19        f2n=(grad(x+d)-grad(x-d))/(2*d)
20        # Since grad=f', f''' is the second derivative of grad.
21        f3n=(grad(x+d)-2*grad(x)+grad(x-d))/(d*d)
22        varn=float(np.var(residuals)); r=self.ema
23        self.f2=(1-r)*self.f2+r*f2n; self.f3=(1-r)*self.f3+r*f3n; self.var=(1-r)*self.var+r*varn
24        aa=max(self.f2,.05); hh=(h-1)*(5*h-1)/h
25        b=-self.f3*self.var/(4*aa*aa)*gamma/n-self.f3*self.var/(12*aa)*hh*gamma*gamma
26        return float(np.clip(b,-self.clip,self.clip))
27
28def simulate(n,h,gamma,mode,seed,rounds=8000,burn=2000):
29    rng=np.random.default_rng(seed); x=0.; xi=np.zeros(n); est=BiasEstimator(); next_c=0.; xs=[]; reported=[]; losses=[]
30    for t in range(rounds):
31        theta=np.full(n,x); residuals=[]
32        for _ in range(h):
33            eps=rng.normal(0,SIGMA,n); residuals.append(eps)
34            theta-=gamma*(grad(theta)+xi+ (next_c if mode=='gradient_correction' and t>=1000 else 0.) + eps)
35        x=float(theta.mean()); xi += (theta-x)/(gamma*h)
36        b=est.update(x,np.concatenate(residuals),gamma,n,h)
37        if mode=='gradient_correction' and t>=1000: next_c=est.f2*b
38        if t>=burn:
39            xs.append(x); reported.append(x-b if mode=='output_correction' else x)
40            losses.append(A*x*x/2+Q*x*x*x/6)
41    xs=np.asarray(xs); reported=np.asarray(reported)
42    return {'mean_x':float(xs.mean()),'abs_bias':float(abs(xs.mean())),'std_x':float(xs.std()),
43            'mean_loss':float(np.mean(losses)),'mean_reported_x':float(reported.mean()),
44            'theory':float(theoretical_bias(gamma,n,h))}
45
46def math_check():
47    est=BiasEstimator(delta=.01,ema=1.)
48    b=est.update(0.,np.random.default_rng(1).normal(0,SIGMA,100000),.01,64,8)
49    return {'f2_hat':est.f2,'f3_hat':est.f3,'sigma2_hat':est.var,'expected_f2':A,'expected_f3':Q,
50            'formula_b':b,'formula_exact':theoretical_bias(.01,64,8)}
51
52def main():
53    out={'math_check':math_check(),'experiments':{}}
54    for h in (8,):
55      for n in (8,64):
56        k=f'n{n}_h{h}'; seed=100+n+h
57        out['experiments'][k]={'vanilla':simulate(n,h,.01,'vanilla',seed),
58          'output_correction':simulate(n,h,.01,'output_correction',seed),
59          'gradient_correction':simulate(n,h,.01,'gradient_correction',seed)}
60    out['scaling_large_N']=[]
61    for gamma in (.012,.020,.028):
62      r=simulate(64,8,gamma,'vanilla',777+int(gamma*10000),rounds=8000,burn=2000)
63      out['scaling_large_N'].append({'gamma':gamma,'empirical_bias':r['mean_x'],'theory':r['theory']})
64    with open('results.json','w') as f: json.dump(out,f,indent=2)
65    print(json.dumps(out,indent=2))
66if __name__=='__main__': main()