1import json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import sys
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import sweep_baseline, make_report
  9
 10META = {"name": "bidirectional_conditional_compatibility", "domain": "paired categorical conditionals", "description": "Synthetic positive joint with two neural conditional directions; evaluates compatibility regularization on held-out quadruples."}
 11SEEDS = tuple(range(8))
 12SWEEP_SEEDS = (0, 1, 2, 3)
 13NX = NY = 4
 14RAW = np.array([[1.0, 2.0, .7, 1.4], [2.1, .8, 1.8, .5], [.6, 1.7, 2.4, 1.1], [1.3, .9, 1.5, 2.2]], dtype=np.float64)
 15JOINT = RAW / RAW.sum()
 16
 17
 18def get_dataset(seed, n_train=400, n_test=400):
 19    rng = np.random.default_rng(1009 + seed)
 20    z = rng.choice(NX * NY, n_train + n_test, p=JOINT.ravel())
 21    x, y = z // NY, z % NY
 22    return {"xtr": x[:n_train].astype(np.int64), "ytr": y[:n_train].astype(np.int64),
 23            "xte": x[n_train:].astype(np.int64), "yte": y[n_train:].astype(np.int64),
 24            "task": "classification", "metric": "bidirectional_test_nll", "input_shape": (2,), "out_dim": 4}
 25
 26
 27class TwoConditionalMLP(nn.Module):
 28    def __init__(self, hidden=16):
 29        super().__init__()
 30        self.q = nn.Sequential(nn.Linear(NY, hidden), nn.Tanh(), nn.Linear(hidden, NX))
 31        self.r = nn.Sequential(nn.Linear(NX, hidden), nn.Tanh(), nn.Linear(hidden, NY))
 32
 33    def log_probs(self):
 34        ey = torch.eye(NY, device=next(self.parameters()).device)
 35        ex = torch.eye(NX, device=next(self.parameters()).device)
 36        return torch.log_softmax(self.q(ey), 1), torch.log_softmax(self.r(ex), 1)
 37
 38
 39def cycle_delta(lq, lr, x1, x2, y1, y2):
 40    return (lq[y1, x1] + lr[x2, y1] + lq[y2, x2] + lr[x1, y2]
 41            - lr[x1, y1] - lq[y2, x1] - lr[x2, y2] - lq[y1, x2])
 42
 43
 44def quadruples(seed):
 45    rows = [(x1,x2,y1,y2) for x1 in range(NX) for x2 in range(NX) if x1 != x2
 46            for y1 in range(NY) for y2 in range(NY) if y1 != y2]
 47    rng = np.random.default_rng(7001 + seed)
 48    p = rng.permutation(len(rows)); cut = len(rows)//2
 49    a = np.asarray(rows, dtype=np.int64)
 50    return torch.tensor(a[p[:cut]]), torch.tensor(a[p[cut:]])
 51
 52
 53def fit(seed, lr, lam, return_model=False):
 54    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 55    ds = get_dataset(seed)
 56    # CPU is intentional here: the custom categorical benchmark is tiny and deterministic.
 57    dev = torch.device('cpu')
 58    model = TwoConditionalMLP().to(dev)
 59    opt = torch.optim.Adam(model.parameters(), lr=lr)
 60    x = torch.tensor(ds['xtr'], dtype=torch.long); y = torch.tensor(ds['ytr'], dtype=torch.long)
 61    trainq, heldq = quadruples(seed)
 62    for ep in range(70):
 63        perm = torch.randperm(len(x))
 64        for i in range(0, len(x), 128):
 65            ix = perm[i:i+128]
 66            lq, lrlog = model.log_probs()
 67            task = -.5 * (lq[y[ix], x[ix]].mean() + lrlog[x[ix], y[ix]].mean())
 68            d = cycle_delta(lq, lrlog, trainq[:,0], trainq[:,1], trainq[:,2], trainq[:,3])
 69            loss = task + lam * .5 * d.square().mean()
 70            opt.zero_grad(); loss.backward(); opt.step()
 71    with torch.no_grad():
 72        lq, lrlog = model.log_probs()
 73        xt = torch.tensor(ds['xte']); yt = torch.tensor(ds['yte'])
 74        nll = float(-.5 * (lq[yt, xt].mean() + lrlog[xt, yt].mean()))
 75        allq = torch.tensor([(x1,x2,y1,y2) for x1 in range(NX) for x2 in range(NX) if x1 != x2
 76                             for y1 in range(NY) for y2 in range(NY) if y1 != y2])
 77        dheld = cycle_delta(lq, lrlog, heldq[:,0], heldq[:,1], heldq[:,2], heldq[:,3]).numpy()
 78        dall = cycle_delta(lq, lrlog, allq[:,0], allq[:,1], allq[:,2], allq[:,3]).numpy()
 79        # Empirical observed conditionals provide an independent data-derived reference.
 80        counts = np.ones((NX, NY), dtype=np.float64) * 1e-3
 81        for xx, yy in zip(ds['xtr'], ds['ytr']): counts[xx, yy] += 1
 82        oq = torch.tensor((counts / counts.sum(0, keepdims=True)).T, dtype=torch.float32).log()
 83        orun = torch.tensor((counts / counts.sum(1, keepdims=True)), dtype=torch.float32).log()
 84        dob = cycle_delta(oq, orun, allq[:,0], allq[:,1], allq[:,2], allq[:,3]).numpy()
 85    out = {'metric': nll, 'heldout_p95_abs_delta': float(np.percentile(np.abs(dheld),95)),
 86           'heldout_mean_abs_delta': float(np.mean(np.abs(dheld))), 'all_mean_abs_delta': float(np.mean(np.abs(dall))),
 87           'observed_empirical_mean_abs_delta': float(np.mean(np.abs(dob)))}
 88    if return_model: out['model'] = model
 89    return out
 90
 91
 92def base_train(cfg, seed): return fit(seed, cfg['lr'], 0.0)['metric']
 93
 94def idea_train(cfg, seed): return fit(seed, cfg['lr'], cfg['lambda'])['metric']
 95
 96def aggregate(rows):
 97    vals = [r if isinstance(r, (float, int)) else r['metric'] for r in rows]
 98    return {'per_seed': [float(v) for v in vals], 'mean': float(np.mean(vals)),
 99            'std': float(np.std(vals, ddof=1))}
100
101
102def main():
103    # The union is shared: all idea learning rates are evaluated by baseline too.
104    lrs = [0.01, 0.03, 0.1]
105    base_grid = [{'lr': lr} for lr in lrs]
106    sweep = sweep_baseline(lambda cfg: (lambda seed: base_train(cfg, seed)), base_grid, seeds=SWEEP_SEEDS)
107    # Select by the harness's baseline sweep result, then use the same 3-config budget for idea.
108    best_lr = sweep['best_cfg']['lr']
109    idea_grid = [{'lr': lr, 'lambda': lam} for lr, lam in [(best_lr, .1), (best_lr, .3), (best_lr, .6)]]
110    idea_cfg_results = []
111    for cfg in idea_grid:
112        rows = [fit(s, cfg['lr'], cfg['lambda']) for s in SEEDS]
113        idea_cfg_results.append((cfg, aggregate(rows)))
114    best_cfg, idea = min(idea_cfg_results, key=lambda z: z[1]['mean'])
115    idea_details = [fit(s, best_cfg['lr'], best_cfg['lambda']) for s in SEEDS]
116    idea = aggregate(idea_details)
117    # Full paired baseline at the selected common learning rate.
118    base_details = [fit(s, best_lr, 0.0) for s in SEEDS]
119    base_full = aggregate(base_details)
120    base_block = {'sweep': sweep, 'best_config': {'lr': best_lr}, 'full': base_full}
121    sig = {'prediction': 'cycle penalty lowers held-out 95th-percentile |Delta| while retaining task NLL',
122           'baseline_heldout_p95': float(np.mean([r['heldout_p95_abs_delta'] for r in base_details])),
123           'idea_heldout_p95': float(np.mean([r['heldout_p95_abs_delta'] for r in idea_details])),
124           'baseline_test_nll': base_full['mean'], 'idea_test_nll': idea['mean'],
125           'observed_empirical_cycle_mean': float(np.mean([r['observed_empirical_mean_abs_delta'] for r in base_details])),
126           'confirmed': bool(np.mean([r['heldout_p95_abs_delta'] for r in idea_details]) <
127                             .8*np.mean([r['heldout_p95_abs_delta'] for r in base_details]))}
128    report = make_report('custom:bidirectional_conditional_compatibility', 'two_conditional_mlp', base_block, idea,
129                         {'custom_track': {'name': META['name'], 'file': 'cycle_bench.py', 'domain': META['domain']}, **sig})
130    Path('bench_report.json').write_text(json.dumps(report, indent=2))
131    print(json.dumps(report, indent=2))
132
133if __name__ == '__main__': main()