File size: 3,788 Bytes
9860743 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | import torch
import cvxpy as cp
import numpy as np
from ffolayer import FFOLayer
def test_example():
# ============================================================
# Problem:
# minimize 0.5 * ||Q_sqrt x||^2 + q^T x
# subject to A x == b
#
# ============================================================
n = 5
m = 2
# CVXPY parameters
Q_sqrt_cp = cp.Parameter((n, n))
q_cp = cp.Parameter(n)
A_cp = cp.Parameter((m, n))
b_cp = cp.Parameter(m)
x = cp.Variable(n)
objective = cp.Minimize(0.5 * cp.sum_squares(Q_sqrt_cp @ x) + q_cp.T @ x)
constraints = [A_cp @ x == b_cp]
problem = cp.Problem(objective, constraints)
assert problem.is_dpp(), "Problem must be DPP-compliant"
# ============================================================
# Torch data
# ============================================================
torch.manual_seed(1)
np.random.seed(1)
torch.set_default_dtype(torch.float64)
M = torch.randn(n, n, dtype=torch.float64)
eps = 0.2
Q_sqrt_tch = (M + eps * torch.eye(n, dtype=torch.float64)).detach().clone().requires_grad_(True)
q_tch = torch.randn(n, dtype=torch.float64).detach().clone().requires_grad_(True)
A_tch = torch.randn(m, n, dtype=torch.float64).detach().clone().requires_grad_(True)
x0 = torch.randn(n, dtype=torch.float64)
b_tch = (A_tch.detach() @ x0).detach().clone().requires_grad_(True)
# ============================================================
# Solve using CVXPY directly (ground-truth x*)
# ============================================================
Q_sqrt_cp.value = Q_sqrt_tch.detach().cpu().numpy()
q_cp.value = q_tch.detach().cpu().numpy()
A_cp.value = A_tch.detach().cpu().numpy()
b_cp.value = b_tch.detach().cpu().numpy()
problem.solve(solver=cp.OSQP, eps_abs=1e-10, eps_rel=1e-10, verbose=False)
print(f"CVXPY solve status: {problem.status}")
print(f"Optimal value: {problem.value:.6f}")
print(f"Optimal x: {x.value}")
# ============================================================
# Compute true gradient wrt x (compare dloss/dx at solution)
# ============================================================
x_tch_true = torch.tensor(x.value, dtype=torch.float64, requires_grad=True)
loss_true = 0.5 * torch.sum((Q_sqrt_tch.detach() @ x_tch_true) ** 2) + (q_tch.detach() @ x_tch_true)
loss_true.backward()
grad_true = x_tch_true.grad.detach().clone()
# ============================================================
# Solve using FFOLayer
# ============================================================
ffo = FFOLayer(problem, parameters=[Q_sqrt_cp, q_cp, A_cp, b_cp], variables=[x])
x_tch_ffo, = ffo(Q_sqrt_tch, q_tch, A_tch, b_tch)
x_tch_ffo = x_tch_ffo.reshape(-1)
x_tch_ffo.retain_grad()
loss_ffo = 0.5 * torch.sum((Q_sqrt_tch @ x_tch_ffo) ** 2) + (q_tch @ x_tch_ffo)
loss_ffo.backward()
grad_ffo = x_tch_ffo.grad.detach().clone()
# ============================================================
# Compare
# ============================================================
print("\n--- Compare x ---")
print("x_tch_true:", x_tch_true.detach().cpu().numpy())
print("x_tch_ffo :", x_tch_ffo.detach().cpu().numpy())
print("||x_true - x_ffo||_2:", torch.norm(x_tch_true.detach() - x_tch_ffo.detach(), p=2).item())
print("\n--- Compare d(loss)/d(x) ---")
cos = torch.nn.functional.cosine_similarity(grad_true, grad_ffo, dim=0).item()
diff = torch.norm(grad_true - grad_ffo, p=2).item()
print("cosine similarity:", cos)
print("L2 gradient difference:", diff)
print("\nTest completed successfully!")
if __name__ == "__main__":
test_example()
|