Merge pull request #864 from HERA-Team/fix-rephase

Temporarily restore old array shape compatibility in HERAData for the sake of pspec
HERA-Team · Jan 20, 2023 · f146720 · f146720
2 parents 541d065 + eff4762
commit f146720
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 22 deletions.
diff --git a/cache_temp/025b413f8a9a021ea648a7dd06839eb9.filter_cache b/cache_temp/025b413f8a9a021ea648a7dd06839eb9.filter_cache
diff --git a/hera_cal/io.py b/hera_cal/io.py
@@ -619,29 +619,55 @@ def _get_slice(self, data_array, key):
             return {pol: self._get_slice(data_array, key + (pol,)) for pol in pols}
         elif len(key) == 3:  # asking for bl-pol
             try:
-                return np.array(
-                    data_array[
-                        self._blt_slices[tuple(key[:2])], :,
-                        self._polstr_indices.get(
-                            key[2],
-                            self._polnum_indices[
-                                polstr2num(key[2], x_orientation=self.x_orientation)
-                            ]
-                        )
-                    ]
-                )
+                if data_array.ndim == 4: # old shapes
+                    return np.array(
+                        data_array[
+                            self._blt_slices[tuple(key[:2])], 0, :,
+                            self._polstr_indices.get(
+                                key[2],
+                                self._polnum_indices[
+                                    polstr2num(key[2], x_orientation=self.x_orientation)
+                                ]
+                            )
+                        ]
+                    )
+                else:
+                    return np.array(
+                        data_array[
+                            self._blt_slices[tuple(key[:2])], :,
+                            self._polstr_indices.get(
+                                key[2],
+                                self._polnum_indices[
+                                    polstr2num(key[2], x_orientation=self.x_orientation)
+                                ]
+                            )
+                        ]
+                    )
             except KeyError:
-                return np.conj(
-                    data_array[
-                        self._blt_slices[tuple(key[1::-1])], :,
-                        self._polstr_indices.get(
-                            conj_pol(key[2]),
-                            self._polnum_indices[
-                                polstr2num(conj_pol(key[2]), x_orientation=self.x_orientation)
-                            ]
-                        )
-                    ]
-                )
+                if data_array.ndim == 4:
+                    return np.conj(
+                        data_array[
+                            self._blt_slices[tuple(key[1::-1])], 0, :,
+                            self._polstr_indices.get(
+                                conj_pol(key[2]),
+                                self._polnum_indices[
+                                    polstr2num(conj_pol(key[2]), x_orientation=self.x_orientation)
+                                ]
+                            )
+                        ]
+                    )
+                else:
+                    return np.conj(
+                        data_array[
+                            self._blt_slices[tuple(key[1::-1])], :,
+                            self._polstr_indices.get(
+                                conj_pol(key[2]),
+                                self._polnum_indices[
+                                    polstr2num(conj_pol(key[2]), x_orientation=self.x_orientation)
+                                ]
+                            )
+                        ]
+                    )
         else:
             raise KeyError('Unrecognized key type for slicing data.')