import json, math, os, random import numpy as np import torch from torch import nn from custom_stochastic_intervention_track import get_dataset, META SEEDS = list(range(8)) LRS = [0.003, 0.01, 0.03] EPOCHS = 90 BATCH = 128 def math_check(seed=123): rng = np.random.RandomState(seed) n = 500000 out = {} vals = [] for sx in [1.0, 2.0]: x = sx*rng.randn(n); u = rng.randn(n); v = rng.randn(n) m = x + u; y = m + x + v keep = np.abs(m-1.0) < .012 empirical = float(y[keep].mean()) analytic = 1.0 + sx*sx/(sx*sx+1.0) vals.append((empirical, analytic, int(keep.sum()))) out['selected_m_equals_1'] = {'sigma1': vals[0], 'sigma2': vals[1]} out['shift_observed'] = abs(vals[1][0]-vals[0][0]) > 0.15 return out def normal_nll(z, pars): mu = pars[:, 0] logsd = pars[:, 1].clamp(-4.0, 3.0) return 0.5*((z-mu)/logsd.exp())**2 + logsd + 0.5*math.log(2*math.pi) class SharedMLP(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential(nn.Linear(3, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 2)) def forward(self, z): return self.net(z) class JointMLP(nn.Module): def __init__(self): super().__init__() self.q = nn.Sequential(nn.Linear(1,32), nn.Tanh(), nn.Linear(32,2)) self.m = nn.Sequential(nn.Linear(2,32), nn.Tanh(), nn.Linear(32,2)) self.y = SharedMLP() def forward(self, c, x, m): return self.q(c), self.m(torch.cat([c,x],1)), self.y(torch.cat([c,x,m],1)) def train_one(seed, lr, idea, device): torch.manual_seed(seed); np.random.seed(seed); random.seed(seed) d = get_dataset(seed+1000, 800, 800) c = torch.tensor(d['xtr'][:,0:1], device=device); x = torch.tensor(d['xtr'][:,1:2], device=device) m = torch.tensor(d['xtr'][:,2:3], device=device); y = torch.tensor(d['ytr'][:,None], device=device) if idea: model = JointMLP().to(device) else: model = SharedMLP().to(device) opt = torch.optim.Adam(model.parameters(), lr=lr) g = torch.Generator(device=device); g.manual_seed(seed+55) n = len(y) model.train() for _ in range(EPOCHS): ix = torch.randperm(n, generator=g, device=device) for start in range(0,n,BATCH): j = ix[start:start+BATCH] if idea: q, mp, yp = model(c[j],x[j],m[j]) loss = (normal_nll(x[j].squeeze(1),q) + normal_nll(m[j].squeeze(1),mp) + normal_nll(y[j].squeeze(1),yp)).mean() else: # Standard mediator-only practice: the identical outcome MLP is trained without X. z = torch.cat([c[j], torch.zeros_like(x[j]), m[j]], 1) loss = normal_nll(y[j].squeeze(1), model(z)).mean() opt.zero_grad(); loss.backward(); opt.step() d = get_dataset(seed+2000, 12000, 12000) ct = torch.tensor(d['xte'][:,0:1], device=device); xt = torch.tensor(d['xte'][:,1:2], device=device) mt = torch.tensor(d['xte'][:,2:3], device=device); yt = torch.tensor(d['yte'], device=device) model.eval() with torch.no_grad(): if idea: _,_,pars = model(ct,xt,mt) pred = pars[:,0] else: pars = model(torch.cat([ct,torch.zeros_like(xt),mt],1)); pred = pars[:,0] mse = float(((pred-yt)**2).mean().cpu()) # Re-test the claimed selection mechanism on model outputs. Selection is on observed mediator; # compare predicted E[Y|M approximately 1] with the empirical selected outcome. keep = (mt[:,0]-1.0).abs() < 0.035 selected_pred = float(pred[keep].mean().cpu()) selected_obs = float(yt[keep].mean().cpu()) return {'seed':seed,'lr':lr,'mse':mse,'selected_pred':selected_pred,'selected_obs':selected_obs,'n_selected':int(keep.sum().cpu())}, model def pvalue(deltas): deltas=np.asarray(deltas); count=0; total=1<>i)&1 else -1 for i in range(len(deltas))]) if (signs*deltas).mean() <= 0: count += 1 return count/total def main(): try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') except Exception: device=torch.device('cpu') try: # A CUDA allocation can fail in the shared environment; retry all work on CPU. torch.zeros(1,device=device) except Exception: device=torch.device('cpu') math_result=math_check() baseline={}; idea={} for lr in LRS: baseline[str(lr)] = [train_one(s,lr,False,device)[0] for s in SEEDS] idea[str(lr)] = [train_one(s,lr,True,device)[0] for s in SEEDS] bmeans={k:float(np.mean([r['mse'] for r in v])) for k,v in baseline.items()} imeans={k:float(np.mean([r['mse'] for r in v])) for k,v in idea.items()} best_lr=min(bmeans,key=bmeans.get); best_idea_lr=min(imeans,key=imeans.get) br=baseline[best_lr]; ir=idea[best_idea_lr] deltas=[ir[i]['mse']-br[i]['mse'] for i in range(8)] sig={ 'selected_m':1.0, 'baseline_predicted_mean':float(np.mean([r['selected_pred'] for r in br])), 'idea_predicted_mean':float(np.mean([r['selected_pred'] for r in ir])), 'observed_mean':float(np.mean([r['selected_obs'] for r in ir])), 'baseline_abs_error':float(abs(np.mean([r['selected_pred'] for r in br])-np.mean([r['selected_obs'] for r in br]))), 'idea_abs_error':float(abs(np.mean([r['selected_pred'] for r in ir])-np.mean([r['selected_obs'] for r in ir]))), 'confirmed': bool(math_result['shift_observed'] and abs(np.mean([r['selected_pred'] for r in ir])-np.mean([r['selected_obs'] for r in ir])) < 0.35) } report={'track':META,'official_bench_available':False,'infrastructure_note':'Specified /home/maxwelhelp/all/math2nn/bench and README.md were absent; this is the required-contract local fallback, not bench.make_report output.','math_check':math_result,'budget':{'seeds':SEEDS,'epochs':EPOCHS,'batch':BATCH,'lr_union':LRS},'baseline_sweep':bmeans,'idea_sweep':imeans,'best_baseline_lr':float(best_lr),'best_idea_lr':float(best_idea_lr),'baseline_per_seed':br,'idea_per_seed':ir,'paired_delta_mean':float(np.mean(deltas)),'paired_deltas':deltas,'permutation_p':pvalue(deltas),'mechanism_signature':sig,'bench_report':{'custom_track':{'name':META['name'],'file':'custom_stochastic_intervention_track.py','domain':META['domain']},'verdict':'idea better (significant)' if np.mean(deltas)<0 and pvalue(deltas)<.05 else 'no significant win'}} with open('custom_bench_results.json','w') as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__=='__main__': main()