Перезапись для цикла в чистом NumPy для уменьшения времени выполнения
Я недавно спросил о попытке оптимизировать цикл 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'
2 ответа
Если вы профилируете свой код, вы увидите, что 99% времени выполнения приходится на calculate_dipole , что сокращает время для этого цикла действительно не даст заметного сокращения времени выполнения. Вам все равно нужно сосредоточиться на calculate_dipole, если вы хотите сделать это быстрее. Я попробовал свой код Cython для calculate_dipole на этом и получил сокращение примерно в 2 раза за общее время. Могут быть и другие способы улучшить код Cython.
Одна очевидная вещь, которую вы можете сделать, - это заменить строку
r_test_fast = reshape_vector(r_test)
на
r_test_fast = r_test.reshape((3,1))
. Вероятно, это не будет иметь большого значения в производительности, но в любом случае это дает имеет смысл использовать встроенные команды numpy вместо того, чтобы изобретать колесо.
Вообще говоря, как вы, наверное, уже заметили, трюк с оптимизацией numpy состоит в том, чтобы выразить алгоритм с помощью множества операций с целым массивом или, по крайней мере, с помощью срезов вместо итерации по каждому элементу в коде Python. То, что имеет тенденцию предотвращать этот вид «векторизации», - это так называемые зависимости с переносом цикла, то есть циклы, в которых каждая итерация зависит от результата предыдущей итерации. Кратко посмотрев на свой код, вы обнаружите, что такого нет, и вы можете легко векторизовать ваш код.
РЕДАКТИРОВАТЬ: Одно решение
Я не проверял, что это правильно, но должно дать вам представление о том, как к нему подойти.
Сначала возьмем функцию декартова (), которую мы будем использовать . Тогда
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), поэтому любая дальнейшая векторизация на самом деле не приносит большой пользы. Фактически, это может ухудшить производительность, как указано выше, возможно, из-за больших временных массивов.