From 4f62c447fbdd1be1ec466c1eb366ac59f9122688 Mon Sep 17 00:00:00 2001
From: Alejandro Montanez <alejomonbar@gmail.com>
Date: Thu, 29 Jun 2023 14:15:19 +0200
Subject: [PATCH] Adding the unbalanced penalization method

Following recent thread about the unbalanced penalization equation to encode inequality constraints. #1339
---
 dimod/binary/binary_quadratic_model.py | 71 +++++++++++++++-----------
 1 file changed, 41 insertions(+), 30 deletions(-)

diff --git a/dimod/binary/binary_quadratic_model.py b/dimod/binary/binary_quadratic_model.py
index 628ba674b..8cadb322d 100644
--- a/dimod/binary/binary_quadratic_model.py
+++ b/dimod/binary/binary_quadratic_model.py
@@ -703,12 +703,13 @@ def add_linear_equality_constraint(
 
     def add_linear_inequality_constraint(
             self, terms: Iterable[Tuple[Variable, int]],
-            lagrange_multiplier: Bias,
+            lagrange_multiplier: Union[Bias, Iterable],
             label: str,
             constant: int = 0,
             lb: int = np.iinfo(np.int64).min,
             ub: int = 0,
-            cross_zero: bool = False
+            cross_zero: bool = False,
+            penalization_method: str = "slack"
     ) -> Iterable[Tuple[Variable, int]]:
         """Add a linear inequality constraint as a quadratic objective.
 
@@ -739,6 +740,10 @@ def add_linear_inequality_constraint(
                 Upper bound for the constraint.
             cross_zero:
                 When True, adds zero to the domain of constraint.
+            penalization_method:
+                Whether to use slack variables or the unbalanced penalization method [1].
+                ("slack", "unbalanced")
+                [1] https://arxiv.org/abs/2211.13914
 
         Returns:
             slack_terms:  Values of :math:`\sum_{i} b_{i} slack_{i}` as an
@@ -772,35 +777,41 @@ def add_linear_inequality_constraint(
             raise ValueError(
                 f'The given constraint ({label}) is infeasible with any value'
                 ' for state variables.')
-
-        slack_upper_bound = int(ub_c - lb_c)
-        if slack_upper_bound == 0:
-            self.add_linear_equality_constraint(terms, lagrange_multiplier, -ub_c)
+        if penalization_method == "slack":
+            slack_upper_bound = int(ub_c - lb_c)
+            if slack_upper_bound == 0:
+                self.add_linear_equality_constraint(terms, lagrange_multiplier, -ub_c)
+                return []
+            else:
+                slack_terms = []
+                zero_constraint = False
+                if cross_zero:
+                    if lb_c > 0 or ub_c < 0:
+                        if ub_c-slack_upper_bound > 0:
+                            zero_constraint = True
+    
+                num_slack = int(np.floor(np.log2(slack_upper_bound)))
+                slack_coefficients = [2 ** j for j in range(num_slack)]
+                if slack_upper_bound - 2 ** num_slack >= 0:
+                    slack_coefficients.append(slack_upper_bound - 2 ** num_slack + 1)
+    
+                for j, s in enumerate(slack_coefficients):
+                    sv = self.add_variable(f'slack_{label}_{j}')
+                    slack_terms.append((sv, s))
+    
+                if zero_constraint:
+                    sv = self.add_variable(f'slack_{label}_{num_slack + 1}')
+                    slack_terms.append((sv, ub_c - slack_upper_bound))
+    
+            self.add_linear_equality_constraint(terms,lagrange_multiplier[0], -ub_c)
+
+            return slack_terms
+
+        elif penalization_method == "unbalanced":
+            self.add_linear_equality_constraint(terms,lagrange_multiplier[0], -ub_c)
+            for v, bias in terms:
+                self.add_linear(v, -lagrange_multiplier[1] * bias)
             return []
-        else:
-            slack_terms = []
-            zero_constraint = False
-            if cross_zero:
-                if lb_c > 0 or ub_c < 0:
-                    if ub_c-slack_upper_bound > 0:
-                        zero_constraint = True
-
-            num_slack = int(np.floor(np.log2(slack_upper_bound)))
-            slack_coefficients = [2 ** j for j in range(num_slack)]
-            if slack_upper_bound - 2 ** num_slack >= 0:
-                slack_coefficients.append(slack_upper_bound - 2 ** num_slack + 1)
-
-            for j, s in enumerate(slack_coefficients):
-                sv = self.add_variable(f'slack_{label}_{j}')
-                slack_terms.append((sv, s))
-
-            if zero_constraint:
-                sv = self.add_variable(f'slack_{label}_{num_slack + 1}')
-                slack_terms.append((sv, ub_c - slack_upper_bound))
-
-        self.add_linear_equality_constraint(terms + slack_terms,
-                                            lagrange_multiplier, -ub_c)
-        return slack_terms
 
     def add_linear_from_array(self, linear: Sequence):
         """Add linear biases from an array-like to a binary quadratic model.