Search code examples
pythonfunctionnumpysub-array

How to efficiently operate on sub-arrays like calculating the determinants, inverse,


I have to to multiple operations on sub-arrays like matrix inversions or building determinants. Since for-loops are not very fast in Python I wonder what is the best way to do this.

import numpy as np
n = 8
a = np.random.rand(3,3,n)
b = np.empty(n)
c = np.zeros_like(a)

for i in range(n):
    b[i] = np.linalg.det(a[:,:,i])
    c[:,:,i] = np.linalg.inv(a[:,:,i])

Solution

  • Those numpy.linalg functions would accept n-dim arrays as long as the last two axes are the ones that form the 2D slices along which functions are intended to be operated upon. Hence, to solve our cases, permute axes to bring-up the axis of iteration as the first one, perform the required operation and if needed push-back that axis back to it's original place.

    Hence, we could get those outputs, like so -

    b = np.linalg.det(np.moveaxis(a,2,0))
    c = np.moveaxis(np.linalg.inv(np.moveaxis(a,2,0)),0,2)