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

* Add test for matrix sum return type

Adds a test to ensure that summing a matrix variable along an axis returns a MatrixExpr type instead of MatrixVariable, addressing issue #1117.

* return `MatrixExpr` type

* Update CHANGELOG.md

* Speed up `MatrixExpr.sum(axis=...)`

Enhanced the MatrixExpr.sum method to accept axis as int or tuple, handle keepdims, and provide better error checking for axis bounds. This improves compatibility with numpy's sum behavior and allows more flexible summation over matrix expressions.

* Test `MatrixExpr.sum(axis=...)`

Removed n=200 from the sum performance test to limit test size. Added additional performance assertions comparing np.ndarray.sum and the optimized sum method for matrix variables.

* Test `MatrixExpr.sum(axis=tuple(range(ndim))`

Renamed test_matrix_sum_argument to test_matrix_sum_axis and updated tests to use explicit axis arguments in sum operations. This clarifies the behavior when summing over all axes and improves test coverage for axis handling.

* Refactor MatrixExpr.sum axis handling with numpy utility

Replaces manual axis validation and normalization in MatrixExpr.sum with numpy's normalize_axis_tuple for improved reliability and code clarity. Updates type hints and simplifies logic for summing across all axes.

* Add tests for matrix sum error

Added tests to verify error handling in matrix variable sum operations for invalid axis types, out-of-range values, and duplicate axes.

* Add tests for matrix sum with keepdims parameter

Introduces test cases to verify the shape of matrix variable sums when using the keepdims argument, ensuring correct behavior for both full and axis-specific summation.

* call `.sum` via positional argument

Replaces the use of the 'axis' keyword argument with a positional argument in the z.sum() method call to align with the expected function signature.

* Refactor sum method to use np.apply_along_axis

Replaces np.fromiter with np.apply_along_axis for summing along specified axes in MatrixExpr. This simplifies the code and improves readability.

* Directly test the `.sum` result

Added a comment to clarify the purpose of the test_matrix_sum_axis function, indicating it compares the result of summing a matrix variable after optimization.

* Update CHANGELOG.md

* Expand docstring for MatrixExpr.sum method

The docstring for the MatrixExpr.sum method was updated to provide detailed information about parameters, return values, and behavior, improving clarity and alignment with numpy conventions.

* Clarify MatrixExpr.sum docstring and note quicksum usage

Updated the docstring for MatrixExpr.sum to specify that it uses quicksum for speed optimization instead of numpy.ndarray.sum. Added a detailed note explaining the difference between quicksum (using __iadd__) and numpy's sum (using __add__).

* Split up two costing time test cases

Renamed the existing performance test to clarify it tests the case where axis is None. Added a new test to measure performance when summing along a specific axis.

* Supports numpy 1.x

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

* Supports Python 3.8

* Simplify a bit

Refactored the sum method in MatrixExpr to clarify axis typing and simplify the application of np.apply_along_axis. The new implementation improves readability and maintains the intended behavior.

* suggestion

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Joao-Dionisio <up201606210@up.pt>
Co-authored-by: João Dionísio <57299939+Joao-Dionisio@users.noreply.github.com>

Author: 3 people

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()