Skip to content

Speed up np.mean(MatrixExpr) and more via MatrixExpr.sum #1153

@Zeroto521

Description

@Zeroto521

Is your feature request related to a problem? Please describe.

Since there is already speed up MatrixExpr.sum
Now we can speed up other high-level operations, such as np.mean, np.trace, and part of np.dot.

Describe the solution you'd like

np.mean will call np.reduce.add in inner, and np.reduce.add will use __add__.
Use numpy __array_ufunc__ protocol to call MatrixExpr.sum instead of np.ndarray.sum.

Additional context

It sppeds up 834x than before in the shape of (200, 200).

from time import time

import numpy as np
from pyscipopt import MatrixVariable, Model


class SpeedMatrix(MatrixVariable):
    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
        if method == "reduce" and ufunc is np.add:
            return args[0].sum(**kwargs)
        return super().__array_ufunc__(ufunc, method, *args, **kwargs)


if __name__ == "__main__":
    model = Model()

    n = 200
    x = model.addMatrixVar((n, n))
    x_optimized = x.view(SpeedMatrix)

    start = time()
    np.mean(x)
    end = time()
    print(f"Standard mean time: {end - start:.6f} seconds")
    # Standard mean time: 16.185321 seconds

    start = time()
    np.mean(x_optimized)
    end = time()
    print(f"Optimized mean time: {end - start:.6f} seconds")
    # Optimized mean time: 0.019406 seconds

    start = time()
    np.sum(x) / x.size
    end = time()
    print(f"Standard mean (sum/size) time: {end - start:.6f} seconds")
    # Standard mean (sum/size) time: 0.017425 seconds

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions