Более быстрое экспонирование сложных массивов в Python с использованием Arrayfire

Question

В соответствии с arrayfire pow documentation, af.pow() в настоящее время поддерживает только реальные (и корни ...) реальных массивов. Ошибка не возникает, но я обнаружил, что использование af.pow() со сложным входом может привести к огромной утечке памяти, особенно если в качестве входных данных используются другие функции (например, af.pow(af.ifft(array), 2)).

Чтобы обойти это, я написал функцию complexPow ниже. Это похоже на сложные массивы без утечки памяти, и быстрое сравнение показало, что моя функция complexPow возвращает те же значения, что и 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 

Tetsted за 4000 итераций по dtype=complex массив с размеры (1, 2**18) с использованием NVidia Quadro K4200, Spyder 3, Python 2.7, Windows 7:

Использование af.ParallelRange: 7,64 сек (1,91 мс на итерации).

Метод выше: 5.94 с (1,49 мсек на итерацию).

Увеличение скорости: 28%.