Feature: Matrix sum method supports arguments inputting (#1025)

* Remove custom sum method from MatrixExpr

Deleted the overridden sum method in MatrixExpr to rely on the default numpy ndarray behavior.

* Revert "Remove custom sum method from MatrixExpr"

This reverts commit 72c42f4.

* Improve MatrixExpr.sum to handle scalar results

Updated the MatrixExpr.sum method to return a scalar when the result is a single value, preserving array output otherwise. This enhances usability for matrix expressions that may reduce to a single value.

* Add test for matrix sum argument in Model

* Update CHANGELOG for MatrixExpr.sum() change

* add case to test 3d array

* Fix axis argument in matrix sum test

Replaces positional argument with explicit 'axis=2' in y.sum() to improve clarity and correctness in test_matrix_sum_argument.

* Correct method name

* Add more complex test cases

* correct function name

* Fix test to constrain matrix variable values

* Fix z value and get x and y values

* Update tests/test_matrix_variable.py

* Update CHANGELOG.md

Co-authored-by: Copilot <[email protected]>

* typo

---------

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

Author: 3 people

CHANGELOG.md

@@ -10,6 +10,7 @@
 ### Fixed
 - Raised an error when an expression is used when a variable is required
 ### Changed
+- MatrixExpr.sum() now supports axis arguments and can return either a scalar or MatrixExpr depending on the result dimensions
 ### Removed
 
 ## 5.5.0 - 2025.05.06

src/pyscipopt/matrix.pxi

@@ -17,8 +17,13 @@ def _is_number(e):
 
 class MatrixExpr(np.ndarray):
     def sum(self, **kwargs):
-        return super().sum(**kwargs).item()
-    
+        """
+        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.
+        """
+        res = super().sum(**kwargs)
+        return res if res.size > 1 else res.item()
+
     def __le__(self, other: Union[float, int, Variable, np.ndarray, 'MatrixExpr']) -> np.ndarray:
 
         expr_cons_matrix = np.empty(self.shape, dtype=object)

tests/test_matrix_variable.py

@@ -164,6 +164,30 @@ def test_expr_from_matrix_vars():
         for term, coeff in expr_list:
             assert len(term) == 3
 
+def test_matrix_sum_argument():
+    m = Model()
+
+    # compare the result of summing 2d array to a scalar with a scalar
+    x = m.addMatrixVar((2, 3), "x", "I", ub=4)
+    m.addCons(x.sum() == 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(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.optimize()
+
+    assert (m.getVal(x) == np.full((2, 3), 4)).all().all()
+    assert (m.getVal(y) == np.full((2, 4), 3)).all().all()
 
 def test_add_cons_matrixVar():
     m = Model()