Elastic sub problem and fixes

.gitignore

@@ -108,3 +108,6 @@ venv.bak/
 .mypy_cache/
 /mypy.ini
 .venv/
+
+# Editors
+.idea/

README.md

@@ -12,11 +12,12 @@ Let's solve an unconstrained SuperSudoku puzzle. A SuperSudoku is a Sudoku
 with the additional requirement that all digits having box coordinates (x, y)
 be distinct for all (x, y).
 
-```import pulp as pp
-from lparray import lparray
+```
+import pulp as pp
+from pulp_lparray import lparray
 
 #                       name      R, C, r, c, n   lb ub type
-X = lparray.create_anon("Board", (3, 3, 3, 3, 9), 0, 1, pp.LpBinary)
+X = lparray.create_anon("Board", (3, 3, 3, 3, 9), cat=pp.LpBinary)
 prob = pp.LpProblem("SuperSudoku", pp.LpMinimize)
 (X.sum(axis=-1) == 1).constrain(prob, "OneDigitPerCell")
 (X.sum(axis=(1, 3)) == 1).constrain(prob, "MaxOnePerRow")

pulp_lparray/lparray.py

@@ -11,7 +11,9 @@
     Protocol,
     TypeVar,
     Union,
+    Tuple
 )
+from functools import partial
 
 import numpy as np  # type: ignore
 from numpy import ndarray
@@ -154,7 +156,7 @@ def recursive_worker(
                     )
 
         arr = np.zeros(
-            tuple(len(ixset) for ixset in index_sets), dtype=np.object
+            tuple(len(ixset) for ixset in index_sets), dtype=object
         )
         recursive_worker(name, arr, index_sets)
 
@@ -223,6 +225,22 @@ def values(self: lparray[LPV]) -> np.ndarray:
 
         return np.vectorize(value)(self).view(np.ndarray)
 
+    @staticmethod
+    def _check_is_constraint(const):
+        if not isinstance(const, LpConstraint):
+            raise TypeError(
+                "Attempting to constrain problem with "
+                f"non-constraint {const}"
+            )
+
+    def _check_is_eq_constraint(self, const):
+        self._check_is_constraint(const)
+        if const.sense != 0:
+            raise TypeError(
+                "Requires an equality constraint (sense=0) "
+                f"sense: {const.sense}"
+            )            
+
     def constrain(self, prob: LpProblem, name: str) -> None:
         """
         Applies the constraints contained in self to the problem.
@@ -268,6 +286,78 @@ def recursive_worker(
 
         recursive_worker(prob, self, name)
 
+    @staticmethod
+    def _convert_to_elastic_subproblem(constraint, proportionFreeBoundList, penalty):
+        elastic_constraint = constraint.makeElasticSubProblem(
+            proportionFreeBoundList=proportionFreeBoundList,
+            penalty=penalty
+        )
+        return elastic_constraint
+
+    def elastically_constrain(
+            self,
+            prob: LpProblem,
+            name: str,
+            proportion_unpenalised_bounds: Tuple[float, float],
+            penalty: float
+    ) -> None:
+        """
+        Applies the constraints contained in self to the problem
+        as elastic constraints. Which are sub problems added to 
+        the objective where an additional symetric penalty term
+        for distance from the constraint target. 
+
+        Preconditions:
+            all entries of self are `LpConstraints` with sense LpConstraintEQ (0).
+            Main objectives are already added to the problem (using +=)
+
+        Arguments:
+            prob: `LpProblem` which to apply constraints to.
+            name: base name to use for the applied constraints.
+            proportion_unpenalised_bounds: proportional distance either 
+                side of the equality constraint within which there 
+                is no penalty.
+            penalty
+
+        Usage:
+            (array_of_lp_vars == array_of_numbers).elasticly_constrain(
+                problem, "SomeElasticConstraint")
+        """
+        elastify = partial(
+            self._convert_to_elastic_subproblem,
+            proportionFreeBoundList=list(proportion_unpenalised_bounds),
+            penalty=penalty
+        )
+
+        if self.ndim == 0:
+            constraint = self.item()
+            # ignore
+            self._check_is_eq_constraint(constraint)
+            constraint.name = name
+            prob.extend(elastify(constraint))
+            return
+
+        if name and self.ndim == 1:
+            name += "("
+
+        def recursive_worker(
+            r_prob: LpProblem, plane: np.ndarray, r_name: str
+        ) -> None:
+            if plane.ndim == 1:
+                close_paren = r_name and (")" if "(" in r_name else "")
+                for cx, constraint in enumerate(plane):
+                    self._check_is_eq_constraint(constraint)
+                    constraint.name = r_name and f"{r_name}{cx}{close_paren}"
+                    r_prob.extend(elastify(constraint))
+            else:
+                open_paren = r_name and ("(" if "(" not in r_name else "")
+                for px, subplane in enumerate(plane):
+                    subname = r_name and f"{r_name}{open_paren}{px},"
+                    recursive_worker(r_prob, subplane, subname)
+
+        recursive_worker(prob, self, name)
+
+
     def abs_decompose(
         self: lparray[LPV],
         prob: LpProblem,
@@ -277,9 +367,9 @@ def abs_decompose(
         **kwargs: Any,
     ) -> tuple[lparray[LpVariable], lparray[LpVariable]]:
         """
-        Generates two arrays, xp and xm, that sum to |self|, with the following
+        Generates two arrays, xp and xm, that difference to self, with the following
         properties:
-
+            xp - xm == self
             xp >= 0
             xm >= 0
             xp == 0 XOR xm == 0

tests/test_lparray.py

@@ -1,7 +1,9 @@
 import numpy as np
 import numpy.random as npr
 import pulp as pp
+import pytest
 from pulp import LpBinary, LpMaximize
+from numpy.testing import assert_allclose
 
 from pulp_lparray import lparray
 
@@ -37,14 +39,35 @@ def check_super_sudoku(arr: np.ndarray) -> bool:
     (X.sum(axis=(1, 3)) == 1).constrain(prob, "MaxOnePerRow")
     (X.sum(axis=(0, 2)) == 1).constrain(prob, "MaxOnePerCol")
     (X.sum(axis=(2, 3)) == 1).constrain(prob, "MaxOnePerBox")
-    (X.sum(axis=(0, 1)) == 1).constrain(prob, "MaxOnePerDust")
+    (X.sum(axis=(0, 1)) == 1).constrain(prob, "MaxOnePerXY")
     prob.solve()
     board = X.values.argmax(axis=-1)
     print(board)
 
     assert check_super_sudoku(X.values)
 
 
+def test_elastically_constrain() -> None:
+    prob = pp.LpProblem("elastically_constrain", pp.LpMinimize)
+    x = lparray.create_anon(
+        "arr", shape=(5,), cat=pp.LpInteger, lowBound=0, upBound=5
+    )
+    target_vals = np.array([1, 2.1, 3, 4, 5])
+    prob += x.sumit()
+    constraints = (x == target_vals)
+    # add 2 to the objective for each 1 away from the constraint
+    constraints.elastically_constrain(
+        prob,
+        "elast",
+        (0.0, 0.0),
+        2
+    )
+    prob.solve()
+    assert_allclose(x.values, np.array([1, 2, 3, 4, 5]))
+    # Sum of X is 15, = 0.2 from the soft constraint.
+    assert prob.objective.value() == 15.2
+
+
 # noinspection PyArgumentList
 def test_logical_clip() -> None: