API: use __array_ufunc__ as numpy ufunc enter (#1163)

* Refactor matrix comparison logic using __array_ufunc__

Replaces custom rich comparison methods for MatrixExpr and MatrixExprCons with __array_ufunc__ implementations, leveraging numpy's ufuncs for elementwise comparisons. Removes legacy helper functions and streamlines operator handling, improving maintainability and performance.

* Remove unused matrix.pxi include from expr.pxi

Deleted the 'include "matrix.pxi"' line from expr.pxi as it is no longer needed. This helps clean up the code and avoid unnecessary dependencies.

* Replace MatrixExpr checks with np.ndarray in expr.pxi

Updated type checks from MatrixExpr to np.ndarray throughout expr.pxi to improve compatibility with NumPy arrays and simplify matrix expression handling.

* Fix logic error in MatrixExpr __array_ufunc__ handling

Corrects the condition to check for NotImplemented in MatrixExpr's __array_ufunc__ method, ensuring proper delegation to the superclass implementation.

* Fix matmul return type assertion in test

Update the test to assert that the result of 1D @ 1D matrix multiplication is of type Expr instead of MatrixExpr.

* Restrict supported ufuncs in MatrixExprCons

MatrixExprCons now only supports the '<=' and '>=' ufuncs, raising NotImplementedError for all others. This clarifies and enforces the intended usage of the class.

* Update matrix sum test to use MatrixExpr view

Refactored test_matrix_sum_result to use a view of MatrixExpr for sum operation, aligning the test with the expected usage pattern and simplifying result comparison.

* Refactor __array_ufunc__ return handling in MatrixExpr

Simplifies the __array_ufunc__ method by directly returning a MatrixExpr view when the result is a numpy ndarray, improving code clarity.

* Replace np.ndarray checks with MatrixExpr in unary ops

Updated type checks in exp, log, sqrt, sin, and cos functions to use MatrixExpr instead of np.ndarray. This change ensures that matrix-specific expression handling is applied only to MatrixExpr instances.

* Remove method stubs from MatrixExpr and MatrixExprCons

Eliminated explicit method declarations from the MatrixExpr and MatrixExprCons classes in the type stub, leaving only ellipses. This simplifies the type hints and may reflect a change in how these classes are intended to be used or documented.

* Update CHANGELOG for MatrixExpr __array_ufunc__ support

Added entry noting that MatrixExpr and MatrixExprCons now use the `__array_ufunc__` protocol to control all numpy.ufunc inputs and outputs.

* Import Expr from pyscipopt.scip in matrix.pxi

Added a cimport for Expr from pyscipopt.scip to enable usage of the Expr type in this module.

* Fix error message grammar

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix error message grammar

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix argument unpacking in ufunc handling for MatrixExpr

Replaces use of argument unpacking (*args) with explicit indexing (args[0], args[1]) in calls to _vec_le, _vec_ge, and _vec_eq within MatrixExpr and MatrixExprCons. This change clarifies argument passing and may prevent potential issues with argument order or count.

* Move _vec_evaluate definition above _ensure_array

Relocated the definition of _vec_evaluate to group all np.frompyfunc vectorized operator functions together for better code organization and readability.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: Zeroto521

CHANGELOG.md

@@ -20,6 +20,7 @@
 - Speed up MatrixExpr.add.reduce via quicksum
 - Speed up np.ndarray(..., dtype=np.float64) @ MatrixExpr
 - Minimum numpy version increased from 1.16.0 to 1.19.0
+- MatrixExpr and MatrixExprCons use `__array_ufunc__` protocol to control all numpy.ufunc inputs and outputs
 ### Removed
 
 ## 6.0.0 - 2025.xx.yy

src/pyscipopt/matrix.pxi

@@ -1,7 +1,4 @@
-"""
-# TODO Cythonize things. Improve performance.
-# TODO Add tests
-"""
+import operator
 from typing import Literal, Optional, Tuple, Union
 import numpy as np
 from numpy.typing import NDArray
@@ -18,42 +15,6 @@ from pyscipopt.scip cimport Expr, Solution
 cnp.import_array()
 
 
-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 _matrixexpr_richcmp(self, other, op):
-    def _richcmp(self, other, op):
-        if op == 1: # <=
-            return self.__le__(other)
-        elif op == 5: # >=
-            return self.__ge__(other)
-        elif op == 2: # ==
-            return self.__eq__(other)
-        else:
-            raise NotImplementedError("Can only support constraints with '<=', '>=', or '=='.")
-
-    if _is_number(other) or isinstance(other, Expr):
-        res = np.empty(self.shape, dtype=object)
-        res.flat = [_richcmp(i, other, op) for i in self.flat]
-
-    elif isinstance(other, np.ndarray):
-        out = np.broadcast(self, other)
-        res = np.empty(out.shape, dtype=object)
-        res.flat = [_richcmp(i, j, op) for i, j in out]
-
-    else:
-        raise TypeError(f"Unsupported type {type(other)}")
-
-    return res.view(MatrixExprCons)
-
-
 class MatrixExpr(np.ndarray):
 
     def __array_ufunc__(
@@ -97,52 +58,21 @@ class MatrixExpr(np.ndarray):
         if method == "__call__":  # Standard ufunc call, e.g., np.add(a, b)
             if ufunc in {np.matmul, np.dot}:
                 res = _core_dot(args[0], args[1])
+            elif ufunc is np.less_equal:
+                return _vec_le(args[0], args[1]).view(MatrixExprCons)
+            elif ufunc is np.greater_equal:
+                return _vec_ge(args[0], args[1]).view(MatrixExprCons)
+            elif ufunc is np.equal:
+                return _vec_eq(args[0], args[1]).view(MatrixExprCons)
+            elif ufunc in {np.less, np.greater, np.not_equal}:
+                raise NotImplementedError("can only support '<=', '>=', or '=='")
 
         if res is NotImplemented:
             if "out" in kwargs: # Unboxing MatrixExpr to stop __array_ufunc__ recursion
                 kwargs["out"] = tuple(_ensure_array(arg, False) for arg in kwargs["out"])
             res = super().__array_ufunc__(ufunc, method, *args, **kwargs)
         return res.view(MatrixExpr) if isinstance(res, np.ndarray) else res
 
-    def __le__(self, other: Union[float, int, "Expr", np.ndarray, "MatrixExpr"]) -> MatrixExprCons:
-        return _matrixexpr_richcmp(self, other, 1)
-
-    def __ge__(self, other: Union[float, int, "Expr", np.ndarray, "MatrixExpr"]) -> MatrixExprCons:
-        return _matrixexpr_richcmp(self, other, 5)
-
-    def __eq__(self, other: Union[float, int, "Expr", np.ndarray, "MatrixExpr"]) -> MatrixExprCons:
-        return _matrixexpr_richcmp(self, other, 2)
-
-    def __add__(self, other):
-        return super().__add__(other).view(MatrixExpr)
-    
-    def __iadd__(self, other):
-        return super().__iadd__(other).view(MatrixExpr)
-
-    def __mul__(self, other):
-        return super().__mul__(other).view(MatrixExpr)
-
-    def __truediv__(self, other):
-        return super().__truediv__(other).view(MatrixExpr)
-    
-    def __rtruediv__(self, other):
-        return super().__rtruediv__(other).view(MatrixExpr)
-        
-    def __pow__(self, other):
-        return super().__pow__(other).view(MatrixExpr)
-    
-    def __sub__(self, other):
-        return super().__sub__(other).view(MatrixExpr)
-    
-    def __radd__(self, other):
-        return super().__radd__(other).view(MatrixExpr)
-    
-    def __rmul__(self, other):
-        return super().__rmul__(other).view(MatrixExpr)
-    
-    def __rsub__(self, other):
-        return super().__rsub__(other).view(MatrixExpr)
-
 
     def _evaluate(self, Solution sol) -> NDArray[np.float64]:
         return _vec_evaluate(self, sol).view(np.ndarray)
@@ -154,14 +84,20 @@ class MatrixGenExpr(MatrixExpr):
 
 class MatrixExprCons(np.ndarray):
 
-    def __le__(self, other: Union[float, int, np.ndarray]) -> MatrixExprCons:
-        return _matrixexpr_richcmp(self, other, 1)
+    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
+        if method == "__call__":
+            args = tuple(_ensure_array(arg) for arg in args)
+            if ufunc is np.less_equal:
+                return _vec_le(args[0], args[1]).view(MatrixExprCons)
+            elif ufunc is np.greater_equal:
+                return _vec_ge(args[0], args[1]).view(MatrixExprCons)
+        raise NotImplementedError("can only support '<=' or '>='")
 
-    def __ge__(self, other: Union[float, int, np.ndarray]) -> MatrixExprCons:
-        return _matrixexpr_richcmp(self, other, 5)
 
-    def __eq__(self, other):
-        raise NotImplementedError("Cannot compare MatrixExprCons with '=='.")
+_vec_le = np.frompyfunc(operator.le, 2, 1)
+_vec_ge = np.frompyfunc(operator.ge, 2, 1)
+_vec_eq = np.frompyfunc(operator.eq, 2, 1)
+_vec_evaluate = np.frompyfunc(lambda expr, sol: expr._evaluate(sol), 2, 1)
 
 
 cdef inline _ensure_array(arg, bool convert_scalar = True):
@@ -172,9 +108,6 @@ cdef inline _ensure_array(arg, bool convert_scalar = True):
     return np.array(arg, dtype=object) if convert_scalar else arg
 
 
-_vec_evaluate = np.frompyfunc(lambda expr, sol: expr._evaluate(sol), 2, 1)
-
-
 def _core_dot(cnp.ndarray a, cnp.ndarray b) -> Union[Expr, np.ndarray]:
     """
     Perform matrix multiplication between a N-Demension constant array and a N-Demension
@@ -261,7 +194,7 @@ def _core_sum(
 
     Parameters
     ----------
-    a : cnp.ndarray
+    a : np.ndarray
         A `np.ndarray` of type `object` and containing `Expr` objects.
 
     axis : None or int or tuple of ints, optional

src/pyscipopt/scip.pyi

@@ -510,25 +510,10 @@ class MatrixConstraint(numpy.ndarray):
     def isStickingAtNode(self) -> Incomplete: ...
 
 class MatrixExpr(numpy.ndarray):
-    def __add__(self, other: Incomplete) -> Incomplete: ...
-    def __eq__(self, other: Incomplete)-> Incomplete: ...
-    def __ge__(self, other: Incomplete) -> MatrixExprCons: ...
-    def __iadd__(self, other: Incomplete) -> Incomplete: ...  # noqa: PYI034
-    def __le__(self, other: Incomplete) -> MatrixExprCons: ...
-    def __matmul__(self, other: Incomplete) -> Incomplete: ...
-    def __mul__(self, other: Incomplete) -> Incomplete: ...
-    def __pow__(self, other: Incomplete) -> Incomplete: ...  # type: ignore[override]
-    def __radd__(self, other: Incomplete) -> Incomplete: ...
-    def __rmul__(self, other: Incomplete) -> Incomplete: ...
-    def __rsub__(self, other: Incomplete) -> Incomplete: ...
-    def __rtruediv__(self, other: Incomplete) -> Incomplete: ...
-    def __sub__(self, other: Incomplete) -> Incomplete: ...
-    def __truediv__(self, other: Incomplete) -> Incomplete: ...
+    ...
 
 class MatrixExprCons(numpy.ndarray):
-    def __eq__(self, other: Incomplete)-> Incomplete: ...
-    def __ge__(self, other: Incomplete) -> MatrixExprCons: ...
-    def __le__(self, other: Incomplete) -> MatrixExprCons: ...
+    ...
 
 class MatrixGenExpr(MatrixExpr):
     ...