Feature: Expr and GenExpr support numpy unary func like np.sin

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Added pre-commit hook for automatic stub regeneration (see .pre-commit-config.yaml)
 - Wrapped isObjIntegral() and test
 - Added structured_optimization_trace recipe for structured optimization progress tracking
+- Expr and GenExpr support numpy unary func (`np.sin`, `np.cos`, `np.sqrt`, `np.exp`, `np.log`, `np.absolute`)
 ### Fixed
 - getBestSol() now returns None for infeasible problems instead of a Solution with NULL pointer
 - all fundamental callbacks now raise an error if not implemented

src/pyscipopt/expr.pxi

@@ -43,13 +43,13 @@
 # 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, Literal
+
+import numpy as np
 
-from pyscipopt.scip cimport Variable, Solution
 from cpython.dict cimport PyDict_Next
 from cpython.ref cimport PyObject
-
-import numpy as np
+from pyscipopt.scip cimport Variable, Solution
 
 
 if TYPE_CHECKING:
@@ -146,6 +146,7 @@ cdef class Term:
 
 CONST = Term()
 
+
 # helper function
 def buildGenExprObj(expr):
     """helper function to generate an object of type GenExpr"""
@@ -181,10 +182,45 @@ def buildGenExprObj(expr):
         assert isinstance(expr, GenExpr)
         return expr
 
+
+cdef class ExprLike:
+
+    def __array_ufunc__(
+        self,
+        ufunc: np.ufunc,
+        method: Literal["__call__", "reduce", "reduceat", "accumulate", "outer", "at"],
+        *args,
+        **kwargs,
+    ):
+        if method == "__call__":
+            if ufunc in UNARY_MAPPER:
+                return getattr(args[0], UNARY_MAPPER[ufunc])()
+
+        return NotImplemented
+
+    def __abs__(self):
+        return UnaryExpr(Operator.fabs, buildGenExprObj(self))
+
+    def exp(self):
+        return UnaryExpr(Operator.exp, buildGenExprObj(self))
+
+    def log(self):
+        return UnaryExpr(Operator.log, buildGenExprObj(self))
+
+    def sqrt(self):
+        return UnaryExpr(Operator.sqrt, buildGenExprObj(self))
+
+    def sin(self):
+        return UnaryExpr(Operator.sin, buildGenExprObj(self))
+
+    def cos(self):
+        return UnaryExpr(Operator.cos, buildGenExprObj(self))
+
+
 ##@details Polynomial expressions of variables with operator overloading. \n
 #See also the @ref ExprDetails "description" in the expr.pxi. 
-cdef class Expr:
-
+cdef class Expr(ExprLike):
+    
     def __init__(self, terms=None):
         '''terms is a dict of variables to coefficients.
 
@@ -202,9 +238,6 @@ cdef class Expr:
     def __iter__(self):
         return iter(self.terms)
 
-    def __abs__(self):
-        return abs(buildGenExprObj(self))
-
     def __add__(self, other):
         left = self
         right = other
@@ -463,17 +496,13 @@ Operator = Op()
 #     so expr[x] will generate an error instead of returning the coefficient of x </pre>
 #
 #See also the @ref ExprDetails "description" in the expr.pxi. 
-cdef class GenExpr:
-
+cdef class GenExpr(ExprLike):
     cdef public _op
     cdef public children
 
     def __init__(self): # do we need it
         ''' '''
 
-    def __abs__(self):
-        return UnaryExpr(Operator.fabs, self)
-
     def __add__(self, other):
         if isinstance(other, np.ndarray):
             return other + self
@@ -758,55 +787,20 @@ cdef class Constant(GenExpr):
         return self.number
 
 
-def exp(expr):
-    """returns expression with exp-function"""
-    if isinstance(expr, MatrixExpr):   
-        unary_exprs = np.empty(shape=expr.shape, dtype=object)
-        for idx in np.ndindex(expr.shape):
-            unary_exprs[idx] = UnaryExpr(Operator.exp, buildGenExprObj(expr[idx]))
-        return unary_exprs.view(MatrixGenExpr)
-    else:
-        return UnaryExpr(Operator.exp, buildGenExprObj(expr))
+exp = np.exp
+log = np.log
+sqrt = np.sqrt
+sin = np.sin
+cos = np.cos
+UNARY_MAPPER = {
+    np.absolute: "__abs__",
+    np.exp: "exp",
+    np.log: "log",
+    np.sqrt: "sqrt",
+    np.sin: "sin",
+    np.cos: "cos",
+}
 
-def log(expr):
-    """returns expression with log-function"""
-    if isinstance(expr, MatrixExpr):
-        unary_exprs = np.empty(shape=expr.shape, dtype=object)
-        for idx in np.ndindex(expr.shape):
-            unary_exprs[idx] = UnaryExpr(Operator.log, buildGenExprObj(expr[idx]))
-        return unary_exprs.view(MatrixGenExpr)
-    else:
-        return UnaryExpr(Operator.log, buildGenExprObj(expr))
-
-def sqrt(expr):
-    """returns expression with sqrt-function"""
-    if isinstance(expr, MatrixExpr):
-        unary_exprs = np.empty(shape=expr.shape, dtype=object)
-        for idx in np.ndindex(expr.shape):
-            unary_exprs[idx] = UnaryExpr(Operator.sqrt, buildGenExprObj(expr[idx]))
-        return unary_exprs.view(MatrixGenExpr)
-    else:
-        return UnaryExpr(Operator.sqrt, buildGenExprObj(expr))
-
-def sin(expr):
-    """returns expression with sin-function"""
-    if isinstance(expr, MatrixExpr):
-        unary_exprs = np.empty(shape=expr.shape, dtype=object)
-        for idx in np.ndindex(expr.shape):
-            unary_exprs[idx] = UnaryExpr(Operator.sin, buildGenExprObj(expr[idx]))
-        return unary_exprs.view(MatrixGenExpr)
-    else:
-        return UnaryExpr(Operator.sin, buildGenExprObj(expr))
-
-def cos(expr):
-    """returns expression with cos-function"""
-    if isinstance(expr, MatrixExpr):   
-        unary_exprs = np.empty(shape=expr.shape, dtype=object)
-        for idx in np.ndindex(expr.shape):
-            unary_exprs[idx] = UnaryExpr(Operator.cos, buildGenExprObj(expr[idx]))
-        return unary_exprs.view(MatrixGenExpr)
-    else:
-        return UnaryExpr(Operator.cos, buildGenExprObj(expr))
 
 def expr_to_nodes(expr):
     '''transforms tree to an array of nodes. each node is an operator and the position of the 

src/pyscipopt/scip.pxd

@@ -2108,7 +2108,10 @@ cdef extern from "scip/scip_var.h":
 cdef extern from "tpi/tpi.h":
     int SCIPtpiGetNumThreads()
 
-cdef class Expr:
+cdef class ExprLike:
+    pass
+
+cdef class Expr(ExprLike):
     cdef public terms
 
     cpdef double _evaluate(self, Solution sol)

src/pyscipopt/scip.pyi

@@ -1,15 +1,16 @@
-from typing import ClassVar
+from typing import ClassVar, Dict
 
-import numpy
+import numpy as np
 from _typeshed import Incomplete
 from typing_extensions import disjoint_base
 
 CONST: Term
+UNARY_MAPPER: Dict[np.ufunc, str]
 EventNames: dict
 MAJOR: int
 MINOR: int
-Operator: Op
 PATCH: int
+Operator: Op
 PY_SCIP_CALL: Incomplete
 StageNames: dict
 TYPE_CHECKING: bool
@@ -20,18 +21,18 @@ _core_sum: Incomplete
 _expr_richcmp: Incomplete
 _is_number: Incomplete
 buildGenExprObj: Incomplete
-cos: Incomplete
 exp: Incomplete
+log: Incomplete
+sin: Incomplete
+cos: Incomplete
+sqrt: Incomplete
 expr_to_array: Incomplete
 expr_to_nodes: Incomplete
 is_memory_freed: Incomplete
-log: Incomplete
 print_memory_in_use: Incomplete
 quickprod: Incomplete
 quicksum: Incomplete
 readStatistics: Incomplete
-sin: Incomplete
-sqrt: Incomplete
 str_conversion: Incomplete
 value_to_array: Incomplete
 
@@ -325,13 +326,27 @@ class Eventhdlr:
     def eventinit(self) -> Incomplete: ...
     def eventinitsol(self) -> Incomplete: ...
 
+class ExprLike:
+    def __array_ufunc__(
+        self,
+        ufunc: np.ufunc,
+        method: str,
+        *args: Incomplete,
+        **kwargs: Incomplete,
+    ) -> Incomplete: ...
+    def __abs__(self) -> Incomplete: ...
+    def exp(self) -> Incomplete: ...
+    def log(self) -> Incomplete: ...
+    def sqrt(self) -> Incomplete: ...
+    def sin(self) -> Incomplete: ...
+    def cos(self) -> Incomplete: ...
+
 @disjoint_base
-class Expr:
+class Expr(ExprLike):
     terms: Incomplete
     def __init__(self, terms: Incomplete = ...) -> None: ...
     def degree(self) -> Incomplete: ...
     def normalize(self) -> Incomplete: ...
-    def __abs__(self) -> Incomplete: ...
     def __add__(self, other: Incomplete) -> Incomplete: ...
     def __eq__(self, other: object) -> bool: ...
     def __ge__(self, other: object) -> bool: ...
@@ -371,13 +386,12 @@ class ExprCons:
     def __ne__(self, other: object) -> bool: ...
 
 @disjoint_base
-class GenExpr:
+class GenExpr(ExprLike):
     _op: Incomplete
     children: Incomplete
     def __init__(self) -> None: ...
     def degree(self) -> Incomplete: ...
     def getOp(self) -> Incomplete: ...
-    def __abs__(self) -> Incomplete: ...
     def __add__(self, other: Incomplete) -> Incomplete: ...
     def __eq__(self, other: object) -> bool: ...
     def __ge__(self, other: object) -> bool: ...
@@ -496,7 +510,7 @@ class LP:
     def solve(self, dual: Incomplete = ...) -> Incomplete: ...
     def writeLP(self, filename: Incomplete) -> Incomplete: ...
 
-class MatrixConstraint(numpy.ndarray):
+class MatrixConstraint(np.ndarray):
     def getConshdlrName(self) -> Incomplete: ...
     def isActive(self) -> Incomplete: ...
     def isChecked(self) -> Incomplete: ...
@@ -512,21 +526,21 @@ class MatrixConstraint(numpy.ndarray):
     def isSeparated(self) -> Incomplete: ...
     def isStickingAtNode(self) -> Incomplete: ...
 
-class MatrixExpr(numpy.ndarray):
+class MatrixExpr(np.ndarray):
     def _evaluate(self, sol: Incomplete) -> Incomplete: ...
     def __array_ufunc__(
         self,
-        ufunc: Incomplete,
-        method: Incomplete,
+        ufunc: np.ufunc,
+        method: str,
         *args: Incomplete,
         **kwargs: Incomplete,
     ) -> Incomplete: ...
 
-class MatrixExprCons(numpy.ndarray):
+class MatrixExprCons(np.ndarray):
     def __array_ufunc__(
         self,
-        ufunc: Incomplete,
-        method: Incomplete,
+        ufunc: np.ufunc,
+        method: str,
         *args: Incomplete,
         **kwargs: Incomplete,
     ) -> Incomplete: ...

tests/test_expr.py

@@ -1,9 +1,9 @@
 import math
 
+import numpy as np
 import pytest
-
-from pyscipopt import Model, sqrt, log, exp, sin, cos
-from pyscipopt.scip import Expr, GenExpr, ExprCons, Term
+from pyscipopt import Model, cos, exp, log, sin, sqrt
+from pyscipopt.scip import Expr, ExprCons, GenExpr, Term
 
 
 @pytest.fixture(scope="module")
@@ -118,10 +118,12 @@ def test_genexpr_op_genexpr(model):
     assert isinstance(1/x + genexpr, GenExpr)
     assert isinstance(1/x**1.5 - genexpr, GenExpr)
     assert isinstance(y/x - exp(genexpr), GenExpr)
-    # sqrt(2) is not a constant expression and
-    # we can only power to constant expressions!
-    with pytest.raises(NotImplementedError):
-        genexpr **= sqrt(2)
+
+    genexpr **= sqrt(2)
+    assert isinstance(genexpr, GenExpr)
+
+    with pytest.raises(TypeError):
+        genexpr **= sqrt("2")
 
 def test_degree(model):
     m, x, y, z = model
@@ -218,3 +220,21 @@ def test_getVal_with_GenExpr():
 
     with pytest.raises(ZeroDivisionError):
         m.getVal(1 / z)
+
+
+def test_unary(model):
+    m, x, y, z = model
+
+    assert str(abs(x)) == "abs(sum(0.0,prod(1.0,x)))"
+    assert str(np.absolute(x)) == "abs(sum(0.0,prod(1.0,x)))"
+    assert str(sin([x, y])) == "[sin(sum(0.0,prod(1.0,x))) sin(sum(0.0,prod(1.0,y)))]"
+    assert (
+        str(np.sin([x, y])) == "[sin(sum(0.0,prod(1.0,x))) sin(sum(0.0,prod(1.0,y)))]"
+    )
+    assert (
+        str(sqrt([x, y])) == "[sqrt(sum(0.0,prod(1.0,x))) sqrt(sum(0.0,prod(1.0,y)))]"
+    )
+    assert (
+        str(np.sqrt([x, y]))
+        == "[sqrt(sum(0.0,prod(1.0,x))) sqrt(sum(0.0,prod(1.0,y)))]"
+    )