labels improved, png filename now has also stokes component and scrunching factors

l2com/plot-bf-hdf5-spectrum.py

@@ -42,6 +42,7 @@ def read_hdf5_file(fname, tscrunch, fscrunch, memory_usage_limit):
             group_name = f"{sap_id}/{beam_id}"
             beamno = int(beam_id[5:])
 
+            chanpersub = h5[group_name].attrs["CHANNELS_PER_SUBBAND"]
             chanbw = h5[group_name].attrs["CHANNEL_WIDTH"] # in Hz
             freq = h5[group_name + "/COORDINATES/COORDINATE_1"].attrs["AXIS_VALUES_WORLD"]
             tsamp = h5[group_name + "/COORDINATES/COORDINATE_0"].attrs["INCREMENT"]
@@ -51,10 +52,12 @@ def read_hdf5_file(fname, tscrunch, fscrunch, memory_usage_limit):
             nof = h5[group_name].attrs["NOF_STOKES"]
             if nof>1:
                 print (f"\nWarning: there are more than 1 Stokes component in the file, only the first one will be read!")
-            #stokes_ids = [key for key in h5[group_name].keys() if "STOKES" in key]
-            (nsamp, nchan) = np.shape(h5[group_name + "/STOKES_0"])
-            dtype = np.dtype(h5[group_name + "/STOKES_0"])
-            stokes_comp = h5[group_name + "/STOKES_0"].attrs["STOKES_COMPONENT"]
+            stokes_ids = [key for key in h5[group_name].keys() if "STOKES_" in key]
+            stokes_group = group_name + "/" + stokes_ids[0]
+            (nsamp, nchan) = np.shape(h5[stokes_group])
+            dtype = np.dtype(h5[stokes_group])
+            stokes_comp = h5[stokes_group].attrs["STOKES_COMPONENT"]
+            nsubs = h5[stokes_group].attrs["NOF_SUBBANDS"]
 
     itemsize = np.dtype(dtype).itemsize # size of an element in bytes
     nsamp_fromfile = int(os.path.getsize(rawfile) / nchan / itemsize)
@@ -73,7 +76,7 @@ def read_hdf5_file(fname, tscrunch, fscrunch, memory_usage_limit):
         print (f"Expected RAM usage is > {memory_usage_limit}% of the total RAM ({expected_ram_usage*100./ramsize:.1f}%, {expected_ram_usage_GB} GB)")
         print ("(based on lofar-default.lua AOFlagger strategy)")
         print (f"Reading of the input file will be done in {nsteps} steps")
-        print ("You can also decrease the RAM usage by tscrunching the data")
+        print ("You can also decrease the RAM usage by scrunching the data")
         print ("")
     else: nsteps = 1
 
@@ -119,7 +122,7 @@ def read_hdf5_file(fname, tscrunch, fscrunch, memory_usage_limit):
     # f-scrunching the freq values as well
     freq = np.mean(np.reshape(freq, (int(nchan/fscrunch), fscrunch)), axis=1)
     
-    return (tsamp, freq, chanbw, data, stations, nstations, band, antenna_set, stokes_comp, obsid, targets)
+    return (tsamp, freq, chanbw, chanpersub, nsubs, data, stations, nstations, band, antenna_set, stokes_comp, obsid, targets)
 
 ###
 ###--- m a i n ---
@@ -143,7 +146,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # Read data
-    (tsamp, freq, chanbw, data, stations, nstations, band, antenna_set, stokes, obsid, targets) = \
+    (tsamp, freq, chanbw, chanpersub, nsubs, data, stations, nstations, band, antenna_set, stokes, obsid, targets) = \
         read_hdf5_file(args.h5file, args.tscrunch, args.fscrunch, args.memory_usage_limit)
     nsamp, nchan = data.shape # here nsamp is already t-scrunched
     freq *= 1e-6 # in MHz
@@ -205,7 +208,7 @@ if __name__ == "__main__":
     fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(18, 8), sharex=True, gridspec_kw={"height_ratios": [0.1, 0.2, 0.7], "hspace": 0.02})
     
     # Top RFI fraction subplot
-    ax1.set_title(f"L{obsid} | {antenna_set} | {band} | Nstations = {nstations} | {stations}", loc='left')
+    ax1.set_title(f"L{obsid} | {antenna_set} | {band} | Nsub: {nsubs} | Nchan/sub = {chanpersub} | Nstations = {nstations} | {stations}", loc='left')
     ax1.set_title(f"{targets} | Stokes: {stokes}", loc='right')
     ax1.set_ylabel("RFI (%)")
     ax1.set_ylim(-10, 110)
@@ -214,7 +217,6 @@ if __name__ == "__main__":
         ax1.legend(loc=1)
     
     # Average spectrum
-    print (freq)
     ax2.plot(freq, z, color="k", linewidth=0.5)
     ax2.set_ylabel("Power (dB)")
     ax2.set_ylim(vmin, vmax)
@@ -232,8 +234,11 @@ if __name__ == "__main__":
     img = ax3.imshow(v, origin="lower", aspect="auto", interpolation="None", extent=[fmin, fmax, tmin, tmax], vmin=vmin, vmax=vmax)
     ax3.set_xlabel("Frequency (MHz)")
     ax3.set_ylabel("Time (s)")
+    ax3label=f"Tscrunch: {args.tscrunch}  Fscrunch: {args.fscrunch}"
+    props = dict(boxstyle='round', facecolor='white', alpha=0.7)
+    ax3.text(0.01, 0.95, ax3label, transform=ax3.transAxes, fontsize=10, verticalalignment='center', bbox=props)
     fig.colorbar(img, ax=ax3, location="bottom", orientation="horizontal", label="Power (dB)", shrink=1, pad=0.12, fraction=0.05)
     
-    plt.savefig(f"L{obsid}_bf_spectrum.png", bbox_inches="tight")
+    plt.savefig(f"L{obsid}_bf_spectrum_stokes{stokes}_tscr{args.tscrunch}_fscr{args.fscrunch}.png", bbox_inches="tight")
     plt.show()
     plt.close()