from complex import ComplexSIMD
from math import sin, cos
alias PI = 3.14159265359
alias width = DType.float64
fn fft[N: Int](x: ComplexSIMD[width, N]) -> ComplexSIMD[width, N]:
alias coeff = 2*PI/N
var rng = SIMD[width, N//2](0.)
for i in range(N):
rng[i] = i
var exp_real = sin(coeff * rng)
var exp_imag = cos(coeff * rng)
var exp_complex = ComplexSIMD(exp_real, exp_imag)
@parameter
if N == 1:
return x
var real = x.re
var imag = x.im
var real_deinterleave = real.deinterleave()
var imag_deinterleave = imag.deinterleave()
var real_even = real_deinterleave[0]
var real_odd = real_deinterleave[1]
var imag_even = imag_deinterleave[0]
var imag_odd = imag_deinterleave[1]
var even_fft = fft(ComplexSIMD(real_even, imag_even))
var odd_fft = fft(ComplexSIMD(real_odd, imag_odd))
var T = odd_fft * exp_complex
var first_half = even_fft + T
var second_half = even_fft + (-T)
var final_real = first_half.re.join(second_half.re)
var final_imag = first_half.im.join(second_half.im)
var final = ComplexSIMD(final_real, final_imag)
return final
问题来自于final声明。编译器说:error: cannot implicitly convert 'ComplexSIMD[f64, mul(2, floordiv(N, 2))]' value to 'ComplexSIMD[f64, N]' in return value。我知道,一般来说,N // 2 *2 != N ,但是有没有办法告诉编译器我只会给出正确的情况?
我知道这段代码不正确,但即使它是正确的,我想我仍然会遇到这个问题。
ASIDE :功能不如.因此,我必须对 FFT 的实部和虚部进行大量不雅的分离。如果有人能想出更好的方法来做到这一点,那将不胜感激。 ComplexSIMDSIMD