Fast fourier transform.

fft.go

package fft

import (
	"math"
	"math/cmplx"
)

func FFT(input []complex128) []complex128 {

	fftNums, fftLen, NPow, WN, fftOrder := prepair(input)

	for count := 0; count < int(NPow); count++ {
		pair := int(math.Pow(2, float64(count)))
		step := int(math.Pow(2, float64(count+1)))
		for i := 0; i < fftLen; i++ {
			if i%step < pair {
				inow := fftOrder[i]
				ipair := fftOrder[i+pair]
				k := fftLen / step * (i % step)

				tmpi := fftNums[inow] + WN[k]*fftNums[ipair]
				fftNums[ipair] = fftNums[inow] - WN[k]*fftNums[ipair]
				fftNums[inow] = tmpi
			}
		}
	}
	fftre := make([]complex128, fftLen)
	for k := range fftNums {
		fftre[k] = fftNums[fftOrder[k]]
	}
	return fftre

}

func cvrt(size, x int) int {
	re := 0
	for i := 0; i < size; i++ {
		bit := (x >> uint(i)) & 1
		if bit == 0 {
			continue
		}
		re += bit << uint(size-i-1)
	}
	return re
}

func prepair(input []complex128) (fftInputs []complex128, fftLen int, NPow float64, WN []complex128, fftOrder []int) {

	rawLen := len(input)
	NPow = math.Floor(math.Log2(float64(rawLen))-0.001) + 1
	fftLen = int(math.Pow(2, NPow))
	fftLenD2 := fftLen / 2
	WN = make([]complex128, fftLen/2)
	fftOrder = make([]int, fftLen)

	for i := 0; i < fftLen; i++ {
		if i >= rawLen {
			input = append(input, 0)
		}
		if i < fftLenD2 {
			WN[i] = cmplx.Rect(1, -2*math.Pi/float64(fftLen)*float64(i))
		}
		fftOrder[i] = cvrt(int(NPow), i)
	}
	fftInputs = input
	return
}