Numpy : matrix of vectors, inversion

Question

I have a matrix A of dimensions N x N, each matrix element is a vector of size M. I want to inverse A matrix. In other terms I want to compute A^-1 that is the inverted A matrix composed of NxN vectors of size M.

Here is code to implement what I want to do, I'm just computing M times inverse matrix to compute C = A^-1 x B and then I'm checking A x C = B. But each time I'm iterating over the M element of A, B elements to do a matrix inversion. I'm quite sure my code does what I need but not in a clever way...

a = np.array([[[3, 4, 8], [1,8,3]],
     [[2, 1, 2], [6, 5, 0]]])

b = np.array([[2, 0, 6],
     [5, 2, 5]])

c = []
# compute c = a^-1 x b
for i in range(a.shape[-1]):
    c.append(np.linalg.inv(a[:,:,i])@b[:,i])
c = np.asarray(c)

# check inversion compute a x c and checks a x c = b 
for i in range(a.shape[-1]):
    if not np.allclose(a[:,:,i]@c[i,:], b[:,i]):
        raise Exception('Inversion ko')
        break
print('inversion ok')

I need a mix of matrix operations and element wise operations. But I do not like my implementations. I'm quite a straightforward implementation with less code exists. Let me know your suggestions.

Answer 1

We can use np.linalg.inv on the 3D array a after pushing the last axis to the front. Then, we can leverage einsum for the final output for a vectorized way -

p = np.linalg.inv(a.transpose(2,0,1))
c = np.einsum('ijk,kli->ij',p,b)[...,None]

A simpler alternative to get the final output c would be with np.matmul/@-operator -

c = p@b.transpose(2,0,1)

Thus, the whole process could be converted to a one-liner -

c = np.linalg.inv(a.transpose(2,0,1))@b.transpose(2,0,1)

Answer 2

If you stacked your matrices along the first dimension, instead of the last, you can invert them all in one single np.linalg.pinv() call. Afterwards use np.einsum() to do all dot products in one go:

a = a.transpose(2, 0, 1)
b = b.transpose()

np.einsum('fmn,fm->fn', np.linalg.inv(a), b)