[Python] How can I optimize the performance of this numpy function

Is there any way optimizing the performance speed of this function?

def func(X):
n, p = X.shape
R = np.eye(p)
delta = 0.0
for i in range(100):
delta_old = delta
Y = X @ R

alpha = 1. / n
Y2 = Y**2
Y3 = Y2 * Y
W = np.sum(Y2, axis=0)

transformed = X.T @ (Y3 - (alpha * Y * W))
U, svals, VT = np.linalg.svd(transformed, full_matrices=False)
R = U @ VT # is used as a stopping criterion
delta = np.sum(svals)

return R


Naively, I thought using numba would help because of the loop (the actual number of loops is higher),

from numba import jit

@jit(nopython=True, parallel=True)
def func_numba(X):
n, p = X.shape
R = np.eye(p)
delta = 0.0
for i in range(100):
delta_old = delta
Y = X @ R

alpha = 1. / n
Y2 = Y**2
Y3 = Y2 * Y
W = np.sum(Y2, axis=0)

transformed = X.T @ (Y3 - (alpha * Y * W))
U, svals, VT = np.linalg.svd(transformed, full_matrices=False)
R = U @ VT
delta = np.sum(svals) # is used as a stopping criterion

return R



but to my surprise the numbaized function is actually slower. Why is numba not more effective in this case? Is there another option for me (preferably using numpy)?

Note: You can assume X to be a "tall-and-skinny" matrix.

MWE


import numpy as np

size = (10_000, 15)
X = np.random.normal(size=size)
%timeit func(X) # 1.28 s
%timeit func_numba(X) # 2.05 s

Continue reading...