根据索引初始化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

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

Answer 3

我认为有几件事需要考虑：

• 是否有pos_data必须列出的原因？
• 没有其他变量（d），你必须硬编码，而它总是应该是其他元组的长度。

考虑到这些，您可以使用itertools.product解决可变数量的for循环问题，这基本上只是嵌套for循环的简写。 product的位置参数是循环的范围。

我的实施是：

from itertools import product

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

pos_data_new = np.zeros((len(ratio), *vol_ext))

# now loop over each dimension in `vol_ext`. Since `product` expects
# positional arguments, we have to unpack a tuple of `range(vol)`s.
for inds in product(*(range(vol) for vol in vol_ext)):
    # inds is now a tuple, and we have to combine it with a slice in 
    # in the first dimension, and use it as an array on the right hand 
    # side to do the computation. 
    pos_data_new[(slice(None),) + inds] = (np.array(inds) - 1) * ratio

我不认为这会更快，但它看起来确实更好。

请注意，pos_data_new现在是一个数组，根据原始示例，将其作为第一维中的列表，很简单。