根据索引初始化NumPy数组

Question

我正在使用NumPy创建几个多维数组,并根据索引对它们进行初始化,如下所示:

pos_data = []
# Some typical values
d = 2  # Could also be 3
vol_ext = (1000, 500)  # If d = 3, this will have another entry
ratio = [5.0, 8.0]  # Again, if d = 3, it will have another entry

for i in range(d):
    pos_data.append(np.zeros(vol_ext))

if d == 2:
    for y in range(vol_ext[1]):
        for x in range(vol_ext[0]):
            pos_data[0][x, y] = (x - 1.0) * ratio[0]
            pos_data[1][x, y] = (y - 1.0) * ratio[1]
elif d == 3:
    for z in range(vol_ext[2]):
        for y in range(vol_ext[1]):
            for x in range(vol_ext[0]):
                pos_data[0][x, y, z] = (x - 1.0) * ratio[0]
                pos_data[1][x, y, z] = (y - 1.0) * ratio[1]
                pos_data[2][x, y, z] = (z - 1.0) * ratio[2]

循环有点丑陋和缓慢.另外,如果我有一个三维对象,那么我必须有另一个嵌套循环.

我想知道是否有一种Pythonic方法来生成这些值,因为它们只是基于x,y和z索引.我尝试使用组合比特itertools,但我无法使其工作.

Answer 1

这很简单np.meshgrid:

pos_data = np.meshgrid(*(r * (np.arange(s) - 1.0)
                         for s, r in zip(vol_ext, ratio)), indexing='ij')

Answer 2

我将生成二维或三维numpy.meshgrid数据,然后按每个切片的比例缩放每个条目.

对于2D情况:

(X, Y) = np.meshgrid(np.arange(vol_ext[1]), np.arange(vol_ext[0]))
pos_data = [(Y - 1) * ratio[0], (X - 1) * ratio[1]]

对于3D案例:

(X, Y, Z) = np.meshgrid(np.arange(vol_ext[2]), np.arange(vol_ext[1]), np.arange(vol_ext[0]))
pos_data = [(Z - 1) * ratio[0], (Y - 1) * ratio[1], (X - 1) * ratio[2]]

使用2D数据的示例

pos_data已由您的代码生成.我创建了一个新列表pos_data2,使用上面的解决方案存储等效列表:

In [40]: vol_ext = (1000, 500)

In [41]: (X, Y) = np.meshgrid(np.arange(vol_ext[1]), np.arange(vol_ext[0]))

In [42]: pos_data2 = [(Y - 1) * ratio[0], (X - 1) * ratio[1]]

In [43]: np.allclose(pos_data[0], pos_data2[0])
Out[43]: True

In [44]: np.allclose(pos_data[1], pos_data2[1])
Out[44]: True

基于此进行自适应 vol_ext

我们可以将它与列表理解结合起来,我们可以利用输出numpy.meshgrid为元组的事实:

pts = [np.arange(v) for v in reversed(vol_ext)]
pos_data = [(D - 1) * R for (D, R) in zip(reversed(np.meshgrid(*pts)), ratio)]

第一行代码生成每个所需维度的点范围作为列表.然后,我们使用列表推导来计算每个切片所需的计算,方法是迭代所需维度中每个所需的点网格,并结合要应用的正确比率.

示例运行

In [49]: pts = [np.arange(v) for v in reversed(vol_ext)]

In [50]:  pos_data2 = [(D - 1) * R for (D, R) in zip(reversed(np.meshgrid(*pts)), ratio)]

In [51]: np.all([np.allclose(p, p2) for (p, p2) in zip(pos_data, pos_data2)])
Out[51]: True

最后一行遍历每个切片并确保两个列表对齐.