Merge pull request #3423 from jsiirola/numpy-scalar

Resolve errors in mapping ScalarVar to numpy ndarray

Author: blnicho

pyomo/core/base/indexed_component.py

@@ -1201,7 +1201,7 @@ def __array__(self, dtype=None):
         if not self.is_indexed():
             ans = _ndarray.NumericNDArray(shape=(1,), dtype=object)
             ans[0] = self
-            return ans
+            return ans.reshape(())
 
         _dim = self.dim()
         if _dim is None:

pyomo/core/expr/compare.py

@@ -66,6 +66,11 @@ def handle_external_function_expression(node: ExternalFunctionExpression, pn: Li
     return node.args
 
 
+def handle_sequence(node: collections.abc.Sequence, pn: List):
+    pn.append((collections.abc.Sequence, len(node)))
+    return list(node)
+
+
 def _generic_expression_handler():
     return handle_expression
 
@@ -79,6 +84,7 @@ def _generic_expression_handler():
 handler[AbsExpression] = handle_unary_expression
 handler[NPV_AbsExpression] = handle_unary_expression
 handler[RangedExpression] = handle_expression
+handler[list] = handle_sequence
 
 
 class PrefixVisitor(StreamBasedExpressionVisitor):
@@ -97,19 +103,26 @@ def enterNode(self, node):
             self._result.append(node)
             return tuple(), None
 
-        if node.is_expression_type():
-            if node.is_named_expression_type():
-                return (
-                    handle_named_expression(
-                        node, self._result, self._include_named_exprs
-                    ),
-                    None,
-                )
-            else:
-                return handler[ntype](node, self._result), None
-        else:
-            self._result.append(node)
-            return tuple(), None
+        if ntype in handler:
+            return handler[ntype](node, self._result), None
+
+        if hasattr(node, 'is_expression_type'):
+            if node.is_expression_type():
+                if node.is_named_expression_type():
+                    return (
+                        handle_named_expression(
+                            node, self._result, self._include_named_exprs
+                        ),
+                        None,
+                    )
+                else:
+                    return handler[ntype](node, self._result), None
+        elif hasattr(node, '__len__'):
+            handler[ntype] = handle_sequence
+            return handle_sequence(node, self._result), None
+
+        self._result.append(node)
+        return tuple(), None
 
     def finalizeResult(self, result):
         ans = self._result
@@ -161,10 +174,7 @@ def convert_expression_to_prefix_notation(expr, include_named_exprs=True):
 
     """
     visitor = PrefixVisitor(include_named_exprs=include_named_exprs)
-    if isinstance(expr, Sequence):
-        return expr.__class__(visitor.walk_expression(e) for e in expr)
-    else:
-        return visitor.walk_expression(expr)
+    return visitor.walk_expression(expr)
 
 
 def compare_expressions(expr1, expr2, include_named_exprs=True):

pyomo/core/tests/unit/test_expr_numpy.py

@@ -0,0 +1,95 @@
+#  ___________________________________________________________________________
+#
+#  Pyomo: Python Optimization Modeling Objects
+#  Copyright (c) 2008-2024
+#  National Technology and Engineering Solutions of Sandia, LLC
+#  Under the terms of Contract DE-NA0003525 with National Technology and
+#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
+#  rights in this software.
+#  This software is distributed under the 3-clause BSD License.
+#  ___________________________________________________________________________
+
+import pyomo.common.unittest as unittest
+
+from pyomo.common.dependencies import numpy as np, numpy_available
+from pyomo.environ import ConcreteModel, Var, Constraint
+
+
+@unittest.skipUnless(numpy_available, "tests require numpy")
+class TestNumpyExpr(unittest.TestCase):
+    def test_scalar_operations(self):
+        m = ConcreteModel()
+        m.x = Var()
+
+        a = np.array(m.x)
+        self.assertEqual(a.shape, ())
+
+        self.assertExpressionsEqual(5 * a, 5 * m.x)
+        self.assertExpressionsEqual(np.array([2, 3]) * a, [2 * m.x, 3 * m.x])
+        self.assertExpressionsEqual(np.array([5, 6]) * m.x, [5 * m.x, 6 * m.x])
+        self.assertExpressionsEqual(np.array([8, m.x]) * m.x, [8 * m.x, m.x * m.x])
+
+        a = np.array([m.x])
+        self.assertEqual(a.shape, (1,))
+
+        self.assertExpressionsEqual(5 * a, [5 * m.x])
+        self.assertExpressionsEqual(np.array([2, 3]) * a, [2 * m.x, 3 * m.x])
+        self.assertExpressionsEqual(np.array([5, 6]) * m.x, [5 * m.x, 6 * m.x])
+        self.assertExpressionsEqual(np.array([8, m.x]) * m.x, [8 * m.x, m.x * m.x])
+
+    def test_vector_operations(self):
+        m = ConcreteModel()
+        m.x = Var()
+        m.y = Var([0, 1, 2])
+
+        with self.assertRaisesRegex(TypeError, "unsupported operand"):
+            # TODO: when we finally support a true matrix expression
+            # system, this test should work
+            self.assertExpressionsEqual(5 * m.y, [5 * m.y[0], 5 * m.y[1], 5 * m.y[2]])
+
+        a = np.array(5)
+        self.assertExpressionsEqual(a * m.y, [5 * m.y[0], 5 * m.y[1], 5 * m.y[2]])
+        self.assertExpressionsEqual(m.y * a, [5 * m.y[0], 5 * m.y[1], 5 * m.y[2]])
+        a = np.array([5])
+        self.assertExpressionsEqual(a * m.y, [5 * m.y[0], 5 * m.y[1], 5 * m.y[2]])
+        self.assertExpressionsEqual(m.y * a, [5 * m.y[0], 5 * m.y[1], 5 * m.y[2]])
+
+        a = np.array(5)
+        with self.assertRaisesRegex(TypeError, "unsupported operand"):
+            # TODO: when we finally support a true matrix expression
+            # system, this test should work
+            self.assertExpressionsEqual(
+                a * m.x * m.y, [5 * m.x * m.y[0], 5 * m.x * m.y[1], 5 * m.x * m.y[2]]
+            )
+        self.assertExpressionsEqual(
+            a * m.y * m.x, [5 * m.y[0] * m.x, 5 * m.y[1] * m.x, 5 * m.y[2] * m.x]
+        )
+        self.assertExpressionsEqual(
+            a * m.y * m.y,
+            [5 * m.y[0] * m.y[0], 5 * m.y[1] * m.y[1], 5 * m.y[2] * m.y[2]],
+        )
+        self.assertExpressionsEqual(
+            m.y * a * m.x, [5 * m.y[0] * m.x, 5 * m.y[1] * m.x, 5 * m.y[2] * m.x]
+        )
+        with self.assertRaisesRegex(TypeError, "unsupported operand"):
+            # TODO: when we finally support a true matrix expression
+            # system, this test should work
+            self.assertExpressionsEqual(
+                m.y * m.x * a, [5 * m.y[0] * m.x, 5 * m.y[1] * m.x, 5 * m.y[2] * m.x]
+            )
+        with self.assertRaisesRegex(TypeError, "unsupported operand"):
+            # TODO: when we finally support a true matrix expression
+            # system, this test should work
+            self.assertExpressionsEqual(
+                m.x * a * m.y, [5 * m.y[0] * m.x, 5 * m.y[1] * m.x, 5 * m.y[2] * m.x]
+            )
+        with self.assertRaisesRegex(TypeError, "unsupported operand"):
+            # TODO: when we finally support a true matrix expression
+            # system, this test should work
+            self.assertExpressionsEqual(
+                m.x * m.y * a, [5 * m.y[0] * m.x, 5 * m.y[1] * m.x, 5 * m.y[2] * m.x]
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()