Compute inverse of 2D arrays along the third axis in a 3D array without loops

Question

I have an array A whose shape is (N, N, K) and I would like to compute another array B with the same shape where B[:, :, i] = np.linalg.inv(A[:, :, i]).

As solutions, I see map and for loops but I am wondering if numpy provides a function to do this (I have tried np.apply_over_axes but it seems that it can only handle 1D array).

with a for loop:

B = np.zeros(shape=A.shape)
for i in range(A.shape[2]):
    B[:, :, i] = np.linalg.inv(A[:, :, i])

with map:

B = np.asarray(map(np.linalg.inv, np.squeeze(np.dsplit(A, A.shape[2])))).transpose(1, 2, 0)

Answer 1

For an invertible matrix M we have inv(M).T == inv(M.T) (the transpose of the inverse is equal to the inverse of the transpose).

Since np.linalg.inv is broadcastable, your problem can be solved by simply transposing A, calling inv and transposing the result:

B = np.linalg.inv(A.T).T

For example:

>>> N, K = 2, 3
>>> A = np.random.randint(1, 5, (N, N, K))
>>> A
array([[[4, 2, 3],
        [2, 3, 1]],

       [[3, 3, 4],
        [4, 4, 4]]])

>>> B = np.linalg.inv(A.T).T
>>> B
array([[[ 0.4  , -4.   ,  0.5  ],
        [-0.2  ,  3.   , -0.125]],

       [[-0.3  ,  3.   , -0.5  ],
        [ 0.4  , -2.   ,  0.375]]])

You can check the values of B match the inverses of the arrays in A as expected:

>>> all(np.allclose(B[:, :, i], np.linalg.inv(A[:, :, i])) for i in range(K))
True