unpack_data.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from make_gqrx_not_crash_my_os import load_spectrogram, kill_time_and_frequency
import pandas as pd

def find_doppler_peaks(Sxx, f_arr, t_arr, threshold_db = 14.5, distance = 25, freq_range = 20e3):
    """
    Searches each time slice of a cleaned spectrogram for peaks above a threshold and returns
    the strongest peak per time slice as the detected Doppler frequency.
    Multiple detections within the same time slice are resolved by taking the strongest peak.

    :param Sxx: cleaned spectrogram in dB (output of kill_time_and_frequency)
    :param f_arr: frequency array in Hz (not fftshifted)
    :param t_arr: time array in seconds
    :param threshold_db: minimum peak height above noise floor in dB (default: 14.5).
                         Tune this to balance sensitivity vs false detections.
    :param distance: minimum separation between peaks in frequency bins (default: 25).
                     Prevents detecting multiple peaks too close together.
    :param freq_range: half-width of frequency range to search in Hz (default: 20000)
    :return: doppler_times — array of times where peaks were detected in seconds,
             doppler_freqs — array of detected peak frequencies in Hz relative to carrier,
             f_shifted — fftshifted frequency array for plotting,
             Sxx_shifted — fftshifted spectrogram for plotting,
             freq_mask — boolean mask for the frequency range of interest
    """

    f_shifted = np.fft.fftshift(f_arr)
    Sxx_shifted = np.fft.fftshift(Sxx, axes=0)
    noise_floor = np.median(Sxx_shifted)

    freq_mask = np.abs(f_shifted) <= freq_range

    doppler_times = []
    doppler_freqs = []

    for i, t in enumerate(t_arr):
        peaks, properties = find_peaks(Sxx_shifted[:, i],
                                       height=noise_floor + threshold_db,
                                       distance=distance)
        if len(peaks) > 0:
            strongest = peaks[np.argmax(properties['peak_heights'])]
            doppler_times.append(t)
            doppler_freqs.append(f_shifted[strongest])

    doppler_times = np.array(doppler_times)
    doppler_freqs = np.array(doppler_freqs)
    print(f"Found peaks in {len(doppler_times)} out of {len(t_arr)} time slices")
    return doppler_times, doppler_freqs, f_shifted, Sxx_shifted, freq_mask

def plot_doppler(doppler_times, doppler_freqs, f_shifted, Sxx_shifted, t_avg,
                 freq_mask, noise_floor):
    """
    Plots the cleaned spectrogram with detected Doppler peaks overlaid as white dots.

    :param doppler_times: array of peak detection times in seconds
    :param doppler_freqs: array of detected peak frequencies in Hz
    :param f_shifted: fftshifted frequency array in Hz
    :param Sxx_shifted: fftshifted cleaned spectrogram in dB
    :param t_avg: time array in seconds
    :param freq_mask: boolean mask for the frequency range to display
    :param noise_floor: noise floor value in dB, used for color scale reference
    """

    plt.figure(figsize=(12, 6))
    plt.pcolormesh(t_avg,
                   f_shifted[freq_mask] / 1e3,
                   Sxx_shifted[freq_mask, :],
                   shading='nearest',
                   cmap='magma',
                   vmin=noise_floor,
                   vmax=noise_floor + 20)
    plt.plot(doppler_times, doppler_freqs / 1e3,
             'w.', markersize=3, label='detected peaks')
    plt.ylabel("Frequency (kHz relative to center)")
    plt.xlabel("Time (s)")
    plt.colorbar(label="Power (dB)")
    plt.legend()
    plt.tight_layout()
    plt.show()

def save_doppler_csv(doppler_times, doppler_freqs, filepath):
    """
    Saves detected Doppler peak coordinates to a CSV file with columns
    time (seconds) and frequency (Hz relative to carrier).
    Output filename is derived from the input filepath.

    :param doppler_times: array of peak times in seconds
    :param doppler_freqs: array of peak frequencies in Hz
    :param filepath: path to the spectrogram file, used to derive the output CSV filename
    :return: csv_path — path to the saved CSV file
    """

    csv_path = filepath.replace("_spectrogram.npz", "").replace("_decimated.npz", "") + "_doppler.csv"
    df = pd.DataFrame({
        'time': doppler_times,
        'frequency': doppler_freqs
    })
    df.to_csv(csv_path, index=False)
    print(f"Saved {len(doppler_times)} points to {csv_path}")
    return csv_path

def load_doppler_csv(csv_path, recording_start_unix= 0):
    """
    Load doppler CSV and convert time axis to Unix timestamps.

    :param csv_path: path to _doppler.csv
    :param recording_start_unix: Unix timestamp (seconds) of recording start
    :return: timestamps_unix — numpy array of Unix timestamps in seconds,
             doppler_obs_hz — numpy array of observed Doppler shifts in Hz
    """
    df = pd.read_csv(csv_path)
    timestamps_unix = df['time'].values + recording_start_unix
    doppler_obs_hz = df['frequency'].values
    return timestamps_unix, doppler_obs_hz

def extract_doppler(filepath, threshold_db=14.5, distance=25, freq_range=20e3):
    """
    Full pipeline — loads a saved spectrogram, cleans it, finds Doppler peaks,
    plots the result and saves the peaks to CSV.
    This is the main entry point for unpack_data.py.

    :param filepath: path to a _spectrogram.npz file
    :param threshold_db: peak detection threshold above noise floor in dB (default: 14.5)
    :param distance: minimum peak separation in frequency bins (default: 25)
    :param freq_range: half-width of frequency range to search in Hz (default: 20000)
    :return: doppler_times — array of detected peak times in seconds,
             doppler_freqs — array of detected peak frequencies in Hz
    """

    Sxx, f_arr, t_avg, carrier_freq = load_spectrogram(filepath)
    Sxx_clean = kill_time_and_frequency(Sxx)

    doppler_times, doppler_freqs, f_shifted, Sxx_shifted, freq_mask = find_doppler_peaks(
        Sxx_clean, f_arr, t_avg,
        threshold_db=threshold_db,
        distance=distance,
        freq_range=freq_range)

    noise_floor = np.median(Sxx_shifted)
    plot_doppler(doppler_times, doppler_freqs, f_shifted, Sxx_shifted, t_avg, freq_mask, noise_floor)
    save_doppler_csv(doppler_times, doppler_freqs, filepath)
    return doppler_times, doppler_freqs


if __name__ == "__main__":
    fname = "raw_data/gqrx_20260428_134823_437075000_2000000_fc_spectrogram.npz"
    extract_doppler(fname, threshold_db=14.5, distance=25, freq_range=20e3)
    print(load_doppler_csv("raw_data/gqrx_20260428_134823_437075000_2000000_fc_doppler.csv",0))