How can I vectorize this linalg.lstq() operation?

Question

I am trying to implement a multi-frequency phase unwrapping algorithm using Python3 and NumPy. I have 7 single channel (gray scale) images of shape (1080, 1920). After stacking them along the third axis I get (1080, 1920, 7).

I have shift-matrix A (which is fixed for all pixels) of shape (7, 7). For all pixels I have different intensity values (array r of shape (1, 7)).

To solve them for each pixel by minimizing L2-norm as ||r - Au||, I can do:

# Just for illustration
A = np.random.randn(7, 7)
r = np.random.randn(7, 1)

# Solved for each pixel location
u = np.linalg.lstsq(a = A, b = r, rcond = None)

This can be implemented using 2 for loops:

for y in range(0, image_height - 1):
    for x in range(0, image_width - 1):
        # code here

This is inefficient. How can I write it as efficient NumPy code?

Answer 1

IIUC you have one A matrix and 1080*1920 RHSs. NumPy's lstsq function is vectorized for that case.

So we reshape your img with shape (1080,1920, 7) to

import numpy as np
rng = np.random.default_rng(234892359843)
img = rng.random((1080, 1920, 7))
A = rng.random((7, 7))
b = img.reshape((-1, 7)).T  # reshape to (7, 1080*1920)
x, _, _, _ = np.linalg.lstsq(A, b, rcond=None)
x = (x.T).reshape(img.shape)  # return to original shape

Then for any valid indices i, j:

np.testing.assert_allclose(x[i, j], np.linalg.lstsq(A, img[i, j], rcond=None)[0])

I'm assuming your matrix A is singular; otherwise you could use the same code with solve.

Then this should hold:

np.testing.assert_allclose(A @ x[i, j], img[i, j])