Speed up MatrixExpr.sum(axis=...) via quicksum

CHANGELOG.md

@@ -2,8 +2,10 @@
 
 ## Unreleased
 ### Added
+- Speed up MatrixExpr.sum(axis=...) via quicksum
 ### Fixed
 - all fundamental callbacks now raise an error if not implemented
+- Fixed the type of MatrixExpr.sum(axis=...) result from MatrixVariable to MatrixExpr.
 ### Changed
 - changed default value of enablepricing flag to True
 ### Removed

src/pyscipopt/matrix.pxi

@@ -3,8 +3,14 @@
 # TODO Add tests
 """
 
+from typing import Optional, Tuple, Union
 import numpy as np
-from typing import Union
+try:
+    # NumPy 2.x location
+    from numpy.lib.array_utils import normalize_axis_tuple
+except ImportError:
+    # Fallback for NumPy 1.x
+    from numpy.core.numeric import normalize_axis_tuple
 
 
 def _is_number(e):
@@ -44,16 +50,61 @@ def _matrixexpr_richcmp(self, other, op):
 
 
 class MatrixExpr(np.ndarray):
-    def sum(self, **kwargs):
-        """
-        Based on `numpy.ndarray.sum`, but returns a scalar if `axis=None`.
-        This is useful for matrix expressions to compare with a matrix or a scalar.
+
+    def sum(
+        self,
+        axis: Optional[Union[int, Tuple[int, ...]]] = None,
+        keepdims: bool = False,
+        **kwargs,
+    ) -> Union[Expr, MatrixExpr]:
         """
+        Return the sum of the array elements over the given axis.
+
+        Parameters
+        ----------
+        axis : None or int or tuple of ints, optional
+            Axis or axes along which a sum is performed. The default, axis=None, will
+            sum all of the elements of the input array. If axis is negative it counts
+            from the last to the first axis. If axis is a tuple of ints, a sum is
+            performed on all of the axes specified in the tuple instead of a single axis
+            or all the axes as before.
+
+        keepdims : bool, optional
+            If this is set to True, the axes which are reduced are left in the result as
+            dimensions with size one. With this option, the result will broadcast
+            correctly against the input array.
+
+        **kwargs : ignored
+            Additional keyword arguments are ignored. They exist for compatibility
+            with `numpy.ndarray.sum`.
+
+        Returns
+        -------
+        Expr or MatrixExpr
+            If the sum is performed over all axes, return an Expr, otherwise return
+            a MatrixExpr.
 
-        if kwargs.get("axis") is None:
-            # Speed up `.sum()` #1070
-            return quicksum(self.flat)
-        return super().sum(**kwargs)
+        """
+        axis: Tuple[int, ...] = normalize_axis_tuple(
+            range(self.ndim) if axis is None else axis, self.ndim
+        )
+        if len(axis) == self.ndim:
+            res = quicksum(self.flat)
+            return (
+                np.array([res], dtype=object).reshape([1] * self.ndim).view(MatrixExpr)
+                if keepdims
+                else res
+            )
+
+        keep_axes = tuple(i for i in range(self.ndim) if i not in axis)
+        shape = (
+            tuple(1 if i in axis else self.shape[i] for i in range(self.ndim))
+            if keepdims
+            else tuple(self.shape[i] for i in keep_axes)
+        )
+        return np.apply_along_axis(
+            quicksum, -1, self.transpose(keep_axes + axis).reshape(shape + (-1,))
+        ).view(MatrixExpr)
 
     def __le__(self, other: Union[float, int, "Expr", np.ndarray, "MatrixExpr"]) -> MatrixExprCons:
         return _matrixexpr_richcmp(self, other, 1)

tests/test_matrix_variable.py

@@ -19,7 +19,7 @@
     sin,
     sqrt,
 )
-from pyscipopt.scip import GenExpr
+from pyscipopt.scip import CONST, GenExpr
 
 
 def test_catching_errors():
@@ -181,7 +181,30 @@ def test_expr_from_matrix_vars():
         for term, coeff in expr_list:
             assert len(term) == 3
 
-def test_matrix_sum_argument():
+
+def test_matrix_sum_error():
+    m = Model()
+    x = m.addMatrixVar((2, 3), "x", "I", ub=4)
+
+    # test axis type
+    with pytest.raises(TypeError):
+        x.sum("0")
+
+    # test axis value (out of range)
+    with pytest.raises(ValueError):
+        x.sum(2)
+
+    # test axis value (out of range)
+    with pytest.raises(ValueError):
+        x.sum((-3,))
+
+    # test axis value (duplicate)
+    with pytest.raises(ValueError):
+        x.sum((0, 0))
+
+
+def test_matrix_sum_axis():
+    # compare the result of summing matrix variable after optimization
     m = Model()
 
     # Return a array when axis isn't None
@@ -190,29 +213,52 @@ def test_matrix_sum_argument():
 
     # compare the result of summing 2d array to a scalar with a scalar
     x = m.addMatrixVar((2, 3), "x", "I", ub=4)
-    m.addMatrixCons(x.sum() == 24)
+    # `axis=tuple(range(x.ndim))` is `axis=None`
+    m.addMatrixCons(x.sum(axis=tuple(range(x.ndim))) == 24)
 
     # compare the result of summing 2d array to 1d array
     y = m.addMatrixVar((2, 4), "y", "I", ub=4)
     m.addMatrixCons(x.sum(axis=1) == y.sum(axis=1))
 
     # compare the result of summing 3d array to a 2d array with a 2d array
     z = m.addMatrixVar((2, 3, 4), "z", "I", ub=4)
-    m.addMatrixCons(z.sum(axis=2) == x)
+    m.addMatrixCons(z.sum(2) == x)
     m.addMatrixCons(z.sum(axis=1) == y)
 
     # to fix the element values
     m.addMatrixCons(z == np.ones((2, 3, 4)))
 
-    m.setObjective(x.sum() + y.sum() + z.sum(), "maximize")
+    m.setObjective(x.sum() + y.sum() + z.sum(tuple(range(z.ndim))), "maximize")
     m.optimize()
 
     assert (m.getVal(x) == np.full((2, 3), 4)).all().all()
     assert (m.getVal(y) == np.full((2, 4), 3)).all().all()
 
 
-@pytest.mark.parametrize("n", [50, 100, 200])
-def test_sum_performance(n):
+@pytest.mark.parametrize(
+    "axis, keepdims",
+    [
+        (0, False),
+        (0, True),
+        (1, False),
+        (1, True),
+        ((0, 2), False),
+        ((0, 2), True),
+    ],
+)
+def test_matrix_sum_result(axis, keepdims):
+    # directly compare the result of np.sum and MatrixExpr.sum
+    _getVal = np.vectorize(lambda e: e.terms[CONST])
+    a = np.arange(6).reshape((1, 2, 3))
+
+    np_res = a.sum(axis, keepdims=keepdims)
+    scip_res = MatrixExpr.sum(a, axis, keepdims=keepdims)
+    assert (np_res == _getVal(scip_res)).all()
+    assert np_res.shape == _getVal(scip_res).shape
+
+
+@pytest.mark.parametrize("n", [50, 100])
+def test_matrix_sum_axis_is_none_performance(n):
     model = Model()
     x = model.addMatrixVar((n, n))
 
@@ -229,6 +275,24 @@ def test_sum_performance(n):
     assert model.isGT(end_orig - start_orig, end_matrix - start_matrix)
 
 
+@pytest.mark.parametrize("n", [50, 100])
+def test_matrix_sum_axis_not_none_performance(n):
+    model = Model()
+    x = model.addMatrixVar((n, n))
+
+    # Original sum via `np.ndarray.sum`, `np.sum` will call subclass method
+    start_orig = time()
+    np.ndarray.sum(x, axis=0)
+    end_orig = time()
+
+    # Optimized sum via `quicksum`
+    start_matrix = time()
+    x.sum(axis=0)
+    end_matrix = time()
+
+    assert model.isGT(end_orig - start_orig, end_matrix - start_matrix)
+
+
 def test_add_cons_matrixVar():
     m = Model()
     matrix_variable = m.addMatrixVar(shape=(3, 3), vtype="B", name="A", obj=1)
@@ -521,6 +585,14 @@ def test_matrix_matmul_return_type():
     assert type(y @ z) is MatrixExpr
 
 
+def test_matrix_sum_return_type():
+    # test #1117, require returning type is MatrixExpr not MatrixVariable
+    m = Model()
+
+    x = m.addMatrixVar((3, 2))
+    assert type(x.sum(axis=1)) is MatrixExpr
+
+
 def test_broadcast():
     # test #1065
     m = Model()