1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 3146
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10except Exception:
 11    device = torch.device('cpu')
 12
 13def scm(n, sigma_x, seed):
 14    r = np.random.RandomState(seed)
 15    x = r.randn(n) * sigma_x
 16    u = r.randn(n); v = r.randn(n)
 17    m = x + u
 18    y = m + x + v
 19    return x.astype('float32'), m.astype('float32'), y.astype('float32')
 20
 21# Stage 1: exact Gaussian calculations and sampling verify the claimed selection effect.
 22def math_check():
 23    # For M=X+U, Y=M+X+V: E[X|M=m] = s2/(s2+1)m and E[Y|M=m] changes with intervention variance.
 24    m0 = 1.0
 25    analytic = {}
 26    for s in (1.0, 2.0):
 27        analytic[str(s)] = (1.0 + s*s/(s*s+1.0))*m0
 28    x,m,y = scm(1500000, 1.0, 11)
 29    selected = np.abs(m-m0) < .015
 30    empirical_train = float(y[selected].mean())
 31    x,m,y = scm(1500000, 2.0, 12)
 32    selected = np.abs(m-m0) < .015
 33    empirical_shift = float(y[selected].mean())
 34    # Joint marginal variance is Var(Y)=4 Var(X)+2, while mediator-only cannot encode this shift.
 35    var_formula = {str(s): 4*s*s+2 for s in (1.0,2.0)}
 36    return {'analytic_EY_given_M1': analytic, 'empirical_EY_given_M1': {'sigma1': empirical_train, 'sigma2': empirical_shift}, 'marginal_Y_variance_formula': var_formula}
 37
 38class Joint(nn.Module):
 39    def __init__(self):
 40        super().__init__()
 41        self.q = nn.Sequential(nn.Linear(1,32), nn.Tanh(), nn.Linear(32,2)) # mu, log sd of X
 42        self.m = nn.Sequential(nn.Linear(1,32), nn.Tanh(), nn.Linear(32,2)) # conditioned on x
 43        self.y = nn.Sequential(nn.Linear(2,32), nn.Tanh(), nn.Linear(32,2)) # x,m
 44    def forward(self,x,m):
 45        q=self.q(torch.zeros_like(x[:,None])); mp=self.m(x[:,None]); yp=self.y(torch.cat([x[:,None],m[:,None]],1))
 46        return q,mp,yp
 47
 48def normal_nll(z, pars):
 49    mu, logsd = pars[:,0], pars[:,1].clamp(-5,3)
 50    return .5*((z-mu)/logsd.exp())**2 + logsd + .5*math.log(2*math.pi)
 51
 52class Base(nn.Module):
 53    def __init__(self):
 54        super().__init__(); self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,2))
 55    def forward(self,m): return self.net(m[:,None])
 56
 57def train_eval():
 58    x,m,y=scm(24000,1.0,20)
 59    xt,mt,yt=scm(120000,2.0,21)
 60    X=torch.tensor(x,device=device); M=torch.tensor(m,device=device); Y=torch.tensor(y,device=device)
 61    # Baseline deliberately follows the common mediator-only predictor p(y|m).
 62    base=Base().to(device); opt=torch.optim.Adam(base.parameters(),lr=.01)
 63    for _ in range(350):
 64        ix=torch.randint(0,len(X),(128,),device=device)
 65        loss=normal_nll(Y[ix],base(M[ix])).mean(); opt.zero_grad(); loss.backward(); opt.step()
 66    # Joint likelihood q(x) + p(m|x) + p(y|x,m).
 67    joint=Joint().to(device); opt=torch.optim.Adam(joint.parameters(),lr=.01)
 68    for _ in range(450):
 69        ix=torch.randint(0,len(X),(128,),device=device)
 70        q,mp,yp=joint(X[ix],M[ix])
 71        loss=(normal_nll(X[ix],q)+normal_nll(M[ix],mp)+normal_nll(Y[ix],yp)).mean()
 72        opt.zero_grad(); loss.backward(); opt.step()
 73    with torch.no_grad():
 74        MT=torch.tensor(mt,device=device); XT=torch.tensor(xt,device=device); YT=torch.tensor(yt,device=device)
 75        base_nll=float(normal_nll(YT,base(MT)).mean().cpu())
 76        _,_,yp=joint(XT,MT)
 77        joint_nll=float(normal_nll(YT,yp).mean().cpu())
 78        # Estimated conditional mean at M=1 under the learned conditional model, versus baseline.
 79        one=torch.ones(2000,device=device); xx=torch.randn(2000,device=device)*2
 80        # draw mediator noise, then retain a narrow mediator selection
 81        mm=xx+torch.randn_like(xx)
 82        keep=(mm-1).abs()<.02
 83        _,_,pred=joint(xx[keep],mm[keep])
 84        joint_sel=float(pred[:,0].mean().cpu())
 85        bsel=float(base(torch.ones(2000,device=device))[:,0].mean().cpu())
 86        # Marginal predictive variance under target X~N(0,2), via structural heads.
 87        xs=torch.randn(20000,device=device)*2
 88        ms=joint.m(xs,)[0] if False else None
 89        mp=joint.m(xs[:,None]); mdraw=mp[:,0]+torch.randn_like(xs)*mp[:,1].clamp(-5,3).exp()
 90        yp=joint.y(torch.stack([xs,mdraw],1)); ydraw=yp[:,0]+torch.randn_like(xs)*yp[:,1].clamp(-5,3).exp()
 91        pred_var=float(ydraw.var().cpu())
 92    return {'test_shift_sigma2_nll':{'baseline_m_only':base_nll,'joint':joint_nll,'joint_minus_baseline':joint_nll-base_nll},'selected_M1_predicted_mean':{'baseline':bsel,'joint_conditional':joint_sel,'true_sigma2':1.8},'joint_marginal_predicted_Y_variance_sigma2':pred_var,'device':str(device)}
 93
 94if __name__=='__main__':
 95    out={'math_check':math_check()}
 96    try: out['neural_experiment']=train_eval()
 97    except Exception as e:
 98        # Explicit CPU fallback as required for shared GPU environments.
 99        device=torch.device('cpu'); out['neural_experiment']=train_eval(); out['cuda_error']=repr(e)
100    with open('results.json','w') as f: json.dump(out,f,indent=2)
101    print(json.dumps(out,indent=2))