获取 3D numpy 数组中沿轴的连续非 nans 值总和的最大值

Posted 2023-04-18

【中文标题】获取 3D numpy 数组中沿轴的连续非 nans 值总和的最大值

【英文标题】：Get the max of the sum of consecutive non-nans values along an axis in 3D numpy array

【发布时间】：2021-02-27 20:39:08

【问题描述】：

我想对一个操作进行矢量化，但我不知道如何优化它。

假设我有一个形状为 (4, 6, 3) 的 3D 数组 a。我想对轴 0 上的连续非 nan 值求和，并计算轴 0 上新形成的数组 b 的最大值。

# a is of shape (6, 4, 3)
a = np.array(
[[[       np.nan,        np.nan,    0.21145476],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [   0.11141416,        np.nan,    0.00345888]],

 [[       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,            1.,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,            1.,        np.nan]],
  
 [[       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,    0.43558095],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [   0.72669291,        np.nan,        np.nan]],
  
 [[       np.nan,        np.nan,    0.11763977],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,    0.01386668,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [   0.41135376,        np.nan,        np.nan]]])
  
# Calculate b to obtain consecutive non-nan sums over axis 0
b = VECTORIZED_FUNCTION(a)

# This is what b should end up being:
b = np.array(
[[[       np.nan,        np.nan,    0.21145476],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [   0.11141416,        np.nan,    0.00345888]],

 [[       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,            1.,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,            1.,        np.nan]],
  
 [[       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,    0.43558095],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [   0.72669291,        np.nan,        np.nan]],
  
 [[       np.nan,        np.nan,    0.11763977],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,    0.01386668,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [   1.13804667,        np.nan,        np.nan]]])

c = np.nanmax(b, axis=0)

# This is what c should end up being:
c = np.array(
  [[      np.nan,        np.nan,    0.21145476],
  [       np.nan,        np.nan,        np.nan],
  [       np.nan,            1.,    0.43558095],
  [       np.nan,    0.01386668,        np.nan],
  [       np.nan,        np.nan,        np.nan],
  [   1.13804667,            1.,    0.00345888]])

我可以让它在轴 1 和轴 2 上使用循环，但这对于我需要处理的数组（1000 个维度 ~(500, 1500, 800) 的数组）来说太慢了。

EDIT1：

第一个（非循环）试验性方法导致：

def pad_series(m, n=3):
    # Function adapted from @Swenzel's answer https://***.com/questions/32706135/extend-numpy-mask-by-n-cells-to-the-right-for-each-bad-value-efficiently
    k = m.copy()
    lenM, lenK = 1, 1
    while lenM+lenK < n:
        m[lenM:] |= k[:-lenM]
        m, k = k, m
        lenM, lenK = lenK, lenM+lenK
    k[n-lenM:] |= m[:-n+lenM]
    return k
    
# Select the non nan value in a mask
mask = ~np.isnan(a)
# Get the start of each "series" greater than 1 element along axis 0
runs = np.zeros_like(mask, dtype=bool)
runs[:-1] = mask[:-1] & mask[1:]
# Pad so that you get all the value in the series in an updated mask
runs_2 = pad_series(runs, n=2)
b = np.copy(a)
# Set everything in a copy of a to nan if they're not part of a serie
b[~runs_2] = np.nan
# Sum cumulatively
b = np.nancumsum(b, axis=0)
# Big issue: works only if one "multi-subheatwave" heatwave
c = np.nanmax(np.maximum(a, b), axis=0)

因此，在我的（简化）示例中，此 edit1 解决方案按预期工作。但是，如果您有两个或多个系列，它并没有被概括，np.nancumsum 函数会将第二个系列添加到第一个系列中，等等。

【问题讨论】：

VECTORIZED_FUNCTION 应该在这里做什么？在给定的示例中，a 和 b 是相同的。也许您有一个未矢量化的版本作为示例？

对不起，问题中可能没有说清楚。 a 和 b 不相同，请参见“左下”值。基本上，目标是“如果轴 = 0 上的 2 个或多个连续非 nan，取它们的总和”，这给你 b，然后“沿轴 0 取 b 的最大值”以获得答案

抱歉，我之前错过了左下角的不同之处……我不确定我能否轻松得出一个完全矢量化的答案，但我发现将numba 与派对混合通常会有所帮助很多

您应该始终制作一个 MCVE，重点放在 M 上。您不需要 (4, 6, 3) 数组。它对解释没有任何帮助，而且更难想象你在做什么。

【参考方案1】：

据我所知，如果不增加相当大的复杂性，矢量化操作无法直接执行您想要的操作。然而，我可以用一个相当简单的函数摆脱除一个循环之外的所有循环（假设我已经正确解释了你的意图）。

def segmented_cumsum(a):
    b = a.copy() #just use `a` if you want in-place operation
    for i in range(1,len(b)):
        mask = ~np.isnan(b[i-1])
        b[i][mask] += b[i-1][mask]
    return b

通过仅从前一个切片中获取非nan 数字，如果我们尝试与当前切片中的nan 值求和，则无关紧要。当前切片中的任何nan 值都将保持nan，如果它不是nan，则任何非nan 值都将与前一个切片中的值相加。

这也应该有利于使用 numba 将循环编译为 LLVM 代码的一些加速。 numba 不支持所有 numpy 的高级数组切片，不幸的是，我们必须将数组扁平化一个维度进行计算，然后在最后重新整形：

from numba import njit
@njit("float64[:,:,:](float64[:,:,:])")
def segmented_cumsum(a):
    shape = a.shape
    b = a.copy().reshape(shape[0], -1) #just use `a` if you want in-place operation
    for i in range(1,len(b)):
        mask = ~np.isnan(b[i-1])
        b[i][mask] += b[i-1][mask] #cumsum
    return b.reshape(shape)

第一次调用 jit'ed 函数时，编译需要一两秒钟，但之后会非常快（在我的机器上使用您的示例 a 作为输入大约是 10 倍）。

【讨论】：

是的，这看起来确实可以，使用 nan 属性，太棒了。使用 numba 也是一个非常好的建议，我还没有对此做任何事情，但这似乎是一个完美的机会

稍微了解 IEEE 754 是您知道如何寻找这些类型的解决方案（如何在算术中处理 nan 和 inf）。幸运的是 numpy 在这种情况下使用了非信令 nan。矢量化只不过是用 c 而不是 python 编写的一个非常有效的循环。在某些情况下，可以使用 simd 指令之类的东西（并且可能不会与 numba 一起使用），但最大的胜利是摆脱 python 来进行繁重的计算。