Optimizing an array mapping operation in python

问题

I am trying to get rid of an inefficient set of nested for loops in python. I have an array that I will call S(f_k,f_q) that needs to be mapped onto a different array that I will call Z(f_i,α_j). The arguments are all sampling frequencies. Both arrays have the same dimensions, which are user-selected. The mapping rule is fairly straightforward:

f_i = 0.5 · (f_k - f_q)
α_j = f_k + f_q

Currently I'm performing this via a series of nested for loops:

import numpy as np
nrows = 64
ncolumns = 16384
fk = np.fft.fftfreq(nrows)
fq = np.fft.fftfreq(ncolumns)
# using random numbers here to simplify the example
# in practice S is the result of several FFTs and complex multiplications
S = np.random.random(size=(nrows,ncolumns)) + 1j*np.random.random(size=(nrows,ncolumns))

fi = []
alphaj = []
Z = []
for k in range(-nrows//2,nrows//2):
    for q in range(-ncolumns//2,ncolumns//2):
        fi.append(0.5*(fk[k] - fq[q]))
        alphaj.append(fk[k] + fq[q])
        Z.append(S[k,q])

Obviously this is highly inefficient -- with this approach the mapping operation takes longer than the actual calculation of S (which in practice is the result of several FFT's and complex multiplications). I would like to find a way to vectorize this, but I'm having trouble coming up with the right approach. Any suggestions would be greatly appreciated.

Note: This is related to another question about how to store the results. Since this is about optimization I thought it would be better to create two separate questions.

回答1:

This doesn’t use the negative indexing of your original function but by returning arrays you can use normal indexing to map values

def weirdMath():
    nrows = 64
    ncolumns = 16384
    fk = np.fft.fftfreq(nrows)
    fq = np.fft.fftfreq(ncolumns)
    S = np.random.random(size=(nrows,ncolumns)) + 1j*np.random.random(size=(nrows,ncolumns))

    fi = .5*fk[:,np.newaxis] - fq
    alphaj = fk[:,np.newaxis] + fq
    return fi, alphaj, S


>>> f1,a1=weirdMath()
>>> f1.size
1048576
>>> f1[32,:10]
array([ 0.25      ,  0.24993896,  0.24987793,  0.24981689,  0.24975586,
        0.24969482,  0.24963379,  0.24957275,  0.24951172,  0.24945068])

Proof with rolling of axes added to match order of output in original code. Note: S was modified to np.arange() so that value comparison between functions could be directly matched:

def origCode():
    nrows = 64
    ncolumns = 16384
    fk = np.fft.fftfreq(nrows)
    fq = np.fft.fftfreq(ncolumns)
# using random numbers here to simplify the example
# in practice S is the result of several FFTs and complex multiplications
    #S = np.random.random(size=(nrows,ncolumns)) + 1j*np.random.random(size=(nrows,ncolumns))
    S = np.arange(nrows*ncolumns).reshape(nrows, ncolumns)
    fi = []
    alphaj = []
    Z = []

    for k in range(-nrows//2,nrows//2):
        for q in range(-ncolumns//2,ncolumns//2):
            fi.append(0.5*fk[k] - fq[q])
            alphaj.append(fk[k] + fq[q])
            Z.append(S[k,q])
    return fi, alphaj,Z 


def weirdMathWithRoll():
    nrows = 64
    ncolumns = 16384
    rowRollAdj = nrows%2
    colRollAdj = ncolumns%2

    fk = np.roll(np.fft.fftfreq(nrows), shift=(-nrows//2) + rowRollAdj, axis=0)

    fq = np.roll(np.fft.fftfreq(ncolumns), (-ncolumns//2) + colRollAdj)
    S = np.random.random(size=(nrows,ncolumns)) + 1j*np.random.random(size=(nrows,ncolumns))
    S = np.arange(nrows*ncolumns).reshape(nrows, ncolumns)
    s2 = np.roll(S,ncolumns//2 + colRollAdj, axis=1)
    s3 = np.roll(s2,nrows//2 + rowRollAdj, axis=0)

    fi = .5*fk[:,np.newaxis] - fq
    alphaj = fk[:,np.newaxis] + fq

    return fi, alphaj, s3

def testMath():
    f,a,z = origCode()
    f1,a1,s1 = weirdMathWithRoll()

    fMatch = f==f1.flatten()
    aMatch = a==a1.flatten()
    sMatch = z==s1.flatten()
    return (np.all(fMatch), np.all(aMatch), np.all(sMatch))

Output of proof:

>>> testMath()
(True, True, True)

Performance improvement:

>>> timeit.timeit(origCode, number = 1)
0.984715332997439
>>> timeit.timeit(weirdMathWithRoll, number = 1)
0.051891374998376705

回答2:

Does indexing with negative k values do what you want? In Python/numpy fk[-1] means last, fk[-2] means second to the last, etc.

In [90]: S = np.arange(1,11)                                                                           
In [91]: Z = []                                                                                        
In [92]: for k in range(-5,5): 
    ...:     Z.append(S[k]) 
    ...:                                                                                               

In [94]: S                                                                                             
Out[94]: array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])
In [95]: Z                                                                                             
Out[95]: [6, 7, 8, 9, 10, 1, 2, 3, 4, 5]

Or with slicing:

In [96]: np.concatenate([S[5:],S[:5]])                                                                 
Out[96]: array([ 6,  7,  8,  9, 10,  1,  2,  3,  4,  5])

来源：https://stackoverflow.com/questions/61104413/optimizing-an-array-mapping-operation-in-python