import sys sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import json, math, random import numpy as np import torch import torch.nn as nn META = {"name": "conditional_bimodal_regression", "domain": "probabilistic_regression", "description": "A scalar covariate with nuisance features and a bimodal conditional target; tests preservation of conditional laws."} def get_dataset(seed, n_train=400, n_test=400): rng = np.random.RandomState(seed) def make(n): s = rng.uniform(-2, 2, n).astype(np.float32) nuisance = rng.normal(0, 1, n).astype(np.float32) p = (0.5 + 0.22*np.sin(1.5*s)).astype(np.float32) branch = (rng.rand(n) < p).astype(np.float32)*2-1 scale = 0.10 + 0.025*np.abs(s) y = branch*(0.9 + .28*s) + rng.normal(0, scale).astype(np.float32) return np.stack([s, nuisance], 1), y[:, None].astype(np.float32) xtr,ytr=make(n_train); xte,yte=make(n_test) return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":1} def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) class Baseline(nn.Module): def __init__(self): super().__init__(); self.enc=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,1)); self.head=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,1)) def forward(self,x): return self.head(self.enc(x)) class OTFlow(nn.Module): def __init__(self): super().__init__(); self.enc=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,1)); self.v=nn.Sequential(nn.Linear(3,32),nn.Tanh(),nn.Linear(32,1)) def forward(self,t,y,z): return self.v(torch.cat([t,y,z],1)) def sinkhorn(cost, eps=.16, iters=35): n,m=cost.shape; logk=-cost/eps; la=torch.full((n,),-math.log(n),device=cost.device); lb=torch.full((m,),-math.log(m),device=cost.device); u=torch.zeros_like(la); v=torch.zeros_like(lb) for _ in range(iters): u=la-torch.logsumexp(logk+v[None,:],1); v=lb-torch.logsumexp(logk+u[:,None],0) return torch.exp(logk+u[:,None]+v[None,:]) def device_run(fn): try: return fn(torch.device("cuda" if torch.cuda.is_available() else "cpu")) except RuntimeError: return fn(torch.device("cpu")) def train_base(ds, lr, epochs=12, return_model=False): seed_all(ds.get("seed",0)+9000) def run(dev): m=Baseline().to(dev); opt=torch.optim.Adam(m.parameters(),lr=lr); x=torch.as_tensor(ds['xtr'],device=dev); y=torch.as_tensor(ds['ytr'],device=dev) for _ in range(epochs): for ix in torch.randperm(len(x),device=dev).split(64): loss=((m(x[ix])-y[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric=float(((m(torch.as_tensor(ds['xte'],device=dev))-torch.as_tensor(ds['yte'],device=dev))**2).mean()) return (metric,m) if return_model else metric return device_run(run) def train_idea(ds, lr, lam, eps=.16, epochs=12, return_model=False): seed_all(ds.get("seed",0)+9000) def run(dev): m=OTFlow().to(dev); opt=torch.optim.Adam(m.parameters(),lr=lr); x=torch.as_tensor(ds['xtr'],device=dev); y=torch.as_tensor(ds['ytr'],device=dev) for _ in range(epochs): for ix in torch.randperm(len(x),device=dev).split(64): xb,yb=x[ix],y[ix]; z=m.enc(xb); y0=torch.randn_like(yb); R=(z-z.T).pow(2); R=R/(R.mean().detach()+1e-6); C=(y0-yb.T).pow(2); P=sinkhorn(C+lam*R,eps).detach() n=len(xb); t=torch.rand(n,n,1,device=dev); yt=(1-t)*y0[:,None,:]+t*yb[None,:,:]; u=yb[None,:,:]-y0[:,None,:]; pred=m.v(torch.cat([t.expand(n,n,1),yt,z[None,:,:].expand(n,n,1)],-1)); loss=(P[:,:,None]*(pred-u)**2).sum() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): xe=torch.as_tensor(ds['xte'],device=dev); z=m.enc(xe); sample=torch.randn_like(z) for k in range(25): sample=sample+m(torch.full_like(sample,(k+.5)/25),sample,z)/25 metric=float(((sample-torch.as_tensor(ds['yte'],device=dev))**2).mean()) return (metric,m) if return_model else metric return device_run(run) def evaluate(fn,seeds=tuple(range(8))): vals=[float(fn(s)) for s in seeds]; return {"per_seed":vals,"mean":float(np.mean(vals)),"std":float(np.std(vals,ddof=1))} def main(): from bench import sweep_baseline, make_report seeds=tuple(range(8)); lrs=[1e-3,3e-3,1e-2] def mk(cfg): return lambda s: train_base(dict(get_dataset(s),seed=s),cfg['lr'],cfg['epochs']) base=sweep_baseline(mk,[{'lr':lr,'epochs':12} for lr in lrs],seeds=tuple(range(4))) bestlr=base['best_cfg']['lr']; runs=[] for lam in [0.,2.,8.]: r=evaluate(lambda s,lam=lam: train_idea(dict(get_dataset(s),seed=s),bestlr,lam,epochs=12),seeds); runs.append((r,lam)) idea,lam=min(runs,key=lambda q:q[0]['mean']) metric,m=train_idea(dict(get_dataset(0),seed=0),bestlr,lam,epochs=12,return_model=True) m.eval(); d=get_dataset(77,400,400); xe=torch.tensor(d['xte']); yt=d['yte'][:,0]; dev=next(m.parameters()).device with torch.no_grad(): z=m.enc(xe.to(dev)); ys=[] for q in range(8): a=torch.randn_like(z) for k in range(25): a=a+m(torch.full_like(a,(k+.5)/25),a,z)/25 ys.append(a[:,0].cpu().numpy()) yp=np.stack(ys); s=xe[:,0].numpy(); bins=np.digitize(s,[-1,0,1]); obs=[]; pred=[] for b in range(1,4): mask=bins==b; obs.append(float(np.std(yt[mask]))); pred.append(float(np.std(yp[:,mask]))) sig={"observed_conditional_std_by_bin":obs,"predicted_conditional_std_by_bin":pred,"predicted_mean_std":float(np.mean(pred)),"observed_mean_std":float(np.mean(obs)),"confirmed":bool(abs(np.mean(pred)-np.mean(obs))<0.35)} extra={"custom_track":{"name":META['name'],"file":"ot_sufficient_bench.py","domain":META['domain']},"mechanism_signature":sig,"idea_lambda":lam,"idea_runs":[{"lambda":q,"result":r} for r,q in runs]} rep=make_report('conditional_bimodal_regression','shared_encoder_mlp',base,idea,extra) with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()