Get last x non-nan elements from numpy array

user_12819237 0
Viewed 439

I have a 2D numpy array with random nan values scattered throughout. My goal now is to get only the last x non-nan elements along dimension 0 from this array.

import numpy as np
np.random.seed(0)
M, N, c = 10, 5, 15
A = np.random.randn(M,N)
A.ravel()[np.random.choice(A.size, c, replace=False)] = np.nan

A
[[ 1.76405235         nan  0.97873798  2.2408932          nan]
 [-0.97727788         nan         nan         nan  0.4105985 ]
 [ 0.14404357  1.45427351         nan  0.12167502  0.44386323]
 [ 0.33367433         nan -0.20515826  0.3130677  -0.85409574]
 [        nan  0.6536186   0.8644362  -0.74216502  2.26975462]
 [        nan  0.04575852         nan  1.53277921  1.46935877]
 [ 0.15494743  0.37816252 -0.88778575         nan -0.34791215]
 [ 0.15634897  1.23029068  1.20237985 -0.38732682         nan]
 [        nan -1.42001794 -1.70627019         nan -0.50965218]
 [-0.4380743  -1.25279536  0.77749036         nan -0.21274028]]

So in this example, for x=3 in the end I would like an array the size of 3x5 containing only the last 3 non nan values. Like this:

[[ 0.15494743  1.23029068  1.20237985 -0.74216502 -0.34791215]
 [ 0.15634897 -1.42001794 -1.70627019  1.53277921 -0.50965218]
 [-0.4380743  -1.25279536  0.77749036 -0.38732682 -0.21274028]]

I know that I can use idx = np.where(~np.isnan(A)) to get the indices of non-nan values but I am not sure on how to continue from there.

Edit: This can easily be done using a for loop, however I am ideally looking for a vectorized solution.

numpy arrays python
1 Answers

So I figured out one way to do it. I don't know if it's the fastest but it is vectorized, so it shouldn't be too slow. If anyone has a solution that might be faster I'm definitely interested.

def last_x_non_nan(a, x):
    # end points
    non_count = np.cumsum(np.count_nonzero(~np.isnan(a), axis=0))
    rel_idx = np.arange(-x, 0)
    abs_idx = (rel_idx[:,None] + non_count).ravel('F')

    # non-nan array ordered in column-major
    arr_raveled = a.ravel('F')
    arr_clear = arr_raveled[~np.isnan(arr_raveled)]

    # Subset and convert to desired shape
    return arr_clear[abs_idx].reshape((x, a.shape[1]), order='F')

This returns the expected output:

print(last_x_non_nan(A, 3))
[[ 0.15494743  1.23029068  1.20237985 -0.74216502 -0.34791215]
 [ 0.15634897 -1.42001794 -1.70627019  1.53277921 -0.50965218]
 [-0.4380743  -1.25279536  0.77749036 -0.38732682 -0.21274028]]
user_128192370
Related