1import json, math, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED = 1344
  7np.random.seed(SEED); torch.manual_seed(SEED)
  8device = "cuda" if torch.cuda.is_available() else "cpu"
  9try:
 10    if device == "cuda":
 11        torch.zeros(1, device="cuda").sum().item()
 12except Exception:
 13    device = "cpu"
 14
 15torch.set_default_dtype(torch.float64)
 16
 17
 18def simplex_map(z, eps):
 19    return torch.softmax(-z / eps, dim=-1)
 20
 21
 22def simplex_jacobian(w, eps):
 23    # derivative of softmax(-z/eps), evaluated at the corresponding w
 24    return -(torch.diag(w) - w[:, None] * w[None, :]) / eps
 25
 26
 27def verify_math():
 28    # Prediction 1: at z=0, tangent singular values are exactly 1/(n eps).
 29    n = 7
 30    z = torch.zeros(n)
 31    rows = []
 32    for eps in [0.05, 0.1, 0.2, 0.4]:
 33        w = simplex_map(z, eps)
 34        J = simplex_jacobian(w, eps).numpy()
 35        sv = np.linalg.svd(J, compute_uv=False)
 36        measured = float(sv[0]); predicted = 1.0 / (n * eps)
 37        rows.append({"eps": eps, "predicted_tangent_gain": predicted,
 38                     "measured_tangent_gain": measured,
 39                     "relative_error": abs(measured-predicted)/predicted})
 40    # Prediction 2: probability ratios obey log(w_i/w_j)=- (z_i-z_j)/eps.
 41    ratio_rows = []
 42    z = torch.tensor([0.0, 0.37, -0.21, 0.8])
 43    for eps in [0.05, 0.1, 0.25, 0.5]:
 44        w = simplex_map(z, eps)
 45        observed = float(torch.log(w[0]/w[1]))
 46        predicted = float(-(z[0]-z[1])/eps)
 47        ratio_rows.append({"eps": eps, "predicted_log_ratio": predicted,
 48                           "measured_log_ratio": observed,
 49                           "absolute_error": abs(observed-predicted)})
 50    # Prediction 3: scaling equivariance T_eps(c z)=T_{eps/c}(z).
 51    z = torch.tensor([0.2, -0.4, 0.9, -0.1, 0.5])
 52    equiv = []
 53    for eps, c in [(0.1, 2.), (0.2, 3.), (0.4, .5), (.07, 5.)]:
 54        lhs = simplex_map(c*z, eps)
 55        rhs = simplex_map(z, eps/c)
 56        equiv.append({"eps": eps, "c": c,
 57                      "max_abs_error": float((lhs-rhs).abs().max())})
 58    return {"jacobian_scaling": rows, "log_ratio_scaling": ratio_rows,
 59            "epsilon_score_equivariance": equiv}
 60
 61
 62class CapacityBarrier(nn.Module):
 63    """Entropy Legendre regularizer on the simplex plus log barriers Aw<b.
 64    Equality sum(w)=1 is handled by a KKT Newton solve; iterations are unrolled.
 65    """
 66    def __init__(self, A, b, eps=0.12, steps=14):
 67        super().__init__()
 68        self.register_buffer("A", A)
 69        self.register_buffer("b", b)
 70        self.eps, self.steps = eps, steps
 71
 72    def forward(self, z):
 73        B, n = z.shape; m = self.A.shape[0]
 74        w = torch.full((B,n), 1.0/n, dtype=z.dtype, device=z.device)
 75        # fixed damped Newton with a feasibility-preserving step chosen per row
 76        for _ in range(self.steps):
 77            slack = self.b[None,:] - w @ self.A.T
 78            # strict interior is guaranteed by construction of the experiment
 79            g = z + self.eps * (torch.log(w)+1.0) + self.eps * (1.0/slack) @ self.A
 80            H = self.eps * (torch.diag_embed(1.0/w) +
 81                torch.einsum('mi,bm,mj->bij', self.A, 1.0/(slack**2), self.A))
 82            C = torch.ones((B,1,n), dtype=z.dtype, device=z.device)
 83            K = torch.cat([torch.cat([H,C.transpose(1,2)], dim=2),
 84                           torch.cat([C, torch.zeros(B,1,1,dtype=z.dtype,device=z.device)], dim=2)], dim=1)
 85            rhs = torch.cat([g, (w.sum(1)-1.0)[:,None]], dim=1)
 86            sol = torch.linalg.solve(K, rhs[...,None])[...,0]
 87            delta = sol[:,:n]
 88            # Backtracking only determines a scalar; detach it to retain a stable graph.
 89            alpha = torch.ones(B, dtype=z.dtype, device=z.device)
 90            neg = delta > 0
 91            alpha = torch.minimum(alpha, torch.where(neg, -0.9*w/(delta-1e-30), alpha[:,None]).min(1).values)
 92            dslack = delta @ self.A.T
 93            alpha = torch.minimum(alpha, torch.where(dslack < 0, -0.9*slack/(dslack-1e-30), alpha[:,None]).min(1).values)
 94            alpha = alpha.clamp(max=1.0, min=1e-4).detach()
 95            w = w - alpha[:,None] * delta
 96        return w
 97
 98
 99def make_data(N=768, n=8, m=3):
100    rng = np.random.default_rng(SEED+4)
101    A = rng.uniform(.05, 1.0, size=(m,n)); A /= A.sum(1,keepdims=True)
102    b = np.full(m, .70) # uniform slack .30, while random policies often violate
103    X = rng.normal(size=(N, 6))
104    # Feasible targets are mixtures of uniform and random simplex points.
105    raw = rng.dirichlet(np.ones(n), size=N)
106    target = .72/n + .28*raw
107    # use a fixed strictly-feasible target independent of rare random violations
108    target = target / target.sum(1, keepdims=True)
109    return map(torch.tensor,(X,target,A,b))
110
111
112def train_compare():
113    X, target, A, b = make_data(); X=X.to(device); target=target.to(device)
114    A=A.to(device); b=b.to(device)
115    split=600; n=target.shape[1]
116    def model(): return nn.Sequential(nn.Linear(6,24), nn.Tanh(), nn.Linear(24,n)).to(device)
117    def run(kind):
118        torch.manual_seed(SEED+ (0 if kind=='penalty' else 1))
119        net=model(); opt=torch.optim.Adam(net.parameters(), lr=.015)
120        layer=CapacityBarrier(A,b).to(device)
121        t0=time.perf_counter()
122        for step in range(180):
123            ix=torch.arange((step*32)%split, (step*32)%split+32, device=device) % split
124            z=net(X[ix]); w=torch.softmax(z, -1) if kind=='penalty' else layer(z)
125            viol=torch.relu(w@A.T-b)
126            loss=((w-target[ix])**2).mean() + (50*viol.square().mean() if kind=='penalty' else 0)
127            opt.zero_grad(); loss.backward(); opt.step()
128        with torch.no_grad():
129            z=net(X[split:]); w=torch.softmax(z,-1) if kind=='penalty' else layer(z)
130            v=torch.relu(w@A.T-b)
131            mse=((w-target[split:])**2).mean().item()
132            maxv=v.max().item(); meanv=v.mean().item(); minsl=(b-w@A.T).min().item()
133        return {"test_mse":mse, "max_capacity_violation":maxv,
134                "mean_capacity_violation":meanv, "minimum_slack":minsl,
135                "seconds":time.perf_counter()-t0}
136    return {"penalty":run('penalty'), "legendre_barrier":run('legendre')}
137
138if __name__ == '__main__':
139    checks=verify_math()
140    comparison=train_compare()
141    out={"device":device, "checks":checks, "comparison":comparison}
142    with open('results.json','w') as f: json.dump(out,f,indent=2)
143    print(json.dumps(out,indent=2))