Implement spectrogram strategy

Author: jcavar

README.md

@@ -49,17 +49,24 @@ class MyAudioTests: XCTestCase {
 }
 ```
 
+Snapshot audio tests are snapshot tested itself. Please find many examples in: [AudioSnapshotTestingTests.swift](Tests/AudioSnapshotTestingTests/AudioSnapshotTestingTests.swift)
+
 ### Features
 
 - [x] `AVAudioPCMBuffer` waveform snapshots
 - [x] `AVAudioPCMBuffer` overlayed waveform snapshots
-- [ ] Spectrogram
+- [x] Spectrogram
 - [x] Spectra
 - [x] Different waveform rendering strategies
-- [ ] Test against other reference implementations and with known audio files
+- [x] Test against other reference implementations and with known audio files
 - [ ] Documentation
 - [ ] Mention JUCE
 - [ ] Link blog post and talk
+- [ ] review stashes
+- [x] Add a link to swift snapshot testing
+- [ ] Multi level comparison (first hash, then data, then image)
+- [ ] Use accelerate in downsampling
+- [ ] Add file strategy
 
 ## Contributing
 

Sources/AudioSnapshotTesting/AVAudioPCMBufferExtensions.swift

@@ -88,6 +88,57 @@ extension AVAudioPCMBuffer {
             sampleRate: Float(mono.format.sampleRate)
         )
     }
+
+    func spectrogram(fftSize: Int, hopSize: Int, width: Int, window: [Float], amplitudeScale: AmplitudeScale) -> [Float] {
+        var spectrogramData = [Float]()
+
+        let mono = mixToMono()
+        let data = mono.floatChannelData![0]
+
+        for iteration in (0..<width) {
+            let frequencyAmplitudes = fft(
+                n: fftSize,
+                signal: data.advanced(by: iteration * hopSize),
+                window: window,
+                sampleRate: Float(mono.format.sampleRate)
+            )
+            .map(\.amplitude)
+            spectrogramData.append(contentsOf: frequencyAmplitudes)
+        }
+
+        // Apply amplitude scaling
+        return applyAmplitudeScale(spectrogramData, scale: amplitudeScale)
+    }
+}
+
+/// Applies amplitude scaling to spectrogram data
+/// - Parameters:
+///   - data: Raw amplitude values from FFT
+///   - scale: The scaling method to apply
+/// - Returns: Scaled amplitude values
+private func applyAmplitudeScale(_ data: [Float], scale: AmplitudeScale) -> [Float] {
+    switch scale {
+    case .linear:
+        // Return raw values unchanged (backward compatible)
+        return data
+
+    case .logarithmic(let range):
+        // Convert to dB scale: 20 * log10(amplitude)
+        // Use a small threshold to avoid log(0)
+        let minThreshold: Float = 1e-10
+        let minDB = -range
+
+        let dbValues = data.map { amplitude -> Float in
+            let clampedAmplitude = max(amplitude, minThreshold)
+            return 20.0 * log10(clampedAmplitude)
+        }
+
+        // Normalize to 0...1 based on dB range
+        return dbValues.map { db in
+            let normalized = (db - minDB) / range
+            return max(0, min(1, normalized)) // Clamp to 0...1
+        }
+    }
 }
 
 private func fft(
@@ -193,9 +244,10 @@ private func computeAmplitudes(
     sampleRate: Float,
     n: Float
 ) -> [FrequencyAmplitude] {
+    // Return all frequency bins without filtering
+    // Filtering is applied at a higher level (e.g., via threshold parameter in spectrum())
     autospectrum
         .enumerated()
-        .filter { $0.element > 1 }
         .map { index, element in
             return FrequencyAmplitude(
                 frequency: Float(index) * sampleRate / n,

Sources/AudioSnapshotTesting/AudioSnapshotTesting.swift

@@ -4,13 +4,24 @@ import Accelerate
 
 @_exported import SnapshotTesting
 
+/// Specifies the amplitude scaling method for spectrograms.
+public enum AmplitudeScale {
+    /// Linear amplitude scale (raw FFT magnitudes normalized to 0...1)
+    case linear
+    /// Logarithmic amplitude scale in decibels with specified dynamic range
+    /// - Parameter range: The dB range from minimum to maximum (e.g., 120 means -120dB to 0dB)
+    case logarithmic(range: Float)
+}
+
 #if os(macOS)
 public typealias PlatformImage = NSImage
 typealias PlatformView = NSView
+typealias PlatformColor = NSColor
 typealias PlatformHostingView = NSHostingView
 #elseif os(iOS)
 public typealias PlatformImage = UIImage
 typealias PlatformView = UIView
+typealias PlatformColor = UIColor
 typealias PlatformHostingView = _UIHostingView
 #endif
 
@@ -90,17 +101,66 @@ public extension Snapshotting where Format == PlatformImage, Value == AVAudioPCM
     ) -> 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 effectiveWindow = window ?? createHannWindow(size: Int(buffer.frameLength))
                 let data = buffer
-                    .spectrum(window: window ?? fallbackWindow)
+                    .spectrum(window: effectiveWindow)
                     .filter { $0.amplitude > threshold }
                 let spectrum = SpectrumView(data: data, height: CGFloat(height))
                 return PlatformHostingView(rootView: spectrum.environment(\.colorScheme, .light))
             }
     }
+
+    /// Generates a spectrogram of the given `AVAudioPCMBuffer`.
+    /// - Parameters:
+    ///    - hopSize: The number of audio frames between successive spectral frames.
+    ///    - frequencyCount: The number of frequency bins to include in each spectral frame.
+    ///    - 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.
+    ///    - amplitudeScale: The amplitude scaling method. Defaults to `.logarithmic(range: 120)` for standard 120 dB dynamic range. Use `.linear` for raw FFT magnitudes.
+    ///    - imageWidth: The width of the resulting snapshot image in pixels. Defaults to 1000.
+    @available(iOS 16, macOS 13, *)
+    static func spectrogram(
+        hopSize: Int,
+        frequencyCount: Int,
+        window: [Float]? = nil,
+        amplitudeScale: AmplitudeScale = .logarithmic(range: 120),
+        imageWidth: Int = 1000
+    ) -> Snapshotting {
+        let height = frequencyCount
+        return Snapshotting<PlatformView, PlatformImage>.image(size: .init(width: imageWidth, height: height))
+            .pullback { buffer in
+                let fftSize = frequencyCount * 2
+                let lastBucketStart = Int(buffer.frameLength) - fftSize
+
+                // Calculate number of FFT windows that fit in the buffer
+                // We need at least fftSize samples for each window, starting at position 0
+                // The last window starts at (frameLength - fftSize), and we hop by hopSize
+                // +1 accounts for the initial window at position 0
+                let width = 1 + (lastBucketStart / hopSize)
+
+                let effectiveWindow = window ?? createHannWindow(size: fftSize)
+                let data = buffer.spectrogram(fftSize: fftSize, hopSize: hopSize, width: width, window: effectiveWindow, amplitudeScale: amplitudeScale)
+
+                // Calculate bin size and max frequency based on sample rate and frequency count
+                let binSize = buffer.format.sampleRate / Double(fftSize)
+                let maxFrequency = Int(binSize * Double(frequencyCount))
+
+                // Calculate duration from buffer length and sample rate
+                let duration = Double(buffer.frameLength) / buffer.format.sampleRate
+
+                let spectrogram = SpectrogramView(data: data, width: width, height: height, maxFrequency: maxFrequency, duration: duration)
+                return PlatformHostingView(rootView: spectrogram.environment(\.colorScheme, .light))
+            }
+    }
+}
+
+/// Creates a normalized Hann window of the specified size
+/// - Parameter size: The number of samples in the window
+/// - Returns: An array of Float values representing the Hann window function
+private func createHannWindow(size: Int) -> [Float] {
+    vDSP.window(
+        ofType: Float.self,
+        usingSequence: .hanningNormalized,
+        count: size,
+        isHalfWindow: false
+    )
 }

Sources/AudioSnapshotTesting/CGImageExtensions.swift

@@ -0,0 +1,105 @@
+import CoreGraphics
+import Accelerate
+
+@available(iOS 16, macOS 13, *)
+extension CGImage {
+    static func spectrogramImage(data: [Float], width: Int, height: Int) -> CGImage {
+        let rgbImageFormat = vImage_CGImageFormat(
+            bitsPerComponent: 32,
+            bitsPerPixel: 32 * 3,
+            colorSpace: CGColorSpaceCreateDeviceRGB(),
+            bitmapInfo: CGBitmapInfo(
+                rawValue: kCGBitmapByteOrder32Host.rawValue |
+                CGBitmapInfo.floatComponents.rawValue |
+                CGImageAlphaInfo.none.rawValue))!
+
+        /// RGB vImage buffer that contains a vertical representation of the audio spectrogram.
+        let redBuffer = vImage.PixelBuffer<vImage.PlanarF>(width: height, height: width)
+        let greenBuffer = vImage.PixelBuffer<vImage.PlanarF>(width: height, height: width)
+        let blueBuffer = vImage.PixelBuffer<vImage.PlanarF>(width: height, height: width)
+        let rgbImageBuffer = vImage.PixelBuffer<vImage.InterleavedFx3>(width: height, height: width)
+        var data = data
+        data.withUnsafeMutableBufferPointer {
+            let planarImageBuffer = vImage.PixelBuffer(
+                data: $0.baseAddress!,
+                width: height,
+                height: width,
+                byteCountPerRow: height * MemoryLayout<Float>.stride,
+                pixelFormat: vImage.PlanarF.self
+            )
+
+            multidimensionalLookupTable.apply(
+                sources: [planarImageBuffer],
+                destinations: [redBuffer, greenBuffer, blueBuffer],
+                interpolation: .half
+            )
+
+            rgbImageBuffer.interleave(
+                planarSourceBuffers: [redBuffer, greenBuffer, blueBuffer]
+            )
+        }
+        return rgbImageBuffer.makeCGImage(cgImageFormat: rgbImageFormat)!
+    }
+}
+
+/// Returns the RGB values from a blue -> red -> green color map for a specified value.
+///
+/// Values near zero return dark blue, `0.5` returns red, and `1.0` returns full-brightness green.
+@available(macOS 13, iOS 16, *)
+nonisolated(unsafe) var multidimensionalLookupTable: vImage.MultidimensionalLookupTable = {
+    // Lookup table resolution: 32 entries provides good color accuracy while keeping memory usage low
+    // Higher values (64, 128) would provide smoother gradients but increase memory and initialization time
+    // Lower values (16, 8) would be faster but produce visible banding in spectrograms
+    let entriesPerChannel = UInt8(32)
+    let srcChannelCount = 1
+    let destChannelCount = 3
+
+    let lookupTableElementCount = Int(pow(Float(entriesPerChannel), Float(srcChannelCount))) * Int(destChannelCount)
+
+    let tableData = [UInt16](unsafeUninitializedCapacity: lookupTableElementCount) {
+        buffer,
+        count in
+        /// Supply the samples in the range `0...65535`. The transform function
+        /// interpolates these to the range `0...1`.
+        let multiplier = CGFloat(UInt16.max)
+        var bufferIndex = 0
+
+        for gray in ( 0 ..< entriesPerChannel) {
+            /// Create normalized red, green, and blue values in the range `0...1`.
+            let normalizedValue = CGFloat(gray) / CGFloat(entriesPerChannel - 1)
+
+            // Define `hue` that's blue at `0.0` to red at `1.0`.
+            let hue = 0.6666 - (0.6666 * normalizedValue)
+            let brightness = sqrt(normalizedValue)
+
+            let color = PlatformColor(
+                hue: hue,
+                saturation: 1,
+                brightness: brightness,
+                alpha: 1
+            )
+
+            var r = CGFloat()
+            var g = CGFloat()
+            var b = CGFloat()
+
+            color.getRed(&r, green: &g, blue: &b, alpha: nil)
+
+            buffer[ bufferIndex ] = UInt16(g * multiplier)
+            bufferIndex += 1
+            buffer[ bufferIndex ] = UInt16(r * multiplier)
+            bufferIndex += 1
+            buffer[ bufferIndex ] = UInt16(b * multiplier)
+            bufferIndex += 1
+        }
+        count = lookupTableElementCount
+    }
+
+    let entryCountPerSourceChannel = [UInt8](repeating: entriesPerChannel, count: srcChannelCount)
+
+    return vImage.MultidimensionalLookupTable(
+        entryCountPerSourceChannel: entryCountPerSourceChannel,
+        destinationChannelCount: destChannelCount,
+        data: tableData
+    )
+}()