Python：在网格节点上构建一个迭代器

Question

我正在使用Python 3。我想要从一个三维节点列表开始构建一个网格。我希望避免该结构。

import numpy as np
l = np.zeros(len(xv)*len(yv)*len(zv))
for (i,x) in zip(range(len(xv)),xv):
   for (j,y) in zip(range(len(yv)),yv):
       for (k,z) in zip(range(len(zv)),zv):
           l[i,j,k] = func(x,y,z)

我正在寻找一种更紧凑的版本，类似于迭代器zip，但它将在网格中迭代所有可能的元组。

- simona

---

itertools.product? - juanpa.arrivillaga

一个更Pythonic的迭代表达式是： for i，x in enumerate（xv）： - hpaulj

@juanpa.arrivillaga 不错，但我如何获得适当的i（和j、k）？ - simona

告诉我们关于 func 的信息。它只能处理标量 x，y，z 值吗？还是可以接受数组？如果可以接受数组，那么有多少维度？ - hpaulj

Answer 1

你可以使用类似于np.meshgrid的工具来构建你的网格。假设func已经正确向量化，那么这应该足以构建l。

X, Y, Z = np.meshgrid(xv, yv, zv)
l = func(X, Y, Z)

如果 func 没有向量化，您可以使用 np.vectorize 构建一个向量化的版本。

另外请注意，通过巧妙地使用 np.newaxis，甚至可以避免使用 np.meshgrid：

>>> x
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> y
array([0, 1, 2])
>>> z
array([0, 1])
>>> def func(x, y, z):
...     return x + y + z
... 
>>> vfunc = np.vectorize(func)
>>> vfunc(x[:, np.newaxis, np.newaxis], y[np.newaxis, :, np.newaxis], z[np.newaxis, np.newaxis, :])
array([[[ 0,  1],
        [ 1,  2],
        [ 2,  3]],

       [[ 1,  2],
        [ 2,  3],
        [ 3,  4]],

       [[ 2,  3],
        [ 3,  4],
        [ 4,  5]],

       [[ 3,  4],
        [ 4,  5],
        [ 5,  6]],

       [[ 4,  5],
        [ 5,  6],
        [ 6,  7]],

       [[ 5,  6],
        [ 6,  7],
        [ 7,  8]],

       [[ 6,  7],
        [ 7,  8],
        [ 8,  9]],

       [[ 7,  8],
        [ 8,  9],
        [ 9, 10]],

       [[ 8,  9],
        [ 9, 10],
        [10, 11]],

       [[ 9, 10],
        [10, 11],
        [11, 12]]])

如评论中所指出的，np.ix_ 可以作为 np.newaxis 的快捷方式：

vfunc(*np.ix_(xv, yv, zv))

请注意，使用这个简单的函数时，不需要使用np.vectorize，实际上会严重影响我们的性能...

- mgilson

---

您IP地址为143.198.54.68，由于运营成本限制，当前对于免费用户的使用频率限制为每个IP每72小时10次对话，如需解除限制，请点击左下角设置图标按钮（手机用户先点击左上角菜单按钮）。 - Divakar

同样地，np.ix_可以用来创建这些扩展数组：func(*np.ix_(xv,yv,zv))。 - Divakar

@Divakar -- 啊，是的。我总是不知道为什么会忘记np.ix_... - mgilson

这是我们一遍又一遍看到的问题，“如何“向量化”这个函数？”或者“摆脱循环”。我们可以建议各种生成网格/网格或广播的方法，但它真正取决于func的性质。如果它处理数组，那么问题就很简单。如果它只适用于标量，则除了表面上的方式外，它是不可能的。如果问题本质上是串行的，你必须拿出cumsum枪，或者诉诸cython。 - hpaulj

@hpaulj -- 是的，但这个问题并不一定是关于性能的 - 至少我是这样理解的 - 更多的是关于“如何以更好的方式编写代码”。因此，“表面”处理（例如np.vectorize）很可能已经足够回答这个问题了。如果我们关心解决方案执行所需的时间，那么我同意这是一个几乎不可能回答的问题，除非我们知道func的具体情况。 - mgilson

Answer 2

假设你的函数是这样的：

def func(x,y,z,indices):
    xv, yv, zv = [i[j] for i,j in zip((x,y,z),indices)]
    #do a calc with the value for the specific x,y,z points

用partial将您想要连接的列表挂钩到它上面，方法如下：

from functools import partial
f = partial(func, x=xv, y=yv, z=zv)

现在只需提供索引并进行映射，你就完成了！

l = list(map(lambda x: f(indices=x), itertools.product(x,y,z)))

- themistoklik

Answer 3

通过一个简单的函数：

def foo(x,y,z):
   return x**2 + y*2 + z

并由空格定义：

In [328]: xv, yv, zv = [np.arange(i) for i in [2,3,4]]

这个迭代速度和任何一个一样快，尽管有点啰嗦：

In [329]: res = np.zeros((xv.shape[0], yv.shape[0], zv.shape[0]), dtype=int)
In [330]: for i,x in enumerate(xv):
     ...:     for j,y in enumerate(yv):
     ...:         for k,z in enumerate(zv):
     ...:             res[i,j,k] = foo(x,y,z)

In [331]: res
Out[331]: 
array([[[0, 1, 2, 3],
        [2, 3, 4, 5],
        [4, 5, 6, 7]],

       [[1, 2, 3, 4],
        [3, 4, 5, 6],
        [5, 6, 7, 8]]])

正如 @mgilson 所解释的那样，您可以生成3个数组来定义3D空间：

In [332]: I,J,K = np.meshgrid(xv,yv,zv,indexing='ij',sparse=True)
In [333]: I.shape
Out[333]: (2, 1, 1)
In [334]: J.shape
Out[334]: (1, 3, 1)
In [335]: I,J,K = np.ix_(xv,yv,zv)    # equivalently
In [336]: I.shape
Out[336]: (2, 1, 1)

foo 被编写成可以与数组一样良好地工作，因此：

In [337]: res1 = foo(I,J,K)
In [338]: res1
Out[338]: 
array([[[0, 1, 2, 3],
      ...
        [5, 6, 7, 8]]])

如果您的函数符合这种模式，请使用它。看看那些带有和不带有稀疏属性的 I，J，K 数组。

还有其他工具可以生成 i，j，k 集。例如：

for i,j,k in np.ndindex(res.shape):
    res[i,j,k] = foo(xv[i], yv[j], zv[k])

for i,j,k in itertools.product(range(2),range(3),range(4)):
    res[i,j,k] = foo(xv[i], yv[j], zv[k])

itertools.product 是快速的，特别是当用作 list(product(...)) 时。但迭代机制并不那么重要。最耗时间的是对 foo 的重复调用。

ndindex 实际上使用了 nditer，可以直接在以下情况下使用：

it = np.nditer([I,J,K,None],flags=['external_loop','buffered'])
for x,y,z,r in it:
    r[...] = foo(x,y,z)
it.operands[-1]

nditer最好的描述在于：https://docs.scipy.org/doc/numpy/reference/arrays.nditer.html。它最好用作迈向Cython版本的垫脚石。否则它没有任何速度优势。(虽然使用这个foo和'external_loop'与foo(I,J,K)一样快)。请注意，这不需要索引(但请参见'multi_index')。

是的，还有vectorize。方便，但不是一个快速的解决方案。

vfoo=np.vectorize(foo, otypes=['int'])
vfoo(I,J,K)

- hpaulj