使用 Arrayfire 在 Python 中更快地对复杂数组求幂

Question

根据 arrayfire pow documentation，af.pow() 目前仅支持实数组的幂（和根...）。没有抛出任何错误，但我发现对复杂输入使用 af.pow() 会导致巨大的内存泄漏，特别是如果使用其他函数作为输入（例如，af.pow(af.ifft(array), 2)）。

为了解决这个问题，我在下面编写了函数 complexPow。这似乎运行对于没有内存泄漏的复杂数组，快速比较表明我的 complexPow 函数 returns 与 numpy.sqrt() 和 ** 的值相同例如运算符。

def complexPow(inData, power):
    for i in af.ParallelRange(inData.shape[0]):
        theta = af.atan(af.imag(inData[i])/af.real(inData[i]))
        rSquared = af.pow(af.real(inData[i]), 2.0) + \
                    af.pow(af.imag(inData[i]), 2.0)
        r = af.pow(rSquared, .5)
        inData[i] = af.pow(r, power) * (af.cos(theta*power) + \
                1j*af.sin(theta*power))
    return inData

有没有比这更快的并行逐元素求幂的方法？我还没有找到，但害怕我在这里错过了一个技巧......

Answer 1

如果没有并行 for 循环，这会快一点：

def complexPow(inData, power):
    theta = af.atan(af.imag(inData)/af.real(inData))
    r = af.pow(af.pow(af.real(inData), 2.0) + 
                af.pow(af.imag(inData), 2.0), .5)
    inData = af.pow(r, power) * (af.cos(theta*power) + \
                1j*af.sin(theta*power))
    return inData

使用 nvidia Quadro K4200、Spyder 3、Python 2.7、Windows 7:[=15= 在维度 (1, 2**18) 的 dtype=complex 数组上测试了 4000 次迭代]

使用 af.ParallelRange: 7.64 秒（每次迭代 1.91 毫秒）。

上述方法： 5.94 秒（每次迭代 1.49 毫秒）。

速度提升： 28%.

使用 Arrayfire 在 Python 中更快地对复杂数组求幂

Faster exponentiation of complex arrays in Python using Arrayfire

python

gpu

gpgpu

pow

arrayfire