54 multiplier models

cooper/formulation/__init__.py

@@ -1,3 +1,4 @@
 from .augmented_lagrangian import AugmentedLagrangianFormulation
 from .formulation import Formulation, UnconstrainedFormulation
 from .lagrangian import LagrangianFormulation, ProxyLagrangianFormulation
+from .lagrangian_model import LagrangianModelFormulation

cooper/formulation/augmented_lagrangian.py

@@ -90,7 +90,7 @@ def weighted_violation(
             # to update the value of the multipliers by lazily filling the
             # multiplier gradients in `backward`
 
-            # TODO (JGP): Verify that call to backward is general enough for
+            # TODO (gallego-posada): Verify that call to backward is general enough for
             # Lagrange Multiplier models
             violation_for_update = torch.sum(multipliers * defect.detach())
             self.update_accumulated_violation(update=violation_for_update)

cooper/formulation/formulation.py

@@ -1,10 +1,13 @@
 import abc
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Union
 
 import torch
 
 from cooper.problem import CMPState, ConstrainedMinimizationProblem
 
+# from .lagrangian_model import CMPModelState
+
+
 # Formulation, and some other classes below, are heavily inspired by the design
 # of the TensorFlow Constrained Optimization (TFCO) library :
 # https://github.com/google-research/tensorflow_constrained_optimization
@@ -34,6 +37,11 @@ def state(self):
         """Returns the internal state of formulation (e.g. multipliers)."""
         pass
 
+    @abc.abstractmethod
+    def flip_dual_gradients(self):
+        """Flips the sign of the dual gradients."""
+        pass
+
     @property
     @abc.abstractmethod
     def is_state_created(self):
@@ -59,7 +67,8 @@ def backward(self, *args, **kwargs):
         formulation."""
         pass
 
-    def write_cmp_state(self, cmp_state: CMPState):
+    # TODO(daoterog): fix circular import type hint can be correct
+    def write_cmp_state(self, cmp_state: CMPState):  # Union[CMPState, CMPModelState]):
         """Provided that the formulation is linked to a
         `ConstrainedMinimizationProblem`, writes a CMPState to the CMP."""
 
@@ -119,6 +128,11 @@ def load_state_dict(self, state_dict: dict):
         """
         pass
 
+    def flip_dual_gradients(self):
+        """Flips the sign of the dual gradients. This is a no-op for
+        unconstrained formulations."""
+        pass
+
     def compute_lagrangian(
         self,
         closure: Callable[..., CMPState],

cooper/formulation/lagrangian.py

@@ -42,7 +42,7 @@ def __init__(
         # Store user-provided initializations for dual variables
         self.ineq_init = ineq_init
         self.eq_init = eq_init
-
+        # TODO(gallego-posada): comment what is the meaning of this object
         self.accumulated_violation_dot_prod: torch.Tensor = None
 
     @property
@@ -193,11 +193,11 @@ def state_dict(self) -> Dict[str, Any]:
         """
         Generates the state dictionary for a Lagrangian formulation.
         """
-
+        # TODO(gallego-posada): fix in next PR
         state_dict = {
             "ineq_multipliers": self.ineq_multipliers,
             "eq_multipliers": self.eq_multipliers,
-            "accumulated_violation_dot_prod": self.accumulated_violation_dot_prod,
+            # "accumulated_violation_dot_prod": self.accumulated_violation_dot_prod,
         }
         return state_dict
 
@@ -220,6 +220,14 @@ def load_state_dict(self, state_dict: Dict[str, Any]):
             ), "LagrangianFormulation received unknown key: {}".format(key)
             setattr(self, key, val)
 
+    def flip_dual_gradients(self):
+        """
+        Flip the sign of the gradients of the dual variables.
+        """
+        for multiplier in self.state():
+            if multiplier is not None:
+                multiplier.grad.mul_(-1.0)
+
 
 class LagrangianFormulation(BaseLagrangianFormulation):
     """
@@ -353,7 +361,7 @@ def weighted_violation(
             # on `accumulated_violation_dot_prod`. This enables easy
             # extensibility to multiplier classes beyond DenseMultiplier.
 
-            # TODO (JGP): Verify that call to backward is general enough for
+            # TODO (gallego-posada): Verify that call to backward is general enough for
             # Lagrange Multiplier models
             violation_for_update = torch.sum(multipliers * defect.detach())
             self.update_accumulated_violation(update=violation_for_update)