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数组进行矢量化处理。例如，对于第一个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

这是实现该目标的一种方法：

import numpy as np
X = np.arange(12).reshape(2, 2, 3)
out = np.stack(np.broadcast_arrays(X[:, 0, :, None], X[:, 1, None, :]), -1).reshape(len(X), -1, 2)
print(out)
# [[[ 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]]]

Answer 2

这是一个通用解决方案，它使用一个循环并缩放为通用形状。它分配给一个初始化的数组并进行广播以复制值，从而实现内存效率。它可以沿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]]])

从3D阵列的2D切片创建网格网格

2 个答案: