Use one convention in get_frop_wrapper

HERA-Team · Nov 6, 2024 · ee6f8bd · ee6f8bd
1 parent bb9f198
commit ee6f8bd
Showing 1 changed file with 21 additions and 12 deletions.
diff --git a/notebooks/single_baseline_postprocessing_and_pspec.ipynb b/notebooks/single_baseline_postprocessing_and_pspec.ipynb
@@ -1419,7 +1419,8 @@
    "source": [
     "def get_frop_wrapper(mode=\"main_lobe\", pol=\"nn\", stream_ind=0, band_ind=0, t_avg=AVERAGING_TIME, \n",
     "                     rephase=True, wgt_tavg_by_nsample=True, bl_vec=None,\n",
-    "                     dlst=None, coherent_avg=True, times=None, band_slice=None):\n",
+    "                     dlst=None, coherent_avg=True, times=None,\n",
+    "                     freq_decimation=CORR_MATRIX_FREQ_DECIMATION):\n",
     "    \"\"\"\n",
     "    This wraps hera_cal.frf.get_frop_for_noise using specific information from this notebook so that we can do\n",
     "    post FRF time-time visibility covariance calculation. It returns an operator that can be used on a dynamic \n",
@@ -1428,15 +1429,15 @@
     "    bl=(ANTPAIR[0], ANTPAIR[1], pol)\n",
     "    \n",
     "    band = bands[band_ind]\n",
-    "    if band_slice is None:\n",
-    "        band_slice = band_slices[band_ind]\n",
-    "    Nfreqs = band_slice.stop - band_slice.start\n",
+    "    band_slice = band_slices[band_ind]\n",
+    "    if freq_decimation > 1:\n",
+    "        band_slice = slice(band_slice.start, band_slice.stop, freq_decimation)\n",
     "\n",
     "    weights=deint_wgts[stream_ind][bl][:, band_slice]\n",
     "    nsamples = deint_nsamples[stream_ind][bl][:, band_slice]\n",
     "    \n",
     "    if times is None:\n",
-    "        times = np.modf(single_bl_times[tslice][stream_ind::4])[0] * 24 * 3600\n",
+    "        times = np.modf(single_bl_times[tslice][stream_ind::NINTERLEAVE])[0] * 24 * 3600\n",
     "        times = times - times[0]\n",
     "\n",
     "    if mode == \"main_lobe\":\n",
@@ -1470,8 +1471,8 @@
     "coherent_avg_correction_factors = []\n",
     "for spw, band in enumerate(bands):\n",
     "    freq_slice = band_slices[spw]\n",
-    "    if CORR_MATRIX_FREQ_DECIMATION > 1:\n",
-    "        freq_slice = slice(freq_slice.start, freq_slice.stop, CORR_MATRIX_FREQ_DECIMATION)\n",
+    "    decimated_freq_slice = slice(freq_slice.start, freq_slice.stop, CORR_MATRIX_FREQ_DECIMATION)\n",
+    "        \n",
     "    if USE_CORR_MATRIX:\n",
     "        corr_factors = []\n",
     "        for stream_ind in range(NINTERLEAVE):\n",
@@ -1484,15 +1485,15 @@
     "                                            deint_nsamples[stream_ind],\n",
     "                                            deint_wgts[stream_ind],\n",
     "                                            cross_antpairpol,\n",
-    "                                            freq_slice,\n",
+    "                                            decimated_freq_slice,\n",
     "                                            auto_ant=0)\n",
     "                main_lobe_frop = get_frop_wrapper(pol=pol, stream_ind=stream_ind, band_ind=spw,\n",
     "                                                  dlst=dlst, bl_vec=bl_vec[cross_antpairpol],\n",
-    "                                                  times=times, band_slice=freq_slice)\n",
+    "                                                  times=times)\n",
     "                xtalk_frop = get_frop_wrapper(pol=pol, stream_ind=stream_ind, band_ind=spw,\n",
     "                                              dlst=dlst, bl_vec=bl_vec[cross_antpairpol],\n",
     "                                              times=times, rephase=False, mode=\"xtalk\", coherent_avg=True,\n",
-    "                                              t_avg=times[1] - times[0], band_slice=freq_slice)\n",
+    "                                              t_avg=times[1] - times[0])\n",
     "                \n",
     "\n",
     "                frop = np.zeros_like(main_lobe_frop)\n",
@@ -1503,8 +1504,6 @@
     "                    frop[:, :, freq_ind] = np.tensordot(main_lobe_frop[:, :, freq_ind],\n",
     "                                                        xtalk_frop[:, :, freq_ind],\n",
     "                                                        axes=1)\n",
-    "\n",
-    "\n",
     "                cov_here += frf.get_FRF_cov(frop, var)  # computes covariance for pI by adding ee and nn\n",
     "            if hd.pol_convention == \"avg\":\n",
     "                cov_here /= 4\n",
@@ -1521,6 +1520,16 @@
     "        "
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2f462068",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(coherent_avg_correction_factors)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "51813f91",