Separate Linear Terms in Nonlinear to Piecewise Linear Transformation

pyomo/contrib/piecewise/tests/test_nonlinear_to_pwl.py

@@ -9,6 +9,7 @@
 
 from io import StringIO
 import logging
+import math
 
 from pyomo.common.dependencies import attempt_import, scipy_available, numpy_available
 from pyomo.common.log import LoggingIntercept
@@ -33,6 +34,7 @@
     Constraint,
     Integers,
     TransformationFactory,
+    exp,
     log,
     Objective,
     Reals,
@@ -292,7 +294,7 @@ def test_error_for_non_separable_exceeding_max_dimension(self):
     def test_do_not_additively_decompose_below_min_dimension(self):
         m = ConcreteModel()
         m.x = Var([0, 1, 2, 3, 4], bounds=(-4, 5))
-        m.c = Constraint(expr=m.x[0] * m.x[1] + m.x[3] <= 4)
+        m.c = Constraint(expr=m.x[0] * m.x[1] + m.x[3] ** 3 <= 4)
 
         n_to_pwl = TransformationFactory('contrib.piecewise.nonlinear_to_pwl')
         n_to_pwl.apply_to(
@@ -345,7 +347,7 @@ def test_uniform_sampling_discrete_vars(self):
         m = ConcreteModel()
         m.x = Var(['rocky', 'bullwinkle'], domain=Binary)
         m.y = Var(domain=Integers, bounds=(0, 5))
-        m.c = Constraint(expr=m.x['rocky'] * m.x['bullwinkle'] + m.y <= 4)
+        m.c = Constraint(expr=m.x['rocky'] * m.x['bullwinkle'] + m.y**2 <= 4)
 
         n_to_pwl = TransformationFactory('contrib.piecewise.nonlinear_to_pwl')
         output = StringIO()
@@ -378,7 +380,7 @@ def test_random_sampling_discrete_vars(self):
         m = ConcreteModel()
         m.x = Var(['rocky', 'bullwinkle'], domain=Binary)
         m.y = Var(domain=Integers, bounds=(0, 5))
-        m.c = Constraint(expr=m.x['rocky'] * m.x['bullwinkle'] + m.y <= 4)
+        m.c = Constraint(expr=m.x['rocky'] * m.x['bullwinkle'] + m.y**2 <= 4)
 
         n_to_pwl = TransformationFactory('contrib.piecewise.nonlinear_to_pwl')
         output = StringIO()
@@ -405,6 +407,34 @@ def test_random_sampling_discrete_vars(self):
                 for z in [0, 1, 5]:
                     self.assertIn((x, y, z), points)
 
+    @unittest.skipUnless(numpy_available, "Numpy is not available")
+    @unittest.skipUnless(scipy_available, "Scipy is not available")
+    def test_do_not_pwl_linear_part(self):
+        m = ConcreteModel()
+        m.x = Var(bounds=(-3, 4))
+        m.y = Var(initialize=7)
+        m.c = Constraint(expr=m.y >= exp(m.x) + m.x**3)
+
+        # we want to make sure m.y does not show up in the PWL function
+        # even if additively_decompose is False
+        n_to_pwl = TransformationFactory('contrib.piecewise.nonlinear_to_pwl')
+        num_points = 5
+        n_to_pwl.apply_to(
+            m,
+            num_points=num_points,
+            additively_decompose=False,
+            domain_partitioning_method=DomainPartitioningMethod.UNIFORM_GRID,
+        )
+
+        transformed_c = n_to_pwl.get_transformed_component(m.c)
+        self.assertIsInstance(transformed_c.body, SumExpression)
+        assertExpressionsEqual(self, transformed_c.body.args[0], -m.y)
+        pwlf = transformed_c.body.args[1].expr.pw_linear_function
+        for tup in pwlf._points:
+            self.assertEqual(len(tup), 1)
+            x = tup[0]
+            self.assertAlmostEqual(pwlf(x), math.exp(x) + x**3)
+
 
 class TestNonlinearToPWL_2D(unittest.TestCase):
     def make_paraboloid_model(self):
@@ -716,17 +746,17 @@ def test_transform_and_solve_additively_decomposes_model(self):
         # two terms
         self.assertIsInstance(new_obj.expr, SumExpression)
         self.assertEqual(len(new_obj.expr.args), 2)
-        first = new_obj.expr.args[0]
-        pwlf = first.expr.pw_linear_function
+        second = new_obj.expr.args[1]
+        pwlf = second.expr.pw_linear_function
         all_pwlf = list(
             xm.component_data_objects(PiecewiseLinearFunction, descend_into=True)
         )
         self.assertEqual(len(all_pwlf), 1)
         # It is on the active tree.
         self.assertIs(pwlf, all_pwlf[0])
 
-        second = new_obj.expr.args[1]
-        assertExpressionsEqual(self, second, 5.04 * xm.x1)
+        first = new_obj.expr.args[0]
+        assertExpressionsEqual(self, first, 5.04 * xm.x1)
 
         objs = n_to_pwl.get_transformed_nonlinear_objectives(xm)
         self.assertEqual(len(objs), 0)

pyomo/contrib/piecewise/transform/nonlinear_to_pwl.py

@@ -40,14 +40,14 @@
 from pyomo.common.dependencies import numpy as np, packaging
 from pyomo.common.enums import IntEnum
 from pyomo.common.modeling import unique_component_name
-from pyomo.core.expr.numeric_expr import SumExpression
+from pyomo.core.expr.numeric_expr import SumExpression, mutable_expression
 from pyomo.core.expr import identify_variables
 from pyomo.core.expr import SumExpression
 from pyomo.core.util import target_list
 from pyomo.contrib.piecewise import PiecewiseLinearExpression, PiecewiseLinearFunction
 from pyomo.gdp import Disjunct, Disjunction
 from pyomo.network import Port
-from pyomo.repn.quadratic import QuadraticRepnVisitor
+from pyomo.repn.quadratic import QuadraticRepnVisitor, QuadraticRepn
 from pyomo.repn.util import ExprType, OrderedVarRecorder
 
 
@@ -646,8 +646,9 @@ def _get_bounds_list(self, var_list, obj):
                 bounds.append((v.bounds, v.is_integer()))
         return bounds
 
-    def _needs_approximating(self, expr, approximate_quadratic):
-        repn = self._quadratic_repn_visitor.walk_expression(expr)
+    def _needs_approximating(self, expr, approximate_quadratic, repn=None):
+        if repn is None:
+            repn = self._quadratic_repn_visitor.walk_expression(expr)
         if repn.nonlinear is None:
             if repn.quadratic is None:
                 # Linear constraint. Always skip.
@@ -659,24 +660,82 @@ def _needs_approximating(self, expr, approximate_quadratic):
                 return ExprType.QUADRATIC, True
         return ExprType.GENERAL, True
 
+    def _separate_linear_parts(
+        self, repn: QuadraticRepn, visitor: QuadraticRepnVisitor
+    ):
+        """
+        The idea here is to ensure that linear parts of constraints
+        always get separated from the nonlinear parts, even if
+        additively_decompose if False. The idea is that
+
+        y >= exp(x) + x**3
+
+        should become
+
+        y >= PWL(exp(x) + x**3)
+
+        and not
+
+        0 >= PWL(exp(x) + x**3 - y)
+        """
+        assert repn.multiplier == 1
+        linear_repn = QuadraticRepn()
+        linear_repn.constant = repn.constant
+        linear_repn.linear = repn.linear
+        linear = linear_repn.to_expression(visitor)
+
+        return linear, repn.quadratic, repn.nonlinear
+
+    def _get_quadratic_part_of_repn(
+        self, repn: QuadraticRepn, visitor: QuadraticRepnVisitor
+    ):
+        assert repn.multiplier == 1
+        r = QuadraticRepn()
+        r.quadratic = repn.quadratic
+        q = r.to_expression(visitor)
+        return q
+
     def _approximate_expression(
         self, expr, obj, trans_block, config, approximate_quadratic
     ):
+        repn = self._quadratic_repn_visitor.walk_expression(expr)
         expr_type, needs_approximating = self._needs_approximating(
-            expr, approximate_quadratic
+            expr, approximate_quadratic, repn
         )
         if not needs_approximating:
             return None, expr_type
 
+        linear_part, quadratic_part, nonlinear_part = self._separate_linear_parts(
+            repn, self._quadratic_repn_visitor
+        )
+        if quadratic_part is None:
+            quadratic_part = {}
+
         # Additively decompose expr and work on the pieces
-        pwl_summands = []
-        for k, subexpr in enumerate(
-            _additively_decompose_expr(
-                expr, config.min_dimension_to_additively_decompose
-            )
-            if config.additively_decompose
-            else (expr,)
-        ):
+        pwl_summands = [linear_part]
+        Nmin = config.min_dimension_to_additively_decompose
+        subexpr_list = []
+        if config.additively_decompose:
+            # dimension = len(list(identify_variables(input_expr)))
+            vset = set()
+            vmap = self._quadratic_repn_visitor.var_map
+            for x12 in quadratic_part.keys():
+                vset.update(x12)
+            if len(vset) < Nmin and nonlinear_part is not None:
+                vset.update(id(v) for v in identify_variables(nonlinear_part))
+            if len(vset) >= Nmin:
+                for (x1k, x2k), coef in quadratic_part.items():
+                    x1 = vmap[x1k]
+                    x2 = vmap[x2k]
+                    subexpr_list.append(coef * (x1 * x2))
+                if nonlinear_part is not None:
+                    subexpr_list.extend(_additively_decompose_expr(nonlinear_part, 0))
+        if not subexpr_list:
+            e = self._get_quadratic_part_of_repn(repn, self._quadratic_repn_visitor)
+            if nonlinear_part is not None:
+                e += nonlinear_part
+            subexpr_list = [e]
+        for k, subexpr in enumerate(subexpr_list):
             # First check if this is a good idea
             expr_vars = list(identify_variables(subexpr, include_fixed=False))
             orig_values = ComponentMap((v, v.value) for v in expr_vars)