Delete all zeros slices from 4d numpy array

I pretend to remove slices from the third dimension of a 4d numpy array if it’s contains only zeros.

I have a 4d numpy array of dimensions [256,256,336,6] and I need to delete the slices in the third dimension that only contains zeros. So the result would have a shape like this , e.g. [256,256,300,6] if 36 slices are fully zeros. I have tried multiple approaches including for loops, np.delete and all(), any() functions without success.

Answer

I’m not an afficionado with numpy, but does this do what you want?

I take the following small example matrix with 4 dimensions all full of 1s and then I set some slices to zero:

import numpy as np
a=np.ones((4,4,5,2))

The shape of a is:

>>> a.shape
(4, 4, 5, 2)

I will artificially set some of the slices in dimension 3 to zero:

a[:,:,0,:]=0
a[:,:,3,:]=0

I can find the indices of the slices with not all zeros by calculating sums (not very efficient, perhaps!)

indices = [i for i in range(a.shape[2]) if a[:,:,i,:].sum() != 0]
>>> indices
[1, 2, 4]

So, in your general case you could do this:

indices = [i for i in range(a.shape[2]) if a[:,:,i,:].sum() != 0]
a_new = a[:, :, indices, :].copy()

Then the shape of a_new is:

>>> anew.shape
(4, 4, 3, 2)