拆分 3D numpy 数组以获取具有满足条件的连续值的所有系列

Question

我想获得 3D 数组中沿轴 0 满足条件的所有数字序列。我正在尝试尽可能矢量化所有内容以避免极其低效的循环。

下面的代码是有效的，它以一个定义的数组为例，形状为(6, 2, 3)。它在其上应用一个掩码，由数组 b（相同维度）定义。

然后，我沿轴 0 对 3D 数组进行切片，这样我就有 2*3 个 1D 切片，从而获得 6 个形状为 (6,) 的 1D 数组。为此，我使用了一个循环，这显然会成为较大数组的效率问题。

然后我根据得到的掩码拆分我的数组，select（现在简单打印）至少有3个连续值满足b数组给定条件的系列。

import numpy as np

a = np.array([[[0.57337127, 0.7626088, 0.26965987],
               [0.66987041, 0.2914202, 0.62678441]],
              [[0.97442524, 0.61656519, 0.10544983],
               [0.05780219, 0.00381356, 0.57118615]],
              [[0.47069657, 0.36802822, 0.67483419],
               [0.32773146, 0.99773064, 0.56042508]],
              [[0.70984651, 0.25093198, 0.71911127],
               [0.05182876, 0.9463291, 0.7222756]],
              [[0.56736192, 0.62692889, 0.33814278],
               [0.72362855, 0.12885637, 0.44096788]],
              [[0.12706838, 0.90640269, 0.5126569],
               [0.62920448, 0.24502599, 0.26754067]]])

b = np.array([[[0.4, 0.4, 0.4],
               [0.4, 0.4, 0.4]],
              [[0.4, 0.4, 0.4],
               [0.4, 0.4, 0.4]],
              [[0.4, 0.4, 0.4],
               [0.4, 0.4, 0.4]],
              [[0.4, 0.4, 0.4],
               [0.4, 0.4, 0.4]],
              [[0.4, 0.4, 0.4],
               [0.4, 0.4, 0.4]],
              [[0.4, 0.4, 0.4],
               [0.4, 0.4, 0.4]]])

# these two loops i and j are very inefficient 
for i in range(a.shape[1]):
    for j in range(a.shape[2]):
        print(i, j)
        aij = a[:, i, j]
        bij = b[:, i, j]
        mask = aij <= bij
        split_indices = np.where(mask)[0]
        for subarray in np.split(aij, split_indices + 1):
            if len(subarray) > 3:
                print(subarray[:-1])

现在，这行得通了。但是，我的实际数据（数组 a 和 b）的形状为 ~(500, 800, 1500)，这意味着循环变得有问题（相当昂贵）。

你能想出一种方法来进一步矢量化它吗？我试图获得一个 3D 蒙版并按 3D 方式进行拆分，但这会导致沿轴 1 和轴 2 的拆分大小不相等，这是一个问题（也是 np.split 仅采用 0 或 1-D 的原因索引列表。

Answer 1

您可以遍历轴 0，而不是在轴 (1,2) 上循环，以迭代检查掩码是否适用于连续的元素。根据您最终想要实现的目标，最多需要 n 次迭代（其中 n 是您拥有的最长序列）。

例如，如果您只想确定满足 mask 的至少 3 个连续元素的任何有效序列的起始元素，您可以这样做：

mask = a > b
runs = np.zeros_like(mask, dtype=bool)
runs[:-2] = mask[:-2] & mask[1:-1] & mask[2:]

在您的示例中，这会产生：

>>> runs
array([[[ True, False, False],
    [False, False,  True]],
   [[ True, False, False],
    [False, False,  True]],
   [[ True, False, False],
    [False, False,  True]],
   [[False, False, False],
    [False, False, False]],
   [[False, False, False],
    [False, False, False]],
   [[False, False, False],
    [False, False, False]]])

在这里，对于长度至少为 3 的有效序列的所有起始元素，runs 的计算结果为 True。因为我不确定一旦确定它们后您想做什么，我是留在这里。但希望很清楚如何从这里开始概括您要尝试做的事情。

拆分 3D numpy 数组以获取具有满足条件的连续值的所有系列

Split a 3D numpy array to get all series with consecutive values that meet a condition

python

arrays

numpy

vectorization

python-3.x