多维度的numpy linspace和网格网格

时间:2016-03-15 10:22:52

标签: python matlab numpy

我使用numpy将一些matlab代码移植到python并且我有以下matlab命令:

[xgrid,ygrid]=meshgrid(linspace(-0.5,0.5, GridSize-1), ...
                       linspace(-0.5,0.5, GridSize-1));

现在,这在2D中很好,但我想将其扩展到n维。因此,根据输入数据,GridSize可以是2,3或4维向量。因此,在2D中,这将是:

[xgrid, grid] = np.meshgrid(np.linspace(-0.5,0.5, GridSize[0]), 
                            np.linspace(-0.5,0.5, GridSize[1]));

但是,之前我不知道输入的维数,所以可以重写这个表达式,以便它可以生成任意数量的维度的网格吗?

1 个答案:

答案 0 :(得分:5)

您可以使用循环理解生成所有1D数组,然后对*内部执行np.meshgrid的运算符使用unpacking of argument lists,这相当于MATLAB's comma separated lists,像这样 -

allG = [np.linspace(-0.5,0.5, G) for G in GridSize]
out = np.meshgrid(*allG)

示例运行

1)2D案例:

In [27]: GridSize = [3,4]

In [28]: allG = [np.linspace(-0.5,0.5, G) for G in GridSize]
    ...: out = np.meshgrid(*allG)
    ...: 

In [29]: out[0]
Out[29]: 
array([[-0.5,  0. ,  0.5],
       [-0.5,  0. ,  0.5],
       [-0.5,  0. ,  0.5],
       [-0.5,  0. ,  0.5]])

In [30]: out[1]
Out[30]: 
array([[-0.5       , -0.5       , -0.5       ],
       [-0.16666667, -0.16666667, -0.16666667],
       [ 0.16666667,  0.16666667,  0.16666667],
       [ 0.5       ,  0.5       ,  0.5       ]])

2)3D案例:

In [51]: GridSize = [3,4,2]

In [52]: allG = [np.linspace(-0.5,0.5, G) for G in GridSize]
    ...: out = np.meshgrid(*allG)
    ...: 

In [53]: out[0]
Out[53]: 
array([[[-0.5, -0.5],
        [ 0. ,  0. ],
        [ 0.5,  0.5]], ...

       [[-0.5, -0.5],
        [ 0. ,  0. ],
        [ 0.5,  0.5]]])

In [54]: out[1]
Out[54]: 
array([[[-0.5       , -0.5       ], ...

       [[ 0.16666667,  0.16666667],
        [ 0.16666667,  0.16666667],
        [ 0.16666667,  0.16666667]],

       [[ 0.5       ,  0.5       ],
        [ 0.5       ,  0.5       ],
        [ 0.5       ,  0.5       ]]])

In [55]: out[2]
Out[55]: 
array([[[-0.5,  0.5], ....

       [[-0.5,  0.5],
        [-0.5,  0.5],
        [-0.5,  0.5]]])