Skip to content

Commit

Permalink
pspec plotting updates
Browse files Browse the repository at this point in the history
  • Loading branch information
@jsdillon
jsdillon committed Feb 1, 2024
1 parent ec91457 commit 6c2d339
Showing 1 changed file with 54 additions and 26 deletions.
80 changes: 54 additions & 26 deletions notebooks/single_baseline_postprocessing_and_pspec.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1604,19 +1604,17 @@
"cell_type": "code",
"execution_count": null,
"id": "6b27cdb7",
"metadata": {
"scrolled": false
},
"metadata": {},
"outputs": [],
"source": [
"def plot_Pk_vs_LST(clim = None, xlim = [-1999, 1999]):\n",
"def plot_Pk_vs_LST(clim=None, xlim=[-2999, 2999], pol='pI'):\n",
" '''Plots the real part of the power spectrum from each band as a function of LST and delay for each band.'''\n",
" \n",
" lsts = np.where(avg_pspec_lsts > avg_pspec_lsts[-1], avg_pspec_lsts - 2 * np.pi, avg_pspec_lsts) * 12 / np.pi\n",
" fig, axes = plt.subplots(1, len(bands), figsize=(28, 12), sharex=True, sharey=True, gridspec_kw={'wspace': .03}, dpi=100)\n",
" for spw, (ax, band, band_slice) in enumerate(zip(axes, bands, band_slices)):\n",
"\n",
" key = (spw, (ANTPAIR, ANTPAIR), ('pI', 'pI')) \n",
" key = (spw, (ANTPAIR, ANTPAIR), (pol, pol)) \n",
" pk_avg = np.mean([uvp.get_data(key) for uvp in uvps], axis=0).real\n",
" delays = uvps[0].get_dlys(key[0]) * 1e9\n",
"\n",
Expand All @@ -1627,16 +1625,17 @@
" vmax=(clim[1] if clim is not None else np.max(np.abs(_to_plot)))),\n",
" extent=[delays[0], delays[-1], lsts[-1], lsts[0]])\n",
"\n",
" for dly in filter_half_widths[0] * 1e9 * np.array([1, -1]):\n",
" ax.axvline(dly, ls='--', color='k', lw=.5)\n",
" for multiple in [1, -1]:\n",
" ax.axvline(multiple * dly_filter_half_widths[0] * 1e9, ls='--', color='k')\n",
" ax.axvline(multiple * inpaint_filter_half_widths[0] * 1e9, ls=':', color='k')\n",
" ax.set_xlim(xlim)\n",
" ax.set_title(f'Band {spw+1}:\\n{band[0]}—{band[1]} MHz', fontsize=10)\n",
" ax.set_xlabel('Delay (ns)')\n",
" if spw == 0:\n",
" ax.set_ylabel('LST (Hours)')\n",
" ax.set_yticklabels([f'{(int(val) if np.isclose(val, int(val)) else val) % 24}' for val in ax.get_yticks()])\n",
"\n",
" plt.colorbar(im, ax=axes, pad=.02, aspect=40, extend='both', label=r'|Re[$P(k)$]| (mK$^2$ $h^{-3}$ Mpc$^3$)')"
" plt.colorbar(im, ax=axes, pad=.02, aspect=40, extend='both', label=f'{pol} ' + r'|Re[$P(k)$]| (mK$^2$ $h^{-3}$ Mpc$^3$)')"
]
},
{
Expand All @@ -1646,15 +1645,15 @@
"metadata": {},
"outputs": [],
"source": [
"def plot_Pk_SNR_vs_LST(clim=[-5, 5], xlim = [-1999, 1999]):\n",
"def plot_Pk_SNR_vs_LST(clim=[-5, 5], xlim = [-2999, 2999], func=np.real):\n",
" '''Plots the real power spectrum SNR (normalized by P_N, not P_SN) as a function of LST and delay for each band.'''\n",
" \n",
" lsts = np.where(data.lsts[tslice] > data.lsts[tslice][-1], data.lsts[tslice] - 2 * np.pi, data.lsts[tslice]) * 12 / np.pi\n",
" fig, axes = plt.subplots(1, len(bands), figsize=(28, 12), sharex=True, sharey=True, gridspec_kw={'wspace': .03}, dpi=100)\n",
" for spw, (ax, band, band_slice) in enumerate(zip(axes, bands, band_slices)):\n",
"\n",
" key = (spw, (ANTPAIR, ANTPAIR), ('pI', 'pI')) \n",
" pk_avg = np.mean([uvp.get_data(key) for uvp in uvps], axis=0).real\n",
" pk_avg = func(np.mean([uvp.get_data(key) for uvp in uvps], axis=0))\n",
" delays = uvps[0].get_dlys(key[0]) * 1e9\n",
" P_N = np.mean([np.abs(uvp.get_stats('P_N', key)) for uvp in uvps], axis=0)\n",
"\n",
Expand All @@ -1663,16 +1662,17 @@
" im = ax.imshow(SNR, interpolation='none', aspect='auto', cmap='bwr', \n",
" vmin=clim[0], vmax=clim[1],\n",
" extent=[delays[0], delays[-1], lsts[-1], lsts[0]])\n",
" for dly in filter_half_widths[0] * 1e9 * np.array([1, -1]):\n",
" ax.axvline(dly, ls='--', color='k', lw=.5)\n",
" for multiple in [1, -1]:\n",
" ax.axvline(multiple * dly_filter_half_widths[0] * 1e9, ls='--', color='k')\n",
" ax.axvline(multiple * inpaint_filter_half_widths[0] * 1e9, ls=':', color='k')\n",
" ax.set_xlim(xlim)\n",
" ax.set_title(f'Band {spw+1}:\\n{band[0]}—{band[1]} MHz', fontsize=10)\n",
" ax.set_xlabel('Delay (ns)')\n",
" if spw == 0:\n",
" ax.set_ylabel('LST (Hours)')\n",
" ax.set_yticklabels([f'{(int(val) if np.isclose(val, int(val)) else val) % 24}' for val in ax.get_yticks()])\n",
"\n",
" plt.colorbar(im, ax=axes, pad=.02, aspect=40, extend='both', label=r'Re[$P(k) / P_N(k)$] (unitless)')"
" plt.colorbar(im, ax=axes, pad=.02, aspect=40, extend='both', label=f'{\"Re\" if func == np.real else \"Im\"}' + r'[$P(k) / P_N(k)$] (unitless)')"
]
},
{
Expand Down Expand Up @@ -1729,7 +1729,7 @@
"def plot_tavg_pspec():\n",
" '''This plots the time-averaged power spectrum over the whole range of LSTs, including 2 sigma errors'''\n",
" \n",
" fig, axes = plt.subplots(len(bands) // 2, 2, figsize=(18, 12), sharex=True, sharey=True, gridspec_kw={'wspace': .03, 'hspace': .0}, dpi=100)\n",
" fig, axes = plt.subplots(int(np.ceil(len(bands) / 2)), 2, figsize=(18, 12), sharex=True, sharey=True, gridspec_kw={'wspace': .03, 'hspace': .0}, dpi=100)\n",
" for spw, (ax, band, band_slice) in enumerate(zip(np.ravel(axes), bands, band_slices)):\n",
"\n",
" key = (spw, (ANTPAIR, ANTPAIR), ('pI', 'pI'))\n",
Expand All @@ -1743,9 +1743,12 @@
" ax.plot(delays, P_SN, 'k-', label='$P_{SN}$')\n",
" ax.set_yscale('log') \n",
" ax.set_xlim([-2500, 2500])\n",
" ax.set_ylim([1e3, 1e14])\n",
" ax.set_xlabel('Delay (ns)')\n",
" for dly in filter_half_widths[0] * 1e9 * np.array([1, -1]):\n",
" ax.axvline(dly, ls='--', color='k', lw=.5)\n",
" ax.tick_params(axis='x', direction='in')\n",
" for multiple in [1, -1]:\n",
" ax.axvline(multiple * dly_filter_half_widths[0] * 1e9, ls='--', color='k', lw=.5, label=(r'Filtering $\\tau_{max}$' if multiple == 1 else None))\n",
" ax.axvline(multiple * inpaint_filter_half_widths[0] * 1e9, ls=':', color='k', lw=.5, label=(r'Inpainting $\\tau_{max}$' if multiple == 1 else None))\n",
" if spw % 2 == 0:\n",
" ax.set_ylabel('Re[$P(k)$]\\n(mK$^2$ $h^{-3}$ Mpc$^3$)') \n",
"\n",
Expand All @@ -1763,28 +1766,45 @@
"id": "d7ab9c3c",
"metadata": {},
"source": [
"# *Figure 8: Interleave-Averaged Power Spectrum and SNR vs. LST*"
"# *Figure 8: Interleave-Averaged Power Spectra (Psuedo-Stokes I and Q) vs. LST*"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f4ae4cc3",
"metadata": {
"scrolled": false
},
"metadata": {},
"outputs": [],
"source": [
"plot_Pk_vs_LST()\n",
"plot_Pk_SNR_vs_LST()"
"plot_Pk_vs_LST(pol='pI')\n",
"plot_Pk_vs_LST(pol='pQ')"
]
},
{
"cell_type": "markdown",
"id": "2cbb558c",
"metadata": {},
"source": [
"# *Figure 9: Interleave-Averaged Power Spectrum SNR vs. LST (Real, Imaginary, and pQ)*"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "93240fe7",
"metadata": {},
"outputs": [],
"source": [
"plot_Pk_SNR_vs_LST(func=np.real)\n",
"plot_Pk_SNR_vs_LST(func=np.imag)"
]
},
{
"cell_type": "markdown",
"id": "638078a7",
"metadata": {},
"source": [
"# *Figure 9: High Delay Power Spectrum SNR Histograms Before and After Incoherent Averaging*"
"# *Figure 10: High Delay Power Spectrum SNR Histograms Before and After Incoherent Averaging*"
]
},
{
Expand All @@ -1804,23 +1824,31 @@
"id": "0d97725a",
"metadata": {},
"source": [
"# *Figure 10: Incoherently Averaged Power Spectrum with Error Bars*"
"# *Figure 11: Incoherently Averaged Power Spectrum with Error Bars*"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0a636cca",
"id": "85ace6b2",
"metadata": {},
"outputs": [],
"source": [
"plot_tavg_pspec()"
]
},
{
"cell_type": "markdown",
"id": "75ef3628",
"metadata": {},
"source": [
"## Save Results"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8089ed50",
"id": "9362190a",
"metadata": {},
"outputs": [],
"source": [
Expand Down

0 comments on commit 6c2d339

Please sign in to comment.