Search code examples
numpyranking

Ignore nan values when ranking using argsort

I have a below function that ranks 3rd column of 2D arrays inside of a three-dimensional array.

arr[:, :, 3] = arr[:, :, 3].argsort(axis=1)[:, ::-1].argsort(axis=1) + 1

The problem is that nan values are also being ranked which should not happen. Is there a way to skip/ignore nan values when ranking using argsort? if not, what other numpy methods could be used?

Thanks.

Solution

  • I found a workaround as per below:

    # Replace NaNs with Negative Infinities
np.nan_to_num(arr, copy=False, nan=-np.inf)
# Rank 
arr[:, :, 3] = arr[:, :, 3].argsort(axis=1)[:, ::-1].argsort(axis=1) + 1
# Return NaNs back if needed (replace negative infinities with NaNs)
arr[arr== -np.inf] = np.nan