1import sys, json, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7
  8sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  9from bench import train_model, evaluate, sweep_baseline, make_report
 10
 11META = {'name':'spd_covariance_classification','domain':'structured_spd','description':'Classification of noisy SPD covariance matrices with Cholesky factors.'}
 12D, C = 5, 3
 13
 14def get_dataset(seed, n_train=400, n_test=200):
 15    rng=np.random.default_rng(seed)
 16    centers=[]
 17    for k in range(C):
 18        q,_=np.linalg.qr(rng.normal(size=(D,D)))
 19        vals=np.exp(np.linspace(-.45,.55,D)+.28*k)
 20        centers.append(q@np.diag(vals)@q.T)
 21    def sample(n):
 22        xs=[]; ys=[]
 23        for i in range(n):
 24            k=i%C
 25            e=rng.normal(size=(D,D)); e=(e+e.T)/2
 26            s=centers[k]+.11*(e@e.T)+1e-3*np.eye(D)
 27            xs.append(s.astype('float32')); ys.append(k)
 28        p=rng.permutation(n)
 29        return np.asarray(xs)[p], np.asarray(ys,dtype='int64')[p]
 30    xtr,ytr=sample(n_train); xte,yte=sample(n_test)
 31    return {'xtr':torch.tensor(xtr.reshape(n_train,-1)), 'ytr':torch.tensor(ytr),
 32            'xte':torch.tensor(xte.reshape(n_test,-1)), 'yte':torch.tensor(yte),
 33            'task':'classification','metric':'cross_entropy','input_shape':(D*D,), 'out_dim':C,
 34            'xte_spd':xte}
 35
 36class Shared(nn.Module):
 37    def __init__(self, head):
 38        super().__init__(); self.feat=nn.Sequential(nn.Linear(D*D,48),nn.ReLU(),nn.Linear(48,24),nn.ReLU()); self.head=head
 39    def forward(self,x): return self.head(self.feat(x))
 40
 41class EuclideanHead(nn.Module):
 42    def __init__(self): super().__init__(); self.fc=nn.Linear(24,C)
 43    def forward(self,z): return self.fc(z)
 44
 45def low(x): return torch.tril(x,-1)
 46def power_spd(s,p):
 47    s=(s+s.transpose(-1,-2))/2
 48    w,v=torch.linalg.eigh(s); return (v*w.clamp_min(1e-7).pow(p).unsqueeze(-2))@v.transpose(-1,-2)
 49
 50class CholeskyHead(nn.Module):
 51    def __init__(self, theta=1.0):
 52        super().__init__(); self.theta=theta
 53        self.rawL=nn.Parameter(torch.randn(C,D,D)*.08); self.rawA=nn.Parameter(torch.randn(C,D,D)*.05)
 54    def factors(self):
 55        L=torch.tril(self.rawL); diag=F.softplus(torch.diagonal(self.rawL,dim1=-2,dim2=-1))+0.15
 56        return L-torch.diag_embed(torch.diagonal(L,dim1=-2,dim2=-1))+torch.diag_embed(diag)
 57    def forward(self,z):
 58        # Map shared representation to a lower-triangular factor, guaranteeing SPD.
 59        b=z.shape[0]; raw=z.new_zeros(b,D,D)
 60        raw[:,:D,:D]=raw
 61        inds=torch.tril_indices(D,D,device=z.device)
 62        raw[:,inds[0],inds[1]]=z[:, :len(inds[0])]
 63        diag=F.softplus(torch.diagonal(raw,dim1=-2,dim2=-1))+0.05
 64        K=raw-torch.diag_embed(torch.diagonal(raw,dim1=-2,dim2=-1))+torch.diag_embed(diag)
 65        L=self.factors(); A=torch.tril(self.rawA,-1)
 66        S=K@K.transpose(-1,-2); Kp=power_spd(S,self.theta/2)
 67        P=L@L.transpose(-1,-2); Lp=power_spd(P,self.theta/2)
 68        q=A # M=I, solve Mq=A
 69        scores=[]
 70        for k in range(C):
 71            t1=((low(K)-low(L[k]))*A[k]).sum((-1,-2))
 72            t2=((Kp-Lp[k])*q[k]).sum((-1,-2))/(4*self.theta)
 73            scores.append(t1+t2)
 74        return torch.stack(scores,-1)
 75
 76# Same base architecture, with the sole intervention being the SPD head.
 77def make_baseline():
 78    return Shared(EuclideanHead())
 79def make_idea(theta=1.0):
 80    return Shared(CholeskyHead(theta))
 81
 82def run_one(factory, seed, epochs, lr, weight_decay=1e-4):
 83    torch.manual_seed(seed); np.random.seed(seed)
 84    ds=get_dataset(seed)
 85    # train_model is the canonical loop; idea changes representation/readout, not training.
 86    net, metric, hist=train_model(factory(), ds, epochs=epochs, lr=lr, batch=64, weight_decay=weight_decay)
 87    with torch.no_grad():
 88        dev=next(net.parameters()).device
 89        pred=net(ds['xte'].to(dev)); acc=float((pred.argmax(1)==ds['yte'].to(dev)).float().mean())
 90    return float(metric), acc, net
 91
 92def main():
 93    # Union parity: each idea lr is also evaluated in the baseline grid.
 94    grid=[{'lr':1e-3,'epochs':18},{'lr':3e-3,'epochs':18},{'lr':1e-2,'epochs':18}]
 95    base=sweep_baseline(lambda cfg: lambda seed: run_one(make_baseline,seed,**cfg)[0],grid)
 96    best_lr=base['best_cfg']['lr']; idea_cfgs=[{'lr':best_lr,'epochs':18},{'lr':1e-3 if best_lr!=1e-3 else 3e-3,'epochs':18},{'lr':1e-2,'epochs':18}]
 97    # evaluate idea settings on the same full paired seeds; report the best mean.
 98    ir=[]
 99    for cfg in idea_cfgs:
100        r=evaluate(lambda seed: run_one(lambda: make_idea(1.0),seed,**cfg)[0])
101        ir.append((r,cfg))
102    idea,cfg=min(ir,key=lambda x:x[0]['mean'])
103    report=make_report('spd_covariance_classification','mlp_tiny',base,idea,extra={})
104    # Re-test stage-1 mechanism at NN scale using trained systems: predicted factor SPD and observed logits.
105    m,a,_=run_one(make_idea,0,18,best_lr)
106    with torch.no_grad():
107        ds=get_dataset(0); dev=next(_.parameters()).device; z=_.feat(ds['xte'].to(dev)); h=_.head; out=h(z)
108        # Quantitative observed invariant from trained model: all generated Cholesky matrices are SPD.
109        L=h.factors(); mineig=float(torch.linalg.eigvalsh(L@L.transpose(-1,-2)).min())
110        observed=float(torch.isfinite(out).all())
111    sig={'prediction':'factor-generated prototypes remain SPD and logits finite after NN training',
112         'predicted_min_eigenvalue_bound':0.0,'observed_min_prototype_eigenvalue':mineig,
113         'observed_finite_logit_fraction':observed,'confirmed':bool(mineig>0 and observed==1.0),
114         'idea_settings_tried':idea_cfgs}
115    report['mechanism_signature']=sig; report['idea_selected_cfg']=cfg
116    Path('bench_report.json').write_text(json.dumps(report,indent=2))
117    print(json.dumps(report,indent=2))
118if __name__=='__main__': main()