Flatten numpy array with list comprehension, alternative solutions-CodePudding

I have an array a with a.shape = (14, 3200, 2500)

I'm trying to extract non nan values from it to a (-1, 14) array.

This is my current code

nans = ~(np.isnan(a)).any(axis=0)
indices = np.where(nans)
vals = np.asarray([a[:, i, j] for i, j in zip(indices[0], indices[1])])

But I think I should be changing the list comprehension for a numpy built-in function but cant find a function that does this as this doesn't seem to be that fast. Anyone has a good suggestions?

Edit: I'm also trying to get the indices where these values are so my outputs should be arrays with shapes (-1, 14), (-1, 2). The second array is made with

np.stack((indices[0],indices[1])).transpose()

So the list comprehension should preserve the order

CodePudding user response：

Using numpy indexing followed by a .T (transpose) solves the problem:

vals = a[:, indices[0], indices[1]].T

It can be checked that it returns the same array as your current approach:

nans = ~(np.isnan(a)).any(axis=0)
indices = np.where(nans)
vals = np.asarray([a[:, i, j] for i, j in zip(indices[0], indices[1])])

# changed the name for comparison purposes
vals2 = a[:, indices[0], indices[1]].T
print(np.isclose(vals, vals2).all())

Output

True

Note

The array a, (used for testing), was generated by:

a = np.random.random((14, 3200, 2500))
a[a < 0.3] = np.nan