在 python 中优化数组映射操作

Question

我正试图摆脱 python 中一组低效的嵌套 for 循环。我有一个我将调用 S(f_k,f_q) 的数组，它需要映射到另一个我将调用 Z 的数组(f_i,α_j)。参数都是采样频率。两个阵列具有相同的尺寸，由用户选择。映射规则相当简单：

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

目前我正在通过一系列嵌套的 for 循环执行此操作：

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])

显然，这是非常低效的——使用这种方法，映射操作比 S 的实际计算（实际上是多个 FFT 和复数乘法的结果）花费的时间更长。我想找到一种对此进行矢量化的方法，但我在想出正确的方法时遇到了麻烦。任何建议将不胜感激。

注意：这与关于如何存储结果有关。由于这是关于优化的，所以我认为创建两个单独的问题会更好。

Answer 1

使用负 k 值进行索引是否符合您的要求？在Python/numpy中fk[-1]表示最后一个，fk[-2]表示倒数第二个，依此类推

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]

或切片：

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

Answer 2

这不使用原始函数的负索引，但通过返回数组，您可以使用正常索引来映射值

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])

添加轴滚动的证明以匹配原始代码中的输出顺序。注：S修改为np.arange()，可以直接匹配函数间的值比较：

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))

证明输出：

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

性能提升：

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

在 python 中优化数组映射操作

Optimizing an array mapping operation in python

python

optimization

numpy

vectorization