more updates for 4pol

HERA-Team · Aug 27, 2024 · 0c61f3f · 0c61f3f
1 parent 5cd82b0
commit 0c61f3f
Showing 1 changed file with 8 additions and 10 deletions.
diff --git a/notebooks/single_baseline_postprocessing_and_pspec.ipynb b/notebooks/single_baseline_postprocessing_and_pspec.ipynb
@@ -467,9 +467,9 @@
    "source": [
     "# FOR DIAGNOSTICS/DEBUGGING ONLY: unflag everything and set nsamples to the median\n",
     "if NO_FLAGS_FLAT_NSAMPLES:\n",
-    "    for bl in [ANTPAIR + ('ee',), ANTPAIR + ('nn',)]:\n",
+    "    for bl in cross_bls:\n",
     "        flags[bl][tslice, :] = False\n",
-    "        nsamples[bl][tslice, :] = np.median([nsamples[bl][tslice, :] for bl in [ANTPAIR + ('ee',), ANTPAIR + ('nn',)]])\n",
+    "        nsamples[bl][tslice, :] = np.median([nsamples[bl][tslice, :] for bl in cross_bls])\n",
     "    \n",
     "# FOR DIAGNOSTICS/DEBUGGING ONLY: make all the autos a flat 10,000 Jy \n",
     "if FLAT_AUTOS:\n",
@@ -794,7 +794,7 @@
     "    dly_filt_data = copy.deepcopy(data)\n",
     "    inpainted_data = copy.deepcopy(data)\n",
     "    if bls is None:\n",
-    "        bls = [ANTPAIR + (pol,) for pol in ['ee', 'nn']]\n",
+    "        bls = cross_bls\n",
     "    \n",
     "    for bl in bls:\n",
     "        d_mdl = np.zeros_like(dly_filt_data[bl])\n",
@@ -823,7 +823,7 @@
     "    '''This function performs a high-pass filter in fringe rate, removing some small range around the 0 FR mode.'''\n",
     "    xtalk_filt_data = copy.deepcopy(data)\n",
     "    if bls is None:\n",
-    "        bls = [ANTPAIR + (pol,) for pol in ['ee', 'nn']]    \n",
+    "        bls = cross_bls\n",
     "    \n",
     "    for bl in bls:\n",
     "        if tslice is None:\n",
@@ -850,7 +850,7 @@
     "    '''This function performs fringe rate filtering, keeping a range determined by fr_ranges for each band.'''\n",
     "    frf_data = copy.deepcopy(data)\n",
     "    if bls is None:\n",
-    "        bls = [ANTPAIR + (pol,) for pol in ['ee', 'nn']]\n",
+    "        bls = cross_bls\n",
     "    \n",
     "    for bl in bls:    \n",
     "        if tslice is None:\n",
@@ -1130,11 +1130,10 @@
     "if USE_SIMULATED_NOISE:\n",
     "    np.random.seed(21)\n",
     "\n",
-    "    reds = redcal.get_reds(filt_data.data_antpos, pols=['ee', 'nn'], include_autos=True, bl_error_tol=2.0)\n",
+    "    reds = redcal.get_reds(filt_data.data_antpos, pols=filt_data.pols(), include_autos=True, bl_error_tol=2.0)\n",
     "    red_inpainted = datacontainer.RedDataContainer(data, reds=reds)\n",
     "\n",
-    "    for pol in ['ee', 'nn']:\n",
-    "        bl = ANTPAIR + (pol,)\n",
+    "    for bl in cross_bls:\n",
     "        predicted_var = noise.predict_noise_variance_from_autos(bl, red_inpainted, nsamples=nsamples)\n",
     "        predicted_var = np.where(~np.isfinite(predicted_var), 0, predicted_var)\n",
     "        filt_data[bl] = np.sqrt(predicted_var) / 2**.5 * (np.random.randn(hd.Ntimes, hd.Nfreqs) + 1.0j * np.random.randn(hd.Ntimes, hd.Nfreqs))"
@@ -1150,8 +1149,7 @@
     "# Inpaint crosses\n",
     "if PERFORM_INPAINT:\n",
     "    _, filt_data = delay_filter(data, freq_filt_wgts, inpaint_filter_centers, inpaint_filter_half_widths, INPAINT_EIGENVAL_CUTOFF)\n",
-    "    for pol in ['ee', 'nn']:\n",
-    "        bl = ANTPAIR + (pol,)\n",
+    "    for bl in cross_bls:\n",
     "        for band in [high_band, low_band]:\n",
     "            for i in range(filt_flags[bl].shape[0]):\n",
     "                if not np.all(filt_flags[bl][i, band]):\n",