Ускорение анализа массивов в numpy

Question

У меня есть код python, который импортирует 4 столбца txt-файла с цифрами Первые три столбца - это координаты x, y, z, а четвертый столбец - плотность в этой координате.

Ниже приведен код, который читает, преобразует в ndarray, преобразовывает это поле в Фурье, вычисляет расстояние от начала координат (k = (0,0,0)) и преобразует координату, принимает среднее значение и записывает их. Благодаря pandas (библиотека python для анализа данных) и Python FFT загрузка 256^3 строк и преобразование Фурье очень быстрые и сделаны за несколько секунд.

Однако преобразование загруженного txt в numpy ndarray, вычисление средней плотности (средних значений каждой координаты) и вычисление расстояния от начала координат (k = (0,0,0)) занимает очень много времени.

Я думаю, что проблема в части np.around в конце, но я не могу понять, как ее оптимизировать.

У меня есть ресурс 32-х основных машин.

Может ли кто-нибудь научить меня, как ускорить, сделать его многопроцессорным кодом или что-то в этом роде, чтобы это можно было сделать очень быстро? Благодарю.

(Если вы космолог и когда-нибудь понадобится этот код, вы можете использовать его, но, пожалуйста, свяжитесь со мной, если вы можете. Спасибо)

from __future__ import division 
import numpy as np 

ngridx = 128 
ngridy = 128  
ngridz = 128 

maxK = max(ngridx,ngridy,ngridz) 

#making input file 
f = np.zeros((ngridx*ngridy*ngridz,4)) 

i = 0 
for i in np.arange(len(f)): 
    f[i][0] = int(i/(ngridy*ngridz)) 
    f[i][1] = int((i/ngridz))%ngridy 
    f[i][2] = int(i%ngridz) 
    f[i][3] = np.random.rand(1) 
    if i%1000000 ==0: 
     print i 
#This takes forever 
#end making input file 

#Thanks to Mike, 
a = f[:,3].reshape(ngridx,ngridy,ngridz) 

avg =np.sum(f[:,3])/len(f) 
a /= avg 
p = np.fft.fftn(a) 
#This part is much much faster than before (Original Post). 

#Keeping track of corresponding wavenumbers (k_x, k_y,k_z) for each element in p 
#This is just a convension on fourier transformation so you can ignore this part 
kValx = np.fft.fftfreq(ngridx , (1.0/ngridx)) 
kValy = np.fft.fftfreq(ngridy , (1.0/ngridy)) 
kValz = np.fft.fftfreq(ngridz , (1.0/ngridz)) 
kx = np.zeros((ngridx,ngridy,ngridz)) 
ky = np.zeros((ngridx,ngridy,ngridz)) 
kz = np.zeros((ngridx,ngridy,ngridz)) 
rangecolx = np.arange(ngridx) 
rangecoly = np.arange(ngridy) 
rangecolz = np.arange(ngridz) 
for row in np.arange(ngridx): 
    for column in np.arange(ngridy): 
     for height in np.arange(ngridz): 
      kx[row][column][height] = (kValx[row]) 
      ky[row][column][height] = (kValy[column]) 
      kz[row][column][height] = (kValz[height]) 
    if row%10 == 0: 
     print row 
print 'wavenumber generate complete!' 

#Calculating the average powerspectrum in terms of fixed K (Distance from origin to a point in fourier space) 
#by taking the spherical shell of thickness 1 and averaging out the values inside it. 
#I am sure that this process can be optimised somehow, but I gave up. 

qlen = maxK/2 #Nyquist frequency 
q = np.zeros(((qlen),4),dtype=complex) 
#q is a four column array with length maxK/2. 
#q[:,0] is integer wavenumber (K, which is the distance from the origin = sqrt(kx^2+ky^2+kz^2)) 
#q[:,1] is the sum of square of the fourier transformed value 
#q[:,2] is the sum of the fourier transformed value, 
#and q[:,3] is the total number of samples with K=q[:,0] 

for i in np.arange(len(q)): 
    q[i][0] = i 
i = 0 
for i in np.arange(len(p)): 
    for r in np.arange(len(p[0])): 
     for s in np.arange(len(p[0,0])): 
      K = np.around(np.sqrt(kx[i,r,s]**2+ky[i,r,s]**2+kz[i,r,s]**2)) 
      if K < qlen: 
       q[K][1]=q[K][1]+np.abs(p[i,r,s])**2 
       q[K][2]=q[K][2]+p[i,r,s] 
       q[K][3]=q[K][3]+1 
    if i%10 ==0: 
     print 'i = ',i,' !' 
print q 

Answer 1

Numpy может обычно делать вещи сотни раз быстрее, чем обычный питона, с очень небольшим дополнительным усилием. Вам просто нужно знать правильные способы написания кода. Просто назвать первые вещи, которые я думаю, из:

Indexing

Plain питон часто очень медленно на то, что компьютер должен быть большим на. Один пример - с индексированием, поэтому строка, подобная

a[f[i,0]][f[i,1]][f[i,2]]=f[i,3] 

вызывает у меня подозрение. Это тот, о котором вы говорите, когда вы говорите, что «преобразование загруженного txt в numpy ndarray» занимает очень много времени? Это меня не удивит, потому что каждый раз, когда python видит a[f[i,0]], он должен сначала индексировать f, убедившись, что i является целым числом, и вы не сбежали с края f; то он должен убедиться, что f[i,0] является целым числом, и вы не сбегаете с края a. Затем он должен повторить это еще два раза, прежде чем он даже узнает, какой элемент вы хотите установить.

Одним из улучшений является использование a[f[i,0],f[i,1],f[i,2]], потому что numpy быстрее с таким индексированием.

Но я предполагаю, что ваши данные на самом деле находятся в каком-то порядке. Например, делает f[i,2] цикл от 0 до 256, затем f[i,1] увеличивается на 1, а f [i, 2] начинается с 0? Если да, то все, что вам действительно нужно сделать, это сказать что-то вроде

a = f[:,3].reshape(ngridx,ngridy,ngridz) 

Это смехотворно быстрая операция, принимая доли миллисекунды. Форма может быть неправильной, поэтому вам может потребоваться изменить порядок аргументов, сделать что-то с транспозицией, но основная идея определенно существует. Вы можете прочитать об этом в the documentation.

Копирование данных плохо

Вам не нужно копировать все, и когда вам нужно скопировать массив (или часть массива), вы должны позволить NumPy сделать это для вас. Например, вместо функции Firstdel просто используйте a[1:]. Или, если вам действительно нужно, чтобы сделать копию данных (которые вы не просто для черчения) используйте:

def Firstdel(a): 
    return numpy.copy(a[1:]) 

Но в целом, вы можете просто использовать «кусочки» из Numpy массивов, а не копируя их. Читайте об этом here.

Loops

Loops также пресловутые испорченные время. Прежде всего, while не распространен в python для простых циклов. Таким образом, вместо

while i < len(f): 
    # do stuff 
    i = i+1 

вы, вероятно, следует использовать

for i in range(len(f)): 
    # do stuff 

Избавиться как много циклов, как вы можете. Чтобы установить kx, ky и kz, этот код примерно в 10 раз быстрее, чем вложенные циклы, но шкалы как N вместо N^3 (где N = ngridx ngridy ngridz):

for row in range(ngridx): 
    kx[row,:,:] = kValx[row] 
for column in range(ngridy): 
    ky[:,column,:] = kValy[column] 
for height in range(ngridz): 
    kz[:,:,height] = kValz[height] 

нарезка может также могут быть полезны при настройке значений, потому что цикл находится внутри numpy. Вместо этого кода

i = 0 
while i < len(q): 
    q[i][0] = i 
    i = i + 1 

просто использовать

q[:,0] = range(len(q)) 

Здесь, я просто установив «срез» q равного другой массив.

Вложенные петли после этого цикла также могут ускоряться, но они будут немного сложнее.

Но вы также должны избегать петель, где это возможно. Что приводит нас к ...

Использование встроенных функций Numpy

Причина NumPy существует, чтобы перевести эти медленные циклы питона и такие в быстрый C код (что нам не нужно знать, существует). Таким образом, есть много функций, которые делают то, что вы хотите сделать уже встроенным в numpy.

Вместо

while i < len(f): 
    masstot = masstot + f[i,3] 
    i = i+1 

вы должны использовать что-то вроде

masstot = np.sum(f[:,3]) 

Это проще читать, но и, вероятно, будет путь быстрее, потому что NumPy имеет прямой доступ к этим данным в памяти компьютера и может быстро использовать функции C, чтобы найти сумму, а не использовать медленные функции python. (Опять же, вам не нужно ничего знать о функциях C, они просто сделают эту работу.)

Вместо того, что больших вложенных значения петли рассчитать по K каждый раз через петлю, просто сделать K массив с соответствующими значениями:

K = np.around(np.sqrt(kx**2+ky**2+kz**2)) 

K будет иметь тот же размер, как kx и т.д. Затем , вы можете использовать advanced indexing для установки значений q. Это, как я хотел бы сделать это последний раздел:

# Again, we get rid of nested loops, to get a large improvement in speed and scaling 
K = np.around(np.sqrt(kx**2+ky**2+kz**2)).astype(int) 
for i in range(qlen): 
    indices = (K==i) # This will be an array of True/False values, 
        # which will be used for "advanced indexing" of p 
    q[i,0] = i 
    q[i,1] = sum(np.abs(p[indices])**2) 
    q[i,2] = sum(p[indices]) 
    q[i,3] = sum(indices) 
print q 

Собираю это все вместе, я получаю около 35 раз над улучшением коды в настоящее время в вашем вопросе.

Answer 2

Может также ускорить создание в входного файла, а также:

size = ngridx*ngridy*ngridz 
f = np.zeros((size,4)) 
a = np.arange(size) 
f[:, 0] = np.floor_divide(a, ngridy*ngridz) 
f[:, 1] = np.fmod(np.floor_divide(a, ngridz), ngridy) 
f[:, 2] = np.fmod(a, ngridz) 
f[:, 3] = np.random.rand(size) 

Чтобы kx, ky и kz, вы можете избавиться от петель с помощью broadcasting:

kx += kValx[:, np.newaxis, np.newaxis] 
ky += kValy[np.newaxis, :, np.newaxis] 
kz += kValz[np.newaxis, np.newaxis, :]