fft over a rolling sample window for greater bin precision

scripts/selcal-fft.py

@@ -0,0 +1,147 @@
+#! /usr/bin/env python3
+
+
+import numpy as np
+from scipy.io import wavfile
+from scipy.fft import fft
+from scipy.signal import butter, lfilter, decimate
+import matplotlib.pyplot as plt
+import sys
+
+frequencies = np.array([10 ** ((n + 22) * 0.0225 + 2) for n in range(32)])
+widths = np.array([6.25 * 10 ** (n * 0.0225) for n in range(32)])
+widths = np.array([6.25 for n in range(32)])
+
+def butter_bandpass(lowcut, highcut, fs, order=5):
+    nyquist = 0.5 * fs
+    low = lowcut / nyquist
+    high = highcut / nyquist
+    b, a = butter(order, [low, high], btype='band')
+    return b, a
+
+def bandpass_filter(data, lowcut, highcut, fs, order=5):
+    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
+    y = lfilter(b, a, data)
+    return y
+
+def butter_lowpass(cutoff, fs, order=5):
+    nyquist = 0.5 * fs
+    normal_cutoff = cutoff / nyquist
+    b, a = butter(order, normal_cutoff, btype='low', analog=False)
+    return b, a
+
+def lowpass_filter(data, cutoff, fs, order=5):
+    b, a = butter_lowpass(cutoff, fs, order=order)
+    y = lfilter(b, a, data)
+    return y
+
+def get_largest_two_indices(numbers, threshold):
+    # Check if the list has at least three numbers
+    if len(numbers) < 3:
+        return None
+
+    # Find the indices of the three largest numbers in the list
+    indices = sorted(range(len(numbers)), key=lambda i: numbers[i], reverse=True)
+    largest_index = indices[0]
+    second_largest_index = indices[1]
+    third_largest_index = indices[2]
+
+    # Check if the largest and second largest numbers are at least threshold larger than the third largest
+    if numbers[largest_index] - numbers[third_largest_index] >= threshold and \
+       numbers[second_largest_index] - numbers[third_largest_index] >= threshold:
+        return largest_index, second_largest_index
+    else:
+        return None
+
+# Step 1: Read the WAV file
+sample_rate, data = wavfile.read(sys.argv[1])
+
+# Handle stereo audio by converting to mono if needed
+if len(data.shape) == 2:
+    data = data.mean(axis=1)
+
+# Define the maximum frequency of interest and Nyquist rate
+max_freq = 1600.0  # 2000 Hz
+nyquist_rate = 2 * max_freq  # Nyquist rate to prevent aliasing
+
+# If the sample rate is higher than the Nyquist rate, downsample
+if sample_rate > nyquist_rate:
+    # Apply a lowpass filter before downsampling
+    cutoff = max_freq #+ 500  # Lowpass filter cutoff slightly above max_freq to prevent aliasing
+    filtered_data = lowpass_filter(data, cutoff, sample_rate)
+
+    # Calculate the downsampling factor
+    downsample_factor = int(sample_rate / nyquist_rate)
+
+    # Downsample the filtered signal
+    filtered_data = decimate(filtered_data, downsample_factor)
+    sample_rate = sample_rate // downsample_factor
+else:
+    filtered_data = data
+
+# Step 2: Define the increment (0.1 seconds) and segment size
+#increment = 0.2  # in seconds
+#segment_size = int(sample_rate * increment)
+segment_size = 1024
+increment = segment_size / sample_rate
+print(f"Segment size: {segment_size}")
+
+# Step 3: Process each segment
+num_segments = len(filtered_data) // segment_size
+
+# Calculate the frequency resolution
+delta_f = sample_rate / segment_size
+
+# Determine the bin range for desired frequency range (100 Hz to 2000 Hz)
+high_bin = int(max_freq / delta_f)
+
+seg_off = int(sample_rate * 0.1)
+act_segs = len(filtered_data) // seg_off
+
+# Initialize a 2D array to store DFT results (magnitude spectrum)
+# Only store the bins within the desired frequency range
+dft_results = np.zeros((act_segs, high_bin))
+
+for i in range(act_segs):
+    end = (i+1) * seg_off
+    start = end - segment_size
+
+    try:
+        segment = filtered_data[start:end]
+
+        # Step 4: Apply the DFT
+        dft_result = fft(segment)
+
+        magnitudes = np.abs(dft_result)
+        total_energy = np.sum(magnitudes ** 2)
+        normalized_magnitudes = magnitudes / np.sqrt(total_energy)
+        normalized_magnitude = np.mean(normalized_magnitudes)
+
+        # Store the magnitude spectrum in the 2D array, only for the desired frequency range
+        dft_results[i, :] = normalized_magnitudes[:high_bin]
+
+        scores = [
+            10 * np.log10(np.sum(normalized_magnitudes[int((f-w)/delta_f):int((f+w)/delta_f)]))
+            for f,w in zip(frequencies, widths)
+        ]
+
+        codes = get_largest_two_indices(scores, 3.0)
+        if codes:
+            print([frequencies[code] for code in sorted(codes)])
+    except:
+        pass
+
+
+# Step 5: Plot the spectrogram
+plt.figure(figsize=(12, 8))
+extent = [0, num_segments * increment, 0, max_freq]
+plt.imshow(dft_results.T, aspect='auto', origin='lower', extent=extent, cmap='viridis')
+plt.colorbar(label='Magnitude')
+plt.title('Spectrogram (100 Hz to 2000 Hz)')
+plt.xlabel('Time (s)')
+plt.ylabel('Frequency (Hz)')
+
+
+for freq in frequencies:
+    plt.axhline(y=freq, color='r', linestyle='--', linewidth=0.5)
+plt.show()