fix: inpainted_regions handling

HERA-Team · May 16, 2024 · ee27771 · ee27771
1 parent 35354ac
commit ee27771
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/hera_notebook_templates/notebooks/lststack.ipynb b/hera_notebook_templates/notebooks/lststack.ipynb
@@ -1488,7 +1488,8 @@
     "allbad = {}\n",
     "inpainted_regions = {}\n",
     "\n",
-    "for lstidx, zuv in enumerate(zsquare):\n",
+    "for lstidx, (zuv, stack) in enumerate(zip(zsquare, cross_stacks)):\n",
+    "    inpaint_flags = stack.inpainted()\n",
     "    for iap, (a, b) in enumerate(zuv.antpairs):\n",
     "        for ipol, pol in enumerate(zuv.pols):\n",
     "            pol = {'ee': 0, 'nn': 1}[pol]\n",
@@ -1505,15 +1506,24 @@
     "                    for rng in ranges:\n",
     "                        allbad[(lstidx, a, b, pol, jdint, rng[0], rng[1])] = zsqn[rng[0]:rng[1]]\n",
     "\n",
-    "                # Get contiguous regions of inpainted data (anything that is flagged outside FM)\n",
-    "                for band in inpaint_bands:\n",
-    "                    badfreqs = zuv.flags[night, iap, band, ipol]\n",
-    "\n",
+    "                if stackconf.inpaint_files is not None:\n",
+    "                    badfreqs = np.nonzero(inpaint_flags[night, iap, :, ipol])\n",
     "                    if len(badfreqs) > 0:\n",
     "                        ranges = consecutive(badfreqs)\n",
     "\n",
     "                        for rng in ranges:\n",
-    "                            inpainted_regions[(lstidx, a, b, pol, jdint, rng[0], rng[1])] = zsqn[band][rng[0]:rng[1]]\n",
+    "                            allbad[(lstidx, a, b, pol, jdint, rng[0], rng[1])] = zsqn[rng[0]:rng[1]]\n",
+    "    \n",
+    "                else:\n",
+    "                    # Get contiguous regions of inpainted data (anything that is flagged outside FM)\n",
+    "                    for band in inpaint_bands:\n",
+    "                        badfreqs = np.nonzero(stack.flags[night, iap, band, ipol])[0]\n",
+    "\n",
+    "                        if len(badfreqs) > 0:\n",
+    "                            ranges = consecutive(badfreqs)\n",
+    "\n",
+    "                            for rng in ranges:\n",
+    "                                inpainted_regions[(lstidx, a, b, pol, jdint, rng[0], rng[1])] = zsqn[band][rng[0]:rng[1]]\n",
     "    "
    ]
   },