Speed up Constant @ MatrixExpr (#1159)

* Add numpy as build dependency and enhance matrix operations

Added numpy to build requirements in pyproject.toml and setup.py, ensuring numpy headers are included during compilation. Refactored matrix.pxi to improve matrix operation support, including custom __array_ufunc__ handling for dot/matmul, utility functions for type checking and array conversion, and a vectorized _core_dot implementation for efficient matrix multiplication.

* Add and update matrix dot product tests

Introduces a new parameterized test for matrix dot product performance and updates an assertion in test_matrix_matmul_return_type to expect Expr instead of MatrixExpr for 1D @ 1D operations.

* Replace MatrixExpr type checks with np.ndarray

Updated type checks throughout expr.pxi to use np.ndarray instead of MatrixExpr, improving compatibility with numpy arrays. Also adjusted matrix.pxi to ensure ufunc results are returned as MatrixExpr views when appropriate.

* Remove redundant 'out' handling in __array_ufunc__

Deleted unnecessary conversion of the 'out' keyword argument to an array in MatrixExpr.__array_ufunc__, as it is not required for correct operation.

* Update CHANGELOG.md

* Remove custom __matmul__ from MatrixExpr

Deleted the overridden __matmul__ method in MatrixExpr, reverting to the default numpy ndarray behavior for matrix multiplication. This simplifies the class and avoids unnecessary type casting.

* Refactor matrix multiplication logic in MatrixExpr

Introduces a new _core_dot function to handle matrix multiplication between constant arrays and arrays of Expr objects, supporting both 1-D and N-D cases. The original _core_dot is renamed to _core_dot_2d and improved for clarity and efficiency. Updates __array_ufunc__ to use the new logic, ensuring correct handling of mixed-type matrix operations.

* Update matrix matmul return type tests

Adjusted tests in test_matrix_matmul_return_type to check the return type when performing matrix multiplication with numpy arrays and matrix variables, including new cases for ND arrays. Ensures correct type inference for various matrix multiplication scenarios.

* Remove `_is_num_dt`

Replaces the _is_num_dt helper function with direct dtype.kind checks for numeric types in _core_dot. This simplifies the code and removes an unnecessary inline function.

* Enhance MatrixExpr __array_ufunc__ with type hints and docs

Added detailed type annotations and a docstring to the MatrixExpr.__array_ufunc__ method for improved clarity and maintainability. Also clarified handling of ufunc methods and argument unboxing.

* Rename test_matrix_dot to test_matrix_dot_performance

Renamed the test function to better reflect its purpose of testing performance for matrix dot operations.

* Add test for matrix dot value retrieval

Introduces test_matrix_dot_value to verify correct value retrieval from matrix variable dot products using getVal. Ensures expected results for both 1D and higher-dimensional dot operations.

Author: Zeroto521

CHANGELOG.md

@@ -14,6 +14,7 @@
 - Fixed incorrect getVal() result when _bestSol.sol was outdated
 ### Changed
 - changed default value of enablepricing flag to True
+- Speed up np.ndarray(..., dtype=np.float64) @ MatrixExpr
 ### Removed
 
 ## 6.0.0 - 2025.xx.yy

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ['setuptools', 'cython >=0.21']
+requires = ["setuptools", "cython >=0.21", "numpy"]
 build-backend = "setuptools.build_meta"
 
 [project]

setup.py

@@ -1,5 +1,9 @@
-from setuptools import find_packages, setup, Extension
-import os, platform, sys
+import os
+import platform
+import sys
+
+import numpy as np
+from setuptools import Extension, find_packages, setup
 
 # look for environment variable that specifies path to SCIP
 scipoptdir = os.environ.get("SCIPOPTDIR", "").strip('"')
@@ -112,7 +116,7 @@
     Extension(
         "pyscipopt.scip",
         [os.path.join(packagedir, "scip%s" % ext)],
-        include_dirs=includedirs,
+        include_dirs=includedirs + [np.get_include()],
         library_dirs=[libdir],
         libraries=[libname],
         extra_compile_args=extra_compile_args,

src/pyscipopt/expr.pxi

@@ -42,7 +42,7 @@
 # which should, in princple, modify the expr. However, since we do not implement __isub__, __sub__
 # 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>
-include "matrix.pxi"
+import numpy as np
 
 
 def _is_number(e):
@@ -61,7 +61,7 @@ def _expr_richcmp(self, other, op):
             return (self - other) <= 0.0
         elif _is_number(other):
             return ExprCons(self, rhs=float(other))
-        elif isinstance(other, MatrixExpr):
+        elif isinstance(other, np.ndarray):
             return _expr_richcmp(other, self, 5)
         else:
             raise TypeError(f"Unsupported type {type(other)}")
@@ -70,7 +70,7 @@ def _expr_richcmp(self, other, op):
             return (self - other) >= 0.0
         elif _is_number(other):
             return ExprCons(self, lhs=float(other))
-        elif isinstance(other, MatrixExpr):
+        elif isinstance(other, np.ndarray):
             return _expr_richcmp(other, self, 1)
         else:
             raise TypeError(f"Unsupported type {type(other)}")
@@ -79,7 +79,7 @@ def _expr_richcmp(self, other, op):
             return (self - other) == 0.0
         elif _is_number(other):
             return ExprCons(self, lhs=float(other), rhs=float(other))
-        elif isinstance(other, MatrixExpr):
+        elif isinstance(other, np.ndarray):
             return _expr_richcmp(other, self, 2)
         else:
             raise TypeError(f"Unsupported type {type(other)}")
@@ -144,7 +144,7 @@ def buildGenExprObj(expr):
                 sumexpr += coef * prodexpr
         return sumexpr
 
-    elif isinstance(expr, MatrixExpr):   
+    elif isinstance(expr, np.ndarray):   
         GenExprs = np.empty(expr.shape, dtype=object)
         for idx in np.ndindex(expr.shape):
             GenExprs[idx] = buildGenExprObj(expr[idx])
@@ -200,7 +200,7 @@ cdef class Expr:
             terms[CONST] = terms.get(CONST, 0.0) + c
         elif isinstance(right, GenExpr):
             return buildGenExprObj(left) + right
-        elif isinstance(right, MatrixExpr):
+        elif isinstance(right, np.ndarray):
             return right + left
         else:
             raise TypeError(f"Unsupported type {type(right)}")
@@ -225,7 +225,7 @@ cdef class Expr:
         return self
 
     def __mul__(self, other):
-        if isinstance(other, MatrixExpr):
+        if isinstance(other, np.ndarray):
             return other * self
 
         if _is_number(other):
@@ -438,7 +438,7 @@ cdef class GenExpr:
         return UnaryExpr(Operator.fabs, self)
 
     def __add__(self, other):
-        if isinstance(other, MatrixExpr):
+        if isinstance(other, np.ndarray):
             return other + self
 
         left = buildGenExprObj(self)
@@ -496,7 +496,7 @@ cdef class GenExpr:
     #    return self
 
     def __mul__(self, other):
-        if isinstance(other, MatrixExpr):
+        if isinstance(other, np.ndarray):
             return other * self
 
         left = buildGenExprObj(self)

src/pyscipopt/matrix.pxi

@@ -3,7 +3,7 @@
 # TODO Add tests
 """
 
-from typing import Optional, Tuple, Union
+from typing import Literal, Optional, Tuple, Union
 import numpy as np
 try:
     # NumPy 2.x location
@@ -12,6 +12,10 @@ except ImportError:
     # Fallback for NumPy 1.x
     from numpy.core.numeric import normalize_axis_tuple
 
+cimport numpy as cnp
+
+cnp.import_array()
+
 
 def _is_number(e):
     try:
@@ -51,6 +55,48 @@ def _matrixexpr_richcmp(self, other, op):
 
 class MatrixExpr(np.ndarray):
 
+    def __array_ufunc__(
+        self,
+        ufunc: np.ufunc,
+        method: Literal["__call__", "reduce", "reduceat", "accumulate", "outer", "at"],
+        *args,
+        **kwargs,
+    ):
+        """
+        Customizes the behavior of NumPy ufuncs for MatrixExpr.
+
+        Parameters
+        ----------
+        ufunc : numpy.ufunc
+            The ufunc object that was called.
+
+        method : {"__call__", "reduce", "reduceat", "accumulate", "outer", "at"}
+            A string indicating which UFunc method was called.
+
+        *args : tuple
+            The input arguments to the ufunc.
+
+        **kwargs : dict
+            Additional keyword arguments to the ufunc.
+
+        Returns
+        -------
+        Expr, GenExpr, MatrixExpr
+            The result of the ufunc operation is wrapped back into a MatrixExpr if
+            applicable.
+
+        """
+        res = NotImplemented
+        if method == "__call__":  # Standard ufunc call, e.g., np.add(a, b)
+            if ufunc in {np.matmul, np.dot}:
+                res = _core_dot(_ensure_array(args[0]), _ensure_array(args[1]))
+
+        if res is NotImplemented:
+            # Unboxing MatrixExpr to stop __array_ufunc__ recursion
+            args = tuple(_ensure_array(arg) for arg in args)
+            res = super().__array_ufunc__(ufunc, method, *args, **kwargs)
+        return res.view(MatrixExpr) if isinstance(res, np.ndarray) else res
+
     def sum(
         self,
         axis: Optional[Union[int, Tuple[int, ...]]] = None,
@@ -145,8 +191,6 @@ class MatrixExpr(np.ndarray):
     def __rsub__(self, other):
         return super().__rsub__(other).view(MatrixExpr)
 
-    def __matmul__(self, other):
-        return super().__matmul__(other).view(MatrixExpr)
 
 class MatrixGenExpr(MatrixExpr):
     pass
@@ -161,3 +205,86 @@ class MatrixExprCons(np.ndarray):
 
     def __eq__(self, other):
         raise NotImplementedError("Cannot compare MatrixExprCons with '=='.")
+
+
+cdef inline _ensure_array(arg, bool convert_scalar = True):
+    if isinstance(arg, np.ndarray):
+        return arg.view(np.ndarray)
+    elif isinstance(arg, (list, tuple)):
+        return np.asarray(arg)
+    return np.array(arg, dtype=object) if convert_scalar else arg
+
+
+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
+    `np.ndarray` of type `object` and containing `Expr` objects.
+
+    Parameters
+    ----------
+    a : np.ndarray
+        A constant n-d `np.ndarray` of type `np.float64`.
+
+    b : np.ndarray
+        A n-d `np.ndarray` of type `object` and containing `Expr` objects.
+
+    Returns
+    -------
+    Expr or np.ndarray
+        If both `a` and `b` are 1-D arrays, return an `Expr`, otherwise return a
+        `np.ndarray` of type `object` and containing `Expr` objects.
+    """
+    cdef bool a_is_1d = a.ndim == 1
+    cdef bool b_is_1d = b.ndim == 1
+    cdef cnp.ndarray a_nd = a[..., np.newaxis, :] if a_is_1d else a
+    cdef cnp.ndarray b_nd = b[..., :, np.newaxis] if b_is_1d else b
+    cdef bool a_is_num = a_nd.dtype.kind in "fiub"
+
+    if a_is_num ^ (b_nd.dtype.kind in "fiub"):
+        res = _core_dot_2d(a_nd, b_nd) if a_is_num else _core_dot_2d(b_nd.T, a_nd.T).T
+        if a_is_1d and b_is_1d:
+            return res.item()
+        if a_is_1d:
+            return res.reshape(np.delete(res.shape, -2))
+        if b_is_1d:
+            return res.reshape(np.delete(res.shape, -1))
+        return res
+    return NotImplemented
+
+
+@np.vectorize(otypes=[object], signature="(m,n),(n,p)->(m,p)")
+def _core_dot_2d(cnp.ndarray a, cnp.ndarray x) -> np.ndarray:
+    """
+    Perform matrix multiplication between a 2-Demension constant array and a 2-Demension
+    `np.ndarray` of type `object` and containing `Expr` objects.
+
+    Parameters
+    ----------
+    a : np.ndarray
+        A 2-D `np.ndarray` of type `np.float64`.
+
+    x : np.ndarray
+        A 2-D `np.ndarray` of type `object` and containing `Expr` objects.
+
+    Returns
+    -------
+    np.ndarray
+        A 2-D `np.ndarray` of type `object` and containing `Expr` objects.
+    """
+    if not a.flags.c_contiguous or a.dtype != np.float64:
+        a = np.ascontiguousarray(a, dtype=np.float64)
+
+    cdef const double[:, :] a_view = a
+    cdef int m = a.shape[0], k = x.shape[1]
+    cdef cnp.ndarray[object, ndim=2] res = np.zeros((m, k), dtype=object)
+    cdef Py_ssize_t[:] nonzero
+    cdef int i, j, idx
+
+    for i in range(m):
+        if (nonzero := np.flatnonzero(a_view[i, :])).size == 0:
+            continue
+
+        for j in range(k):
+            res[i, j] = quicksum(a_view[i, idx] * x[idx, j] for idx in nonzero)
+
+    return res

tests/test_matrix_variable.py

@@ -293,6 +293,33 @@ def test_matrix_sum_axis_not_none_performance(n):
     assert model.isGT(end_orig - start_orig, end_matrix - start_matrix)
 
 
+@pytest.mark.parametrize("n", [50, 100])
+def test_matrix_dot_performance(n):
+    model = Model()
+    x = model.addMatrixVar((n, n))
+    a = np.random.rand(n, n)
+
+    start = time()
+    a @ x.view(np.ndarray)
+    orig = time() - start
+
+    start = time()
+    a @ x
+    matrix = time() - start
+
+    assert model.isGT(orig, matrix)
+
+
+def test_matrix_dot_value():
+    model = Model()
+    x = model.addMatrixVar(3, lb=[1, 2, 3], ub=[1, 2, 3])
+    y = model.addMatrixVar((3, 2), lb=1, ub=1)
+    model.optimize()
+
+    assert model.getVal(np.ones(3) @ x) == 6
+    assert (model.getVal(np.ones((2, 2, 3)) @ y) == np.full((2, 2, 2), 3)).all()
+
+
 def test_add_cons_matrixVar():
     m = Model()
     matrix_variable = m.addMatrixVar(shape=(3, 3), vtype="B", name="A", obj=1)
@@ -574,7 +601,7 @@ def test_matrix_matmul_return_type():
 
     # test 1D @ 1D → 0D
     x = m.addMatrixVar(3)
-    assert type(x @ x) is MatrixExpr
+    assert type(np.ones(3) @ x) is Expr
 
     # test 1D @ 1D → 2D
     assert type(x[:, None] @ x[None, :]) is MatrixExpr
@@ -584,6 +611,9 @@ def test_matrix_matmul_return_type():
     z = m.addMatrixVar((3, 4))
     assert type(y @ z) is MatrixExpr
 
+    # test ND @ 2D → ND
+    assert type(np.ones((2, 4, 3)) @ z) is MatrixExpr
+
 
 def test_matrix_sum_return_type():
     # test #1117, require returning type is MatrixExpr not MatrixVariable