1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 2953
  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"
 11
 12# q(x|y) and r(y|x) are represented by two small categorical tables.
 13class TwoConditionals(nn.Module):
 14    def __init__(self, nx, ny):
 15        super().__init__()
 16        self.q_logits = nn.Parameter(torch.zeros(ny, nx))
 17        self.r_logits = nn.Parameter(torch.zeros(nx, ny))
 18
 19    def log_q(self): return torch.log_softmax(self.q_logits, dim=1)
 20    def log_r(self): return torch.log_softmax(self.r_logits, dim=1)
 21
 22
 23def cycle_delta(logq, logr, x1, x2, y1, y2):
 24    # Indexing follows q[x|y] = logq[y,x], r[y|x] = logr[x,y].
 25    return (logq[y1, x1] + logr[x2, y1] + logq[y2, x2] + logr[x1, y2]
 26            - logr[x1, y1] - logq[y2, x1] - logr[x2, y2] - logq[y1, x2])
 27
 28
 29def all_quadruples(nx, ny, device):
 30    # Distinct pairs make the certificate nontrivial; all held-out evaluation
 31    # quadruples are deterministic rather than sampled from the training data.
 32    rows = []
 33    for x1 in range(nx):
 34      for x2 in range(nx):
 35       if x1 == x2: continue
 36       for y1 in range(ny):
 37        for y2 in range(ny):
 38         if y1 == y2: continue
 39         rows.append((x1,x2,y1,y2))
 40    return [torch.tensor([z[i] for z in rows], device=device) for i in range(4)]
 41
 42
 43def residuals(model, quads):
 44    lq, lr = model.log_q(), model.log_r()
 45    return cycle_delta(lq, lr, *quads).detach().cpu().numpy()
 46
 47
 48def train(joint, lam, steps=900):
 49    nx, ny = joint.shape
 50    model = TwoConditionals(nx, ny).to(device)
 51    opt = torch.optim.Adam(model.parameters(), lr=0.08)
 52    p = torch.tensor(joint.ravel(), dtype=torch.float32, device=device)
 53    # Fixed synthetic population gives exactly matched data across conditions.
 54    rng = np.random.default_rng(SEED + 17)
 55    samples = rng.choice(nx*ny, size=5000, p=joint.ravel())
 56    xs = torch.tensor(samples // ny, dtype=torch.long, device=device)
 57    ys = torch.tensor(samples % ny, dtype=torch.long, device=device)
 58    train_x, train_y = xs[:4000], ys[:4000]
 59    qs = all_quadruples(nx, ny, device)
 60    # Hold out a deterministic quarter of quadruples for generalization testing.
 61    nq = qs[0].numel()
 62    perm = torch.randperm(nq, generator=torch.Generator(device=device).manual_seed(SEED+91), device=device)
 63    cut = int(0.75*nq)
 64    train_q = tuple(z[perm[:cut]] for z in qs)
 65    heldout_q = tuple(z[perm[cut:]] for z in qs)
 66    for step in range(steps):
 67        ix = torch.arange((step*64) % 3900, (step*64) % 3900 + 64, device=device) % 4000
 68        x, y = train_x[ix], train_y[ix]
 69        lq, lr = model.log_q(), model.log_r()
 70        task = -0.5 * (lq[y, x].mean() + lr[x, y].mean())
 71        # Use all small-support quadruples, which is cheap and lower variance.
 72        d = cycle_delta(lq, lr, *train_q)
 73        cycle = 0.5 * d.square().mean()
 74        loss = task + lam * cycle
 75        opt.zero_grad(); loss.backward(); opt.step()
 76    # Recover a joint from q and r using the least-squares log-ratio identity.
 77    with torch.no_grad():
 78        lq, lr = model.log_q(), model.log_r()
 79        # For a compatible pair, log p(x,y) differs from log q(x|y)+log p(y),
 80        # and p(y) can be estimated by averaging the two directional joints.
 81        joint_q = torch.exp(lq).T * torch.tensor(joint.sum(0), device=device)
 82        joint_r = torch.exp(lr) * torch.tensor(joint.sum(1), device=device)[:,None]
 83        recovered = 0.5*(joint_q + joint_r)
 84        recovered = recovered / recovered.sum()
 85        true = torch.tensor(joint, dtype=torch.float32, device=device)
 86        joint_l1 = torch.abs(recovered-true).sum().item()
 87        task_final = float((-0.5*(lq[train_y,train_x].mean()+lr[train_x,train_y].mean())).item())
 88    rr = residuals(model, qs)
 89    rh = residuals(model, heldout_q)
 90    return {"model": model, "p95_abs_residual": float(np.percentile(np.abs(rr),95)),
 91            "heldout_p95_abs_residual": float(np.percentile(np.abs(rh),95)),
 92            "mean_abs_residual": float(np.mean(np.abs(rr))),
 93            "heldout_mean_abs_residual": float(np.mean(np.abs(rh))), "joint_l1": joint_l1,
 94            "task_nll": task_final}
 95
 96
 97def main():
 98    # Strictly positive non-product joint, so the exact identity is meaningful.
 99    raw = np.array([[1.0, 2.0, 0.7, 1.4], [2.1, 0.8, 1.8, 0.5],
100                    [0.6, 1.7, 2.4, 1.1], [1.3, 0.9, 1.5, 2.2]], dtype=np.float64)
101    joint = raw/raw.sum()
102    nx, ny = joint.shape
103    px, py = joint.sum(1), joint.sum(0)
104    q = joint/py[None,:]       # q[x,y]
105    r = joint/px[:,None]       # r[x,y]
106    # Core numerical verification: exact conditionals satisfy Delta=0.
107    exact = []
108    arbitrary = []
109    rng = np.random.default_rng(SEED)
110    aq = rng.dirichlet(np.ones(nx), size=ny)
111    ar = rng.dirichlet(np.ones(ny), size=nx)
112    for x1,x2,y1,y2 in [(0,1,0,1),(1,3,2,0),(2,3,1,3)]:
113        exact.append(math.log(q[x1,y1])+math.log(r[x2,y1])+math.log(q[x2,y2])+math.log(r[x1,y2])
114          -math.log(r[x1,y1])-math.log(q[x1,y2])-math.log(r[x2,y2])-math.log(q[x2,y1]))
115        arbitrary.append(math.log(aq[y1,x1])+math.log(ar[x2,y1])+math.log(aq[y2,x2])+math.log(ar[x1,y2])
116          -math.log(ar[x1,y1])-math.log(aq[y2,x1])-math.log(ar[x2,y2])-math.log(aq[y1,x2]))
117    exact_max = float(np.max(np.abs(exact)))
118    arbitrary_mean = float(np.mean(np.abs(arbitrary)))
119    baseline = train(joint, 0.0)
120    idea = train(joint, 0.5)
121    out = {"seed":SEED,"device":device,"math_check_exact_max_abs_delta":exact_max,
122           "math_check_random_mean_abs_delta":arbitrary_mean,"baseline":{k:v for k,v in baseline.items() if k!="model"},
123           "idea_lambda_0.5":{k:v for k,v in idea.items() if k!="model"}}
124    Path("results.json").write_text(json.dumps(out, indent=2))
125    print(json.dumps(out, indent=2))
126
127if __name__ == "__main__":
128    try:
129        main()
130    except (RuntimeError, torch.cuda.CudaError) as e:
131        if device == "cuda":
132            print("CUDA failed; rerun on CPU", repr(e))
133            torch.cuda.empty_cache(); device = "cpu"; main()
134        else: raise