fix bug when one pol is flagged

notebooks/single_baseline_2D_filtered_SNRs.ipynb

@@ -188,8 +188,10 @@
     "    \n",
     "    axes[0].set_ylabel(f'JD - {int(data.times[0])}')\n",
     "    plt.tight_layout()\n",
-    "    plt.colorbar(im, ax=axes, label='|2D DPSS Filtered SNR|', pad=.02,\n",
-    "                 extend=('max' if np.max([np.abs(filtered_SNR[bl][~flags[bl]]) for bl in filtered_SNR]) > vmax else None))\n",
+    "    largest_pixel = np.max([np.abs(filtered_SNR[bl][~flags[bl]]) \n",
+    "                            for bl in filtered_SNR if not np.all(flags[bl])])\n",
+    "    plt.colorbar(im, ax=axes, label='|2D DPSS Filtered SNR|', pad=.02, \n",
+    "                 extend=('max' if largest_pixel > vmax else None))\n",
     "    plt.close(fig)\n",
     "    return fig"
    ]
@@ -285,6 +287,7 @@
     "\n",
     "for single_bl_file in corner_turn_map['files_to_outfiles_map'][RED_AVG_FILE]:\n",
     "    # Load data\n",
+    "    print(f'Now loading {single_bl_file}')\n",
     "    hd = io.HERAData(single_bl_file)\n",
     "    data, flags, nsamples = hd.read(polarizations=['ee', 'nn'])\n",
     "    dt = np.median(np.diff(hd.times)) * 24 * 3600\n",
@@ -331,6 +334,7 @@
     "\n",
     "    # save results\n",
     "    hd.update(data=filtered_SNR)\n",
+    "    print(f\"Writing results to {single_bl_file.replace('.uvh5', SNR_SUFFIX)}\")\n",
     "    hd.write_uvh5(single_bl_file.replace('.uvh5', SNR_SUFFIX), clobber=True)"
    ]
   },