Flatten numpy array with list comprehension, alternative solutions

Question

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

Does flattening preserve the order? I'm also trying to get coordinates where these values area in the array. I'll add an edit — MesRot
– MesRot, Commented Jul 21, 2022 at 7:18

Dani Mesejo · Accepted Answer · 2022-07-21 07:25:31Z

2

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

edited Jul 21, 2022 at 7:25

answered Jul 21, 2022 at 7:19

Dani Mesejo

62.2k6 gold badges56 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

MesRot Over a year ago

Works like a charm for numpy array. Anyone stumbling over this solution take a note that this does not work for xarray.DataArray and you have to convert to numpy first

Collectives™ on Stack Overflow

Flatten numpy array with list comprehension, alternative solutions

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related