Merge pull request Pyomo#3201 from jsiirola/mixed-standard-form

Add "mixed" standard form representation

pyomo/repn/plugins/standard_form.py

@@ -139,6 +139,15 @@ class LinearStandardFormCompiler(object):
             description='Add slack variables and return `min cTx s.t. Ax == b`',
         ),
     )
+    CONFIG.declare(
+        'mixed_form',
+        ConfigValue(
+            default=False,
+            domain=bool,
+            description='Return A in mixed form (the comparison operator is a '
+            'mix of <=, ==, and >=)',
+        ),
+    )
     CONFIG.declare(
         'show_section_timing',
         ConfigValue(
@@ -332,6 +341,9 @@ def write(self, model):
         # Tabulate constraints
         #
         slack_form = self.config.slack_form
+        mixed_form = self.config.mixed_form
+        if slack_form and mixed_form:
+            raise ValueError("cannot specify both slack_form and mixed_form")
         rows = []
         rhs = []
         con_data = []
@@ -372,7 +384,30 @@ def write(self, model):
                     f"model contains a trivially infeasible constraint, '{con.name}'"
                 )
 
-            if slack_form:
+            if mixed_form:
+                N = len(repn.linear)
+                _data = np.fromiter(repn.linear.values(), float, N)
+                _index = np.fromiter(map(var_order.__getitem__, repn.linear), float, N)
+                if ub == lb:
+                    rows.append(RowEntry(con, 0))
+                    rhs.append(ub - offset)
+                    con_data.append(_data)
+                    con_index.append(_index)
+                    con_index_ptr.append(con_index_ptr[-1] + N)
+                else:
+                    if ub is not None:
+                        rows.append(RowEntry(con, 1))
+                        rhs.append(ub - offset)
+                        con_data.append(_data)
+                        con_index.append(_index)
+                        con_index_ptr.append(con_index_ptr[-1] + N)
+                    if lb is not None:
+                        rows.append(RowEntry(con, -1))
+                        rhs.append(lb - offset)
+                        con_data.append(_data)
+                        con_index.append(_index)
+                        con_index_ptr.append(con_index_ptr[-1] + N)
+            elif slack_form:
                 _data = list(repn.linear.values())
                 _index = list(map(var_order.__getitem__, repn.linear))
                 if lb == ub:  # TODO: add tolerance?
@@ -437,24 +472,22 @@ def write(self, model):
         # at the index pointer list (an O(num_var) operation).
         c_ip = c.indptr
         A_ip = A.indptr
-        active_var_idx = list(
-            filter(
-                lambda i: A_ip[i] != A_ip[i + 1] or c_ip[i] != c_ip[i + 1],
-                range(len(columns)),
-            )
-        )
-        nCol = len(active_var_idx)
+        active_var_mask = (A_ip[1:] > A_ip[:-1]) | (c_ip[1:] > c_ip[:-1])
+
+        # Masks on NumPy arrays are very fast.  Build the reduced A
+        # indptr and then check if we actually have to manipulate the
+        # columns
+        augmented_mask = np.concatenate((active_var_mask, [True]))
+        reduced_A_indptr = A.indptr[augmented_mask]
+        nCol = len(reduced_A_indptr) - 1
         if nCol != len(columns):
-            # Note that the indptr can't just use range() because a var
-            # may only appear in the objectives or the constraints.
-            columns = list(map(columns.__getitem__, active_var_idx))
-            active_var_idx.append(c.indptr[-1])
+            columns = [v for k, v in zip(active_var_mask, columns) if k]
             c = scipy.sparse.csc_array(
-                (c.data, c.indices, c.indptr.take(active_var_idx)), [c.shape[0], nCol]
+                (c.data, c.indices, c.indptr[augmented_mask]), [c.shape[0], nCol]
             )
-            active_var_idx[-1] = A.indptr[-1]
+            # active_var_idx[-1] = len(columns)
             A = scipy.sparse.csc_array(
-                (A.data, A.indices, A.indptr.take(active_var_idx)), [A.shape[0], nCol]
+                (A.data, A.indices, reduced_A_indptr), [A.shape[0], nCol]
             )
 
         if self.config.nonnegative_vars:

pyomo/repn/tests/test_standard_form.py

@@ -42,6 +42,23 @@ def test_linear_model(self):
         self.assertTrue(np.all(repn.c == np.array([0, 0, 0])))
         self.assertTrue(np.all(repn.A == np.array([[-1, -2, 0], [0, 1, 4]])))
         self.assertTrue(np.all(repn.rhs == np.array([-3, 5])))
+        self.assertEqual(repn.rows, [(m.c, -1), (m.d, 1)])
+        self.assertEqual(repn.columns, [m.x, m.y[1], m.y[3]])
+
+    def test_almost_dense_linear_model(self):
+        m = pyo.ConcreteModel()
+        m.x = pyo.Var()
+        m.y = pyo.Var([1, 2, 3])
+        m.c = pyo.Constraint(expr=m.x + 2 * m.y[1] + 4 * m.y[3] >= 10)
+        m.d = pyo.Constraint(expr=5 * m.x + 6 * m.y[1] + 8 * m.y[3] <= 20)
+
+        repn = LinearStandardFormCompiler().write(m)
+
+        self.assertTrue(np.all(repn.c == np.array([0, 0, 0])))
+        self.assertTrue(np.all(repn.A == np.array([[-1, -2, -4], [5, 6, 8]])))
+        self.assertTrue(np.all(repn.rhs == np.array([-10, 20])))
+        self.assertEqual(repn.rows, [(m.c, -1), (m.d, 1)])
+        self.assertEqual(repn.columns, [m.x, m.y[1], m.y[3]])
 
     def test_linear_model_row_col_order(self):
         m = pyo.ConcreteModel()
@@ -57,6 +74,8 @@ def test_linear_model_row_col_order(self):
         self.assertTrue(np.all(repn.c == np.array([0, 0, 0])))
         self.assertTrue(np.all(repn.A == np.array([[4, 0, 1], [0, -1, -2]])))
         self.assertTrue(np.all(repn.rhs == np.array([5, -3])))
+        self.assertEqual(repn.rows, [(m.d, 1), (m.c, -1)])
+        self.assertEqual(repn.columns, [m.y[3], m.x, m.y[1]])
 
     def test_suffix_warning(self):
         m = pyo.ConcreteModel()
@@ -222,6 +241,28 @@ def test_alternative_forms(self):
         )
         self._verify_solution(soln, repn, True)
 
+        repn = LinearStandardFormCompiler().write(
+            m, mixed_form=True, column_order=col_order
+        )
+
+        self.assertEqual(
+            repn.rows, [(m.c, -1), (m.d, 1), (m.e, 1), (m.e, -1), (m.f, 0)]
+        )
+        self.assertEqual(list(map(str, repn.x)), ['x', 'y[0]', 'y[1]', 'y[3]'])
+        self.assertEqual(
+            list(v.bounds for v in repn.x), [(None, None), (0, 10), (-5, 10), (-5, -2)]
+        )
+        ref = np.array(
+            [[1, 0, 2, 0], [0, 0, 1, 4], [0, 1, 6, 0], [0, 1, 6, 0], [1, 1, 0, 0]]
+        )
+        self.assertTrue(np.all(repn.A == ref))
+        self.assertTrue(np.all(repn.b == np.array([3, 5, 6, -3, 8])))
+        self.assertTrue(np.all(repn.c == np.array([[-1, 0, -5, 0], [1, 0, 0, 15]])))
+        # Note that the mixed_form solution is a mix of inequality and
+        # equality constraints, so we cannot (easily) reuse the
+        # _verify_solutions helper (as in the above cases):
+        # self._verify_solution(soln, repn, False)
+
         repn = LinearStandardFormCompiler().write(
             m, slack_form=True, nonnegative_vars=True, column_order=col_order
         )