python - Numpy:将三个 1D 数组组装成 3D(但不完全是简单的坐标网格)

Question

我有三个一维数组，其中两个长度相同，但第三个长度不同，例如

A = np.array([1, 2, 3, 4])
B = np.array([10, 20, 30, 40])
C = np.array([100, 200, 300, 400, 500])

我想将它们组合成一个网格，如下所示:

D = np.array([[[1, 10, 100], [1, 10, 200], [1, 10, 300], [1, 10, 400], [1, 10, 500]],
              [[2, 20, 100], [2, 20, 200], [2, 20, 300], [2, 20, 400], [2, 20, 500]],
              [[3, 30, 100], [3, 30, 200], [3, 30, 300], [3, 30, 400], [3, 30, 500]],
              [[4, 40, 100], [4, 40, 200], [4, 40, 300], [4, 40, 400], [4, 40, 500]]])

它类似于坐标网格，但是 A 和 B 一起变化，所以我假设它们需要先组合，然后将生成的 2D 数组与C 不知何故，但我找不到 stack、concatenate、meshgrid 等的正确组合。

实际的数组会比这个例子大很多(每个都有数千个值)，而且代码会被调用很多次，所以速度很重要。它将存在于 MC​​MC 采样器的主循环中，其中 A 和 B 是在每个循环中变化的参数，而 C 是一个常量.

Answer 1

我们分别称A, B, C a, b的长度>，c，分别。您正在寻找形状为 (a, c, 3)(或 (b, c, 3))的输出数组。正确堆叠 A 和 B 会为您提供形状为 (a, 2):

的数组

zipped = np.stack((A, B), axis=-1)

为了能够将 C 附加到 zipped，您需要沿结果的前两个轴将数组广播到相同的形状:

Aa = np.broadcast_to(zipped[..., None, :], (A.size, C.size, 2))
Cc = np.broadcast_to(C[..., None], (A.size, C.size, 1))
result = np.append(Aa, Cc, axis=-1)

整个操作可以写成一行:

result = np.append(np.broadcast_to(np.stack((A, B), axis=-1)[..., None, :], (A.size, C.size, 2)), np.broadcast_to(C[..., None], (A.size, C.size, 1)), axis=-1)

这个操作相当有效:np.broadcast_to 创建一个 View 而不复制任何数据。只有 append 操作本身会进行复制。