1import json, math, random
  2import numpy as np
  3import torch
  4from sklearn.datasets import make_moons
  5from sklearn.model_selection import train_test_split
  6from scipy.stats import spearmanr
  7
  8SEED = 3016
  9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11
 12class Net(torch.nn.Module):
 13    def __init__(self):
 14        super().__init__()
 15        self.net = torch.nn.Sequential(
 16            torch.nn.Linear(2, 48), torch.nn.ReLU(), torch.nn.Dropout(.35),
 17            torch.nn.Linear(48, 48), torch.nn.ReLU(), torch.nn.Dropout(.35),
 18            torch.nn.Linear(48, 2))
 19    def forward(self, x): return self.net(x)
 20
 21def math_sanity():
 22    # Empirical Phi uses local ball mass. A sparse cluster should have larger Phi.
 23    x = np.r_[np.linspace(-1, 1, 31), np.linspace(5, 5.12, 5)]
 24    d = np.abs(x[:, None] - x[None, :]); anchor = 15
 25    def phi(i):
 26        v = d[i, anchor]
 27        rs = np.linspace(max(1e-5, 4*v/500), 4*v, 500)
 28        masses = np.array([(d[i] <= r).mean() for r in rs])
 29        return np.trapz(np.log(1/masses), rs)
 30    ph = np.array([phi(i) for i in range(len(x))])
 31    return {"dense_cluster_mean_phi": float(ph[:31].mean()),
 32            "sparse_cluster_mean_phi": float(ph[31:].mean()),
 33            "sparse_larger": bool(ph[31:].mean() > ph[:31].mean()),
 34            "integrand_monotonicity": bool(math.log(1/.05) > math.log(1/.5))}
 35
 36def stochastic_logits(model, X, K):
 37    model.train()  # activate dropout, without gradients
 38    vals=[]
 39    with torch.no_grad():
 40        xt=torch.tensor(X, dtype=torch.float32)
 41        for _ in range(K): vals.append(model(xt).cpu().numpy())
 42    return np.stack(vals)
 43
 44def main():
 45    sanity=math_sanity()
 46    X,y=make_moons(n_samples=1800, noise=.25, random_state=SEED)
 47    X=(X-X.mean(0))/(X.std(0)+1e-8)
 48    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.35,random_state=SEED,stratify=y)
 49    model=Net(); opt=torch.optim.Adam(model.parameters(),lr=2e-3,weight_decay=1e-4)
 50    tx=torch.tensor(Xtr,dtype=torch.float32); ty=torch.tensor(ytr,dtype=torch.long)
 51    model.train()
 52    for epoch in range(180):
 53        perm=torch.randperm(len(tx))
 54        for j in range(0,len(tx),128):
 55            ix=perm[j:j+128]; loss=torch.nn.functional.cross_entropy(model(tx[ix]),ty[ix])
 56            opt.zero_grad(); loss.backward(); opt.step()
 57    # Calibration/reference set and held-out test subset.
 58    rng=np.random.default_rng(SEED)
 59    cal_idx=rng.choice(len(Xtr),size=180,replace=False)
 60    cal=Xtr[cal_idx]
 61    test_idx=np.arange(min(360,len(Xte)))
 62    Xt=Xte[test_idx]; yt=yte[test_idx]
 63    # Multiple stochastic outputs on calibration and test points.
 64    A=stochastic_logits(model,cal,24); B=stochastic_logits(model,Xt,80)
 65    # Robust metric: median pairwise same-seed output distance, with calibration medoid anchor.
 66    flat=A.mean(0)
 67    pd=((flat[:,None,:]-flat[None,:,:])**2).sum(2)**.5
 68    medoid=int(np.argmin(pd.mean(1)))
 69    dist_anchor=pd[:,medoid]
 70    # Phi approximation on a log radius grid using empirical ball masses.
 71    def phi_for(z):
 72        dd=np.sqrt(((flat-z[None,:])**2).sum(1)); v=np.sqrt(((z-flat[medoid])**2).sum())
 73        if v < 1e-8: return 0.
 74        lo=max(v*1e-5,1e-7); rs=np.geomspace(lo,4*v,180)
 75        mass=np.maximum((dd[:,None] <= rs[None,:]).mean(0),1/len(flat))
 76        return float(np.trapz(np.log(1/mass),rs))
 77    phi=np.array([phi_for(z) for z in B.mean(0)])
 78    va=np.sqrt(((B.mean(0)-flat[medoid])**2).sum(1))
 79    alpha=1.0; delta=.1
 80    M0=phi+va*math.log(math.e/delta)
 81    # Fit a conservative multiplicative C from calibration-like split of test points,
 82    # using high-sample mean as target stochastic prediction error.
 83    probs=(np.exp(B-B.max(2,keepdims=True))).sum(2)
 84    probs= np.exp(B-B.max(2,keepdims=True)); probs/=probs.sum(2,keepdims=True)
 85    meanp=probs.mean(0); pred=meanp.argmax(1)
 86    # Actual MC error is deviation of K=80 mean from single stochastic draws.
 87    point_err=np.sqrt(((probs-meanp[None,:,:])**2).sum(2)).mean(0)
 88    C=float(np.quantile(point_err/(M0+1e-8),.9))
 89    M=C*M0
 90    # Estimate adaptive K from the proposed inverse-square rule, with practical bounds.
 91    eps=.035; kmin=2; kmax=32
 92    K=np.clip(np.ceil(2*(M/eps)**2),kmin,kmax).astype(int)
 93    # Compare mean compute and accuracy for fixed K chosen at the adaptive mean,
 94    # and a global variance gate with same bounds.
 95    fixed=int(np.clip(round(K.mean()),kmin,kmax))
 96    # Evaluate predictions using prefixes of the same stochastic draws.
 97    def acc_for(ks):
 98        out=[]
 99        for i,k in enumerate(ks): out.append(probs[:k,i].mean(0).argmax())
100        return float(np.mean(np.array(out)==yt))
101    fixed_acc=acc_for(np.full(len(Xt),fixed))
102    adaptive_acc=acc_for(K)
103    var=np.sqrt(((probs-probs.mean(0)[None,:,:])**2).sum(2).mean(0))
104    order=np.argsort(var)
105    # Global variance gating uses a single threshold/rank and equal average K.
106    globalK=np.full(len(Xt),fixed)
107    budget=int(K.sum())
108    globalK[:]=kmin
109    rem=max(0,budget-kmin*len(Xt)); globalK[order[::-1][:min(len(Xt),rem//(kmax-kmin+1))]] = kmax
110    # Fill leftover approximately with intermediate allocation.
111    for i in order[::-1]:
112        if globalK.sum()>=budget: break
113        globalK[i]=min(kmax,globalK[i]+1)
114    global_acc=acc_for(globalK)
115    rho_phi=float(spearmanr(M,point_err).statistic) if np.std(M)>0 else 0.
116    rho_var=float(spearmanr(var,point_err).statistic) if np.std(var)>0 else 0.
117    result={"math_sanity":sanity,"n_test":len(Xt),"C":C,"epsilon":eps,
118      "adaptive_mean_K":float(K.mean()),"fixed_K":fixed,"global_mean_K":float(globalK.mean()),
119      "adaptive_accuracy":adaptive_acc,"fixed_accuracy":fixed_acc,"global_variance_accuracy":global_acc,
120      "phi_error_spearman":rho_phi,"variance_error_spearman":rho_var,
121      "adaptive_failure_rate_at_80_error":float(np.mean(point_err>M)),
122      "M_range":[float(M.min()),float(M.max())]}
123    print(json.dumps(result,indent=2))
124    with open('results.json','w') as f: json.dump(result,f,indent=2)
125
126if __name__=='__main__': main()