I want to take the SVD of an image (treating each of the four RGBa channels separately), and then reconstruct the image using only the first N vectors. I have the code below, which works as expected, but I would ideally like to write the reconstruction without using np.stack. I'd rather like to use some clever broadcasting trick to achieve the same result.
import matplotlib.image as mpimg
from matplotlib import pyplot as plt
# Load image (python logo from python.org)
img = mpimg.imread('https://www.python.org/static/img/python-logo@2x.png')
# img has shape (164, 580, 4)
# Take SVD, using transpose to get the 4 as
# the first dimension, which makes np.linalg.svd
# treat it as 4 separate 580 x 164 matrices
U, S, V = np.linalg.svd(img.T)
# U has shape (4, 580, 580)
# S has shape (4, 164)
# V has shape (4, 164, 164)
# Reconstruct image with the first N vectors,
# using np.stack(..., axis = -1) to get the 4
# as the last dimension, as in the original image
N = 3
img_reconstructed = np.stack((
np.dot(U[0,:,:N], S[0,:N,None]*V[0,:N,:]).T, # Red channel
np.dot(U[1,:,:N], S[1,:N,None]*V[1,:N,:]).T, # Green channel
np.dot(U[2,:,:N], S[2,:N,None]*V[2,:N,:]).T, # Blue channel
np.dot(U[3,:,:N], S[3,:N,None]*V[3,:N,:]).T, # alpha channel
), axis = -1)
# img_reconstructed has shape (164, 580, 4),
# just like the original image
plt.imshow(img_reconstructed)
I thought maybe I could achieve the same with np.tensordot, but I couldn't get it to work. I tried
img_reconstructed = np.tensordot(U[:,:,:n], S[:,:n,None]*V[:,:n,:], axes = ((2), (1))).T
but it comes out with shape (164, 4, 580, 4), instead of (164, 580, 4) as I had hoped.
I managed to get a working solution with np.einsum:
img_reconstructed = np.einsum('...ij,...jk->ki...', U[:,:,:n], S[:,:n,None]*V[:,:n,:])
but when I time it, I find that it's about 25 times slower than the solution above with np.stack.