Calcolo più performante delle forze newtoniane in numpy/scipy

Question

Per un esercizio universitario, abbiamo dovuto implementare uno Leapfrog integrator con forze Newtoniane esatte in Python. Il corso è finito e le nostre soluzioni erano più che sufficienti, ma mi chiedevo se/come si potesse migliorare ulteriormente le prestazioni del calcolo delle forze.

il collo di bottiglia per calcolare tutte le forze (alias accelerazioni):

per un grande (1000 e oltre) numero di particelle N (i, j < N).

Qui r e r sono i vettori 3-dimensionale della posizione delle particelle immagazzinato in un ndarray di forma (N, 3) e Gm è la particella massa per la costante gravitazionale che ho salvato in un narray di forma (N).

la versione più veloce che ho trovato finora è la seguente:

def a(self): 
    sep = self.r[np.newaxis, :] - self.r[:, np.newaxis] 
    # Calculate the distances between all particles with cdist 
    # this is much faster than by hand 
    dists = cdist(self.r, self.r) 
    scale =dists*dists*dists 
    # set diagonal elements of dist to something != 0, to avoid division by 0 
    np.fill_diagonal(scale,1) 
    Fsum = (sep/scale.reshape(self.particlenr,self.particlenr,1))*self.Gm[:,None] 
    return np.add.reduce(Fsum,axis=1) 

ma mi bug che questo non è probabilmente la versione più veloce. La prima riga sembra essere troppo lenta quando si confronta con il cd che sta facendo essenzialmente lo stesso calcolo. Inoltre, questa soluzione non utilizza la simmetria di commutazione r e r nel problema e calcola tutti gli elementi due volte.

Conosci qualche miglioramento delle prestazioni (senza modificare il calcolo della forza in una certa approssimazione o cambiare il linguaggio di programmazione)?

Answer 1

Dubito che Numpy stia effettivamente calcolando due volte le distanze (poiché sarà sempre simmetrica). Probabilmente sta facendo un calcolo e assegnando lo stesso valore in due punti.

Alcune delle idee si verificano a me, però:

1. Si potrebbe seguire la fonte NumPy e scrivere una versione personalizzata di cdist. È possibile che il programma analizzi molte opzioni ogni volta che si esegue l'iterazione di. Non molto, ma forse potrebbe darti una piccola percentuale di urti.
2. Pre-allocazione. Ogni volta che si esegue a(), si sta potenzialmente riallocando la memoria per tutti i valori intermedi della matrice. Potresti fare queste quantità persistenti?

Non ho ancora eseguito i calcoli, ma non sarei sorpreso se potessi ridurre elegantemente i calcoli simmetrici ridondanti in qualche modo.

Answer 2

darò una prova: ho implementato una routine, che determina un unico a_i:

import numpy as np 

GM = .01 # article mass times the gravitation 

def calc_a_i(rr, i): 
    """ Calculate one a_i """ 
    drr = rr - rr[i, :] # r_j - r_i 
    dr3 = np.linalg.norm(drr, axis=1)**3 # |r_j - r_i|**3 
    dr3[i] = 1 # case i==j: drr = [0, 0, 0] 
    # this would be more robust (elimnate small denominators): 
    #dr3 = np.where(np.abs(dr3) > 1e-12, dr3, 1) 
    return np.sum(drr.T/dr3, axis=1) 

n = 4000 # number of particles 
rr = np.random.randn(n, 3) # generate some particles 

# Calculate each a_i separately: 
aa = np.array([calc_a_i(rr, i) for i in range(n)]) * GM # all a_i 

Per provarlo, mi sono imbattuto:

In [1]: %timeit aa = np.array([calc_a_i(rr, i) for i in range(n)]) 
1 loops, best of 3: 2.93 s per loop 

Il modo più semplice per accelerare tale codice, è quello di utilizzare numexpr per la valutazione rapida delle espressioni di matrice:

import numexpr as ne 
ne.set_num_threads(1) # multithreading causes to much overhead 

def ne_calc_a_i(i): 
    """ Use numexpr - here rr is global for easier parallelization""" 
    dr1, dr2, dr3 = (rr - rr[i, :]).T # r_j - r_i 
    drrp3 = ne.evaluate("sqrt(dr1**2 + dr2**2 + dr3**2)**3") 
    drrp3[i] = 1 
    return np.sum(np.vstack([dr1, dr2, dr3])/drrp3, axis=1) 

# Calculate each a_i separately: 
aa_ne = np.array([ne_calc_a_i(i) for i in range(n)]) * GM # all a_i  

Thi s migliora la velocità di un fattore di 2:

In [2]: %timeit aa_ne = np.array([ne_calc_a_i(i) for i in range(n)]) 
    1 loops, best of 3: 1.29 s per loop 

Per accelerare il codice ulteriormente, diciamo eseguirlo IPython Cluster:

# Start local cluster with 4 clients in a shell with: 
# ipcluster start -n 4 

rc = Client() # clients of cluster 
dview = rc[:] # view of clusters 

dview.execute("import numpy as np") # import libraries on clients 
dview.execute("import numexpr as ne") 
dview.execute("ne.set_num_threads(1)") 

def para_calc_a(dview, rr): 
    """ Only in function for %timeit """ 
    # send rr and ne_calc_a_i() to clients: 
    dview.push(dict(rr=rr, ne_calc_a_i=ne_calc_a_i), block=True) 
    return np.array(dview.map_sync(ne_calc_a_i, range(n)))*GM 

L'aumento di velocità è più di un fattore di quattro:

Come già notato @mathdan, non è ovvio come ottimizzare un tale problema: dipende dall'architettura della CPU se il bus di memoria o l'unità in virgola mobile è il fattore limitante, che chiama f o tecniche diverse.

Per ulteriori vantaggi si potrebbe voler guardare Theano: Può generare dinamicamente codice codice GPU da Python.

Answer 3

Il seguente è un po 'più ottimale:

import numpy as np 
from scipy.spatial.distance import pdist, squareform  

def a6(r, Gm): 
    dists = pdist(r) 
    dists *= dists*dists 
    dists = squareform(dists) 
    np.fill_diagonal(dists, 1.) 
    sep = r[np.newaxis, :] - r[:, np.newaxis] 
    return np.einsum('ijk,ij->ik', sep, Gm/dists) 

guadagno velocità è principalmente dovuta alla linea einsum; L'utilizzo di pdist e squareform come questo è solo marginalmente più veloce rispetto al modo originale con cdist.

Puoi ancora un po 'fermarti, ad es. usando threading e Numba (richiesta la versione 0.17.0). Anche se il codice sotto è molto brutto e sicuramente può essere migliorato molto, è piuttosto veloce.

import numpy as np 
import math 
from numba import jit 
from threading import Thread 
NUM_THREADS = 2 # choose wisely 

def a_numba_par(r, Gm): 
    a = np.zeros_like(r) 
    N = r.shape[0] 

    offset = range(0, N+1, N//NUM_THREADS) 
    chunks = zip(offset, offset[1:]) 
    threads = [Thread(target=_numba_loop, args=(r,Gm,a)+c) for c in chunks] 

    for thread in threads: 
     thread.start() 
    for thread in threads: 
     thread.join() 

    return a 

@jit(nopython=True, nogil=True) 
def _numba_loop(r, Gm, a, i1, i2): 
    N = r.shape[0] 
    for i in range(i1, i2): 
     _helper(r, Gm, i, 0 , i, a[i,:]) 
     _helper(r, Gm, i, i+1, N, a[i,:]) 
    return a 

@jit(nopython=True, nogil=True) 
def _helper(r, Gm, i, j1, j2, a): 
    for j in range(j1, j2): 
     dx = r[j,0] - r[i,0] 
     dy = r[j,1] - r[i,1] 
     dz = r[j,2] - r[i,2] 

     sqeuc = dx*dx + dy*dy + dz*dz 
     scale = Gm[j]/(sqeuc * math.sqrt(sqeuc)) 

     a[0] += scale * dx 
     a[1] += scale * dy 
     a[2] += scale * dz