rangeProcessing.py

import argparse
import numpy as np
import matplotlib.pyplot as plt
import pickle
import os
import scipy


def parse_data(x, samples_per_pulse):
    """
    Parse the digital data by looking for 'start of pulse' flags.
    Converts data to voltages.

    Parameters
    ----------
    x : numpy.ndarray
        Raw sample data.
    samples_per_pulse : int
        Number of samples per pulse.

    Returns
    -------
    parse_data : numpy.ndarray
        Parsed data array of shape (num_pulses-1, samples_per_pulse).
    num_pulses : int
        Number of pulses found.
    """
    flag_pulse_start = 5000
    scale_factor = 3.3 / (2 ** 12)

    pulse_starts = np.where(x > flag_pulse_start)[0]
    num_pulses = len(pulse_starts)

    parsed = np.zeros((num_pulses - 1, samples_per_pulse))

    for i in range(num_pulses - 1):
        start_idx = pulse_starts[i] + 1
        end_idx = start_idx + samples_per_pulse
        parsed[i, :] = x[start_idx:end_idx]

    # Convert digital data to voltages (3.3V / 12-bits) centered about 0V
    parsed = (parsed * scale_factor) - 3.3 / 2

    return parsed, num_pulses


def detect_and_load(file_path):
    """
    Detect the file type and load sample data accordingly.

    Supports the following formats:
      - .mat  (MATLAB)
      - .pkl  (pickle)
      - .npy  (NumPy binary)
      - .npz  (NumPy compressed)
      - .csv  (comma-separated values)

    If file_path has no extension, searches for existing files in the
    order: .pkl, .mat, .npy, .npz, .csv

    Parameters
    ----------
    file_path : str
        Path to the data file (with or without extension).

    Returns
    -------
    samples : numpy.ndarray
        1-D array of sample data.
    """
    # Supported extensions in priority order
    supported_extensions = [".pkl", ".mat", ".npy", ".npz", ".csv"]

    # Split into base and extension
    base, ext = os.path.splitext(file_path)

    if not os.path.isfile(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")
    resolved_path = file_path
    resolved_ext = ext.lower()

    print(f"Detected file: {resolved_path} (format: {resolved_ext})")

    # Load based on detected extension
    if resolved_ext == ".mat":
        from scipy.io import loadmat
        a = loadmat(resolved_path)
        # Try common variable names
        for key in ["samples", "data", "s"]:
            if key in a:
                return np.array(a[key]).flatten()
        # Fall back to first non-metadata key
        data_keys = [k for k in a.keys() if not k.startswith("__")]
        if data_keys:
            return np.array(a[data_keys[0]]).flatten()
        raise ValueError(f"No recognizable data found in {resolved_path}")

    elif resolved_ext == ".pkl":
        with open(resolved_path, "rb") as f:
            data = pickle.load(f)
        if isinstance(data, np.ndarray):
            return data.flatten()
        return np.array(data, dtype=np.float64).flatten()

    elif resolved_ext == ".npy":
        return np.load(resolved_path).flatten()

    elif resolved_ext == ".npz":
        npz = np.load(resolved_path)
        # Try common key names
        for key in ["samples", "data", "arr_0"]:
            if key in npz:
                return npz[key].flatten()
        # Fall back to first key
        first_key = list(npz.keys())[0]
        return npz[first_key].flatten()

    elif resolved_ext == ".csv":
        return np.loadtxt(resolved_path, delimiter=",").flatten()

    else:
        raise ValueError(f"Unsupported file format: {resolved_ext}")


def ranging_processing(
    file_path,
    pulse_width,
    sample_rate,
    ovs_rng,
    f_start,
    f_stop,
    n_pulse_cancel,
    max_range,
):
    """
    Produces an RTI (range x time intensity) image of the radar recording.
    Also applies a simple two-pulse canceller to filter out clutter and
    CFAR normalization to improve visual detection.

    Parameters
    ----------
    file_path : str
        Path to the data file (with or without extension).
    pulse_width : float
        Pulse width in seconds.
    sample_rate : float
        Sample rate in Hz.
    ovs_rng : int
        Oversampling factor applied in range.
    f_start : float
        LFM start frequency in Hz.
    f_stop : float
        LFM stop frequency in Hz.
    n_pulse_cancel : int
        Number of pulses to use for canceller.
    max_range : float
        Maximum range to display in meters.
    """
    c = 299792458.0  # (m/s) speed of light

    # Calculate bandwidth
    bw = f_stop - f_start
    print(f"Using {bw * 1e-6:g} MHz bandwidth")

    # Load data (auto-detect format)
    print("Loading data file...")
    s = detect_and_load(file_path)

    # Derived parameters
    Np = int(round(pulse_width * sample_rate)) - 1  # samples per pulse (-1 due to start flag)

    # Parse data and format as [pulse x sample]
    print("Parsing the recording...")
    parsed_data, num_pulses = parse_data(s, Np)
    print(f"Found {num_pulses} pulses")

    # Taper applied to reduce range sidelobes
    # win = np.hanning(Np)
    win = scipy.signal.windows.gaussian(Np, (Np-1)/20)
    rng_win = np.tile(win, (num_pulses - 1, 1))

    # Fourier transform to convert frequency to range
    print("Transforming from frequency to range...")
    nfft = Np * ovs_rng
    range_data = np.fft.fftshift(
        np.fft.fft(parsed_data * rng_win, n=nfft, axis=1), axes=1
    )

    # Compute the range axis
    freq_axis = np.linspace(-sample_rate / 2, sample_rate / 2, nfft)
    range_axis = c * pulse_width * freq_axis / (2 * bw)

    # Compute the time axis
    t = np.arange(num_pulses - 1) * pulse_width * 2

    # Trim data and range axis based on specified max range
    samples_keep = (range_axis < max_range) & (range_axis >= 0)
    range_axis_keep = range_axis[samples_keep]
    range_data = range_data[:, samples_keep]

    # Apply N-pulse canceller
    print("Cancelling stationary background clutter...")
    n_rows = (num_pulses - 1) - n_pulse_cancel + 1
    n_cols = len(range_axis_keep)
    range_cancelled = np.zeros((n_rows, n_cols), dtype=complex)

    for i in range(n_rows):
        pulses = range_data[i : i + n_pulse_cancel, :]
        mask = -np.ones((n_pulse_cancel, n_cols))
        mask[0, :] = 1.0
        mask[1:, :] = mask[1:, :] / (n_pulse_cancel - 1)
        range_cancelled[i, :] = np.sum(pulses * mask, axis=0)

    # Trim time axis to match cancelled data
    t_cancelled = t[:n_rows]

    # Apply the median CFAR normalization
    range_cancelled_dB = 20 * np.log10(np.abs(range_cancelled) + 1e-30)

    # Subtract median over time (column-wise median)
    median_over_time = np.median(range_cancelled_dB, axis=0, keepdims=True)
    range_cancelled_dB = range_cancelled_dB - median_over_time

    # Subtract median over range (row-wise median)
    median_over_range = np.median(range_cancelled_dB, axis=1, keepdims=True)
    range_cancelled_dB = range_cancelled_dB - median_over_range

    # ----------------------------- Plotting ---------------------------------
    print("Plotting...")

    # RTI without clutter rejection
    fig1, ax1 = plt.subplots()
    im1 = ax1.imshow(
        20 * np.log10(np.abs(range_data) + 1e-30),
        aspect="auto",
        extent=[range_axis_keep[0], range_axis_keep[-1], t[-1], t[0]],
        cmap="jet",
    )
    ax1.set_ylabel("Time (s)")
    ax1.set_xlabel("Range (m)")
    ax1.set_title("RTI without Clutter Rejection")
    plt.colorbar(im1, ax=ax1)

    # RTI with MTI clutter rejection
    fig2, ax2 = plt.subplots()
    im2 = ax2.imshow(
        20 * np.log10(np.abs(range_cancelled) + 1e-30),
        aspect="auto",
        extent=[range_axis_keep[0], range_axis_keep[-1], t_cancelled[-1], t_cancelled[0]],
        cmap="jet",
    )
    ax2.set_ylabel("Time (s)")
    ax2.set_xlabel("Range (m)")
    ax2.set_title("RTI with MTI Clutter Rejection")
    plt.colorbar(im2, ax=ax2)

    # RTI with MTI clutter rejection and CFAR
    fig3, ax3 = plt.subplots()
    im3 = ax3.imshow(
        range_cancelled_dB,
        aspect="auto",
        extent=[range_axis_keep[0], range_axis_keep[-1], t_cancelled[-1], t_cancelled[0]],
        vmin=0,
        vmax=35,
        cmap="jet",
    )
    ax3.set_ylabel("Time (s)")
    ax3.set_xlabel("Range (m)")
    ax3.set_title("RTI with MTI Clutter Rejection and CFAR")
    plt.colorbar(im3, ax=ax3)

    plt.show()
    print("Done.")


def main():
    parser = argparse.ArgumentParser(
        description="Radar ranging processor. Produces RTI (range x time intensity) "
                    "images with optional MTI clutter rejection and CFAR normalization. "
                    "Automatically detects input file format (.pkl, .mat, .npy, .npz, .csv)."
    )

    parser.add_argument(
        "-f", "--file",
        type=str,
        default="Test",
        help="Path to data file. Can include extension (.pkl, .mat, .npy, .npz, .csv) "
             "or omit it for auto-detection. Default: Test.",
    )
    parser.add_argument(
        "--pulse-width",
        type=float,
        default=20e-3,
        help="Pulse width in seconds. Default: 20e-3.",
    )
    parser.add_argument(
        "--sample-rate",
        type=float,
        default=40000,
        help="Sample rate in Hz. Default: 40000.",
    )
    parser.add_argument(
        "--oversample",
        type=int,
        default=10,
        help="Oversampling factor applied in range. Default: 10.",
    )
    parser.add_argument(
        "--f-start",
        type=float,
        default=2400e6,
        help="LFM start frequency in Hz. Default: 2400e6.",
    )
    parser.add_argument(
        "--f-stop",
        type=float,
        default=2480e6,
        help="LFM stop frequency in Hz. Default: 2480e6.",
    )
    parser.add_argument(
        "--n-pulse-cancel",
        type=int,
        default=2,
        help="Number of pulses to use for canceller. Default: 2.",
    )
    parser.add_argument(
        "--max-range",
        type=float,
        default=100,
        help="Maximum range to display in meters. Default: 100.",
    )

    args = parser.parse_args()

    ranging_processing(
        file_path=args.file,
        pulse_width=args.pulse_width,
        sample_rate=args.sample_rate,
        ovs_rng=args.oversample,
        f_start=args.f_start,
        f_stop=args.f_stop,
        n_pulse_cancel=args.n_pulse_cancel,
        max_range=args.max_range,
    )


if __name__ == "__main__":
    main()