在python中优化数组映射操作 [英] Optimizing an array mapping operation in python

问题描述

我正试图摆脱python中无效的嵌套for循环集.我有一个数组，我将其称为S(f _k ，f _q )，该数组需要映射到另一个称为Z(f _i ，α _j ).参数是所有采样频率.这两个阵列具有相同的尺寸，这些尺寸是用户选择的.映射规则非常简单:

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

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

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

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

显然，这是非常低效的-使用这种方法，映射操作所需的时间比S的实际计算要长(实际上，这是几次FFT和复数乘法的结果).我想找到一种向量化的方法，但是我在想出正确的方法时遇到了麻烦.任何建议将不胜感激.

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.

推荐答案

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

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

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

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

证明输出:

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

性能改进:

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

这篇关于在python中优化数组映射操作的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！