使用 3 个 2D numpy 数组或 3 个 1D numpy 数组创建一个 3D numpy 数组

Question

我有三个 numpy 二维数组：

A1 with the shape of (x * y)
A2 with the shape of (x * z)
A3 with the shape of (y * z)

三个数组中的值为 True 或 False。现在，我想创建一个形状为 (x * y * z) 的 3D 数组，以便 3D 数组中的每个元素如下所示：

3D_array[x, y, z] = A1[x, y] & A2[x, z] & A3[y, z]

我知道我可以使用循环。但是有没有更快的方法来做到这一点？比如通过矢量化？

另外，三个二维数组实际上只是三个一维数组之间的一些成对交互。因此，提出我的问题的更一般的方式是：

给出三个长度分别为x，y，z的1D numpy数组（X），（Y），（Z），创建3D numpy数组的最佳方法是使3D数组中的每个元素等于：

3D_array[i, j, k] = my_function(X[i], Y[j], Z[k])

其中 my_function 是一个自定义函数，它返回 True/False。同样，我正在寻找比循环更好的东西。

欢迎对第一个问题或更一般的问题提供任何解决方案。非常感谢！

Answer 1

TL;DR：

np.apply_along_axis(sum, axis=0, np.stack(np.meshgrid(x, y, z)))

演练：

让我们看一个具体的例子，从创建三个一维数组开始：

x, y, z = [1, 2], [3, 4], [5, 6]

现在，让我们把这些变成meshgrid：

a, b, c = np.meshgrid(x, y, z)

这次我们得到了三个 3-d 数组：

>>> a
array([[[1, 1],
        [2, 2]],

       [[1, 1],
        [2, 2]]])

这个数组对应我们的第一个参数，一半时间是1，另一半时间是2。这是我们在x 中的两个值。我们总共有八个值，因为我们最终输出的形状应该是2 x 2 x 2 = 8。

同样，我们有b 和c：

>>> b
array([[[3, 3],
        [3, 3]],

       [[4, 4],
        [4, 4]]])
>>> c
array([[[5, 6],
        [5, 6]],

       [[5, 6],
        [5, 6]]])

此时，您已经可以开始以各种方式组合这三个数组，因此对于简单的函数（可以分为两步的函数），您可以简单地运行 a + b + c 或 a & b & c。

在一般情况下，您可能希望定义接受三个数字并基于任意逻辑产生单个输出的函数。要应用这些功能，我们首先需要将stack 这三个数组合并为一个：

>>> np.stack([a, b, c])
array([[[[1, 1],
         [2, 2]],

        [[1, 1],
         [2, 2]]],


       [[[3, 3],
         [3, 3]],

        [[4, 4],
         [4, 4]]],


       [[[5, 6],
         [5, 6]],

        [[5, 6],
         [5, 6]]]])

现在我们有 3 个大小为 8 的数组（实际形状是 3,2,2,2），我们想将一个函数 apply 一次 8 次变成三个数字：

>>> np.apply_along_axis(sum, axis=0, np.stack(np.meshgrid(x, y, z)))
array([[[ 9, 10],
        [10, 11]],

       [[10, 11],
        [11, 12]]])

我们得到 8 个点，结果如预期的那样，每个点是 1 或 2 之一、3 或 4 之一和 5 或 6 之一的总和。

注意numpy 会将一维数组传递给f，而不是一系列参数，所以如果你有这样的函数：

def f(a, b, c):
  return a + b - c

您将需要定义一个包装原始函数的附加函数：

def f2(vals):
  return f(*vals)

# Alternatively:
f2 = lambda x: f(*x)

所以，现在我们可以将f2 应用于我们的数据：

>>> np.apply_along_axis(f2, 0, np.stack(np.meshgrid(x, y, z)))
array([[[-1, -2],
        [ 0, -1]],

       [[ 0, -1],
        [ 1,  0]]])

Answer 2

有广播

3D_array[x, y, z] = A1[x, y] & A2[x, z] & A3[y, z]

可以通过：

B = A1[:,:,None] & A2[:,None,:] & A3[None,:,:]

这实际上将所有 3 个数组转换为与大多数 numpy 运算符一起工作的 3d 数组。

更一般的情况：

3D_array[i, j, k] = my_function(X[i], Y[j], Z[k])

更难。

B = X[:,None,None] + Y[None,:,None] + Z[None,None,:]

这同样适用于使用广播的运算符（和ufunc），但不适用于仅适用于标量的用户定义函数。

我犹豫是否建议这样做，但也许有必要。 np.vectorize 可以像这样创建一个与广播一起使用的函数。但它不是性能工具。

其他答案中建议的apply_along_axis 不是性能工具。它在 other 轴上迭代。对于像这样的多维数组，它的代码比更显式的迭代更简单，但速度并不快。