Skip to content

Commit

Permalink
Merge pull request Pyomo#3423 from jsiirola/numpy-scalar
Browse files Browse the repository at this point in the history
Resolve errors in mapping ScalarVar to numpy ndarray
  • Loading branch information
blnicho authored Nov 18, 2024
2 parents 5669dfb + 05bdaf7 commit 3f21a2d
Show file tree
Hide file tree
Showing 3 changed files with 123 additions and 18 deletions.
2 changes: 1 addition & 1 deletion pyomo/core/base/indexed_component.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
44 changes: 27 additions & 17 deletions pyomo/core/expr/compare.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand Down
95 changes: 95 additions & 0 deletions pyomo/core/tests/unit/test_expr_numpy.py
Original file line number Diff line number Diff line change
@@ -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()

0 comments on commit 3f21a2d

Please sign in to comment.