diff --git a/CHANGELOG.md b/CHANGELOG.md index 5468f187e..75adedf36 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/pyscipopt/matrix.pxi b/src/pyscipopt/matrix.pxi index 9eb0349f2..51225f032 100644 --- a/src/pyscipopt/matrix.pxi +++ b/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: diff --git a/tests/test_matrix_variable.py b/tests/test_matrix_variable.py index 36e6422bb..04f9d0fcc 100644 --- a/tests/test_matrix_variable.py +++ b/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 = {}