Skip to content

Commit

Permalink
use flags and tslice more properly in plot_waterfall and plot_real_de…
Browse files Browse the repository at this point in the history 
…lay_vs_lst
  • Loading branch information
@jsdillon
jsdillon committed Jun 19, 2024
1 parent fc3ecf7 commit f4d3f8b
Showing 1 changed file with 15 additions and 12 deletions.
27 changes: 15 additions & 12 deletions notebooks/single_baseline_postprocessing_and_pspec.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -632,13 +632,13 @@
" extent = [data.freqs[0]/1e6, data.freqs[-1]/1e6, lsts[-1], lsts[0]]\n",
" \n",
" fig, axes = plt.subplots(1, 3, figsize=(20, 12), sharex=True, sharey=True, dpi=100)\n",
" im = axes[0].imshow(np.where(flags[bl], np.nan, np.abs(data[bl])), aspect='auto', norm=matplotlib.colors.LogNorm(), interpolation='none', cmap='inferno', extent=extent)\n",
" im = axes[0].imshow(np.where(flags[bl], np.nan, np.abs(data[bl]))[tslice], aspect='auto', norm=matplotlib.colors.LogNorm(), interpolation='none', cmap='inferno', extent=extent)\n",
" fig.colorbar(im, ax=axes[0], location='top', pad=.02, label=f'{bl}: Amplitude (Jy)')\n",
"\n",
" im = axes[1].imshow(np.where(flags[bl], np.nan, np.angle(data[bl])), aspect='auto', cmap='twilight', interpolation='none', extent=extent)\n",
" im = axes[1].imshow(np.where(flags[bl], np.nan, np.angle(data[bl]))[tslice], aspect='auto', cmap='twilight', interpolation='none', extent=extent)\n",
" fig.colorbar(im, ax=axes[1], location='top', pad=.02, label=f'{bl}: Phase (Radians)')\n",
"\n",
" im = axes[2].imshow(nsamples[bl], aspect='auto', interpolation='none', extent=extent)\n",
" im = axes[2].imshow(nsamples[bl][tslice], aspect='auto', interpolation='none', extent=extent)\n",
" fig.colorbar(im, ax=axes[2], location='top', pad=.02, label=f'{bl}: Number of Samples')\n",
" plt.tight_layout()\n",
"\n",
Expand All @@ -659,17 +659,20 @@
},
"outputs": [],
"source": [
"def plot_real_delay_vs_lst(data, bl=(ANTPAIR + ('ee',)), xlim=[-1999, 1999], clim=None, linthresh=10, taper=TAPER, tslice=tslice):\n",
"def plot_real_delay_vs_lst(data, bl=(ANTPAIR + ('ee',)), flags=None, xlim=[-1999, 1999], clim=None, linthresh=10, taper=TAPER, tslice=tslice):\n",
" '''Plots the real part of the tapered FFT of the data in each power spectrum band as a function of delay and LST.'''\n",
" if tslice is None:\n",
" tslice = slice(0, data[bl].shape[0], 1)\n",
" lsts = np.where(data.lsts > data.lsts[-1], data.lsts - 2 * np.pi, data.lsts)[tslice] * 12 / np.pi\n",
"\n",
" fig, axes = plt.subplots(1, len(bands), figsize=(28, 12), sharex=True, sharey=True, gridspec_kw={'wspace': .03}, dpi=100)\n",
" for i, (ax, band, band_slice) in enumerate(zip(axes, bands, band_slices)):\n",
" dfft = uvtools.utils.FFT(data[bl][:, band_slice], axis=1, taper=taper)\n",
" dfft = uvtools.utils.FFT(data[bl][tslice, band_slice], axis=1, taper=taper)\n",
" delays = uvtools.utils.fourier_freqs(data.freqs[band_slice]) * 1e9\n",
" to_plot = np.real(dfft)\n",
" if flags is not None:\n",
" flagged_times = np.all(flags[bl][tslice, band_slice], axis=1)\n",
" to_plot[flagged_times, :] = np.nan\n",
" if i == 0:\n",
" _to_plot = copy.deepcopy(to_plot)\n",
" ax.set_ylabel('LST (Hours)')\n",
Expand Down Expand Up @@ -1237,8 +1240,8 @@
"outputs": [],
"source": [
"if PLOT:\n",
" plot_waterfall(data)\n",
" plot_real_delay_vs_lst(data, xlim=[-3999, 3999], clim=[-1e4, 1e4])\n",
" plot_waterfall(data, flags=flags)\n",
" plot_real_delay_vs_lst(data, flags=flags, xlim=[-3999, 3999], clim=[-1e4, 1e4])\n",
" plot_dly_vs_fr(data, xlim=[-1999, 1999], clim=[1e0, 1e5])"
]
},
Expand All @@ -1259,7 +1262,7 @@
"source": [
"if PLOT and (PERFORM_DLY_FILT or PERFORM_INPAINT):\n",
" plot_waterfall(filt_data, flags=filt_flags, nsamples=nsamples)\n",
" plot_real_delay_vs_lst(filt_data, xlim=[-3999, 3999], clim=[-1e4, 1e4])\n",
" plot_real_delay_vs_lst(filt_data, flags=filt_flags, xlim=[-3999, 3999], clim=[-1e4, 1e4])\n",
" plot_dly_vs_fr(filt_data, xlim=[-1999, 1999], clim=[1e0, 1e5])"
]
},
Expand All @@ -1280,7 +1283,7 @@
"source": [
"if PLOT and not SKIP_XTALK_AND_FRF:\n",
" plot_waterfall(deint_xtalk_filt_data[0], flags=deint_flags[0], nsamples=deint_nsamples[0], tslice=None)\n",
" plot_real_delay_vs_lst(deint_xtalk_filt_data[0], xlim=[-3999, 3999], clim=[-1e4, 1e4], tslice=None)\n",
" plot_real_delay_vs_lst(deint_xtalk_filt_data[0], flags=deint_flags[0], xlim=[-3999, 3999], clim=[-1e4, 1e4], tslice=None)\n",
" plot_dly_vs_fr(deint_xtalk_filt_data[0], xlim=[-1999, 1999], clim=[1e0, 1e5], tslice=None)"
]
},
Expand All @@ -1301,7 +1304,7 @@
"source": [
"if PLOT and not SKIP_XTALK_AND_FRF:\n",
" plot_waterfall(deint_frf_data[0], flags=deint_flags[0], nsamples=deint_nsamples[0], tslice=None)\n",
" plot_real_delay_vs_lst(deint_frf_data[0], xlim=[-3999, 3999], clim=[-2e2, 2e2], linthresh=1, tslice=None)\n",
" plot_real_delay_vs_lst(deint_frf_data[0], flags=deint_flags[0], xlim=[-3999, 3999], clim=[-2e2, 2e2], linthresh=1, tslice=None)\n",
" plot_dly_vs_fr(deint_frf_data[0], xlim=[-1999, 1999], clim=[1e0, 1e5], tslice=None)"
]
},
Expand All @@ -1322,7 +1325,7 @@
"source": [
"if PLOT:\n",
" plot_waterfall(deint_frf_data[0], bl=(ANTPAIR + ('pI',)), flags=deint_flags[0], nsamples=deint_nsamples[0], tslice=None)\n",
" plot_real_delay_vs_lst(deint_frf_data[0], bl=(ANTPAIR + ('pI',)), xlim=[-3999, 3999], clim=[-2e2, 2e2], linthresh=1, tslice=None)\n",
" plot_real_delay_vs_lst(deint_frf_data[0], bl=(ANTPAIR + ('pI',)), flags=deint_flags[0], xlim=[-3999, 3999], clim=[-2e2, 2e2], linthresh=1, tslice=None)\n",
" plot_dly_vs_fr(deint_frf_data[0], bl=(ANTPAIR + ('pI',)), xlim=[-1999, 1999], clim=[1e0, 1e5], tslice=None)"
]
},
Expand All @@ -1343,7 +1346,7 @@
"source": [
"if PLOT:\n",
" plot_waterfall(deint_avg_data[0], bl=(ANTPAIR + ('pI',)), flags=deint_avg_flags[0], nsamples=deint_avg_nsamples[0], tslice=None)\n",
" plot_real_delay_vs_lst(deint_avg_data[0], bl=(ANTPAIR + ('pI',)), xlim=[-3999, 3999], clim=[-2e2, 2e2], linthresh=1, tslice=None)\n",
" plot_real_delay_vs_lst(deint_avg_data[0], bl=(ANTPAIR + ('pI',)), flags=deint_flags[0], xlim=[-3999, 3999], clim=[-2e2, 2e2], linthresh=1, tslice=None)\n",
" plot_dly_vs_fr(deint_avg_data[0], bl=(ANTPAIR + ('pI',)), xlim=[-1999, 1999], clim=[1e0, 1e5], tslice=None)"
]
},
Expand Down

0 comments on commit f4d3f8b

Please sign in to comment.