1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 19
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9try:
 10    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 11    if device == 'cuda': torch.zeros(1, device='cuda')
 12except Exception:
 13    device = 'cpu'
 14
 15
 16def cube_neighbors(n, dev):
 17    ids = torch.arange(2 ** n, device=dev)
 18    return torch.stack([ids ^ (1 << i) for i in range(n)], 1)
 19
 20
 21def laplacian_target(f, nb):
 22    return 0.5 * (f[:, None] - f[nb]).sum(1)
 23
 24
 25def divergence(v, nb):
 26    n = nb.shape[1]
 27    return (v - v[nb, torch.arange(n, device=v.device)]).sum(1)
 28
 29
 30def solve_beckmann(f, nb, steps=80, rho=2.0, eta=0.002, eps=1e-5):
 31    """Differentiable unrolled minimization of mean(||V||_2)+rho*mean(div(V)-Lf)^2."""
 32    n = nb.shape[1]; N = f.numel()
 33    b = laplacian_target(f, nb)
 34    v = torch.zeros((N, n), device=f.device, dtype=f.dtype)
 35    cols = torch.arange(n, device=f.device)
 36    for _ in range(steps):
 37        r = divergence(v, nb) - b
 38        # div^* r has entries r(x)-r(x^flip); div is self-adjoint here.
 39        adj = r[:, None] - r[nb]
 40        norm = torch.sqrt((v * v).sum(1, keepdim=True) + eps)
 41        # adj has shape [N,n]; this is the gradient of the residual term.
 42        grad = v / norm + 2.0 * rho * adj
 43        v = v - eta * grad
 44    residual = divergence(v, nb) - b
 45    value = torch.sqrt((v * v).sum(1) + eps).mean() + rho * (residual ** 2).mean()
 46    return value, v, residual, b
 47
 48
 49class SmallMLP(nn.Module):
 50    def __init__(self, n):
 51        super().__init__()
 52        self.net = nn.Sequential(nn.Linear(n, 24), nn.Tanh(), nn.Linear(24, 1))
 53    def forward(self, x): return self.net(x).squeeze(1)
 54
 55
 56def metrics(model, x, y, nb, train_idx):
 57    with torch.no_grad():
 58        f = model(x); pred = (f > 0).long()
 59        acc = (pred == y).float().mean().item()
 60        one = (f[:, None] - f[nb]).abs().mean().item()
 61        # random reproducible two-bit perturbation, using coordinates 0 and 1
 62        two_nb = nb[:, 0][nb[:, 1]] if False else None
 63        ids = torch.arange(x.shape[0], device=x.device)
 64        two = ids ^ 1 ^ 2
 65        multi = (f - f[two]).abs().mean().item()
 66        ce = nn.functional.binary_cross_entropy_with_logits(f[train_idx], y[train_idx].float()).item()
 67    return {'accuracy_all': acc, 'one_bit_logit_change': one,
 68            'two_bit_logit_change': multi, 'train_ce': ce}
 69
 70
 71def run_experiment(n=8, train_fraction=0.5, steps=260):
 72    nb = cube_neighbors(n, device)
 73    N = 2 ** n
 74    ids = torch.arange(N, device=device)
 75    # {-1,+1} cube; bit order is consistent with neighbor table.
 76    x = 2.0 * ((ids[:, None] >> torch.arange(n, device=device)) & 1).float() - 1.0
 77    # parity interaction on first three coordinates, with a weak nuisance term.
 78    y = (((x[:, :3].prod(1) > 0)).long())
 79    perm = torch.randperm(N, device=device)
 80    train_idx = perm[:int(train_fraction * N)]
 81    out = {}
 82    for name in ('ce', 'edge', 'beckmann'):
 83        torch.manual_seed(SEED + {'ce': 0, 'edge': 1, 'beckmann': 2}[name])
 84        model = SmallMLP(n).to(device)
 85        opt = torch.optim.Adam(model.parameters(), lr=0.025)
 86        last_flow = None
 87        for t in range(steps):
 88            opt.zero_grad()
 89            f = model(x)
 90            loss = nn.functional.binary_cross_entropy_with_logits(f[train_idx], y[train_idx].float())
 91            if name == 'edge':
 92                edge = ((f[:, None] - f[nb]) ** 2).mean()
 93                loss = loss + 0.018 * edge
 94            elif name == 'beckmann':
 95                br, _, res, _ = solve_beckmann(f, nb)
 96                loss = loss + 0.008 * br
 97                last_flow = res.detach()
 98            loss.backward(); opt.step()
 99        m = metrics(model, x, y, nb, train_idx)
100        if name == 'beckmann':
101            with torch.no_grad():
102                _, _, rr, bb = solve_beckmann(model(x), nb)
103            m['flow_residual_rms'] = float(torch.sqrt((rr * rr).mean()).cpu())
104            m['target_rms'] = float(torch.sqrt((bb * bb).mean()).cpu())
105        m['steps'] = steps
106        out[name] = m
107    return out
108
109
110def verify_math(n=4):
111    nb = cube_neighbors(n, device); N = 2 ** n
112    g = torch.randn(N, device=device)
113    v = torch.randn(N, n, device=device)
114    b = laplacian_target(g, nb)
115    d = divergence(v, nb)
116    # Both operators have zero mean, and divergence is the adjoint of itself.
117    mean_b = abs(float(b.mean()))
118    mean_d = abs(float(d.mean()))
119    lhs = (divergence(v, nb) * g).sum()
120    rhs = (v * (g[:, None] - g[nb])).sum()
121    adjoint_err = abs(float(lhs - rhs))
122    # A convex solver should reduce its penalized objective from V=0.
123    f = torch.sin(torch.arange(N, device=device).float())
124    b_f = laplacian_target(f, nb)
125    # Explicit feasible reference: V_i=D_i f/2 gives div(V)=Lf exactly.
126    vf = 0.25 * (f[:, None] - f[nb])
127    feasible_residual = divergence(vf, nb) - b_f
128    zero_obj = 2.0 * (b_f ** 2).mean()
129    val, vv, rr, _ = solve_beckmann(f, nb)
130    return {'mean_abs_Lf': mean_b, 'mean_abs_divV': mean_d,
131            'adjoint_relative_error': adjoint_err / (abs(float(lhs)) + 1e-8),
132            'explicit_feasible_flow_residual': float(torch.sqrt((feasible_residual ** 2).mean())),
133            'explicit_feasible_flow_norm': float(torch.sqrt((vf * vf).sum(1)).mean()),
134            'zero_flow_penalized_objective': float(zero_obj),
135            'optimized_penalized_objective': float(val.detach()),
136            'solver_residual_rms': float(torch.sqrt((rr * rr).mean()).detach()),
137            'solver_residual_over_target_rms': float(torch.sqrt((rr * rr).mean()) / (torch.sqrt((b_f * b_f).mean()) + 1e-8))}
138
139
140if __name__ == '__main__':
141    checks = verify_math()
142    results = run_experiment()
143    report = {'device': device, 'math_check': checks, 'experiment': results}
144    with open('results.json', 'w') as f: json.dump(report, f, indent=2)
145    print(json.dumps(report, indent=2))