从三维阵列的二维切片创建网格

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

我想矢量化的所有二维阵列在三维阵列以下。例如，对于第一个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]]])

第二种情况：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]]])