wydajne numpy.roll przed numpy.sum() lub mean()

Question

Mam kilka (rzędu 1000) 3D-tablic kształtu (1000, 800, 1024) Chcę się uczyć. Muszę obliczyć średnią wzdłuż osi = 0, ale zanim to zrobię, muszę przetoczyć dane wzdłuż osi 2, aż znajdzie się "we właściwym miejscu".

To brzmi dziwnie, więc spróbuję wyjaśnić. Pod-macierz kształtu 1D (1024,) to dane z fizycznego bufora pierścieniowego. Bufor pierścieniowy jest odczytywany w różnych miejscach, które znam. Tak więc mam kilka tablic pos o kształcie (1000, 800). Mówienie mi, w jakiej pozycji został odczytany bufor pierścieniowy. I moje tablice 3D o rozmiarze (1000, 800, 1024), które muszę przetasować zgodnie z pos.

Dopiero po przewróceniu .. tablice 3D są dla mnie znaczące i mogę zacząć je analizować. W C można napisać kod, który robi to bardzo prosto, więc zastanawiam się, czy mogę rodzaju, "powiedz" numpy mean() lub sum() powinny one zacząć od różnych wskaźników i "roll around" na końcu podobrazia 1D.

Co ja obecnie zrobić to w ten sposób:

rolled = np.zeros_like(data) # shape (1000, 800, 1024) 
for a in range(rolled.shape[0]): 
    for b in range(rolled.shape[1]): 
     rolled[a,b] = np.roll(data[a,b], pos[a,b]) 

ten trwa ~ 60sec I wtedy zrobić np:

m = rolled.mean(axis=0) 
s = rolled.std(axis=0) 

Która zajmuje tylko 15sek lub tak.

Chodzi mi o to, że wykonanie zwiniętej kopii zajmuje dużo miejsca i czasu (w porządku, mogłem zaoszczędzić miejsce, pisząc zwinięty materiał z powrotem do data), podczas gdy zdecydowanie istnieje sposób (w C) do realizacji tego uśredniania i toczą się w jednej pętli, dzięki czemu oszczędzają dużo czasu. Moje pytanie brzmi ... czy istnieje sposób na zrobienie czegoś podobnego z numpy?

Answer 1

Nudzę się i napisałem twoją funkcję w Cython. Działa około 10 razy szybciej niż opublikowany kod, bez przydzielania tablicy pośredniej.

import numpy as np 
cimport numpy as np 
cimport cython 
from libc.math cimport sqrt 

@cython.boundscheck(False) 
@cython.wraparound(False) 
@cython.nonecheck(False) 
@cython.cdivision(True) 
def rolled_mean_std(double[:,:,::1] data, int[:,::1] pos): 
    cdef Py_ssize_t s1,s2,s3,a,b,c 
    s1 = data.shape[0] 
    s2 = data.shape[1] 
    s3 = data.shape[2] 
    cdef double[:,::1] sums = np.zeros((s2,s3)) 
    cdef double[:,::1] sumsq = np.zeros((s2,s3)) 
    cdef double d 
    cdef int p 
    # Compute sums and sum-of-squares. 
    for a in range(s1): 
    for b in range(s2): 
     p = pos[a,b] 
     for c in range(s3): 
     d = data[a,b,(c+s3-p)%s3] 
     sums[b,c] += d 
     sumsq[b,c] += d * d 
    # Calculate mean + std in place. 
    for b in range(s2): 
    for c in range(s3): 
     d = sums[b,c] 
     sums[b,c] /= s1 
     sumsq[b,c] = sqrt((s1*sumsq[b,c] - (d*d)))/s1 
    return sums, sumsq 

Zauważ, że ten wykorzystuje Naive mean+stdv algorithm, więc może napotkasz zmiennoprzecinkowej błędów dokładności. Moje testy nie wykazały jednak większego efektu.