1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=7
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8
  9def psd_check(SY,SX):
 10    d=SY-SX
 11    return float(np.linalg.eigvalsh((d+d.T)/2).min()), float(np.max(np.abs(d-d.T)))
 12
 13def affine_prob(A,b,z,S,n=300000,seed=11):
 14    rng=np.random.default_rng(seed)
 15    L=np.linalg.cholesky(S+1e-12*np.eye(S.shape[0]))
 16    u=rng.standard_normal((n,S.shape[0]))@L.T + z
 17    ok=np.all(u@A.T <= b[None,:],axis=1)
 18    return float(ok.mean()), int(ok.sum())
 19
 20def theorem_sanity():
 21    # A bounded convex rectangle, with deliberately inflated correlated covariance.
 22    SX=np.array([[.20,.08],[.08,.18]])
 23    SY=SX + .05*np.outer([1.,2.],[1.,2.])
 24    lo, asym=psd_check(SY,SX)
 25    A=np.array([[1,0],[-1,0],[0,1],[0,-1.]])
 26    b=np.array([1.55,1.55,1.55,1.55])
 27    # Centered rectangle has probability > 1/2; compare paired Monte Carlo estimates.
 28    py,ny=affine_prob(A,b,np.zeros(2),SY,300000,13)
 29    px,nx=affine_prob(A,b,np.zeros(2),SX,300000,14)
 30    # Also check an affine shifted set at several z values only where inflated p >= .5.
 31    rectangle_py, rectangle_px = py, px
 32    rows=[]
 33    for z in [np.array([0.,0.]),np.array([.25,-.15]),np.array([-.35,.10])]:
 34        py,_=affine_prob(A,b,z,SY,120000,30+len(rows))
 35        px,_=affine_prob(A,b,z,SX,120000,40+len(rows))
 36        rows.append((float(py),float(px)))
 37    return {'min_eigenvalue_SY_minus_SX':lo,'symmetry_error':asym,
 38            'rectangle_p_inflated':rectangle_py,'rectangle_p_smaller':rectangle_px,
 39            'shifted_pairs':rows,'all_observed_ordering':all(y>=.5 and x+0.012>=y for y,x in rows)}
 40
 41class Net(nn.Module):
 42    def __init__(self):
 43        super().__init__(); self.body=nn.Sequential(nn.Linear(2,12),nn.Tanh(),nn.Linear(12,2))
 44    def forward(self,x): return self.body(x)
 45
 46def train(mode, X, y, SY, steps=700, K=32):
 47    torch.manual_seed(SEED); model=Net(); opt=torch.optim.Adam(model.parameters(),lr=.025)
 48    L=torch.tensor(np.linalg.cholesky(SY),dtype=torch.float32)
 49    # acceptance is true-class logit minus false-class logit >= margin
 50    for step in range(steps):
 51        ix=torch.randint(0,len(X),(64,)); xb=X[ix]; yb=y[ix]
 52        z=model(xb); other=1-yb
 53        margin=z[torch.arange(64),yb]-z[torch.arange(64),other]
 54        # same inflated-noise Monte Carlo cost for both methods
 55        eps=torch.randn(K,64,2); noise=torch.einsum('ij,kbj->kbi',L,eps)
 56        zn=model(xb[None,:,:].expand(K,-1,-1).reshape(K*64,2)).reshape(K,64,2)
 57        zn=zn+noise
 58        noisy_margin=zn[:,:,0]-zn[:,:,1] # only valid for class order, signs below
 59        noisy_margin=torch.where(yb[None,:].bool(), noisy_margin, -noisy_margin)
 60        task=nn.functional.softplus(-margin).mean()
 61        if mode=='clean':
 62            loss=task
 63        elif mode=='ordinary':
 64            aug=nn.functional.softplus(-noisy_margin).mean()
 65            loss=task+0.7*aug
 66        else:
 67            # smooth estimate of P(margin >= 1), target p0=.80
 68            p_hat=torch.sigmoid(noisy_margin/.18).mean()
 69            loss=task+0.10*torch.relu(.55-p_hat)**2
 70        opt.zero_grad(); loss.backward(); opt.step()
 71    return model
 72
 73def evaluate(model,X,y,Slist):
 74    with torch.no_grad():
 75        logits=model(X); acc=(logits.argmax(1)==y).float().mean().item()
 76    out={'clean_accuracy':acc}
 77    for name,S in Slist:
 78        L=torch.tensor(np.linalg.cholesky(S),dtype=torch.float32)
 79        n=12000; xx=X[:min(256,len(X))]; yy=y[:len(xx)]
 80        eps=torch.randn(n,len(xx),2); zz=model(xx[None].expand(n,-1,-1).reshape(n*len(xx),2)).reshape(n,len(xx),2)
 81        noise=torch.einsum('ij,kbj->kbi',L,eps)
 82        # representation perturbation is applied after model, matching the affine head.
 83        zz=zz+noise
 84        pred=zz.argmax(2); accept=(pred==yy[None,:]).float().mean().item()
 85        out[name]=accept
 86    return out
 87
 88def mini_experiment():
 89    rng=np.random.default_rng(SEED); n=900
 90    X=np.r_[rng.multivariate_normal([-1.0,-.7],[[.45,.10],[.10,.35]],n//2),rng.multivariate_normal([1.0,.7],[[.45,.10],[.10,.35]],n//2)].astype('float32')
 91    y=np.r_[np.zeros(n//2,dtype=np.int64),np.ones(n//2,dtype=np.int64)]
 92    p=rng.permutation(n); X=X[p]; y=y[p]
 93    Xt=torch.tensor(X); yt=torch.tensor(y)
 94    SX=np.array([[.08,.025],[.025,.06]])
 95    SY=SX+.12*np.outer([1.,.5],[1.,.5])
 96    Slist=[('small_cov',SX),('inflated_cov',SY),('isotropic_small',np.eye(2)*.04)]
 97    results={}
 98    for mode in ['clean','ordinary','chance']:
 99        m=train(mode,Xt,yt,SY); results[mode]=evaluate(m,Xt,yt,Slist)
100    return {'covariance_min_eigenvalue':float(np.linalg.eigvalsh(SY-SX).min()),'results':results}
101
102if __name__=='__main__':
103    print(json.dumps({'sanity':theorem_sanity(),'experiment':mini_experiment()},indent=2))