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

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 ### Added
+- `Expr` and `GenExpr` support NumPy unary functions (`np.sin`, `np.cos`, `np.sqrt`, `np.exp`, `np.log`, `np.absolute`)
 - Added `getMemUsed()`, `getMemTotal()`, and `getMemExternEstim()` methods
 ### Fixed
 - Removed `Py_INCREF`/`Py_DECREF` on `Model` in `catchEvent`/`dropEvent` that caused memory leak for imbalanced usage

src/pyscipopt/expr.pxi

@@ -43,16 +43,16 @@
 # 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 cpython.dict cimport PyDict_Next, PyDict_GetItem
 from cpython.object cimport Py_TYPE
 from cpython.ref cimport PyObject
 from cpython.tuple cimport PyTuple_GET_ITEM
 from pyscipopt.scip cimport Variable, Solution
 
-import numpy as np
-
 
 if TYPE_CHECKING:
     double = float
@@ -180,6 +180,7 @@ cdef class Term:
 
 CONST = Term()
 
+
 # helper function
 def buildGenExprObj(expr):
     """helper function to generate an object of type GenExpr"""
@@ -215,10 +216,60 @@ 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 kwargs.get("out", None) is not None:
+            raise TypeError(
+                f"{self.__class__.__name__} doesn't support the 'out' parameter in __array_ufunc__"
+            )
+
+        if method == "__call__":
+            if ufunc is np.absolute:
+                return args[0].__abs__()
+            elif ufunc is np.exp:
+                return args[0].exp()
+            elif ufunc is np.log:
+                return args[0].log()
+            elif ufunc is np.sqrt:
+                return args[0].sqrt()
+            elif ufunc is np.sin:
+                return args[0].sin()
+            elif ufunc is np.cos:
+                return args[0].cos()
+
+        return NotImplemented
+
+    def __abs__(self) -> GenExpr:
+        return UnaryExpr(Operator.fabs, buildGenExprObj(self))
+
+    def exp(self) -> GenExpr:
+        return UnaryExpr(Operator.exp, buildGenExprObj(self))
+
+    def log(self) -> GenExpr:
+        return UnaryExpr(Operator.log, buildGenExprObj(self))
+
+    def sqrt(self) -> GenExpr:
+        return UnaryExpr(Operator.sqrt, buildGenExprObj(self))
+
+    def sin(self) -> GenExpr:
+        return UnaryExpr(Operator.sin, buildGenExprObj(self))
+
+    def cos(self) -> GenExpr:
+        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.
 
@@ -236,9 +287,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
@@ -502,17 +550,14 @@ 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
@@ -816,55 +861,23 @@ 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 = lambda x: _wrap_ufunc(x, np.exp)
+log = lambda x: _wrap_ufunc(x, np.log)
+sqrt = lambda x: _wrap_ufunc(x, np.sqrt)
+sin = lambda x: _wrap_ufunc(x, np.sin)
+cos = lambda x: _wrap_ufunc(x, np.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))
+cdef inline object _to_const(object x):
+    return Constant(<double>x) if _is_number(x) else x
 
-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))
+cdef object _vec_to_const = np.frompyfunc(_to_const, 1, 1)
 
-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))
+cdef inline object _wrap_ufunc(object x, object ufunc):
+    if isinstance(x, (np.ndarray, list, tuple)):
+        res = ufunc(_vec_to_const(x))
+        return res.view(MatrixGenExpr) if isinstance(res, np.ndarray) else res
+    return ufunc(_to_const(x))
 
-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

@@ -2126,7 +2126,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

@@ -8,8 +8,8 @@ CONST: Term
 EventNames: dict
 MAJOR: int
 MINOR: int
-Operator: Op
 PATCH: int
+Operator: Op
 PY_SCIP_CALL: Incomplete
 StageNames: dict
 TYPE_CHECKING: bool
@@ -20,18 +20,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 +325,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) -> GenExpr: ...
+    def exp(self) -> GenExpr: ...
+    def log(self) -> GenExpr: ...
+    def sqrt(self) -> GenExpr: ...
+    def sin(self) -> GenExpr: ...
+    def cos(self) -> GenExpr: ...
+
 @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,7 +385,7 @@ class ExprCons:
     def __ne__(self, other: object) -> bool: ...
 
 @disjoint_base
-class GenExpr:
+class GenExpr(ExprLike):
     _op: Incomplete
     children: Incomplete
     def __init__(self) -> None: ...
@@ -516,17 +530,17 @@ 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(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,10 @@
 import math
 
+import numpy as np
 import pytest
 
-from pyscipopt import Model, sqrt, log, exp, sin, cos
-from pyscipopt.scip import Expr, GenExpr, ExprCons, CONST
+from pyscipopt import Model, cos, exp, log, sin, sqrt
+from pyscipopt.scip import CONST, Expr, ExprCons, GenExpr, MatrixGenExpr
 
 
 @pytest.fixture(scope="module")
@@ -117,11 +118,13 @@ 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)
 
+
 def test_degree(model):
     m, x, y, z = model
     expr = GenExpr()
@@ -219,6 +222,57 @@ def test_getVal_with_GenExpr():
         m.getVal(1 / z)
 
 
+def test_unary(model):
+    m, x, y, z = model
+
+    res = "abs(sum(0.0,prod(1.0,x)))"
+    assert str(abs(x)) == res
+    assert str(np.absolute(x)) == res
+
+    res = "[sin(sum(0.0,prod(1.0,x))) sin(sum(0.0,prod(1.0,y)))]"
+    assert str(sin([x, y])) == res
+    assert str(np.sin([x, y])) == res
+
+    res = "[cos(sum(0.0,prod(1.0,x))) cos(sum(0.0,prod(1.0,y)))]"
+    assert str(cos([x, y])) == res
+    assert str(np.cos([x, y])) == res
+
+    res = "[sqrt(sum(0.0,prod(1.0,x))) sqrt(sum(0.0,prod(1.0,y)))]"
+    assert str(sqrt([x, y])) == res
+    assert str(np.sqrt([x, y])) == res
+
+    res = "[exp(sum(0.0,prod(1.0,x))) exp(sum(0.0,prod(1.0,y)))]"
+    assert str(exp([x, y])) == res
+    assert str(np.exp([x, y])) == res
+
+    res = "[log(sum(0.0,prod(1.0,x))) log(sum(0.0,prod(1.0,y)))]"
+    assert str(log([x, y])) == res
+    assert str(np.log([x, y])) == res
+
+    assert str(log([1, x])) == "[log(1.0) log(sum(0.0,prod(1.0,x)))]"
+
+    assert str(sqrt(4)) == "sqrt(4.0)"
+    assert str(sqrt([4, 4])) == "[sqrt(4.0) sqrt(4.0)]"
+    assert str(exp(3)) == "exp(3.0)"
+    assert str(exp([3, 3])) == "[exp(3.0) exp(3.0)]"
+    assert str(log(5)) == "log(5.0)"
+    assert str(log([5, 5])) == "[log(5.0) log(5.0)]"
+    assert str(sin(1)) == "sin(1.0)"
+    assert str(sin([[1, 1]])) == "[[sin(1.0) sin(1.0)]]"
+    assert str(cos(1)) == "cos(1.0)"
+    assert str(cos([[1]])) == "[[cos(1.0)]]"
+
+    assert isinstance(sqrt(2), GenExpr)
+    assert isinstance(sqrt([2, 2]), MatrixGenExpr)
+    assert isinstance(sqrt([[2], [2]]), MatrixGenExpr)
+    assert isinstance(sqrt([2, x]), MatrixGenExpr)
+    assert isinstance(sqrt([[2], [x]]), MatrixGenExpr)
+
+    # test invalid unary operations
+    with pytest.raises(TypeError):
+        np.arcsin(x)
+
+
 def test_mul():
     m = Model()
     x = m.addVar(name="x")