1import json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6
  7SEED = 117
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 11try:
 12    if DEVICE == 'cuda': torch.zeros(1, device='cuda')
 13except Exception:
 14    DEVICE = 'cpu'
 15
 16def power_mean(a, weights, p, eps=1e-12):
 17    a = a.clamp_min(eps); w = weights.to(a.device)
 18    if abs(p) < 1e-10: return torch.exp((w * torch.log(a)).sum(-1))
 19    return (w * a.pow(p)).sum(-1).clamp_min(eps).pow(1.0/p)
 20
 21def q_exponent(p, d): return p/(1.0-d*p)
 22
 23def softmin(u, tau): return -tau * torch.logsumexp(-u/tau, dim=0)
 24
 25def bbl_penalty(x1, x2, heads, p=0.0, d=1, tau_factor=.02):
 26    z = .5*x1 + .5*x2
 27    f1 = F.softplus(heads[0](x1)).squeeze(-1) + 1e-6
 28    f2 = F.softplus(heads[1](x2)).squeeze(-1) + 1e-6
 29    g1 = F.softplus(heads[0](z)).squeeze(-1) + 1e-6
 30    g2 = F.softplus(heads[1](z)).squeeze(-1) + 1e-6
 31    r = torch.stack((g1/f1, g2/f2), -1)
 32    # Independent Monte Carlo mass estimates for f and g.
 33    fx1 = f1.mean(); fx2 = f2.mean()
 34    gx1 = (F.softplus(heads[0](x1)).squeeze(-1)+1e-6).mean()
 35    gx2 = (F.softplus(heads[1](x2)).squeeze(-1)+1e-6).mean()
 36    ratios = torch.stack((gx1/fx1, gx2/fx2))
 37    w = torch.tensor([.5, .5], device=x1.device)
 38    u = power_mean(r, w, p)
 39    rhs = power_mean(ratios[None, :], w, q_exponent(p,d)).squeeze()
 40    tau = tau_factor * u.detach().median().clamp_min(1e-4)
 41    sm = softmin(u, tau)
 42    gap = sm-rhs
 43    return F.relu(gap).pow(2), gap, u.detach(), rhs.detach()
 44
 45class Branch(nn.Module):
 46    def __init__(self):
 47        super().__init__(); self.net=nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,1))
 48    def forward(self,x): return self.net(x)
 49
 50def make_data(n, noise=.55, seed=0):
 51    gen=torch.Generator().manual_seed(seed)
 52    y=(torch.randint(0,2,(n,),generator=gen)*2-1).float()
 53    x1=(y+noise*torch.randn(n,generator=gen)).unsqueeze(1)
 54    x2=(y+noise*torch.randn(n,generator=gen)).unsqueeze(1)
 55    return x1,x2,(y>0).float()
 56
 57def ece(prob,y,bins=10):
 58    out=0.; edges=torch.linspace(0,1,bins+1)
 59    for j in range(bins):
 60        mask=(prob>=edges[j])&((prob<edges[j+1]) if j<bins-1 else (prob<=edges[j+1]))
 61        if mask.any(): out += mask.float().mean().item()*abs(prob[mask].mean().item()-y[mask].mean().item())
 62    return out
 63
 64def run(use_bbl, seed=SEED):
 65    torch.manual_seed(seed); x1,x2,y=make_data(1800,seed=seed); tx1,tx2,ty=make_data(1600,seed=seed+100)
 66    heads=[Branch().to(DEVICE),Branch().to(DEVICE)]
 67    opt=torch.optim.Adam([p for h in heads for p in h.parameters()],lr=3e-3)
 68    x1,x2,y,tx1,tx2,ty=[v.to(DEVICE) for v in (x1,x2,y,tx1,tx2,ty)]
 69    max_pen=0.; active=0
 70    for step in range(260):
 71        ix=torch.randint(0,len(y),(96,),device=DEVICE); a,b,t=x1[ix],x2[ix],y[ix]
 72        loss=F.binary_cross_entropy_with_logits((heads[0](a)+heads[1](b)).squeeze(1)/2,t)
 73        pen,gap,_,_=bbl_penalty(a,b,heads)
 74        max_pen=max(max_pen,pen.item()); active += int(gap.item()>0)
 75        if use_bbl: loss=loss+.15*pen
 76        opt.zero_grad(); loss.backward(); opt.step()
 77    with torch.no_grad():
 78        probs=torch.sigmoid((heads[0](tx1)+heads[1](tx2)).squeeze(1)); pred=(probs>.5).float()
 79        clean_acc=(pred==ty).float().mean().item()
 80        disagreement=(torch.sigmoid(heads[0](tx1))-torch.sigmoid(heads[1](tx2))).abs().mean().item()
 81        corrupt=tx1+2*torch.randn_like(tx1)
 82        cp=torch.sigmoid((heads[0](corrupt)+heads[1](tx2)).squeeze(1))
 83        pen,gap,u,rhs=bbl_penalty(tx1[:96],tx2[:96],heads)
 84    return {'acc':clean_acc,'corrupt_acc':((cp>.5).float()==ty).float().mean().item(),
 85            'ece':ece(probs,ty),'disagreement':disagreement,'eval_penalty':pen.item(),
 86            'eval_gap':gap.item(),'ratio_q25':torch.quantile(u,.25).item(),'rhs':rhs.item(),
 87            'train_active_steps':active,'train_max_penalty':max_pen,'device':DEVICE}
 88
 89def math_check():
 90    grid=torch.linspace(-4,4,401); x1=grid[:,None]; x2=grid[None,:]; z=(x1+x2)/2
 91    r1=torch.exp((-z*z+x1*x1)/2); r2=torch.exp((-z*z+x2*x2)/2); u=torch.sqrt(r1*r2)
 92    # For equal Gaussian fields and p=0, inf geometric ratio = mass-ratio RHS = 1.
 93    return {'gaussian_min':u.min().item(),'gaussian_rhs':1.0,'gaussian_max_violation':max(0.,u.min().item()-1.)}
 94
 95def gradient_check():
 96    torch.manual_seed(SEED); h=[Branch().to(DEVICE),Branch().to(DEVICE)]
 97    x1,x2,_=make_data(96,seed=31); x1,x2=x1.to(DEVICE),x2.to(DEVICE)
 98    p,g,_,_=bbl_penalty(x1,x2,h); p.backward()
 99    return {'penalty':p.item(),'gap':g.item(),'gradient_norm':float(sum((v.grad.norm()**2 for q in h for v in q.parameters() if v.grad is not None),torch.tensor(0.,device=DEVICE)).sqrt().cpu())}
100
101def main():
102    print(json.dumps({'math_check':math_check(),'gradient_check':gradient_check(),'baseline':run(False),'idea':run(True)},indent=2))
103if __name__=='__main__': main()