在Python和numpy中列出两个变量循环的理解

时间:2012-06-13 14:41:48

标签: python list numpy

我必须使用matplotlib中的contourf从函数返回到绘图的值x,y创建一个2D numpy数组,到目前为止,我使用" C"就像结构一样,它似乎在Python中非常低效:

    dim_x = np.linspace(self.min_x, self.max_x, self.step)
    dim_y = np.linspace(self.min_y, self.max_y, self.step)
    X, Y = np.meshgrid(dim_x, dim_y)

    len_x = len(dim_x)
    len_y = len(dim_y)


    a = np.zeros([len_x, len_y], dtype=complex)

    for i, y in enumerate(dim_y):
        for j, x in enumerate(dim_x):
            a[i][j] = aux_functions.final_potential(complex(x, y), element_list)

cs = plt.contourf(X, Y, (a.real), 100)

如何以更加pythonic的方式完成?

谢谢!

2 个答案:

答案 0 :(得分:2)

如果您可以将final_potential重写为矢量化函数,那将是理想的选择。一个简单的,也许是太明显的例子:

>>> dim_x = np.linspace(0, 2, 5)
>>> dim_y = np.linspace(0, 2, 5)
>>> X * Y
array([[ 0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
       [ 0.  ,  0.25,  0.5 ,  0.75,  1.  ],
       [ 0.  ,  0.5 ,  1.  ,  1.5 ,  2.  ],
       [ 0.  ,  0.75,  1.5 ,  2.25,  3.  ],
       [ 0.  ,  1.  ,  2.  ,  3.  ,  4.  ]])

但如果你真的不能这样做,你也可以vectorize

>>> np.vectorize(lambda x, y: x * y + 2)(X, Y)
array([[ 2.  ,  2.  ,  2.  ,  2.  ,  2.  ],
       [ 2.  ,  2.25,  2.5 ,  2.75,  3.  ],
       [ 2.  ,  2.5 ,  3.  ,  3.5 ,  4.  ],
       [ 2.  ,  2.75,  3.5 ,  4.25,  5.  ],
       [ 2.  ,  3.  ,  4.  ,  5.  ,  6.  ]])

在你的情况下,它可能看起来像这样:

def wrapper(x, y): 
    return aux_functions.final_potential(complex(x, y), element_list)

a = np.vectorize(wrapper)(X, Y)

这可能比嵌套的for循环快一点,尽管python函数调用的开销会抵消很多numpy的效率。在我过去做过的测试中,使用vectorize提供了适度的5倍加速。 (这与纯粹的numpy操作的加速比为100倍或1000倍,如X * Y示例中所示。)

答案 1 :(得分:1)

我不知道这是否会提高您的效果,但numpy.fromfunction可用于构建数组a。如果没有一个很好的自成一体的例子,那么看看性能有多么不同将会有点棘手。例如运行aux_functions.final_potential需要多长时间?如果该功能很昂贵,那么优化它周围的循环并不重要。

另外,我认为a[i,j]的效率略高于a[i][j],但我还没有测试过。