1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6from torch.utils.data import TensorDataset, DataLoader
  7
  8SEED = 1337
  9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11try:
 12    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 13except Exception:
 14    DEVICE = torch.device('cpu')
 15
 16# Smoothed R-function-like disk support: positive inside, negative outside.
 17def disk_score(x, cx=0.0, cy=0.0, radius=0.65, smooth=0.035):
 18    r = torch.sqrt((x[:, 0]-cx)**2 + (x[:, 1]-cy)**2 + 1e-8)
 19    return torch.sigmoid((radius-r)/smooth)
 20
 21def make_data(kind, n, seed, train=True):
 22    g = torch.Generator().manual_seed(seed)
 23    x = (torch.rand(n, 2, generator=g) * 2.4 - 1.2)
 24    true_g = disk_score(x)
 25    wrong_g = disk_score(x, cx=0.48, cy=-0.30)
 26    if kind == 'correct':
 27        # Geometry is the dominant signal, with a small smooth nuisance term.
 28        y = true_g + 0.08*x[:, 0] - 0.04*x[:, 1]
 29        gp = true_g
 30    elif kind == 'wrong':
 31        y = torch.sin(2.2*x[:,0]) + 0.35*x[:,1]**2
 32        gp = wrong_g
 33    elif kind == 'none':
 34        y = torch.sin(2.2*x[:,0]) + 0.35*x[:,1]**2
 35        gp = true_g
 36    elif kind == 'half':
 37        # The geometry rule applies to the left half only; the right half is unrelated.
 38        obey = (x[:,0] < 0).float()
 39        y = obey*true_g + (1-obey)*(torch.sin(2.2*x[:,0]) + 0.35*x[:,1]**2)
 40        y = y + 0.05*x[:,0]
 41        gp = true_g
 42    else: raise ValueError(kind)
 43    if train:
 44        y = y + 0.035*torch.randn(n, generator=g)
 45    return x, y[:,None], gp[:,None]
 46
 47class Free(nn.Module):
 48    def __init__(self):
 49        super().__init__(); self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
 50    def forward(self,x,g=None): return self.net(x)
 51
 52class Geom(nn.Module):
 53    def __init__(self):
 54        super().__init__(); self.net=nn.Sequential(nn.Linear(3,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
 55    def forward(self,x,g): return self.net(torch.cat([x,g],1))
 56
 57class Concat(Geom):
 58    pass
 59
 60class Mixture(nn.Module):
 61    def __init__(self, input_gate=False):
 62        super().__init__(); self.free=Free(); self.geom=Geom()
 63        self.input_gate=input_gate
 64        self.beta=nn.Parameter(torch.tensor(-4.0))
 65        self.mixer=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,1)) if input_gate else None
 66    def forward(self,x,g):
 67        hf=self.free(x); hg=self.geom(x,g)
 68        a=torch.sigmoid(self.mixer(x) if self.input_gate else self.beta)
 69        return a*hg+(1-a)*hf, hf, hg, a
 70
 71def train_model(kind, mode, epochs=160, n=640, lam=0.0):
 72    x,y,g=make_data(kind,n,100+SEED)
 73    if mode=='free': model=Free().to(DEVICE)
 74    elif mode in ('geom','concat'): model=Geom().to(DEVICE)
 75    elif mode=='mixer': model=Mixture(input_gate=True).to(DEVICE)
 76    elif mode=='global': model=Mixture(input_gate=False).to(DEVICE)
 77    else: raise ValueError(mode)
 78    opt=torch.optim.Adam(model.parameters(),lr=0.012,weight_decay=1e-4)
 79    x,y,g=x.to(DEVICE),y.to(DEVICE),g.to(DEVICE)
 80    for _ in range(epochs):
 81        opt.zero_grad()
 82        if mode=='free': pred=model(x); loss=((pred-y)**2).mean()
 83        elif mode in ('geom','concat'): pred=model(x,g); loss=((pred-y)**2).mean()
 84        else:
 85            pred,hf,hg,a=model(x,g); loss=((pred-y)**2).mean()+lam*a.mean()
 86        loss.backward(); opt.step()
 87    return model
 88
 89def evaluate(model, kind, mode):
 90    x,y,g=make_data(kind,1800,900+SEED,False)
 91    x,y,g=x.to(DEVICE),y.to(DEVICE),g.to(DEVICE)
 92    with torch.no_grad():
 93        if mode=='free': p=model(x); a=torch.zeros_like(p); hf=p; hg=p
 94        elif mode in ('geom','concat'): p=model(x,g); a=torch.ones_like(p); hf=p; hg=p
 95        else: p,hf,hg,a=model(x,g)
 96        rmse=torch.sqrt(((p-y)**2).mean()).item()
 97        branch_f=torch.sqrt(((hf-y)**2).mean()).item(); branch_g=torch.sqrt(((hg-y)**2).mean()).item()
 98        ent=(-(a.clamp(1e-7,1-1e-7)*torch.log(a.clamp(1e-7,1-1e-7))+(1-a).clamp(1e-7,1-1e-7)*torch.log((1-a).clamp(1e-7,1-1e-7))).mean()).item()
 99        return {'rmse':rmse,'free_branch_rmse':branch_f,'geom_branch_rmse':branch_g,'mean_alpha':a.mean().item(),'gate_entropy':ent}
100
101def math_check():
102    torch.manual_seed(7)
103    h1=torch.randn(10000); h2=torch.randn(10000); a=torch.rand(10000)
104    lo=torch.minimum(h1,h2); hi=torch.maximum(h1,h2)
105    inside=bool(((a*h1+(1-a)*h2 >= lo-1e-6)&(a*h1+(1-a)*h2 <= hi+1e-6)).all())
106    beta=torch.tensor(0.,requires_grad=True); alpha=torch.sigmoid(beta); (alpha).backward()
107    deriv=float(beta.grad); beta2=torch.tensor(-4.)
108    return {'convex_interpolation_inside_bounds':inside,'d_mean_alpha_d_beta_at_zero':deriv,'initial_alpha_beta_minus4':float(torch.sigmoid(beta2)),'suppression_gradient_sign':deriv>0}
109
110def main():
111    out={'device':str(DEVICE),'math_check':math_check(),'results':{}}
112    for kind in ['correct','wrong','none','half']:
113        out['results'][kind]={}
114        for mode in ['free','geom','concat','global','mixer']:
115            # concat is architecturally the same geometry-feature baseline; included explicitly.
116            model=train_model(kind,mode)
117            out['results'][kind][mode]=evaluate(model,kind,mode)
118        # Repeat the global gate with mild suppression to test the stated regularizer behavior.
119        m=train_model(kind,'global',lam=0.02)
120        out['results'][kind]['global_lambda_0.02']=evaluate(m,kind,'global')
121    Path('results.json').write_text(json.dumps(out,indent=2))
122    print(json.dumps(out,indent=2))
123
124if __name__=='__main__': main()