Speed up Constant @ MatrixExpr

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