BST-plotter.py

'''
Author: Owen A. Johnson
Date: 2025-05-20
Code Purpose: Plot dynamic spectra from LOFAR 8-bit BST files
Example usage: python BST-plotter.py 20250513_110039_bst_00X.dat --rcumode 7 --integration-time 5 --time-range 100 300 --freq-range 220 250

'''
import argparse
import numpy as np
import matplotlib.pyplot as plt

def read_bst_file(filename, n_beamlets=488):
    """
    Reads a LOFAR 8-bit BST file and returns a 2D NumPy array.

    Parameters:
    - filename: str, path to the .dat file
    - n_beamlets: int, number of beamlets per integration (default: 488 for 8-bit mode)

    Returns:
    - data: np.ndarray of shape (n_integrations, n_beamlets)
    """
    raw_data = np.fromfile(filename, dtype='<f8')
    
    if len(raw_data) % n_beamlets != 0:
        raise ValueError("File size is not a multiple of beamlets per integration")
    
    data = raw_data.reshape((-1, n_beamlets))
    return data

def get_beamlet_frequencies(rcumode=7, n_beamlets=488, subband_offset=0, integration_time_sec=1):
    """
    Computes beamlet center frequencies and prints resolution info.

    Parameters:
    - rcumode: int, receiver mode (only 3, 5, or 7 supported)
    - n_beamlets: int, number of beamlets (default: 488 for 8-bit)
    - subband_offset: int, starting subband index (default: 0)
    - integration_time_sec: float, time resolution per integration (default: 1 sec)
    
    NOTE: I assume clock rate is always 200 MHz for all modes.

    Returns:
    - frequencies: np.ndarray of beamlet center frequencies in Hz
    """
    # Define RCU mode parameters
    rcumode_settings = {
        3: {"start_freq": 110e6, "sampling_rate": 200e6},
        5: {"start_freq": 170e6, "sampling_rate": 200e6},
        7: {"start_freq": 210e6, "sampling_rate": 200e6},
    }

    if rcumode not in rcumode_settings:
        raise ValueError("Unsupported rcumode. Use 3, 5, or 7.")

    settings = rcumode_settings[rcumode]
    subband_width = settings["sampling_rate"] / 1024  # Hz
    start_freq = settings["start_freq"]

    # Compute center frequencies
    frequencies = start_freq + (np.arange(subband_offset, subband_offset + n_beamlets) * subband_width)

    # Print resolutions
    print(f"RCU Mode: {rcumode}")
    print(f"Start frequency: {start_freq / 1e6:.2f} MHz")
    print(f"Subband (beamlet) width: {subband_width:.2f} Hz")
    print(f"Frequency range: {frequencies[0]/1e6:.2f} MHz – {frequencies[-1]/1e6:.2f} MHz")
    print(f"Time resolution: {integration_time_sec} seconds")
    print(f"Spectral resolution: {subband_width:.2f} Hz")

    return frequencies


def plot_bst_dynamic_spectrum_with_freq(
    filename,
    data,
    rcumode=7,
    subband_offset=0,
    integration_time_sec=1,
    time_range=None,
    freq_range_mhz=None
):
    """
    Plots a dynamic spectrum of BST data with frequency on the y-axis.

    Parameters:
    - filename: str, used in the plot title
    - data: np.ndarray, shape (n_integrations, n_beamlets)
    - rcumode: int, RCU mode (3, 5, or 7)
    - subband_offset: int, starting subband number
    - integration_time_sec: float, integration time per step
    - time_range: tuple (t_min, t_max) in seconds, or None to use full range
    - freq_range_mhz: tuple (f_min, f_max) in MHz, or None to use full range
    """
    import matplotlib.pyplot as plt
    import numpy as np

    n_integrations, n_beamlets = data.shape
    time_axis = np.arange(n_integrations) * integration_time_sec

    # Get beamlet frequencies
    freqs_hz = get_beamlet_frequencies(
        rcumode=rcumode,
        n_beamlets=n_beamlets,
        subband_offset=subband_offset,
        integration_time_sec=integration_time_sec
    )
    freqs_mhz = freqs_hz / 1e6

    # Determine time indices
    if time_range is None:
        time_mask = np.full(n_integrations, True)
    else:
        t_min, t_max = time_range
        time_mask = (time_axis >= t_min) & (time_axis <= t_max)

    # Determine frequency indices
    if freq_range_mhz is None:
        freq_mask = np.full(n_beamlets, True)
    else:
        f_min, f_max = freq_range_mhz
        freq_mask = (freqs_mhz >= f_min) & (freqs_mhz <= f_max)

    # Subset data and axes
    data_sub = data[time_mask][:, freq_mask]
    time_sub = time_axis[time_mask]
    freq_sub = freqs_mhz[freq_mask]
    
    vmin, vmax = np.percentile(data_sub, [90, 30])

    plt.figure(figsize=(12, 6))
    plt.imshow(data_sub.T, aspect='auto', origin='lower',
               extent=[time_sub[0], time_sub[-1], freq_sub[0], freq_sub[-1]],
               cmap='magma',
                vmin=vmin, vmax=vmax)
    plt.colorbar(label='Power (arb. units)')
    plt.xlabel('Time (s)')
    plt.ylabel('Frequency (MHz)')
    plt.title(filename)
    plt.tight_layout()
    plt.show()
    
def get_args():
    """
    Parses command-line arguments
    
    Returns:
    - argparse.Namespace with attributes:
        filename, rcumode, integration_time, subband_offset,
        time_range, freq_range
    """
    parser = argparse.ArgumentParser(description="Plot LOFAR BST dynamic spectrum with frequency axis.")

    parser.add_argument("filename", help="Path to the .dat BST file")
    parser.add_argument("--rcumode", type=int, default=5, choices=[3, 5, 7],
                        help="RCU mode (3, 5, or 7). Default: 5")
    parser.add_argument("--integration-time", type=float, default=1,
                        help="Integration time in seconds. Default: 1")
    parser.add_argument("--subband-offset", type=int, default=0,
                        help="Index of the first subband used. Default: 0")

    parser.add_argument("--time-range", nargs=2, type=float, metavar=('T_MIN', 'T_MAX'),
                        help="Time range to plot in seconds, e.g. --time-range 100 200")
    parser.add_argument("--freq-range", nargs=2, type=float, metavar=('F_MIN', 'F_MAX'),
                        help="Frequency range to plot in MHz, e.g. --freq-range 110 190")

    return parser.parse_args()

if __name__ == "__main__":
    args = get_args()
    data = read_bst_file(args.filename)
    plot_bst_dynamic_spectrum_with_freq(
        filename=args.filename,
        data=data,
        rcumode=args.rcumode,
        subband_offset=args.subband_offset,
        integration_time_sec=args.integration_time,
        time_range=tuple(args.time_range) if args.time_range else None,
        freq_range_mhz=tuple(args.freq_range) if args.freq_range else None
    )