Waterfall Plotting

plot_waterfall.py

#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colorbar import Colorbar
import sys
import warnings
with warnings.catch_warnings():
    warnings.filterwarnings("ignore",category=FutureWarning)
    import h5py

if __name__ == "__main__":
    # Filename
    h5file = h5py.File(sys.argv[1], "r")

    # Read observation attributes
    obsid = h5file["/"].attrs["observation_id"]
#    obs_timestamp = h5file.attrs["observation_timestamp"]
#    frequency = h5file.attrs["center_frequency"]
#    script_name = h5file.attrs["script_name"]
#    baud = h5file.attrs["baud"]
#    tle0 = h5file.attrs["tle_line_0"]
#    tle1 = h5file.attrs["tle_line_1"]
#    tle2 = h5file.attrs["tle_line_2"]
#    ground_station_id = h5file.attrs["ground_station_id"]
#    ground_station_lat = h5file.attrs["ground_station_lat"]
#    ground_station_lon = h5file.attrs["ground_station_lon"]
#    ground_station_elev = h5file.attrs["ground_station_elev"]

    # Read waterfall attributes
    timestamp = h5file["/waterfall"].attrs["start_time"]
#    data_min = h5file["/waterfall"].attrs["data_min"]
#    data_max = h5file["/waterfall"].attrs["data_max"]
    offset_in_stds = h5file["/waterfall"].attrs["offset_in_stds"]
    scale_in_stds = h5file["/waterfall"].attrs["scale_in_stds"]

    # Read datasets as numpy arrays
    data = np.array(h5file["/waterfall"]["data"])
    offset = np.array(h5file["/waterfall"]["offset"])
    scale = np.array(h5file["/waterfall"]["scale"])
    trel = np.array(h5file["/waterfall"]["relative_time"])
    tabs = np.array(h5file["/waterfall"]["absolute_time"])
    freq = np.array(h5file["/waterfall"]["frequency"])

    # Apply scale and offset
    data = data.astype("float32")*scale+offset

    data_min = np.min(data)
    data_max = np.max(data)

    h5file.close()

    print(data.shape, trel.shape)
    
    # Axis ranges
    tmin, tmax = np.min(trel), np.max(trel)
    fmin, fmax = np.min(freq), np.max(freq)

    # Create plot
    fig = plt.figure(figsize=(15, 10))
    gs = gridspec.GridSpec(3, 2, height_ratios=[0.05, 1.0, 0.2], width_ratios=[1.0, 0.1])
    gs.update(left=0.05, right=0.95, bottom=0.08, top=0.93, wspace=0.02, hspace=0.03)

    # Plot waterfall
    ax1 = plt.subplot(gs[1, 0])
    img = ax1.imshow(data.T, origin="lower", aspect="auto", interpolation="None",
                     extent=[tmin, tmax, fmin, fmax],
                     vmin=data_min, vmax=data_max,
                     cmap="viridis")
    ax1.set_ylabel("Frequency (kHz)")
    ax1.set_xticklabels([])
    ax1.grid()

    # Plot colorbar
    ax2 = plt.subplot(gs[0, 0])
    cbar = Colorbar(ax=ax2, mappable=img, orientation="horizontal", ticklocation="top")
    cbar.set_label("Power (dB)")

    # Horizontal
    ax3 = plt.subplot(gs[2, 0])
    ax3.plot(trel, np.mean(data, axis=1))
    ax3.set_xlim(tmin, tmax)
    ax3.set_xlabel("Time (seconds) since %s" % timestamp.decode("utf-8"))
    ax3.set_ylabel("Power (dB)")
    ax3.grid()

    # Vertical
    ax4 = plt.subplot(gs[1, 1])
    ax4.plot(np.mean(data, axis=0), freq)
    ax4.set_ylim(fmin, fmax)
    ax4.set_yticklabels([])
    ax4.set_xlabel("Power (dB)")
    ax4.grid()
    
    plt.show()

    # Plot waterfall
    #fig, ax1 = plt.subplots(figsize=(15, 10))
    
    #img = ax1.imshow(data, origin="lower", aspect="auto", interpolation="None",
    #                 extent=[tmin, tmax, fmin, fmax],
    #                 vmin=data_min, vmax=data_max,
    #                 cmap="viridis")
    #ax1.set_title(r"Station %d [%g$^\circ$, %g$^\circ$, %gm], Observation %d_%s [%.3f MHz]" % (
    #    ground_station_id, ground_station_lat, ground_station_lon, ground_station_elev,
    #    obsid, obs_timestamp, frequency*1e-6))
    #ax1.set_xlabel("Time (seconds) since %s" % timestamp.decode("utf-8"))
    #ax1.set_ylabel("Frequency (kHz)")
    #ax1.grid()
    #cbar = plt.colorbar(img, aspect=50, orientation="horizontal")
    #cbar.set_label("Power (dB)")
        
    #plt.tight_layout()
    #plt.show()
    #plt.savefig("plot.png", bbox_inches="tight")