1import json, math, time
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6SEED=2049
 7np.random.seed(SEED); torch.manual_seed(SEED)
 8try:
 9    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10    if device.type=='cuda':
11        torch.zeros(1,device=device).sum().item()
12except Exception:
13    device=torch.device('cpu')
14
15def directions(n,d,rng):
16    q=rng.normal(size=(n,d)); return q/np.linalg.norm(q,axis=1,keepdims=True)
17
18def wos_ball(x, K, M, d, lam=0., delta=1e-4, corrected=False, seed=0):
19    rng=np.random.default_rng(seed); x=np.asarray(x); B=len(x)
20    z=np.repeat(x[:,None,:],K,axis=1); t=np.zeros((B,K)); y=np.zeros((B,K))
21    for j in range(M):
22        r=1.-np.linalg.norm(z,axis=2)
23        active=r>delta
24        dt=r*r/((d if corrected else 2*d))
25        y += np.exp(-lam*t)*dt # h1=1, h0=0
26        u=directions(B*K,d,rng).reshape(B,K,d)
27        z += r[...,None]*u; t += dt
28        if not active.any(): break
29    # zero boundary data, approximate truncation by zero
30    return y.mean(axis=1), y, t
31
32def wos_square(x,K,M,delta=2e-3,seed=0):
33    rng=np.random.default_rng(seed); x=np.asarray(x); B=len(x); d=2
34    z=np.repeat(x[:,None,:],K,axis=1); t=np.zeros((B,K)); y=np.zeros((B,K))
35    for j in range(M):
36        r=np.minimum(z,1-z).min(axis=2); dt=r*r/(2*d)
37        f=2*(z[:,:,0]*(1-z[:,:,0])+z[:,:,1]*(1-z[:,:,1]))
38        y += f*dt
39        u=directions(B*K,d,rng).reshape(B,K,d); z += r[...,None]*u; t += dt
40        if np.max(r)<delta: break
41    return y.mean(axis=1)
42
43def analytic_square(x):
44    return x[:,0]*(1-x[:,0])*x[:,1]*(1-x[:,1])
45
46def analytic_ball(x): return (1-(x*x).sum(axis=1))/(2*x.shape[1])
47
48class MLP(nn.Module):
49    def __init__(self,d):
50        super().__init__(); self.net=nn.Sequential(nn.Linear(d,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
51    def forward(self,x): return self.net(x).squeeze(-1)
52
53def train_wos(idea=True, steps=350, d=2, K=8):
54    torch.manual_seed(SEED); rng=np.random.default_rng(SEED+3)
55    model=MLP(d).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
56    start=time.time()
57    for s in range(steps):
58        x=rng.uniform(.03,.97,size=(64,d))
59        if idea: y=wos_square(x,K,25,seed=SEED+s)
60        else: y=analytic_square(x) # ideal exact target is a favorable standard supervised reference
61        xt=torch.tensor(x,dtype=torch.float32,device=device); yt=torch.tensor(y,dtype=torch.float32,device=device)
62        loss=((model(xt)-yt)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
63    grid=np.stack(np.meshgrid(np.linspace(.03,.97,25),np.linspace(.03,.97,25),indexing='ij'),-1).reshape(-1,d)
64    with torch.no_grad(): pred=model(torch.tensor(grid,dtype=torch.float32,device=device)).cpu().numpy()
65    rel=np.sqrt(np.mean((pred-analytic_square(grid))**2))/np.sqrt(np.mean(analytic_square(grid)**2))
66    return float(rel),time.time()-start
67
68def main():
69    # Prediction 1: exact Brownian ball exit time is r^2/d; supplied formula predicts r^2/(2d).
70    dvals=[1,2,4,8]; r=.37; exit_rows=[]
71    for d in dvals:
72        supplied=r*r/(2*d); exact=r*r/d
73        exit_rows.append({'d':d,'predicted_supplied':supplied,'exact_MC_theory':exact,'ratio_exact/supplied':exact/supplied})
74    # Direct one-step isotropic identity verifies E|z+rU|^2=|z|^2+r^2 and Poisson exit-time equation.
75    rng=np.random.default_rng(SEED); z=np.array([.2,-.1]); rr=.4; u=directions(500000,2,rng)
76    radial_increment=float(np.mean(np.sum((z+rr*u)**2,axis=1)-np.sum(z*z)))
77    # Prediction 2: with h1=1 and zero boundary, infinite-depth WOS converges to half the exact ball solution under supplied dt.
78    x=np.array([[0.,0.],[.3,0.],[.0,.6]])
79    exact=analytic_ball(x); rows=[]
80    for M in [2,4,8,16,32,64]:
81        est,_,_=wos_ball(x,30000,M,2,seed=SEED+M)
82        rows.append({'M':M,'mean_relative_to_exact':float(np.mean(est/exact)),'abs_bias_center':float(abs(est[0]-exact[0]))})
83    # Prediction 3: MC variance of the mean scales as 1/K.
84    x0=np.array([[.2,.15]])
85    _, paths,_=wos_ball(x0,20000,32,2,seed=77)
86    var1=float(np.var(paths[0],ddof=1)); varrows=[]
87    for K in [1,2,4,8,16,32,64,128]:
88        # grouped means from common samples; observed K*Var(mean) should be constant
89        m=paths[0][:len(paths[0])//K*K].reshape(-1,K).mean(1)
90        var=float(np.var(m,ddof=1)); varrows.append({'K':K,'variance':var,'K_times_variance':K*var,'ratio_to_K1':var/var1})
91    # small comparison: WOS stochastic target versus exact-target MLP (same model/steps).
92    idea_rel,idea_sec=train_wos(True); base_rel,base_sec=train_wos(False)
93    out={'device':str(device),'exit_time_check':exit_rows,'sphere_second_moment_increment':radial_increment,'predicted_increment':rr*rr,'finite_depth_source_check':rows,'variance_scaling':{'single_path_variance':var1,'rows':varrows},'training_comparison':{'wos_target_relative_rmse':idea_rel,'exact_target_relative_rmse':base_rel,'wos_seconds':idea_sec,'exact_seconds':base_sec},'note':'The authoritative formula uses r^2/(2d); standard Brownian motion uses r^2/d, so the source estimator is systematically mis-scaled by approximately 1/2.'}
94    print(json.dumps(out,indent=2))
95
96if __name__=='__main__': main()