Speed up subproblem

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
   "numpy",
   "torch",
   "cvxpy-base",
-  "clarabel",
+  "osqp",
 ]
 
 [project.urls]

scripts/example_residual.py

@@ -61,7 +61,7 @@
 gI = [ObjectiveOrConstraint(const, dim_out=1)]
 gE = []
 
-options = {"num_points_obj": 5, "num_points_gI": 5}
+options = {"num_points_obj": 5, "num_points_gI": 5, "qp_solver": "osqp"}
 
 problem = SQPGS(f, gI, gE, x0=None, tol=1e-10, max_iter=500, options=options, verbose=True)
 

src/ncopt/sqpgs/defaults.py

@@ -32,6 +32,7 @@ class Dotdict(dict):
     "num_points_obj": 2,
     "num_points_gI": 3,
     "num_points_gE": 4,
+    "qp_solver": "osqp",
 }
 
 DEFAULT_ARG = Dotdict(_defaults)

src/ncopt/sqpgs/main.py

@@ -141,7 +141,7 @@ def solve(self):
         pE = np.repeat(pE_, self.dimE)
         ###############################################################
 
-        self.SP = SubproblemSQPGS(self.dim, p0, pI, pE, self.assert_tol)
+        self.SP = SubproblemSQPGS(self.dim, p0, pI, pE, self.assert_tol, self.options["qp_solver"])
 
         E_k = np.inf  # for stopping criterion
         x_kmin1 = None  # last iterate
@@ -158,9 +158,8 @@ def solve(self):
         timings = {
             "total": [],
             "sample_and_grad": [],
+            "subproblem": [],
             "step": [],
-            "sp_update": [],
-            "sp_solve": [],
             "other": [],
         }
         metrics = {
@@ -230,17 +229,13 @@ def solve(self):
             else:
                 gE_k = np.array([])
 
-            t01 = time.perf_counter()
             ##############################################
             # Subproblem solve
             ##############################################
+            t1 = time.perf_counter()
+            self.SP.solve(np.linalg.cholesky(H), rho, D_f, D_gI, D_gE, f_k, gI_k, gE_k)
+            t2 = time.perf_counter()
 
-            self.SP.solve(H, rho, D_f, D_gI, D_gE, f_k, gI_k, gE_k)
-
-            timings["sp_update"].append(self.SP.setup_time)
-            timings["sp_solve"].append(self.SP.solve_time)
-
-            t02 = time.perf_counter()
             d_k = self.SP.d.value.copy()
             # compute g_k from paper
             g_k = (
@@ -284,7 +279,7 @@ def solve(self):
             ##############################################
             # Step
             ##############################################
-            t04 = time.perf_counter()
+            t3 = time.perf_counter()
             do_step = delta_q > nu * eps**2  # Flag whether step is taken or not
             if do_step:
                 alpha = 1.0
@@ -346,11 +341,12 @@ def solve(self):
             if self.store_history:
                 x_hist.append(self.x_k)
 
-            t1 = time.perf_counter()
-            timings["total"].append(t1 - t0)
-            timings["sample_and_grad"].append(t01 - t0)
-            timings["other"].append(t04 - t02)
-            timings["step"].append(t1 - t04)
+            t4 = time.perf_counter()
+            timings["total"].append(t4 - t0)
+            timings["sample_and_grad"].append(t1 - t0)
+            timings["subproblem"].append(t2 - t1)
+            timings["other"].append(t3 - t2)
+            timings["step"].append(t4 - t3)
 
         ##############################################
         # End of loop
@@ -511,10 +507,18 @@ def stop_criterion(gI, gE, g_k, SP, gI_k, gE_k, V_gI, V_gE, nI_, nE_):
 
 # %%
 
+CVXPY_SOLVER_DICT = {"osqp": cp.OSQP, "clarabel": cp.CLARABEL}
+
 
 class SubproblemSQPGS:
     def __init__(
-        self, dim: int, p0: np.ndarray, pI: np.ndarray, pE: np.ndarray, assert_tol: float
+        self,
+        dim: int,
+        p0: np.ndarray,
+        pI: np.ndarray,
+        pE: np.ndarray,
+        assert_tol: float,
+        solver: str = DEFAULT_OPTION.qp_solver,
     ) -> None:
         """
         dim : solution space dimension
@@ -532,6 +536,7 @@ def __init__(
 
         self.d = cp.Variable(self.dim)
         self._problem = None
+        self._qp_solver = CVXPY_SOLVER_DICT.get(solver, cp.OSQP)
 
     @property
     def nI(self) -> int:
@@ -572,7 +577,7 @@ def solve_time(self) -> float:
 
     def solve(
         self,
-        H: np.ndarray,
+        L: np.ndarray,
         rho: float,
         D_f: np.ndarray,
         D_gI: List[np.ndarray],
@@ -586,8 +591,8 @@ def solve(
 
         Parameters
 
-        H : array
-            Hessian approximation
+        L : array
+            Cholesky factor of Hessian approximation
         rho : float
             parameter
         D_f : array
@@ -625,7 +630,7 @@ def solve(
         if self.has_eq_constraints:
             r_E = cp.Variable(gE_k.size, nonneg=True)
 
-        objective = rho * z + (1 / 2) * cp.quad_form(d, H)
+        objective = rho * z + (1 / 2) * cp.sum_squares(L.T @ d)
 
         obj_constraint = f_k + D_f @ d <= z
         constraints = [obj_constraint]
@@ -642,7 +647,7 @@ def solve(
             objective = objective + cp.sum(r_E)
 
         problem = cp.Problem(cp.Minimize(objective), constraints)
-        problem.solve(solver=cp.CLARABEL)
+        problem.solve(solver=self._qp_solver, verbose=False)
 
         assert problem.status in {cp.OPTIMAL, cp.OPTIMAL_INACCURATE}
         self._problem = problem

tests/test_rosenbrock.py

@@ -15,11 +15,9 @@
     params=(torch.diag(torch.tensor([torch.sqrt(torch.tensor(2.0)), 2.0])), -torch.ones(2))
 )
 
-np.random.seed(12345)
-torch.manual_seed(12345)
-
 
 def test_rosenbrock_from_zero():
+    torch.manual_seed(1)
     gI = [ObjectiveOrConstraint(g, dim_out=1)]
     gE = []
     xstar = np.array([1 / np.sqrt(2), 0.5])
@@ -29,6 +27,7 @@ def test_rosenbrock_from_zero():
 
 
 def test_rosenbrock_from_rand():
+    torch.manual_seed(1)
     gI = [ObjectiveOrConstraint(g, dim_out=1)]
     gE = []
     xstar = np.array([1 / np.sqrt(2), 0.5])
@@ -40,6 +39,7 @@ def test_rosenbrock_from_rand():
 
 
 def test_rosenbrock_with_eq():
+    torch.manual_seed(12)
     g1 = ObjectiveOrConstraint(torch.nn.Linear(2, 2), dim_out=2)
     g1.model.weight.data = torch.eye(2)
     g1.model.bias.data = -torch.ones(2)

tests/test_subproblem.py

@@ -20,7 +20,7 @@ def test_subproblem_ineq(subproblem_ineq: SubproblemSQPGS):
     D_f = np.array([[-2.0, 1.0], [-2.04236205, -1.0], [-1.92172864, -1.0]])
     D_gI = [np.array([[0.0, 2.0], [0.0, 2.0], [1.41421356, 0.0], [1.41421356, 0.0]])]
     subproblem_ineq.solve(
-        H=np.eye(2, dtype=float),
+        L=np.eye(2, dtype=float),
         rho=0.1,
         D_f=D_f,
         D_gI=D_gI,
@@ -50,7 +50,7 @@ def test_subproblem_eq(subproblem_eq: SubproblemSQPGS):
         np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0], [0.0, 1.0], [0.0, 1.0]]),
     ]
     subproblem_eq.solve(
-        H=np.eye(2, dtype=float),
+        L=np.eye(2, dtype=float),
         rho=0.1,
         D_f=np.array([-2.0, 1.0]),
         D_gI=[],