在python中优化数组映射操作

Question

我正试图摆脱python中一组低效的嵌套for循环。我有一个数组，我将调用S(f_k，f_q，α_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的实际计算花费更长的时间(在实践中，S是几个FFT和复数乘法的结果(。我想找到一种将其矢量化的方法，但我很难找到正确的方法。如有任何建议，我们将不胜感激。

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

Answer 1

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

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

Answer 2

使用负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])