uopi/internal/dsp/fft.go

package dsp

import "math"

// nextPow2 returns the smallest power of two >= n (and at least 1).
func nextPow2(n int) int {
	p := 1
	for p < n {
		p <<= 1
	}
	return p
}

// fftRadix2 computes the in-place iterative radix-2 Cooley-Tukey FFT of the
// complex signal held in re/im. len(re) == len(im) must be a power of two. The
// transform overwrites re/im with the frequency-domain result. This is a small
// self-contained implementation (no external dependency) used by the synthetic
// fft op.
func fftRadix2(re, im []float64) {
	n := len(re)
	if n <= 1 {
		return
	}

	// Bit-reversal permutation.
	for i, j := 1, 0; i < n; i++ {
		bit := n >> 1
		for ; j&bit != 0; bit >>= 1 {
			j ^= bit
		}
		j ^= bit
		if i < j {
			re[i], re[j] = re[j], re[i]
			im[i], im[j] = im[j], im[i]
		}
	}

	// Danielson-Lanczos butterflies.
	for length := 2; length <= n; length <<= 1 {
		ang := -2 * math.Pi / float64(length)
		wReal, wImag := math.Cos(ang), math.Sin(ang)
		for i := 0; i < n; i += length {
			curReal, curImag := 1.0, 0.0
			half := length >> 1
			for k := 0; k < half; k++ {
				a := i + k
				b := i + k + half
				tReal := curReal*re[b] - curImag*im[b]
				tImag := curReal*im[b] + curImag*re[b]
				re[b] = re[a] - tReal
				im[b] = im[a] - tImag
				re[a] += tReal
				im[a] += tImag
				curReal, curImag = curReal*wReal-curImag*wImag, curReal*wImag+curImag*wReal
			}
		}
	}
}