fft.go

// Package gofft provides a fast discrete Fourier transformation algorithm.
//
// Implemented is the 1-dimensional DFT of complex input data
// for with input lengths which are powers of 2.
//
// The algorithm is non-recursive, works in-place overwriting
// the input array, and requires O(1) additional space.
package gofft

import (
	"fmt"
	"sync"
)

var (
	prepareLock sync.RWMutex
	factors     = [64]complex128{
		complex(1, 0),
		complex(-1, 0),
		complex(0, -1),
		complex(0.7071067811865475244008443621048490392845, -0.7071067811865475244008443621048490392845),
		complex(0.9238795325112867561281831893967882868226, -0.3826834323650897717284599840303988667606),
		complex(0.9807852804032304491261822361342390369746, -0.1950903220161282678482848684770222409278),
		complex(0.9951847266721968862448369531094799215744, -0.09801714032956060199419556388864184586113),
		complex(0.9987954562051723927147716047591006944438, -0.04906767432741801425495497694268265831472),
		complex(0.9996988186962042201157656496661721968491, -0.02454122852291228803173452945928292506544),
		complex(0.9999247018391445409216464911963832243506, -0.01227153828571992607940826195100321214037),
		complex(0.9999811752826011426569904377285677161734, -0.006135884649154475359640234590372580917052),
		complex(0.9999952938095761715115801257001198995544, -0.003067956762965976270145365490919842518943),
		complex(0.9999988234517019099290257101715260190497, -0.001533980186284765612303697150264079079954),
		complex(0.9999997058628822191602282177387656771166, -0.0007669903187427045269385683579485766431393),
		complex(0.9999999264657178511447314807073878569491, -0.0003834951875713955890724616811813812633942),
		complex(0.9999999816164292938083469154029097145039, -0.0001917475973107033074399095619890009334687),
		complex(0.9999999954041073128909719331396061489591, -9.587379909597734587051721097647635118683e-05),
		complex(0.9999999988510268275626733077945541083999, -4.793689960306688454900399049465887274681e-05),
		complex(0.9999999997127567068494139722186417760899, -2.396844980841821872918657716502182009474e-05),
		complex(0.9999999999281891767097750958838504902617, -1.198422490506970642152156159698898480473e-05),
		complex(0.9999999999820472941772826241477841073806, -5.992112452642427842879711808890861729988e-06),
		complex(0.9999999999955118235443105841729973244406, -2.996056226334660750454812808357059811825e-06),
		complex(0.9999999999988779558860770165517525365038, -1.49802811316901122885427884615536112069e-06),
		complex(0.9999999999997194889715192147947195844515, -7.490140565847157211304985667306556371544e-07),
		complex(0.9999999999999298722428798012397287367588, -3.745070282923841239031691790846331773967e-07),
		complex(0.9999999999999824680607199501562477367286, -1.872535141461953448688245765935636171201e-07),
		complex(0.9999999999999956170151799875294566562175, -9.362675707309808279906728668088562019313e-08),
		complex(0.9999999999999989042537949968817638341817, -4.681337853654909269511551813854009695942e-08),
		complex(0.9999999999999997260634487492204034379288, -2.340668926827455275950549341903484403787e-08),
		complex(0.9999999999999999315158621873050985144444, -1.170334463413727718124621350323810379807e-08),
		complex(0.9999999999999999828789655468262744820468, -5.851672317068638690809790100834139694376e-09),
		complex(0.9999999999999999957197413867065686113524, -2.92583615853431935792823046906895590201e-09),
		complex(0.999999999999999998929935346676642152265, -1.462918079267159680529532161865963710372e-09),
		complex(0.9999999999999999997324838366691605380295, -7.314590396335798404604431968494175751862e-10),
		complex(0.9999999999999999999331209591672901345066, -3.657295198167899202546812379142632525948e-10),
		complex(0.9999999999999999999832802397918225336267, -1.828647599083949601303980738933259344228e-10),
		complex(0.9999999999999999999958200599479556334074, -9.143237995419748006558121881368725572707e-11),
		complex(0.9999999999999999999989550149869889083504, -4.571618997709874003283838214022166392804e-11),
		complex(0.9999999999999999999997387537467472270869, -2.285809498854937001642516266178308647208e-11),
		complex(0.9999999999999999999999346884366868067732, -1.142904749427468500821332777985057504954e-11),
		complex(0.999999999999999999999983672109171701694, -5.714523747137342504106757196045166501456e-12),
		complex(0.999999999999999999999995918027292925422, -2.857261873568671252053390261287568122819e-12),
		complex(0.9999999999999999999999989795068232313555, -1.428630936784335626026696588551907170421e-12),
		complex(0.9999999999999999999999997448767058078403, -7.143154683921678130133484765144689738369e-13),
		complex(0.9999999999999999999999999362191764519594, -3.571577341960839065066742610370489104958e-13),
		complex(0.9999999999999999999999999840547941129913, -1.785788670980419532533371333660012581954e-13),
		complex(0.9999999999999999999999999960136985282486, -8.928943354902097662666856703893522946613e-14),
		complex(0.9999999999999999999999999990034246320621, -4.464471677451048831333428356395943977908e-14),
		complex(0.9999999999999999999999999997508561580163, -2.232235838725524415666714178754119802035e-14),
		complex(0.9999999999999999999999999999377140395026, -1.116117919362762207833357089446578377652e-14),
		complex(0.9999999999999999999999999999844285098756, -5.580589596813811039166785447319789984036e-15),
		complex(0.9999999999999999999999999999961071274689, -2.790294798406905519583392723670757253998e-15),
		complex(0.9999999999999999999999999999990267818665, -1.395147399203452759791696361836736409743e-15),
		complex(0.9999999999999999999999999999997566954681, -6.975736996017263798958481809185379277168e-16),
		complex(0.9999999999999999999999999999999391738656, -3.487868498008631899479240904592901792128e-16),
		complex(0.9999999999999999999999999999999847934671, -1.74393424900431594973962045229647741526e-16),
		complex(0.999999999999999999999999999999996198366, -8.719671245021579748698102261482420225305e-17),
		complex(0.9999999999999999999999999999999990495922, -4.35983562251078987434905113074121425627e-17),
		complex(0.9999999999999999999999999999999997623973, -2.179917811255394937174525565370607646089e-17),
		complex(0.9999999999999999999999999999999999405993, -1.089958905627697468587262782685303887788e-17),
		complex(0.9999999999999999999999999999999999851506, -5.449794528138487342936313913426519519873e-18),
		complex(0.9999999999999999999999999999999999962884, -2.724897264069243671468156956713259770051e-18),
		complex(0.9999999999999999999999999999999999990714, -1.36244863203462183573407847835662988629e-18),
		complex(0.9999999999999999999999999999999999997678, -6.812243160173109178670392391783149433029e-19),
	}
	permMap = map[int][]int{}
)

// Prepare precomputes values used for FFT on a vector of length N.
// N must be a perfect power of 2, otherwise this will return an error.
func Prepare(N int) error {
	if !IsPow2(N) {
		return fmt.Errorf("Input dimension must be power of 2, is: %d", N)
	}
	prepareLock.RLock()
	if _, ok := permMap[N]; ok {
		prepareLock.RUnlock()
		// Already prepared, no need to do anything
		return nil
	}
	prepareLock.RUnlock()
	prepareLock.Lock()
	defer prepareLock.Unlock()
	permMap[N] = permutationIndex(N)
	return nil
}

// FFT implements the fast Fourier transform.
// This is done in-place (modifying the input array).
// Requires O(1) additional memory.
// len(x) must be a perfect power of 2, otherwise this will return an error.
func FFT(x []complex128) error {
	N, perm, err := getVars(x)
	if err != nil {
		return err
	}
	fft(x, N, perm)
	return nil
}

// IFFT implements the inverse fast Fourier transform.
// This is done in-place (modifying the input array).
// Requires O(1) additional memory.
// len(x) must be a perfect power of 2, otherwise this will return an error.
func IFFT(x []complex128) error {
	N, perm, err := getVars(x)
	if err != nil {
		return err
	}
	ifft(x, N, perm)
	return nil
}

// Pre-load the fft variables for later use.
func getVars(x []complex128) (N int, perm []int, err error) {
	N = len(x)
	err = Prepare(N)
	perm = permMap[N]
	return
}

// fft does the actual work for FFT
func fft(x []complex128, N int, perm []int) {
	// Handle small N quickly
	switch N {
	case 1:
		return
	case 2:
		x[0], x[1] = x[0]+x[1], x[0]-x[1]
		return
	case 4:
		f := complex(imag(x[1])-imag(x[3]), real(x[3])-real(x[1]))
		x[0], x[1], x[2], x[3] = x[0]+x[1]+x[2]+x[3], x[0]-x[2]+f, x[0]-x[1]+x[2]-x[3], x[0]-x[2]-f
		return
	}
	// Reorder the input array.
	permute(x, perm, N)
	s := 0
	// Butterfly
	for n := 1; n < N; n <<= 1 {
		s++
		w := factors[s]
		for o := 0; o < N; o += (n << 1) {
			wj := complex(1, 0)
			for k := 0; k < n; k++ {
				i := k + o
				f := wj * x[i+n]
				x[i], x[i+n] = x[i]+f, x[i]-f
				wj *= w
			}
		}
	}
}

// ifft does the actual work for IFFT
func ifft(x []complex128, N int, perm []int) {
	// Reverse the input vector
	for i := 1; i < N/2; i++ {
		j := N - i
		x[i], x[j] = x[j], x[i]
	}

	// Do the transform.
	fft(x, N, perm)

	// Scale the output by 1/N
	invN := complex(1.0/float64(N), 0)
	for i := 0; i < N; i++ {
		x[i] *= invN
	}
}

// permutationIndex builds the bit-inverted index vector,
// which is needed to permutate the input data.
func permutationIndex(N int) []int {
	index := make([]int, N)
	index[0] = 0 // Initial sequence for N=1
	// For every next power of two, the sequence is multiplied by 2 in-place.
	// Then the result is also appended to the end and increased by one.
	for n := 1; n < N; n <<= 1 {
		for i := 0; i < n; i++ {
			index[i] <<= 1
			index[i+n] = index[i] + 1
		}
	}
	// Re-arrange the permutation to just the necessary swaps
	for i := 1; i < N-1; i++ {
		ind := index[i]
		for ind < i {
			ind = index[ind]
		}
		index[i] = ind
	}
	return index
}

// permutate permutes the input vector according to the permutation vector.
// Uses an in-place algorithm that on FFT permutation vectors runs in O(N) time.
func permute(x []complex128, perm []int, N int) {
	// perm[0] is always 0, and perm[N-1] is always N-1, so skip those
	for i := 1; i < N-1; i++ {
		x[i], x[perm[i]] = x[perm[i]], x[i]
	}
}