-
Notifications
You must be signed in to change notification settings - Fork 275
Open
Description
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 secondsMetadata
Metadata
Assignees
Labels
No labels