在纯NumPy中重写for循环以减少执行时间
我最近询问了如何为一个科学应用程序优化一个Python循环，并且在NumPy中获得了一种优秀的，巧妙的重新编码方式，这为我减少了大约100倍的执行时间！
但是， B值的计算实际上是嵌套在其他几个循环中，因为它是在规则的网格位置进行评估的。 是否有类似智能的NumPy重写来缩短这个过程的时间？
我怀疑这部分的性能增益不会很明显，并且可能的缺点是，不可能向用户报告计算进度，直到结果不能写入输出文件，直到计算的结束，并且可能在一个巨大的步骤中这样做会产生内存影响？ 是否有可能绕过这些？
import numpy as np
import time
def reshape_vector(v):
b = np.empty((3,1))
for i in range(3):
b[i][0] = v[i]
return b
def unit_vectors(r):
return r / np.sqrt((r*r).sum(0))
def calculate_dipole(mu, r_i, mom_i):
relative = mu - r_i
r_unit = unit_vectors(relative)
A = 1e-7
num = A*(3*np.sum(mom_i*r_unit, 0)*r_unit - mom_i)
den = np.sqrt(np.sum(relative*relative, 0))**3
B = np.sum(num/den, 1)
return B
N = 20000 # number of dipoles
r_i = np.random.random((3,N)) # positions of dipoles
mom_i = np.random.random((3,N)) # moments of dipoles
a = np.random.random((3,3)) # three basis vectors for this crystal
n = [10,10,10] # points at which to evaluate sum
gamma_mu = 135.5 # a constant
t_start = time.clock()
for i in range(n[0]):
r_frac_x = np.float(i)/np.float(n[0])
r_test_x = r_frac_x * a[0]
for j in range(n[1]):
r_frac_y = np.float(j)/np.float(n[1])
r_test_y = r_frac_y * a[1]
for k in range(n[2]):
r_frac_z = np.float(k)/np.float(n[2])
r_test = r_test_x +r_test_y + r_frac_z * a[2]
r_test_fast = reshape_vector(r_test)
B = calculate_dipole(r_test_fast, r_i, mom_i)
omega = gamma_mu*np.sqrt(np.dot(B,B))
# write r_test, B and omega to a file
frac_done = np.float(i+1)/(n[0]+1)
t_elapsed = (time.clock()-t_start)
t_remain = (1-frac_done)*t_elapsed/frac_done
print frac_done*100,'% done in',t_elapsed/60.,'minutes...approximately',t_remain/60.,'minutes remaining'
你可以做的一件显而易见的事情就是替换这条线
r_test_fast = reshape_vector(r_test)
同
r_test_fast = r_test.reshape((3,1))
在性能上可能不会有太大的区别，但无论如何，使用numpy内置代替重新发明轮子是有意义的。
一般来说，正如你现在可能已经注意到的那样，优化numpy的技巧是在numpy全阵列操作的帮助下表达算法，或者至少使用切片而不是遍历Python代码中的每个元素。 阻止这种“矢量化”的是所谓的循环承载依赖性，即循环，其中每次迭代取决于先前迭代的结果。 简单地看一下你的代码，你就没有这种东西，应该可以将代码向量化。
编辑：一种解决方案
我没有证实这是正确的，但应该给你一个如何处理它的想法。
首先，使用我们将使用的cartesian（）函数。 然后
def calculate_dipole_vect(mus, r_i, mom_i):
# Treat each mu sequentially
Bs = []
omega = []
for mu in mus:
rel = mu - r_i
r_norm = np.sqrt((rel * rel).sum(1))
r_unit = rel / r_norm[:, np.newaxis]
A = 1e-7
num = A*(3*np.sum(mom_i * r_unit, 0)*r_unit - mom_i)
den = r_norm ** 3
B = np.sum(num / den[:, np.newaxis], 0)
Bs.append(B)
omega.append(gamma_mu * np.sqrt(np.dot(B, B)))
return Bs, omega
# Transpose to get more "natural" ordering with row-major numpy
r_i = r_i.T
mom_i = mom_i.T
t_start = time.clock()
r_frac = cartesian((np.arange(n[0]) / float(n[0]),
np.arange(n[1]) / float(n[1]),
np.arange(n[2]) / float(n[2])))
r_test = np.dot(r_frac, a)
B, omega = calculate_dipole_vect(r_test, r_i, mom_i)
print 'Total time for vectorized: %f s' % (time.clock() - t_start)
那么，在我的测试中，这实际上比我开始的基于循环的方法稍慢。 问题是，在问题的原始版本中，它已经通过整形阵列（20000，3）的全数组操作进行了矢量化，因此任何进一步的矢量化都没有带来太多的好处。 事实上，如上所述，这可能会使性能恶化，这可能是由于大型临时阵列造成的。
如果你分析你的代码，你会发现99％的运行时间在calculate_dipole因此缩短这个循环的时间并不会显着减少执行时间。 如果你想让这个更快，你仍然需要关注calculate_dipole。 我尝试了我的用于calculate_dipole Cython代码，并在总体时间内减少了大约2倍。 也可能有其他方法来改进Cython代码。
上一篇: Rewriting a for loop in pure NumPy to decrease execution time
下一篇: how to use control