Feature: Speed up MatrixVariable.sum(axis=None) via quicksum (#1078)

* Speed up `matrix.sum`

* Add test for sum with axis in matrix variable

Added a test to verify that summing a matrix variable with axis=0 returns a MatrixExpr with the expected shape. This improves coverage for matrix sum operations with axis specified.

* Add matrix sum performance test

Introduced a new test to compare performance of matrix sum versus element-wise sum. Refactored imports for clarity and consistency. Renamed performance test for better description.

* Use a more large size data

* Fix variable naming in test_sum_performance

Renamed the Model instance from 'm' to 'model' to avoid confusion with the integer variable 'm' and improve code clarity in the test_sum_performance function.

* Fix variable name in performance test assertion

Replaces incorrect usage of 'm' with 'model' in the assertion within test_sum_performance to ensure the correct object is referenced.

* Fix expected shape in matrix sum test

Updated the assertion in test_matrix_sum_argument to expect shape (1,) instead of (1, 1) when summing along axis 0. This aligns the test with the actual output of the sum operation.

* Update CHANGELOG.md

* Adjust performance test assertion in matrix variable tests

Modified the assertion in test_sum_performance to compare orig_time + 1 with matrix_time instead of orig_time. This may address timing precision or test flakiness.

* Try a bigger data size

* Compare `np.sum` and `quicksum`

* Try a bigger size

---------

Co-authored-by: João Dionísio <[email protected]>

Author: Zeroto521

CHANGELOG.md

@@ -6,6 +6,7 @@
 ### Fixed
 - Fixed the type of @ matrix operation result from MatrixVariable to MatrixExpr.
 ### Changed
+- Speed up MatrixVariable.sum(axis=None) via quicksum
 ### Removed
 
 ## v5.6.0 - 2025.08.26

src/pyscipopt/matrix.pxi

@@ -18,11 +18,14 @@ def _is_number(e):
 class MatrixExpr(np.ndarray):
     def sum(self, **kwargs):
         """
-        Based on `numpy.ndarray.sum`, but returns a scalar if the result is a single value.
-        This is useful for matrix expressions where the sum might reduce to a single value.
+        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.
         """
-        res = super().sum(**kwargs)
-        return res if res.size > 1 else res.item()
+
+        if kwargs.get("axis") is None:
+            # Speed up `.sum()` #1070
+            return quicksum(self.flat)
+        return super().sum(**kwargs)
 
     def __le__(self, other: Union[float, int, Variable, np.ndarray, 'MatrixExpr']) -> np.ndarray:
 

tests/test_matrix_variable.py

@@ -1,11 +1,24 @@
-import pdb
-import pprint
-import pytest
-from pyscipopt import Model, Variable, log, exp, cos, sin, sqrt
-from pyscipopt import Expr, MatrixExpr, MatrixVariable, MatrixExprCons, MatrixConstraint, ExprCons
 from time import time
 
 import numpy as np
+import pytest
+
+from pyscipopt import (
+    Expr,
+    ExprCons,
+    MatrixConstraint,
+    MatrixExpr,
+    MatrixExprCons,
+    MatrixVariable,
+    Model,
+    Variable,
+    cos,
+    exp,
+    log,
+    quicksum,
+    sin,
+    sqrt,
+)
 
 
 def test_catching_errors():
@@ -170,6 +183,10 @@ def test_expr_from_matrix_vars():
 def test_matrix_sum_argument():
     m = Model()
 
+    # Return a array when axis isn't None
+    res = m.addMatrixVar((3, 1)).sum(axis=0)
+    assert isinstance(res, MatrixExpr) and res.shape == (1,)
+
     # 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)
@@ -192,6 +209,25 @@ def test_matrix_sum_argument():
     assert (m.getVal(x) == np.full((2, 3), 4)).all().all()
     assert (m.getVal(y) == np.full((2, 4), 3)).all().all()
 
+
+def test_sum_performance():
+    n = 1000
+    model = Model()
+    x = model.addMatrixVar((n, n))
+
+    # Original sum via `np.sum`
+    start_orig = time()
+    np.sum(x)
+    end_orig = time()
+
+    # Optimized sum via `quicksum`
+    start_matrix = time()
+    x.sum()
+    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)
@@ -339,7 +375,7 @@ def test_MatrixVariable_attributes():
     assert x.varMayRound().tolist() == [[True, True], [True, True]]
 
 @pytest.mark.skip(reason="Performance test")
-def test_performance():
+def test_add_cons_performance():
     start_orig = time()
     m = Model()
     x = {}