使用 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]

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

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

给出三个一维numpy数组(X), (Y), (Z)，长度分别为x, y, z，创建3D numpy数组的最佳方法是什么，以便3D数组中的每个元素等于：

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

其中 my_function 是 returns 的自定义函数，例如 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)

这次我们得到三个三维数组：

>>> 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]]]])

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

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

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

我们得到八分作为预期的结果，其中每个点是 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 个数组转换为 3d 数组，可与大多数 numpy 运算符一起工作。

更一般的情况：

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 轴上迭代。对于像这样的多维数组，它可以使代码比更明确的迭代更简单，但不会更快。

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

Create a 3D numpy array with 3 2D numpy arrays or 3 1D numpy array

python

arrays

numpy

vectorization