Optimizing numpy least squares by omitting loop and cast

Viewed 116

Currently I am trying to calculate least squares with 2 numpy arrays (X, Y) with n arrays with some same number of values in each. My output that I want is 2 numpy arrays that contain the slope and intercept respectively. Right now I have the following inefficient code:

M = []
C = []
for i in range(len(X)):
    x = X[i]
    y = Y[i]
    A = np.vstack([x, np.ones(len(x))]).T
    m, c = np.linalg.lstsq(A, y, rcond=None)[0]
    M.append(m)
    C.append(c)
return np.array(M), np.array(C)

Since this code relies on a couple conversions and a for loop I think there has to be a more efficient way to solve this problem, but I can't think of one that doesn't rely on the for loop and casts. I need to keep the order of M and C maintained as well.

2 Answers

Actually it should not be that bad, numpy is highly optimized for this kind on stuff (stacking matrices, casting...).

You cannot avoid the for-loop since np.linalg.lstsq does not support so-called "batch-processing", but it is possible to avoid memory allocation. I don't know if it will really improve anything:

M, C = [], []
for i, x, y in enumerate(zip(X, Y)):
    if not i:
        A = np.stack((x.T, np.ones(len(x))), axis=-1)
    else:
        A[:, :-1] = x.T
    m, c = np.linalg.lstsq(A, y, rcond=None)[0]
    M.append(m)
    C.append(c)
return np.array(M), np.array(C)

After thinking twice about it, there is a way to do it by "emulating" batch processing, by reformulating your problem as a block diagonal sparse problem:

import scipy.linalg

A = scipy.linalg.block_diag(*[
    np.stack((x.T, np.ones(len(x))), axis=-1) for x in X])
Y = np.concatenate(Y)
M, C = np.linalg.lstsq(A, y, rcond=None)[0]

Then, doing the inverse of block_diag for M you should be able to recover what you want. C should be fine already.

I haven't tried it myself but I think it would lead to a speedup thanks to efficient matrix algebra. You could probably even further improve efficiency by using spicy space lsqr solver in this case. Good luck with that !

Related