Add spectrum analysis functionality (#2)

Add support for generating frequency spectrum snapshots of audio buffers using FFT analysis.

This includes:

- New spectrum view for visualizing frequency data using SwiftUI Charts
- Extension for AVAudioPCMBuffer to compute spectrum data
- Test coverage for various waveforms including pure tones and noise

Author: jcavar

Sources/AudioSnapshotTesting/AVAudioPCMBufferExtensions.swift

@@ -63,3 +63,132 @@ private extension AVAudioPCMBuffer {
         return buffer
     }
 }
+
+extension AVAudioPCMBuffer {
+    func spectrum(window: [Float]) -> [FrequencyAmplitude] {
+        let mono = mixToMono()
+        guard let data = mono.floatChannelData?[0] else {
+            fatalError("Not a float audio format")
+        }
+        return fft(
+            n: Int(frameLength),
+            signal: data,
+            window: window,
+            sampleRate: Float(mono.format.sampleRate)
+        )
+    }
+}
+
+private func fft(
+    n: Int,
+    signal: UnsafeMutablePointer<Float>,
+    window: [Float],
+    sampleRate: Float
+) -> [FrequencyAmplitude] {
+    let n = vDSP_Length(n)
+    let log2n = vDSP_Length(log2(Float(n)))
+
+    guard let fft = vDSP.FFT(log2n: log2n, radix: .radix2, ofType: DSPSplitComplex.self) else {
+        fatalError("Can't create FFT Setup.")
+    }
+
+    vDSP_vmul(signal, 1, window, 1, signal, 1, n)
+
+    let signal = UnsafeRawBufferPointer(
+        start: signal,
+        count: Int(n) * MemoryLayout<Float>.size
+    )
+
+    let halfN = Int(n / 2)
+
+    var forwardInputReal = [Float](repeating: 0, count: halfN)
+    var forwardInputImag = [Float](repeating: 0, count: halfN)
+    var forwardOutputReal = [Float](repeating: 0, count: halfN)
+    var forwardOutputImag = [Float](repeating: 0, count: halfN)
+
+    forwardInputReal.withUnsafeMutableBufferPointer { forwardInputRealPtr in
+        forwardInputImag.withUnsafeMutableBufferPointer { forwardInputImagPtr in
+            forwardOutputReal.withUnsafeMutableBufferPointer { forwardOutputRealPtr in
+                forwardOutputImag.withUnsafeMutableBufferPointer { forwardOutputImagPtr in
+
+                    // Create a `DSPSplitComplex` to contain the signal.
+                    var forwardInput = DSPSplitComplex(
+                        realp: forwardInputRealPtr.baseAddress!,
+                        imagp: forwardInputImagPtr.baseAddress!
+                    )
+
+                    // Convert the real values in `signal` to complex numbers.
+                    vDSP.convert(
+                        interleavedComplexVector: [DSPComplex](signal.bindMemory(to: DSPComplex.self)),
+                        toSplitComplexVector: &forwardInput
+                    )
+
+                    // Create a `DSPSplitComplex` to receive the FFT result.
+                    var forwardOutput = DSPSplitComplex(
+                        realp: forwardOutputRealPtr.baseAddress!,
+                        imagp: forwardOutputImagPtr.baseAddress!
+                    )
+
+                    // Perform the forward FFT.
+                    fft.forward(input: forwardInput, output: &forwardOutput)
+                }
+            }
+        }
+    }
+
+    let autospectrum = computeAutospectrum(
+        halfN: halfN,
+        forwardOutputReal: &forwardOutputReal,
+        forwardOutputImag: &forwardOutputImag
+    )
+
+    return computeAmplitudes(
+        autospectrum: autospectrum,
+        sampleRate: sampleRate,
+        n: Float(n)
+    )
+}
+
+private func computeAutospectrum(
+    halfN: Int,
+    forwardOutputReal: inout [Float],
+    forwardOutputImag: inout [Float]
+) -> [Float] {
+    [Float](unsafeUninitializedCapacity: halfN) { autospectrumBuffer, initializedCount in
+        // The `vDSP_zaspec` function accumulates its output. Clear the
+        // uninitialized `autospectrumBuffer` before computing the spectrum.
+        vDSP.clear(&autospectrumBuffer)
+
+        forwardOutputReal.withUnsafeMutableBufferPointer { forwardOutputRealPtr in
+            forwardOutputImag.withUnsafeMutableBufferPointer { forwardOutputImagPtr in
+                var frequencyDomain = DSPSplitComplex(
+                    realp: forwardOutputRealPtr.baseAddress!,
+                    imagp: forwardOutputImagPtr.baseAddress!
+                )
+
+                vDSP_zaspec(
+                    &frequencyDomain,
+                    autospectrumBuffer.baseAddress!,
+                    vDSP_Length(halfN)
+                )
+            }
+        }
+        initializedCount = halfN
+    }
+}
+
+private func computeAmplitudes(
+    autospectrum: [Float],
+    sampleRate: Float,
+    n: Float
+) -> [FrequencyAmplitude] {
+    autospectrum
+        .enumerated()
+        .filter { $0.element > 1 }
+        .map { index, element in
+            return FrequencyAmplitude(
+                frequency: Float(index) * sampleRate / n,
+                amplitude: sqrt(element) / n
+            )
+        }
+}

Sources/AudioSnapshotTesting/AudioSnapshotTesting.swift

@@ -17,6 +17,9 @@ typealias PlatformHostingView = _UIHostingView
 @MainActor
 public extension Snapshotting where Format == PlatformImage, Value == (AVAudioPCMBuffer, AVAudioPCMBuffer) {
     /// Generates a overlayed waveform snapshot of the given tuple of `AVAudioPCMBuffer`s.
+    /// - Parameters:
+    ///    - width: The width of the resulting image.
+    ///    - height: The height of the resulting image.
     static func waveform(width: Int, height: Int) -> Snapshotting {
         Snapshotting<PlatformView, PlatformImage>.image(size: .init(width: width, height: height))
             .pullback { buffer1, buffer2 in
@@ -52,6 +55,9 @@ public extension Snapshotting where Format == PlatformImage, Value == (AVAudioPC
 @MainActor
 public extension Snapshotting where Format == PlatformImage, Value == AVAudioPCMBuffer {
     /// Generates a waveform snapshot of the given `AVAudioPCMBuffer`.
+    /// - Parameters:
+    ///    - width: The width of the resulting image.
+    ///    - height: The height of the resulting image.
     static func waveform(width: Int, height: Int) -> Snapshotting {
         Snapshotting<PlatformView, PlatformImage>.image(size: .init(width: width, height: height))
             .pullback { buffer in
@@ -70,4 +76,33 @@ public extension Snapshotting where Format == PlatformImage, Value == AVAudioPCM
                 return PlatformHostingView(rootView: waveform.environment(\.colorScheme, .light))
             }
     }
+
+    /// Generates a frequency spectrum of the given `AVAudioPCMBuffer`.
+    /// - Parameters:
+    ///    - width: The width of the resulting image.
+    ///    - height: The height of the resulting image.
+    ///    - window: An optional array of floats representing the window function to apply before computing the FFT. If not provided, a Hann window will be used by default.
+    ///    - threshold: A float value between 0 and 1 that determines the minimum amplitude required for a frequency bin to be included in the resulting image.
+    @available(iOS 16, macOS 13, *)
+    static func spectrum(
+        width: Int,
+        height: Int,
+        window: [Float]? = nil,
+        threshold: Float = 0.005
+    ) -> Snapshotting {
+        Snapshotting<PlatformView, PlatformImage>.image(size: .init(width: width, height: height))
+            .pullback { buffer in
+                let fallbackWindow = vDSP.window(
+                    ofType: Float.self,
+                    usingSequence: .hanningNormalized,
+                    count: Int(buffer.frameLength),
+                    isHalfWindow: false
+                )
+                let data = buffer
+                    .spectrum(window: window ?? fallbackWindow)
+                    .filter { $0.amplitude > threshold }
+                let spectrum = SpectrumView(data: data, height: CGFloat(height))
+                return PlatformHostingView(rootView: spectrum.environment(\.colorScheme, .light))
+            }
+    }
 }

Sources/AudioSnapshotTesting/FrequencyMagnitude.swift

@@ -0,0 +1,10 @@
+struct FrequencyAmplitude {
+    let frequency: Float
+    let amplitude: Float
+
+    var frequencyString: String { String(format: "%.1f", frequency) }
+}
+
+extension FrequencyAmplitude: Identifiable {
+    var id: Float { frequency }
+}

Sources/AudioSnapshotTesting/SpectrumView.swift

@@ -0,0 +1,51 @@
+import SwiftUI
+import Charts
+
+@available(iOS 16, macOS 13, *)
+struct SpectrumView: View {
+    let data: [FrequencyAmplitude]
+    let height: CGFloat
+
+    var body: some View {
+        Chart(data) { frequency in
+            BarMark(
+                x: .value("Frequency", frequency.frequencyString),
+                y: .value("Amplitude", frequency.amplitude)
+            )
+            .foregroundStyle(Color.red)
+        }
+        .chartYAxis {
+            AxisMarks(position: .leading) { _ in
+                AxisValueLabel()
+                    .foregroundStyle(.white)
+                AxisGridLine()
+                    .foregroundStyle(.white.opacity(0.3))
+            }
+        }
+        .chartYAxisLabel(position: .leading) { Text("Amplitude").foregroundStyle(.white) }
+        .chartXAxis {
+            AxisMarks(preset: .aligned, values: xAxisValues) { _ in
+                AxisValueLabel()
+                    .foregroundStyle(.white)
+                AxisGridLine()
+                    .foregroundStyle(.white.opacity(0.3))
+            }
+        }
+        .chartXAxisLabel(position: .bottom) { Text("Frequency (Hz)").foregroundStyle(.white) }
+        .padding(16)
+        .frame(height: height)
+        .background(Color.black)
+    }
+
+    private var xAxisValues: [String] {
+        let count = data.count
+        let numberOfLabels = min(10, data.count)
+        guard numberOfLabels > 1 else { return [data[0].frequencyString] }
+        let spacing = Float(count - 1) / Float(numberOfLabels - 1)
+        let points = (0..<numberOfLabels).map {
+            let index = Int(Float($0) * spacing)
+            return data[index].frequencyString
+        }
+        return points
+    }
+}

Tests/AudioSnapshotTestingTests/AudioSnapshotTestingTests.swift

@@ -30,6 +30,38 @@ func fileWaveformOverlay() async throws {
     )
 }
 
+@Test(
+    .snapshots(record: false, diffTool: .ksdiff),
+    arguments: ["1hz@32768", "1hz2hz@32768", "2hz@32768"]
+)
+@MainActor
+func perfectSpectrum(wave: String) async throws {
+    let buffer = try AVAudioPCMBuffer.read(wave: wave)
+    assertSnapshot(
+        of: buffer,
+        as: .spectrum(
+            width: 500,
+            height: 200,
+            window: .init(repeating: 1, count: Int(buffer.frameLength))
+        ),
+        named: wave
+    )
+}
+
+@Test(
+    .snapshots(record: false, diffTool: .ksdiff),
+    arguments: ["1hz@44100", "white", "brown", "pink", "square", "triangle", "sawtooth"]
+)
+@MainActor
+func windowedSpectrum(wave: String) async throws {
+    let buffer = try AVAudioPCMBuffer.read(wave: wave)
+    assertSnapshot(
+        of: buffer,
+        as: .spectrum(width: 1500, height: 400),
+        named: wave
+    )
+}
+
 private extension AVAudioPCMBuffer {
     static func read(wave: String) throws -> AVAudioPCMBuffer {
         let file = try AVAudioFile(