69 lines
2 KiB
Swift
69 lines
2 KiB
Swift
import Foundation
|
|
import Numerics
|
|
|
|
typealias Complex = Numerics.Complex<Double>
|
|
|
|
extension Complex {
|
|
var exp: Complex {
|
|
Complex(cos(imaginary), sin(imaginary)) * Complex(cosh(real), sinh(real))
|
|
}
|
|
|
|
var pretty: String {
|
|
let fmt = { String(format: "%1.3f", $0) }
|
|
let re = fmt(real)
|
|
let im = fmt(abs(imaginary))
|
|
|
|
if im == "0.000" {
|
|
return re
|
|
} else if re == "0.000" {
|
|
return im
|
|
} else if imaginary > 0 {
|
|
return re + " + " + im + "i"
|
|
} else {
|
|
return re + " - " + im + "i"
|
|
}
|
|
}
|
|
}
|
|
|
|
func fft(_ array: [Complex]) -> [Complex] { _fft(array, direction: Complex(0.0, 2.0), scalar: 1) }
|
|
func rfft(_ array: [Complex]) -> [Complex] { _fft(array, direction: Complex(0.0, -2.0), scalar: 2) }
|
|
|
|
private func _fft(_ arr: [Complex], direction: Complex, scalar: Double) -> [Complex] {
|
|
guard arr.count > 1 else {
|
|
return arr
|
|
}
|
|
|
|
let n = arr.count
|
|
let cScalar = Complex(scalar, 0)
|
|
|
|
precondition(n % 2 == 0, "The Cooley-Tukey FFT algorithm only works when the length of the input is even.")
|
|
|
|
var (evens, odds) = arr.lazy.enumerated().reduce(into: ([Complex](), [Complex]()), {res, cur in
|
|
if cur.offset & 1 == 0 {
|
|
res.0.append(cur.element)
|
|
} else {
|
|
res.1.append(cur.element)
|
|
}
|
|
})
|
|
|
|
evens = _fft(evens, direction: direction, scalar: scalar)
|
|
odds = _fft(odds, direction: direction, scalar: scalar)
|
|
|
|
let (left, right) = (0 ..< n / 2).map({i -> (Complex, Complex) in
|
|
let offset = (direction * Complex((.pi * Double(i) / Double(n)), 0)).exp * odds[i] / cScalar
|
|
let base = evens[i] / cScalar
|
|
|
|
return (base + offset, base - offset)
|
|
}).reduce(into: ([Complex](), [Complex]()), {res, cur in
|
|
res.0.append(cur.0)
|
|
res.1.append(cur.1)
|
|
})
|
|
|
|
return left + right
|
|
}
|
|
|
|
let dat = [Complex(1.0, 0.0), Complex(1.0, 0.0), Complex(1.0, 0.0), Complex(1.0, 0.0),
|
|
Complex(0.0, 0.0), Complex(0.0, 2.0), Complex(0.0, 0.0), Complex(0.0, 0.0)]
|
|
|
|
print(fft(dat).map({ $0.pretty }))
|
|
print(rfft(f).map({ $0.pretty }))
|