python - 从 3D 数组的 2D 切片创建 meshgrid

Question

我有 3D numpy 数组。

import numpy as np
X = np.arange(12).reshape(2, 2, 3)
print(X)

[[[ 0  1  2]
  [ 3  4  5]]

 [[ 6  7  8]
  [ 9 10 11]]]

我想为 3D 数组中的所有 2D 数组向量化以下内容。例如，对于第一个二维数组:

ss = np.array(np.meshgrid(*X[0]), dtype=object).T.reshape(-1,2)
print(ss)

[[0 3]
 [0 4]
 [0 5]
 [1 3]
 [1 4]
 [1 5]
 [2 3]
 [2 4]
 [2 5]]

我试过以下:

def f(x):
    return np.array(np.meshgrid(*x), dtype=object).T.reshape(-1,2)

ff = np.apply_along_axis(f, 0, X)
print(ff)

Answer 1

这是一个通用解决方案，它使用一个循环并缩放到通用形状。它分配给一个初始化的数组并广播以复制值，从而实现内存效率。它适用于沿 X 第二轴的任何长度。因此，实现将是 -

def meshgrid_2D_blocks(X):
    m,n,r = X.shape
    out_shp = [m]+[r]*n+[n]
    out = np.empty(out_shp,dtype=X.dtype)

    # Assign each block iteratively
    shp = [-1]+[1]*n
    for i in range(n):
        shp[i+1] = r
        out[...,i] = X[:,i].reshape(shp)
        shp[i+1] = 1
    return out.reshape(m,-1,n)

样本运行

案例 #1:length=2 的第二个轴

In [167]: X = np.arange(12).reshape(2, 2, 3)

In [168]: X
Out[168]: 
array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]]])

In [169]: meshgrid_2D_blocks(X)
Out[169]: 
array([[[ 0,  3],
        [ 0,  4],
        [ 0,  5],
        [ 1,  3],
        [ 1,  4],
        [ 1,  5],
        [ 2,  3],
        [ 2,  4],
        [ 2,  5]],

       [[ 6,  9],
        [ 6, 10],
        [ 6, 11],
        [ 7,  9],
        [ 7, 10],
        [ 7, 11],
        [ 8,  9],
        [ 8, 10],
        [ 8, 11]]])

案例 #2:length=3 的第二个轴

In [170]: X = np.arange(12).reshape(2, 3, 2)

In [171]: X
Out[171]: 
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]]])

In [172]: meshgrid_2D_blocks(X)
Out[172]: 
array([[[ 0,  2,  4],
        [ 0,  2,  5],
        [ 0,  3,  4],
        [ 0,  3,  5],
        [ 1,  2,  4],
        [ 1,  2,  5],
        [ 1,  3,  4],
        [ 1,  3,  5]],

       [[ 6,  8, 10],
        [ 6,  8, 11],
        [ 6,  9, 10],
        [ 6,  9, 11],
        [ 7,  8, 10],
        [ 7,  8, 11],
        [ 7,  9, 10],
        [ 7,  9, 11]]])