SciPy: element-wise non-negative least squares using an array of b vectors

Question

I need to solve the linear problem Ax = b, obtaining x using a least squares approach. All elements of x must be non-negative, so I am using scipy.optimize.nnls (documentation here).

The trouble is, I need to solve this problem many times with a single A matrix and many b vectors. I have a 3d numpy ndarray in which vectors along axis 0 are the b vectors, while the other two axes correspond to points in space. I wish to output all x vectors to a corresponding array such that the spatial information for each answer is preserved.

A first pass at the problem looks like this:

A = np.random.rand(5,3)
b_array = B = np.random.rand(5,100,100)
x_array = np.zeros((3,100,100))

for i in range(100):
    for j in range(100):
        x_array[:,i,j] = sp.optimize.nnls(A, b_array[:,i,j])[0]

This code is perfectly functional, but feels completely inelegant. More importantly, it will likely be prohibitively slow (my actual code uses very large datasets and is looped thousands of times with random parameter changes, so efficiency is important).

A little while back I asked this very similar question about element-wise matrix multiplication. I was introduced to np.einsum, which has proven extremely useful in many situations. I had hoped that there would be a similar function for least-squares solutions, but have been unable to find anything. If anyone knows of a function that might work, or an alternate approach to solve this problem efficiently/pythonically, that would be much appreciated!

Answer 1

NNLS does not have a closed-form solution, and apart from sharing memory for the design matrix, there is no algorithmic speedup to be gained from treating the problems together. Although pushing down multi-target capability to C level can cause some speedups, it looks as though the scipy implementation only supports one target at a time, so looping looks like the only option here. The problem is embarassingly parallel, so you can use e.g. joblib to parallelize the loop as follows

from joblib import Parallel, delayed
from itertools import product
from scipy.optimize import nnls
results = Parallel(n_jobs=10)(delayed(nnls)(A, b_array[:,i,j])[0]
             for i, j in product(range(100), range(100)))
x_array = np.array(results).reshape(100, 100, -1).transpose(2, 0, 1)

However, if you use e.g. ridge regression or OLS (maybe not useful in your case), then the solution is closed form obtainable by matrix multiplication and everything can be done in one reshape and matrix multiplication, pushing the multi-target aspect of your problem down to C-level treatment.