python - 选择 NaN 之后出现的 ndarray 中的值

Question

我有一个很大的 2D ndarray of floats，称它为 ar。它包含一些 NaN。我对右侧 NaN 的直接邻居感兴趣(例如，沿 axis=1)。例如，如果我知道 say 点 (3, 7) 是一个 NaN，我想选择 ar[3, 8:8+N]。然后我想对 NaN 的所有位置重复，并 vstack 所有由此获得的切片。

我可以愉快地使用 np.where 定位 NaN，并对这些值执行 for 循环。可悲的是，这有点慢。有没有一种有效的方法来以矢量化方式进行索引？所以我有一个元组列表 (x, y)，我想得到或多或少，

result=np.vstack([ ar[x, y+1:y+1+N] for x, y, in tuples ])

只是没有循环。这可能吗？

非常感谢。

Answer 1

如果 nan 从边缘发生少于 N 列，则您要求的定义不明确，但以下内容应该有效:

rows, cols = np.where(np.isnan(ar))
cols = (cols[:, None] + np.arange(1, N+1)).reshape(-1)
# Handle indices out of range by repeating the last column
cols = np.clip(cols, 0, ar.shape[1] - 1)
rows = np.repeat(rows, N)
result = ar[rows, cols].reshape(-1, 2)

编造一些假数据:

>>> ar = np.random.rand(25)
>>> ar[np.random.randint(25, size=5)] = np.nan
>>> ar = ar.reshape(5, 5)
>>> N = 2

并在其上运行上面的代码会产生:

>>> ar
array([[ 0.96556647,         nan,  0.02934316,  0.82174232,  0.29293098],
       [ 0.34819313,  0.57449136,         nan,         nan,  0.32791866],
       [ 0.14020414,  0.60668458,  0.95613773,  0.09533064,  0.43401037],
       [ 0.83888255,  0.34240687,         nan,  0.02495232,  0.36234979],
       [ 0.21870906,  0.24181006,  0.81447603,  0.24216213,         nan]])
>>> result
array([[ 0.02934316,  0.82174232],
       [        nan,  0.32791866],
       [ 0.32791866,  0.32791866],
       [ 0.02495232,  0.36234979],
       [        nan,         nan]])