1"""MVP for a thermodynamic-confidence controller for minibatch SGD.
  2Run: /home/maxwelhelp/main/bin/python3 thermo_controller.py
  3"""
  4import math, random, json
  5import numpy as np
  6import torch
  7from sklearn.datasets import load_digits
  8from sklearn.model_selection import train_test_split
  9
 10SEED = 3166
 11
 12def seed_all(seed=SEED):
 13    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 14    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 15
 16# h from Eq. (9), evaluated stably for small u.
 17def h(u):
 18    u = np.asarray(u, dtype=float)
 19    return 1.0 + u - np.sqrt(1.0 + 2.0*u)
 20
 21def rate(a, C, D):
 22    """I(a)=(C/D^2)h(Da/C), with the D=0 equilibrium limit."""
 23    if D <= 1e-12: return a*a/(2.0*C)
 24    return float(C/(D*D) * h(D*a/C))
 25
 26def radius(T, delta, C, D, N=1.0):
 27    """Smallest a for 2N exp(-T I(a)) <= delta, by bisection."""
 28    target = math.log(2*N/delta)/T
 29    if target <= 0: return 0.0
 30    lo, hi = 0., max(math.sqrt(2*C*target), C*target + D*target, 1e-8)
 31    while rate(hi,C,D) < target: hi *= 2
 32    for _ in range(70):
 33        mid=(lo+hi)/2
 34        if rate(mid,C,D) >= target: hi=mid
 35        else: lo=mid
 36    return hi
 37
 38def autocorr_time(x):
 39    x=np.asarray(x,float); x=x-x.mean()
 40    if len(x)<4 or np.dot(x,x)<1e-14: return 1.0
 41    ac=np.correlate(x,x,mode='full')[len(x)-1:]/np.dot(x,x)
 42    # Initial-positive-sequence estimate, conservatively clipped.
 43    s=1.0
 44    for z in ac[1:]:
 45        if z <= 0: break
 46        s += 2*z
 47    return float(np.clip(s,1.,len(x)/2))
 48
 49class CurrentController:
 50    def __init__(self, window=32, delta=.05, batch=32, max_batch=128):
 51        self.window=window; self.delta=delta; self.batch=batch; self.max_batch=max_batch
 52        self.values=[]; self.radii=[]; self.means=[]; self.good=0
 53    def observe(self, j):
 54        self.values.append(float(j))
 55        if len(self.values) >= self.window:
 56            x=np.array(self.values[-self.window:]); mean=float(x.mean())
 57            var=float(max(x.var(ddof=1),1e-10)); tau=autocorr_time(x)
 58            # Empirical sub-gamma proxy: effective variance inflated by correlation.
 59            C=2*var*tau; D=max(1e-8, 0.25*math.sqrt(var)*max(tau-1,0))
 60            r=radius(self.window,self.delta,C,D)
 61            self.radii.append(r); self.means.append(mean)
 62            # The prescribed relative trigger, with a tiny noise-scale floor for m ~= 0.
 63            if r > .25*max(abs(mean), .1*math.sqrt(var)) and self.batch < self.max_batch:
 64                self.batch=min(self.max_batch,self.batch*2); self.good=0
 65            elif r < .1*max(abs(mean), .1*math.sqrt(var)):
 66                self.good += 1
 67                if self.good >= 4: self.good=0 # permit LR increase in a full optimizer; not used here
 68            else: self.good=0
 69            return r, mean, tau
 70        return None
 71
 72def math_check():
 73    # Exact bound proxy: radius should fall ~ T^-1/2 in Gaussian regime and rise with tau/variance.
 74    delta=.05; Ts=np.array([32,64,128,256,512])
 75    rs=np.array([radius(int(T),delta,C=2.,D=0.) for T in Ts])
 76    slope=np.polyfit(np.log(Ts),np.log(rs),1)[0]
 77    r_var=radius(128,delta,2.,0.)
 78    r_corr=radius(128,delta,8.,0.)
 79    return {'T':Ts.tolist(),'radii':rs.tolist(),'loglog_slope':float(slope),
 80            'radius_var1':float(r_var),'radius_var4':float(r_corr),
 81            'scaling_ok':bool(-.65 < slope < -.35),'variance_ok':bool(r_corr>r_var)}
 82
 83class Net(torch.nn.Module):
 84    def __init__(self):
 85        super().__init__(); self.net=torch.nn.Sequential(torch.nn.Linear(64,64),torch.nn.ReLU(),torch.nn.Linear(64,10))
 86    def forward(self,x): return self.net(x)
 87
 88def train(controlled, device, steps=300):
 89    seed_all(SEED)
 90    d=load_digits(); X=d.data.astype('float32')/16.; y=d.target.astype('int64')
 91    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
 92    xt=torch.tensor(Xtr,device=device); yt=torch.tensor(ytr,device=device)
 93    xval=torch.tensor(Xte,device=device); yval=torch.tensor(yte,device=device)
 94    model=Net().to(device); opt=torch.optim.SGD(model.parameters(),lr=.08)
 95    ctrl=CurrentController() if controlled else None
 96    v=torch.randn(sum(p.numel() for p in model.parameters()),device=device); v/=v.norm()
 97    rng=np.random.default_rng(SEED); losses=[]; batch_history=[]; n=len(xt)
 98    for step in range(steps):
 99        bs=ctrl.batch if ctrl else 32
100        batch_history.append(bs)
101        ii=torch.tensor(rng.integers(0,n,size=bs),device=device)
102        opt.zero_grad(set_to_none=True); out=model(xt[ii]); loss=torch.nn.functional.cross_entropy(out,yt[ii]); loss.backward()
103        g=torch.cat([p.grad.detach().flatten() for p in model.parameters()]); j=float(torch.dot(v,g).detach().cpu())
104        if ctrl: ctrl.observe(j)
105        opt.step(); losses.append(float(loss.detach().cpu()))
106    with torch.no_grad(): acc=float((model(xval).argmax(1)==yval).float().mean().cpu())
107    return {'final_train_loss':float(np.mean(losses[-20:])),'test_accuracy':acc,
108            'mean_batch':float(np.mean(batch_history)),
109            'total_examples':int(sum(batch_history)),
110            'batch_values':sorted(set(batch_history)),
111            'batches_changed':int(len(set(batch_history))),
112            'radius_count':len(ctrl.radii) if ctrl else 0,
113            'radius_first_last':([float(ctrl.radii[0]),float(ctrl.radii[-1])] if ctrl and ctrl.radii else [])}
114
115def main():
116    try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
117    except Exception: device=torch.device('cpu')
118    try:
119        check=math_check(); base=train(False,device,steps=1102); idea=train(True,device)
120    except Exception as e:
121        if device.type=='cuda':
122            device=torch.device('cpu'); check=math_check(); base=train(False,device,steps=1102); idea=train(True,device)
123        else: raise
124    out={'device':str(device),'math_check':check,'baseline':base,'idea':idea}
125    print(json.dumps(out,indent=2))
126if __name__=='__main__': main()