diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ab002ee6..5a94d66b0 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/pyscipopt/matrix.pxi b/src/pyscipopt/matrix.pxi index 8353ed767..1a6a09cf3 100644 --- a/src/pyscipopt/matrix.pxi +++ b/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) diff --git a/tests/test_matrix_variable.py b/tests/test_matrix_variable.py index 27f549000..e4758f077 100644 --- a/tests/test_matrix_variable.py +++ b/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,7 +213,8 @@ 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) @@ -198,21 +222,43 @@ def test_matrix_sum_argument(): # 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()