added pspecdata.trim_dset_lst for lst aligning of input dsets

hera_pspec/pspecdata.py

@@ -1640,6 +1640,44 @@ def Jy_to_mK(self, beam=None):
                 dset.data_array[:, :, :, j] *= factors[p][None, None, :]
             dset.vis_units = 'mK'
 
+    def trim_dset_lsts(self, lst_tol=6):
+        """
+        Assuming all datasets in self.dsets are locked to the same LST grid (but
+        each may have a constant offset), trim LSTs from each dset that aren't found
+        in all other dsets (within some decimal tolerance specified by lst_tol). 
+
+        Warning: this edits the data in dsets in-place, and is not reversible.
+
+        Parameters
+        ----------
+        lst_tol : float
+            Decimal tolerance [radians] for comparing float-valued LST bins.
+        """
+        # ensure each dset has same dLST within tolerance / Ntimes
+        dlst = np.median(np.diff(np.unique(self.dsets[0].lst_array)))
+        for dset in self.dsets:
+            _dlst = np.median(np.diff(np.unique(dset.lst_array)))
+            if not np.isclose(dlst, _dlst, atol=10**(-lst_tol) / dset.Ntimes):
+                print "not all datasets in self.dsets are on the same LST grid, cannot LST trim."
+                return
+
+        # get lst array of each dataset and turn into string and add to common_lsts
+        lst_arrs = []
+        common_lsts = set()
+        for i, dset in enumerate(self.dsets):
+            lsts = ["{lst:0.{tol}f}".format(lst=l, tol=lst_tol) for l in dset.lst_array]
+            lst_arrs.append(lsts)
+            if i == 0:
+                common_lsts = common_lsts.union(set(lsts))
+            else:
+                common_lsts = common_lsts.intersection(set(lsts))
+
+        # iterate through dsets and trim off integrations whose lst isn't in common_lsts
+        for i, dset in enumerate(self.dsets):
+            trim_inds = np.array([l not in common_lsts for l in lst_arrs[i]])
+            if np.any(trim_inds):
+                self.dsets[i].select(times=dset.time_array[~trim_inds])
+
 
 def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False, group=False, Nblps_per_group=1):
     """

hera_pspec/tests/test_pspecdata.py

@@ -550,6 +550,26 @@ def test_Jy_to_mK(self):
         nt.assert_equal(ds.dsets[1].vis_units, "UNCALIB")
         nt.assert_not_equal(ds.dsets[0].get_data(24, 25, 'xx')[30, 30], ds.dsets[1].get_data(24, 25, 'pI')[30, 30])
 
+    def test_trim_dset_lsts(self):
+        fname = os.path.join(DATA_PATH, "zen.2458042.17772.xx.HH.uvXA")
+        uvd1 = UVData()
+        uvd1.read_miriad(fname)
+        uvd2 = copy.deepcopy(uvd1)
+        uvd2.lst_array = (uvd2.lst_array + 10 * np.median(np.diff(np.unique(uvd2.lst_array)))) % (2*np.pi)
+        # test basic execution
+        ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None])
+        ds.trim_dset_lsts()
+        nt.assert_true(ds.dsets[0].Ntimes, 52)
+        nt.assert_true(ds.dsets[1].Ntimes, 52)
+        nt.assert_true(np.all( (2458042.178948477 < ds.dsets[0].time_array) \
+                        + (ds.dsets[0].time_array < 2458042.1843023109)))
+        # test exception
+        uvd2.lst_array += np.linspace(0, 1e-3, uvd2.Nblts)
+        ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None])
+        ds.trim_dset_lsts()
+        nt.assert_true(ds.dsets[0].Ntimes, 60)
+        nt.assert_true(ds.dsets[1].Ntimes, 60)
+
     def test_units(self):
         ds = pspecdata.PSpecData()
         # test exception