numpy 数字化中的屏蔽值

Masked values in numpy digitize

我希望 numpy digitize 忽略数组中的某些值。为此,我将不需要的值替换为 NaN 并屏蔽了 NaN 值:

import numpy as np
A = np.ma.array(A, mask=np.isnan(A))

尽管如此 np.digitize 将屏蔽值作为 -1 抛出。是否有其他方法可以让 np.digitize 忽略屏蔽值(或 NaN)?

我希望它不是为了性能优化,否则你可以 数字化功能后的掩码:

import numpy as np

A = np.arange(10,dtype=np.float)
A[0] = np.nan
A[-1] = np.nan

bins = np.array([1,2,7])

res = np.digitize(A,bins)

# here np.nan is assigned to the highes bin 
# using numpy '1.17.2'
print(res)

# sp you mask you array after the execution of 
# np.digitize
print(res[~np.isnan(A)])
>>> [3 1 2 2 2 2 2 3 3 3]
>>> [1 2 2 2 2 2 3 3]