1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6# Cross-Partial Nash Compatibility Regularizer MVP.
  7# q1=a*u1*u2 and q2=b*u1*u2 are deliberately incompatible unless a=b.
  8# The supervised target (a,b)=(1,-1) creates a controlled critic-fitting problem.
  9
 10torch.set_default_dtype(torch.float64)
 11SEED = 1167
 12
 13def set_seed(seed=SEED):
 14    np.random.seed(seed); torch.manual_seed(seed)
 15
 16
 17def cross_residual(a, b, u):
 18    """Autograd implementation of d/du2(dq1/du1)-d/du1(dq2/du2)."""
 19    u = u.detach().clone().requires_grad_(True)
 20    q1 = a * u[:, 0] * u[:, 1]
 21    q2 = b * u[:, 0] * u[:, 1]
 22    g1 = torch.autograd.grad(q1.sum(), u, create_graph=True)[0][:, 0]
 23    g2 = torch.autograd.grad(q2.sum(), u, create_graph=True)[0][:, 1]
 24    h12 = torch.autograd.grad(g1.sum(), u, create_graph=True)[0][:, 1]
 25    h21 = torch.autograd.grad(g2.sum(), u, create_graph=True)[0][:, 0]
 26    return h12 - h21
 27
 28
 29def analytic_lambda_prediction(lam):
 30    # argmin .5[(a-1)^2+(b+1)^2] + lam*(a-b)^2 gives d=a-b=2/(1+4 lam).
 31    return 2.0 / (1.0 + 4.0 * lam)
 32
 33
 34def solve_coeff(lam, steps=500, lr=0.04):
 35    set_seed()
 36    ab = torch.tensor([0.0, 0.0], requires_grad=True)
 37    opt = torch.optim.SGD([ab], lr=lr)
 38    u = torch.linspace(-1, 1, 64).reshape(-1, 1).repeat(1, 2)
 39    target = torch.tensor([1.0, -1.0])
 40    for _ in range(steps):
 41        # The coefficient objective is the exact average critic fitting objective.
 42        fit = 0.5 * ((ab - target) ** 2).sum()
 43        compat = (ab[0] - ab[1]) ** 2
 44        loss = fit + lam * compat
 45        opt.zero_grad(); loss.backward(); opt.step()
 46    residual = float((ab[0] - ab[1]).detach().abs())
 47    fit = float((0.5 * ((ab.detach() - target) ** 2).sum()))
 48    return ab.detach().numpy().tolist(), residual, fit
 49
 50
 51def gamma_sweep():
 52    # Gibbs log-conditionals have cross derivative mismatch (a-b)/gamma.
 53    # Therefore the squared mismatch must be residual^2/gamma^2.
 54    d = 2.0
 55    out = []
 56    for gamma in [0.25, 0.5, 1.0, 2.0, 4.0]:
 57        observed = (d / gamma) ** 2
 58        predicted = 4.0 / gamma**2
 59        out.append({"gamma": gamma, "observed": observed, "predicted": predicted,
 60                    "relative_error": abs(observed-predicted)/(predicted+1e-12)})
 61    return out
 62
 63
 64def lambda_sweep():
 65    out = []
 66    for lam in [0.0, 0.05, 0.25, 0.5, 1.0, 2.0, 5.0]:
 67        ab, observed, fit = solve_coeff(lam)
 68        predicted = analytic_lambda_prediction(lam)
 69        out.append({"lambda": lam, "a_b": ab, "observed_abs_residual": observed,
 70                    "predicted_abs_residual": predicted, "fit_loss": fit,
 71                    "relative_error": abs(observed-predicted)/(predicted+1e-12)})
 72    return out
 73
 74
 75def finite_difference_check():
 76    # Independently check the autograd residual against finite differences.
 77    u = torch.tensor([[0.37, -0.21]])
 78    a, b = 1.3, -0.7
 79    auto = float(cross_residual(torch.tensor(a), torch.tensor(b), u)[0])
 80    eps = 1e-5
 81    def q1(x, y): return a*x*y
 82    def q2(x, y): return b*x*y
 83    # d/du2 dq1/du1 and d/du1 dq2/du2 via central differences of first derivatives.
 84    fd12 = (q1(1, u[0,1].item()+eps)-q1(1, u[0,1].item()-eps))/(2*eps)
 85    fd21 = (q2(u[0,0].item()+eps, 1)-q2(u[0,0].item()-eps, 1))/(2*eps)
 86    return {"autograd_residual": auto, "finite_difference_residual": fd12-fd21,
 87            "absolute_error": abs(auto-(fd12-fd21))}
 88
 89
 90def mini_compare():
 91    # Same coefficient critic and optimizer, interpreted as a tiny centralized critic.
 92    rows = []
 93    for lam in [0.0, 1.0]:
 94        ab, residual, fit = solve_coeff(lam, steps=500, lr=0.04)
 95        rows.append({"method": "baseline" if lam == 0 else "compat_regularized",
 96                     "lambda": lam, "final_abs_cross_partial": residual,
 97                     "critic_fit_loss": fit, "coefficients": ab})
 98    return rows
 99
100
101def main():
102    set_seed()
103    device = "cuda" if torch.cuda.is_available() else "cpu"
104    # This MVP is deliberately CPU-sized; no GPU allocation is needed.
105    result = {"device_available": device, "finite_difference": finite_difference_check(),
106              "lambda_sweep": lambda_sweep(), "gamma_sweep": gamma_sweep(),
107              "mini_compare": mini_compare(),
108              "predictions": [
109                "At lambda=0 residual=2; increasing lambda follows 2/(1+4 lambda).",
110                "The Gibbs incompatibility squared scales as 1/gamma^2, so halving gamma quadruples it.",
111                "At large lambda residual tends to zero (compatibility), while fitting error increases because targets are intentionally incompatible."
112              ]}
113    Path("results.json").write_text(json.dumps(result, indent=2))
114    print(json.dumps(result, indent=2))
115
116if __name__ == "__main__":
117    main()