Off the shelf Cooper

cooper/formulation/augmented_lagrangian.py

@@ -25,7 +25,7 @@ class AugmentedLagrangianFormulation(LagrangianFormulation):
 
     def __init__(
         self,
-        cmp: ConstrainedMinimizationProblem,
+        cmp: Optional[ConstrainedMinimizationProblem] = None,
         ineq_init: Optional[torch.Tensor] = None,
         eq_init: Optional[torch.Tensor] = None,
     ):
@@ -104,7 +104,7 @@ def composite_objective(
         closure: Callable[..., CMPState] = None,
         *closure_args,
         pre_computed_state: Optional[CMPState] = None,
-        write_state: bool = True,
+        write_state: Optional[bool] = True,
         **closure_kwargs
     ) -> torch.Tensor:
         """
@@ -146,8 +146,8 @@ def composite_objective(
         else:
             cmp_state = closure(*closure_args, **closure_kwargs)
 
-        if write_state:
-            self.cmp.state = cmp_state
+        if write_state and self.cmp is not None:
+            self.write_cmp_state(cmp_state)
 
         # Extract values from ProblemState object
         loss = cmp_state.loss

cooper/formulation/formulation.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Optional
 
 import torch
 
@@ -13,8 +13,8 @@
 class Formulation(abc.ABC):
     """Base class for formulations of CMPs."""
 
-    def __init__(self):
-        self.cmp = None
+    def __init__(self, cmp: Optional[ConstrainedMinimizationProblem] = None):
+        self.cmp = cmp
 
     @abc.abstractmethod
     def state_dict(self) -> Dict[str, Any]:
@@ -67,6 +67,18 @@ def custom_backward(self, *args, **kwargs):
         """
         self._populate_gradients(*args, **kwargs)
 
+    def write_cmp_state(self, cmp_state: CMPState):
+        """Provided that the formulation is linked to a
+        `ConstrainedMinimizationProblem`, writes a CMPState to the CMP."""
+
+        if self.cmp is None:
+            raise RuntimeError(
+                """Cannot write state to a formulation which is not linked to a
+                ConstrainedMinimizationProblem"""
+            )
+
+        self.cmp.state = cmp_state
+
 
 class UnconstrainedFormulation(Formulation):
     """
@@ -77,7 +89,7 @@ class UnconstrainedFormulation(Formulation):
             to solve and which gives rise to the Lagrangian.
     """
 
-    def __init__(self, cmp: ConstrainedMinimizationProblem):
+    def __init__(self, cmp: Optional[ConstrainedMinimizationProblem] = None):
         """Construct new `UnconstrainedFormulation`"""
 
         self.cmp = cmp
@@ -119,7 +131,7 @@ def composite_objective(
         self,
         closure: Callable[..., CMPState],
         *closure_args,
-        write_state: bool = True,
+        write_state: Optional[bool] = True,
         **closure_kwargs
     ) -> torch.Tensor:
         """
@@ -138,8 +150,9 @@ def composite_objective(
         """
 
         cmp_state = closure(*closure_args, **closure_kwargs)
-        if write_state:
-            self.cmp.state = cmp_state
+
+        if write_state and self.cmp is not None:
+            self.write_cmp_state(cmp_state)
 
         return cmp_state.loss
 

cooper/formulation/lagrangian.py

@@ -28,7 +28,7 @@ class BaseLagrangianFormulation(Formulation, metaclass=abc.ABCMeta):
 
     def __init__(
         self,
-        cmp: ConstrainedMinimizationProblem,
+        cmp: Optional[ConstrainedMinimizationProblem] = None,
         ineq_init: Optional[torch.Tensor] = None,
         eq_init: Optional[torch.Tensor] = None,
     ):
@@ -240,7 +240,7 @@ def composite_objective(
         closure: Callable[..., CMPState] = None,
         *closure_args,
         pre_computed_state: Optional[CMPState] = None,
-        write_state: bool = True,
+        write_state: Optional[bool] = True,
         **closure_kwargs
     ) -> torch.Tensor:
         """
@@ -281,8 +281,8 @@ def composite_objective(
         else:
             cmp_state = closure(*closure_args, **closure_kwargs)
 
-        if write_state:
-            self.cmp.state = cmp_state
+        if write_state and self.cmp is not None:
+            self.write_cmp_state(cmp_state)
 
         # Extract values from ProblemState object
         loss = cmp_state.loss

cooper/optim/constrained_optimizers/alternating_optimizer.py

@@ -23,7 +23,6 @@ def __init__(
         dual_restarts: bool = False,
     ):
         self.formulation = formulation
-        self.cmp = self.formulation.cmp
 
         if isinstance(primal_optimizers, torch.optim.Optimizer):
             self.primal_optimizers = [primal_optimizers]

cooper/optim/constrained_optimizers/extrapolation_optimizer.py

@@ -56,7 +56,6 @@ def __init__(
         dual_restarts: bool = False,
     ):
         self.formulation = formulation
-        self.cmp = self.formulation.cmp
 
         if isinstance(primal_optimizers, ExtragradientOptimizer):
             self.primal_optimizers = [primal_optimizers]

cooper/optim/constrained_optimizers/simultaneous_optimizer.py

@@ -55,7 +55,6 @@ def __init__(
         dual_restarts: bool = False,
     ):
         self.formulation = formulation
-        self.cmp = self.formulation.cmp
 
         if isinstance(primal_optimizers, torch.optim.Optimizer):
             self.primal_optimizers = [primal_optimizers]

cooper/optim/unconstrained_optimizer.py

@@ -36,7 +36,6 @@ def __init__(
             )
 
         self.formulation = formulation
-        self.cmp = self.formulation.cmp
 
         if isinstance(primal_optimizers, torch.optim.Optimizer):
             self.primal_optimizers = [primal_optimizers]

tests/test_simplest_pipeline.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python
+
+"""Tests for Constrained Optimizer class. This test already verifies that the
+code behaves as expected for an unconstrained setting."""
+
+import pytest
+import torch
+
+import cooper
+
+
+@pytest.fixture()
+def params():
+    return torch.nn.Parameter(torch.tensor([0.0, -1.0]))
+
+
+@pytest.fixture()
+def formulation():
+    return cooper.LagrangianFormulation()
+
+
+@pytest.fixture()
+def constrained_optimizer(params, formulation):
+    primal_optim = torch.optim.SGD([params], lr=1e-2, momentum=0.3)
+    dual_optim = cooper.optim.partial_optimizer(torch.optim.SGD, lr=1e-2)
+
+    return cooper.SimultaneousConstrainedOptimizer(
+        formulation, primal_optim, dual_optim, dual_restarts=True
+    )
+
+
+def loss_fn(params):
+    param_x, param_y = params
+
+    return param_x**2 + 2 * param_y**2
+
+
+def defect_fn(params):
+
+    param_x, param_y = params
+
+    # Two inequality constraints
+    defect = torch.stack(
+        [
+            -param_x - param_y + 1.0,  # x + y \ge 1
+            param_x**2 + param_y - 1.0,  # x**2 + y \le 1.0
+        ]
+    )
+
+    return defect
+
+
+def test_simplest_pipeline(params, formulation, constrained_optimizer):
+
+    for step_id in range(1500):
+        constrained_optimizer.zero_grad()
+
+        loss = loss_fn(params)
+        defect = defect_fn(params)
+
+        # Create a CMPState object to hold the loss and defect values
+        cmp_state = cooper.CMPState(loss=loss, ineq_defect=defect)
+
+        lagrangian = formulation.composite_objective(pre_computed_state=cmp_state)
+        formulation.custom_backward(lagrangian)
+
+        constrained_optimizer.step()
+
+    assert torch.allclose(params[0], torch.tensor(2.0 / 3.0))
+    assert torch.allclose(params[1], torch.tensor(1.0 / 3.0))