Return NotImplemented for Expr and GenExpr operators to simplify

CHANGELOG.md

@@ -9,6 +9,7 @@
 - Used `getIndex()` instead of `ptr()` for sorting nonlinear expression terms to avoid nondeterministic behavior
 - Fixed stubtest failures with mypy 1.20 by marking dunder method parameters as positional-only
 ### Changed
+- Return NotImplemented for `Expr` and `GenExpr` operators, if they can't handle input types in the calculation
 - Speed up `constant * Expr` via C-level API
 - Speed up `Term.__eq__` via the C-level API
 ### Removed

src/pyscipopt/expr.pxi

@@ -43,63 +43,27 @@
 # gets called (I guess) and so a copy is returned.
 # Modifying the expression directly would be a bug, given that the expression might be re-used by the user. </pre>
 import math
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
 
 from cpython.dict cimport PyDict_Next, PyDict_GetItem
-from cpython.object cimport Py_TYPE
+from cpython.float cimport PyFloat_Check
+from cpython.long cimport PyLong_Check
+from cpython.number cimport PyNumber_Check
+from cpython.object cimport Py_LE, Py_EQ, Py_GE, Py_TYPE
 from cpython.ref cimport PyObject
 from cpython.tuple cimport PyTuple_GET_ITEM
-from pyscipopt.scip cimport Variable, Solution
 
-import numpy as np
+cimport numpy as cnp
+
+from pyscipopt.scip cimport Variable, Solution
 
 
 if TYPE_CHECKING:
     double = float
 
 
-def _is_number(e):
-    try:
-        f = float(e)
-        return True
-    except ValueError: # for malformed strings
-        return False
-    except TypeError: # for other types (Variable, Expr)
-        return False
-
-
-def _expr_richcmp(self, other, op):
-    if op == 1: # <=
-        if isinstance(other, Expr) or isinstance(other, GenExpr):
-            return (self - other) <= 0.0
-        elif _is_number(other):
-            return ExprCons(self, rhs=float(other))
-        elif isinstance(other, np.ndarray):
-            return _expr_richcmp(other, self, 5)
-        else:
-            raise TypeError(f"Unsupported type {type(other)}")
-    elif op == 5: # >=
-        if isinstance(other, Expr) or isinstance(other, GenExpr):
-            return (self - other) >= 0.0
-        elif _is_number(other):
-            return ExprCons(self, lhs=float(other))
-        elif isinstance(other, np.ndarray):
-            return _expr_richcmp(other, self, 1)
-        else:
-            raise TypeError(f"Unsupported type {type(other)}")
-    elif op == 2: # ==
-        if isinstance(other, Expr) or isinstance(other, GenExpr):
-            return (self - other) == 0.0
-        elif _is_number(other):
-            return ExprCons(self, lhs=float(other), rhs=float(other))
-        elif isinstance(other, np.ndarray):
-            return _expr_richcmp(other, self, 2)
-        else:
-            raise TypeError(f"Unsupported type {type(other)}")
-    else:
-        raise NotImplementedError("Can only support constraints with '<=', '>=', or '=='.")
-
-
 cdef class Term:
     '''This is a monomial term'''
 
@@ -195,8 +159,11 @@ cdef class Term:
 CONST = Term()
 
 # helper function
-def buildGenExprObj(expr):
+def buildGenExprObj(expr: Union[int, float, np.number, Expr, GenExpr]) -> GenExpr:
     """helper function to generate an object of type GenExpr"""
+    if not _is_genexpr_compatible(expr):
+        raise TypeError(f"unsupported type {type(expr).__name__!s}")
+
     if _is_number(expr):
         return Constant(expr)
 
@@ -219,15 +186,7 @@ def buildGenExprObj(expr):
                 sumexpr += coef * prodexpr
         return sumexpr
 
-    elif isinstance(expr, np.ndarray):   
-        GenExprs = np.empty(expr.shape, dtype=object)
-        for idx in np.ndindex(expr.shape):
-            GenExprs[idx] = buildGenExprObj(expr[idx])
-        return GenExprs.view(MatrixExpr)
-
-    else:
-        assert isinstance(expr, GenExpr)
-        return expr
+    return expr
 
 ##@details Polynomial expressions of variables with operator overloading. \n
 #See also the @ref ExprDetails "description" in the expr.pxi. 
@@ -254,6 +213,9 @@ cdef class Expr:
         return abs(buildGenExprObj(self))
 
     def __add__(self, other):
+        if not _is_expr_compatible(other):
+            return NotImplemented
+
         left = self
         right = other
         terms = left.terms.copy()
@@ -265,35 +227,23 @@ cdef class Expr:
         elif _is_number(right):
             c = float(right)
             terms[CONST] = terms.get(CONST, 0.0) + c
-        elif isinstance(right, GenExpr):
-            return buildGenExprObj(left) + right
-        elif isinstance(right, np.ndarray):
-            return right + left
-        else:
-            raise TypeError(f"Unsupported type {type(right)}")
-
         return Expr(terms)
 
     def __iadd__(self, other):
+        if not _is_expr_compatible(other):
+            return NotImplemented
+
         if isinstance(other, Expr):
             for v,c in other.terms.items():
                 self.terms[v] = self.terms.get(v, 0.0) + c
         elif _is_number(other):
             c = float(other)
             self.terms[CONST] = self.terms.get(CONST, 0.0) + c
-        elif isinstance(other, GenExpr):
-            # is no longer in place, might affect performance?
-            # can't do `self = buildGenExprObj(self) + other` since I get
-            # TypeError: Cannot convert pyscipopt.scip.SumExpr to pyscipopt.scip.Expr
-            return buildGenExprObj(self) + other
-        else:
-            raise TypeError(f"Unsupported type {type(other)}")
-
         return self
 
     def __mul__(self, other):
-        if isinstance(other, np.ndarray):
-            return other * self
+        if not _is_expr_compatible(other):
+            return NotImplemented
 
         cdef dict res = {}
         cdef Py_ssize_t pos1 = <Py_ssize_t>0, pos2 = <Py_ssize_t>0
@@ -306,10 +256,9 @@ cdef class Expr:
         cdef double coef
 
         if _is_number(other):
-            coef = float(other)
+            coef = <double>other
             while PyDict_Next(self.terms, &pos1, &k1_ptr, &v1_ptr):
                 res[<Term>k1_ptr] = <double>(<object>v1_ptr) * coef
-            return Expr(res)
 
         elif isinstance(other, Expr):
             while PyDict_Next(self.terms, &pos1, &k1_ptr, &v1_ptr):
@@ -321,22 +270,20 @@ cdef class Expr:
                         res[child] = <double>(<object>old_v_ptr) + coef
                     else:
                         res[child] = coef
-            return Expr(res)
+        return Expr(res)
 
-        elif isinstance(other, GenExpr):
-            return buildGenExprObj(self) * other
-        else:
-            raise NotImplementedError
+    def __truediv__(self, other):
+        if not _is_expr_compatible(other):
+            return NotImplemented
 
-    def __truediv__(self,other):
         if _is_number(other):
-            f = 1.0/float(other)
-            return f * self
-        selfexpr = buildGenExprObj(self)
-        return selfexpr.__truediv__(other)
+            return 1.0 / other * self
+        return buildGenExprObj(self) / other
 
     def __rtruediv__(self, other):
         ''' other / self '''
+        if not _is_expr_compatible(other):
+            return NotImplemented
         return buildGenExprObj(other) / self
 
     def __pow__(self, other, modulo):
@@ -355,13 +302,11 @@ cdef class Expr:
         Implements base**x as scip.exp(x * scip.log(base)).
         Note: base must be positive.
         """
-        if _is_number(other):
-            base = float(other)
-            if base <= 0.0:
-                raise ValueError("Base of a**x must be positive, as expression is reformulated to scip.exp(x * scip.log(a)); got %g" % base)
-            return exp(self * log(base))
-        else:
+        if not _is_number(other):
             raise TypeError(f"Unsupported base type {type(other)} for exponentiation.")
+        if other <= 0.0:
+            raise ValueError("Base of a**x must be positive, as expression is reformulated to scip.exp(x * scip.log(a)); got %g" % other)
+        return exp(self * log(float(other)))
 
     def __neg__(self):
         return Expr({v:-c for v,c in self.terms.items()})
@@ -378,7 +323,7 @@ cdef class Expr:
     def __rsub__(self, other):
         return -1.0 * self + other
 
-    def __richcmp__(self, other, op):
+    def __richcmp__(self, other, int op):
         '''turn it into a constraint'''
         return _expr_richcmp(self, other, op)
 
@@ -443,24 +388,21 @@ cdef class ExprCons:
 
     def __richcmp__(self, other, op):
         '''turn it into a constraint'''
+        if not _is_number(other):
+            raise TypeError('Ranged ExprCons is not well defined!')
+
         if op == 1: # <=
             if not self._rhs is None:
                 raise TypeError('ExprCons already has upper bound')
             assert not self._lhs is None
 
-            if not _is_number(other):
-                raise TypeError('Ranged ExprCons is not well defined!')
-
             return ExprCons(self.expr, lhs=self._lhs, rhs=float(other))
         elif op == 5: # >=
             if not self._lhs is None:
                 raise TypeError('ExprCons already has lower bound')
             assert self._lhs is None
             assert not self._rhs is None
 
-            if not _is_number(other):
-                raise TypeError('Ranged ExprCons is not well defined!')
-
             return ExprCons(self.expr, lhs=float(other), rhs=self._rhs)
         else:
             raise NotImplementedError("Ranged ExprCons can only support with '<=' or '>='.")
@@ -530,8 +472,8 @@ cdef class GenExpr:
         return UnaryExpr(Operator.fabs, self)
 
     def __add__(self, other):
-        if isinstance(other, np.ndarray):
-            return other + self
+        if not _is_genexpr_compatible(other):
+            return NotImplemented
 
         left = buildGenExprObj(self)
         right = buildGenExprObj(other)
@@ -588,8 +530,8 @@ cdef class GenExpr:
     #    return self
 
     def __mul__(self, other):
-        if isinstance(other, np.ndarray):
-            return other * self
+        if not _is_genexpr_compatible(other):
+            return NotImplemented
 
         left = buildGenExprObj(self)
         right = buildGenExprObj(other)
@@ -654,16 +596,17 @@ cdef class GenExpr:
         Implements base**x as scip.exp(x * scip.log(base)). 
         Note: base must be positive.
         """
-        if _is_number(other):
-            base = float(other)
-            if base <= 0.0:
-                raise ValueError("Base of a**x must be positive, as expression is reformulated to scip.exp(x * scip.log(a)); got %g" % base)
-            return exp(self * log(base))
-        else:
+        if not _is_number(other):
             raise TypeError(f"Unsupported base type {type(other)} for exponentiation.")
+        if other <= 0.0:
+            raise ValueError("Base of a**x must be positive, as expression is reformulated to scip.exp(x * scip.log(a)); got %g" % other)
+        return exp(self * log(float(other)))
 
     #TODO: ipow, idiv, etc
     def __truediv__(self,other):
+        if not _is_genexpr_compatible(other):
+            return NotImplemented
+
         divisor = buildGenExprObj(other)
         # we can't divide by 0
         if isinstance(divisor, GenExpr) and divisor.getOp() == Operator.const and divisor.number == 0.0:
@@ -672,8 +615,9 @@ cdef class GenExpr:
 
     def __rtruediv__(self, other):
         ''' other / self '''
-        otherexpr = buildGenExprObj(other)
-        return otherexpr.__truediv__(self)
+        if not _is_genexpr_compatible(other):
+            return NotImplemented
+        return buildGenExprObj(other) / self
 
     def __neg__(self):
         return -1.0 * self
@@ -690,7 +634,7 @@ cdef class GenExpr:
     def __rsub__(self, other):
         return -1.0 * self + other
 
-    def __richcmp__(self, other, op):
+    def __richcmp__(self, other, int op):
         '''turn it into a constraint'''
         return _expr_richcmp(self, other, op)
 
@@ -921,3 +865,41 @@ def expr_to_array(expr, nodes):
     else: # var
         nodes.append( tuple( [op, expr.children] ) )
     return len(nodes) - 1
+
+
+cdef bint _is_number(object x):
+    if PyLong_Check(x) or PyFloat_Check(x):
+        return True
+    if cnp.PyArray_Check(x) or isinstance(x, (Expr, GenExpr, list, tuple)):
+        return False
+    return PyNumber_Check(x)
+
+cdef inline bint _is_expr_compatible(object x):
+    return _is_number(x) or isinstance(x, Expr)
+
+cdef inline bint _is_genexpr_compatible(object x):
+    return _is_expr_compatible(x) or isinstance(x, GenExpr)
+
+cdef object _expr_richcmp(
+    self,
+    other: Union[int, float, np.number, Expr, GenExpr],
+    int op,
+):
+    if isinstance(other, np.ndarray):
+        return NotImplemented
+    if not _is_genexpr_compatible(other):
+        raise TypeError(f"unsupported type {type(other).__name__!s}")
+
+    if op == Py_LE:
+        if _is_number(other):
+            return ExprCons(self, rhs=<double>other)
+        return ExprCons(self - other, rhs=0.0)
+    elif op == Py_GE:
+        if _is_number(other):
+            return ExprCons(self, lhs=<double>other)
+        return ExprCons(self - other, lhs=0.0)
+    elif op == Py_EQ:
+        if _is_number(other):
+            return ExprCons(self, lhs=<double>other, rhs=<double>other)
+        return ExprCons(self - other, lhs=0.0, rhs=0.0)
+    raise NotImplementedError("can only support with '<=', '>=', or '=='")