Allows for multiple polarizations pairs

hera_pspec/pspecdata.py

@@ -999,7 +999,7 @@ def validate_pol(self, dsets, pol_pair):
            return True
 
 
-    def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm='I', 
+    def pspec(self, bls1, bls2, dsets, pols, input_data_weight='identity', norm='I', 
               taper='none', little_h=True, spw_ranges=None, verbose=True, 
               history=''):
         """
@@ -1020,6 +1020,7 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
         If the bl chosen from bls1 is (ant1, ant2) and the bl chosen from bls2 
         is (ant3, ant4), the "baseline-pair" describing their cross 
         multiplication is ((ant1, ant2), (ant3, ant4)).
+
         Parameters
         ----------
         bls1, bls2 : list
@@ -1029,25 +1030,31 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
             Contains indices of self.dsets to use in forming power spectra, 
             where the first index is for the Left-Hand dataset and second index 
             is used for the Right-Hand dataset (see above).
-        pols : length-2 tuple or list of strings or integers
-            contains polarization pairs to use in forming power spectra 
-            e.g. ('XX','XX') or ('XY','YX') or list of polarization pairs. 
+    
+        pols : length-2 tuple of strings or integers or list of length-2 tuples of strings or integers
+            Contains polarization pairs to use in forming power spectra 
+            e.g. ('XX','XX') or [('XX','XX'),('XY','YX')] or list of polarization pairs. 
             Only auto/equal polarization pairs are implemented at the moment. 
             It uses the polarizations of the UVData onjects (specified in dsets)
-            by default only if the UVData object consists of only one polarization.
+            by default only if the UVData object consists of equal polarizations.
+    
         input_data_weight : str, optional
             String specifying which weighting matrix to apply to the input
             data. See the options in the set_R() method for details. 
             Default: 'identity'.
+
         norm : str, optional
             String specifying how to choose the normalization matrix, M. See 
             the 'mode' argument of get_MW() for options. Default: 'I'.
+
         taper : str, optional
             Tapering (window) function to apply to the data. Takes the same
             arguments as aipy.dsp.gen_window(). Default: 'none'.
+
         little_h : boolean, optional
                 Whether to have cosmological length units be h^-1 Mpc or Mpc
                 Default: h^-1 Mpc
+
         spw_ranges : list of tuples, optional
             A list of spectral window channel ranges to select within the total 
             bandwidth of the datasets, each of which forms an independent power 
@@ -1060,10 +1067,12 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
             If True, print progress, warnings and debugging info to stdout.
         history : str, optional
             history string to attach to UVPSpec object
+
         Returns
         -------
         uvp : UVPSpec object
             Instance of UVPSpec that holds the output power spectrum data.
+
         Examples
         --------
         *Example 1:* No grouping; i.e. each baseline is its own group, no 
@@ -1122,15 +1131,14 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
         dset2 = self.dsets[self.dset_idx(dsets[1])]
 
         # check  inputs for pol
-        npols0 = dset1.Npols # number of polarization for zero'th UVData of the first set of UVData objects
-        npols1 = dset2.Npols # number of polarization for zero'th UVData of the second set of UVData objects
-        if pols == None:
-            if npols0 == npols1 == 1:
-                pols = [(dset1.get_pols()[0], dset2.get_pols()[0])]
-                raise_warning("UVData objects have pols {} which are used to estimate the power spectrum.".format(tuple(pols)), verbose=verbose)
-            else:
-                raise ValueError("UVData objects must have only one polarization axis otherwise pols should be specified.")
-
+        #npols0 = dset1.Npols # number of polarization for zero'th UVData of the first set of UVData objects
+        #npols1 = dset2.Npols # number of polarization for zero'th UVData of the second set of UVData objects
+        #if pols == None:
+        #    if npols0 == npols1 == 1:
+        #        pols = [(dset1.get_pols()[0], dset2.get_pols()[0])]
+        #        raise_warning("UVData objects have pols {} which are used to estimate the power spectrum.".format(tuple(pols)), verbose=verbose)
+        #    else:
+        #        raise ValueError("UVData objects must have only one polarization axis otherwise pols should be specified.")
 
         # assert form of bls1 and bls2
         assert len(bls1) == len(bls2), "length of bls1 must equal length of bls2"
@@ -1199,6 +1207,7 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
                 dset1.freq_array.flatten()[spw_ranges[i][0]:spw_ranges[i][1]] )
 
             # Loop over polarizations
+            if isinstance(pols, tuple): pols = [pols]
             for j, ps in enumerate(pols):
                 p = [] 
                 if isinstance(ps[0], (int, np.int)): 
@@ -1214,7 +1223,8 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
                 # validating polarization pair of UVData objects
                 valid = self.validate_pol(dsets, tuple(p))
                 if valid:
-                   spw_pol.append(p)
+                   # storing only one polarization as only equal polarization are allowed at the moment and UVPspec objec also understands one polarization for each UVspec object
+                   spw_pol.append(p[0]) 
                    pass
                 else:
                    print ("Polarization pair: {} failed the validation test".format(p))
@@ -1306,8 +1316,8 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
                     wgts2 = self.w(key2).T
 
                     # get average of nsample across frequency axis, weighted by wgts
-                    nsamp1 = np.sum(dset1.get_nsamples(bl1)[:, self.spw_range[0]:self.spw_range[1]] * wgts1, axis=1) / np.sum(wgts1, axis=1).clip(1, np.inf)
-                    nsamp2 = np.sum(dset2.get_nsamples(bl2)[:, self.spw_range[0]:self.spw_range[1]] * wgts2, axis=1) / np.sum(wgts2, axis=1).clip(1, np.inf)
+                    nsamp1 = np.sum(dset1.get_nsamples(bl1 + (p[0],))[:, self.spw_range[0]:self.spw_range[1]] * wgts1, axis=1) / np.sum(wgts1, axis=1).clip(1, np.inf)
+                    nsamp2 = np.sum(dset2.get_nsamples(bl2 + (p[1],))[:, self.spw_range[0]:self.spw_range[1]] * wgts2, axis=1) / np.sum(wgts2, axis=1).clip(1, np.inf)
 
                     # take average of nsamp1 and nsamp2 and multiply by integration time [seconds] to get total integration
                     pol_ints.extend(np.mean([nsamp1, nsamp2], axis=0) * dset1.integration_time)
@@ -1371,7 +1381,7 @@ def pspec(self, bls1, bls2, dsets, pols=None, input_data_weight='identity', norm
         uvp.Ndlys = len(np.unique(dlys))
         uvp.Nspwdlys = len(spws)
         uvp.Nfreqs = len(np.unique(freqs))
-        uvp.pol_array = np.array(map(lambda p: uvutils.polstr2num(p[0]), spw_pol))
+        uvp.pol_array = np.array(map(lambda p: uvutils.polstr2num(p), spw_pol))
         uvp.Npols = len(spw_pol)
         uvp.scalar_array = np.array(sclr_arr)
         uvp.channel_width = dset1.channel_width

hera_pspec/tests/test_plot.py

@@ -46,7 +46,7 @@ def setUp(self):
                         bls, exclude_permutations=False, exclude_auto_bls=True)
 
         # Calculate the power spectrum
-        self.uvp = self.ds.pspec(self.bls1, self.bls2, (0, 1), 
+        self.uvp = self.ds.pspec(self.bls1, self.bls2, (0, 1), ('xx','xx'),  
                                  spw_ranges=[(300, 400), (600,721)], 
                                  input_data_weight='identity', norm='I', 
                                  taper='blackman-harris', verbose=False)

hera_pspec/tests/test_pspecdata.py

@@ -509,10 +509,10 @@ def test_rephase_to_dst(self):
         ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd1)], wgts=[None, None])
         # get normal pspec
         bls = [(37, 39)]
-        uvp1 = ds.pspec(bls, bls, (0, 1), verbose=False)
+        uvp1 = ds.pspec(bls, bls, (0, 1), pols=('xx','xx'), verbose=False)
         # rephase and get pspec
         ds.rephase_to_dset(0)
-        uvp2 = ds.pspec(bls, bls, (0, 1), verbose=False)
+        uvp2 = ds.pspec(bls, bls, (0, 1), pols=('xx','xx'), verbose=False)
         blp = (0, ((37,39),(37,39)), 'XX')
         nt.assert_true(np.isclose(np.abs(uvp2.get_data(blp)/uvp1.get_data(blp)), 1.0).min())
 
@@ -556,7 +556,7 @@ def test_pspec(self):
 
         # check basic execution with baseline list
         bls = [(24, 25), (37, 38), (38, 39), (52, 53)]
-        uvp = ds.pspec(bls, bls, (0, 1), input_data_weight='identity', norm='I', taper='none',
+        uvp = ds.pspec(bls, bls, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none',
                                 little_h=True, verbose=False)
         nt.assert_equal(len(uvp.bl_array), len(bls))
         nt.assert_true(uvp.antnums_to_blpair(((24, 25), (24, 25))) in uvp.blpair_array)
@@ -568,11 +568,11 @@ def test_pspec(self):
         antpos = dict(zip(ants, antpos))
         red_bls = map(lambda blg: sorted(blg), redcal.get_pos_reds(antpos, low_hi=True))[2]
         bls1, bls2, blps = pspecdata.construct_blpairs(red_bls, exclude_permutations=True)
-        uvp = ds.pspec(bls1, bls2, (0, 1), input_data_weight='identity', norm='I', taper='none',
+        uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none',
                                 little_h=True, verbose=False)
         nt.assert_true(uvp.antnums_to_blpair(((24, 25), (37, 38))) in uvp.blpair_array)
         nt.assert_equal(uvp.Nblpairs, 10)
-        uvp = ds.pspec(bls1, bls2, (0, 1), input_data_weight='identity', norm='I', taper='none',
+        uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none',
                                 little_h=True, verbose=False)
         nt.assert_true(uvp.antnums_to_blpair(((24, 25), (52, 53))) in uvp.blpair_array)
         nt.assert_true(uvp.antnums_to_blpair(((52, 53), (24, 25))) not in uvp.blpair_array)
@@ -581,15 +581,15 @@ def test_pspec(self):
         # test mixed bl group and non blgroup, currently bl grouping of more than 1 blpair doesn't work
         bls1 = [[(24, 25)], (52, 53)]
         bls2 = [[(24, 25)], (52, 53)]
-        uvp = ds.pspec(bls1, bls2, (0, 1), input_data_weight='identity', norm='I', taper='none',
+        uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none',
                                 little_h=True, verbose=False)
 
         # test select
         red_bls = [(24, 25), (37, 38), (38, 39), (52, 53)]
         bls1, bls2, blp = pspecdata.construct_blpairs(red_bls, exclude_permutations=False, exclude_auto_bls=False)
         uvd = copy.deepcopy(self.uvd)
         ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm)
-        uvp = ds.pspec(bls1, bls2, (0, 1), spw_ranges=[(20,30), (30,40)], verbose=False)
+        uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), spw_ranges=[(20,30), (30,40)], verbose=False)
         nt.assert_equal(uvp.Nblpairs, 16)
         nt.assert_equal(uvp.Nspws, 2)
         uvp2 = uvp.select(spws=[0], bls=[(24, 25)], only_pairs_in_bls=False, inplace=False)
@@ -602,7 +602,7 @@ def test_pspec(self):
         # check w/ multiple spectral ranges
         uvd = copy.deepcopy(self.uvd)
         ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm)
-        uvp = ds.pspec(bls, bls, (0, 1), spw_ranges=[(10, 24), (30, 40), (45, 64)], verbose=False)
+        uvp = ds.pspec(bls, bls, (0, 1), ('xx','xx'), spw_ranges=[(10, 24), (30, 40), (45, 64)], verbose=False)
         nt.assert_equal(uvp.Nspws, 3)
         nt.assert_equal(uvp.Nspwdlys, 43)
         nt.assert_equal(uvp.data_array[0].shape, (240, 14, 1))
@@ -617,33 +617,37 @@ def test_pspec(self):
         # test polarization pairs
         uvd = copy.deepcopy(self.uvd)
         ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm)
-        uvp = ds.pspec(bls, bls, (0, 1), pols=[('xx','xx')], spw_ranges=[(10, 24)], verbose=False)
+        uvp = ds.pspec(bls, bls, (0, 1), ('xx','xx'), spw_ranges=[(10, 24)], verbose=False)
         nt.assert_raises(NotImplementedError, ds.pspec, bls, bls, (0, 1), pols=[('xx','yy')])        
         uvd = copy.deepcopy(self.uvd)
         ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm)
-        uvp = ds.pspec(bls, bls, (0, 1), pols=[('xx','xx'), ('yy','yy')], spw_ranges=[(10, 24)], verbose=False)
+        uvp = ds.pspec(bls, bls, (0, 1), [('xx','xx'), ('yy','yy')], spw_ranges=[(10, 24)], verbose=False)
 
         uvd = copy.deepcopy(self.uvd)
         ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm)
-        uvp = ds.pspec(bls, bls, (0, 1), pols=[(-5, -5)], spw_ranges=[(10, 24)], verbose=False)
+        uvp = ds.pspec(bls, bls, (0, 1), (-5, -5), spw_ranges=[(10, 24)], verbose=False)
 
         # test exceptions
-        nt.assert_raises(AssertionError, ds.pspec, bls1[:1], bls2, (0, 1))
-        nt.assert_raises(ValueError, ds.pspec, bls, bls, (0, 1), pols=[('yy','yy')]) 
-        uvd = copy.deepcopy(self.uvd)
+        nt.assert_raises(AssertionError, ds.pspec, bls1[:1], bls2, (0, 1), ('xx','xx'))
+        nt.assert_raises(ValueError, ds.pspec, bls, bls, (0, 1), pols=('yy','yy')) 
         uvd1 = copy.deepcopy(self.uvd)
         uvd1.polarization_array = np.array([-6])
         ds = pspecdata.PSpecData(dsets=[uvd, uvd1], wgts=[None, None], beam=self.bm)
-        nt.assert_raises(ValueError, ds.pspec, bls, bls, (0, 1), pols=[('xx','xx')])
+        nt.assert_raises(ValueError, ds.pspec, bls, bls, (0, 1), ('xx','xx'))
 
         # test files with more than one polarizations
-        uvd = copy.deepcopy(self.uvd)
         uvd1 = copy.deepcopy(self.uvd)
         uvd1.polarization_array = np.array([-6])
-        uvd3 = uvd + uvd1
-        ds = pspecdata.PSpecData(dsets=[uvd, uvd3], wgts=[None, None], beam=self.bm)
-        nt.assert_raises(ValueError, ds.pspec, bls, bls, (0, 1))
+        uvd2 = self.uvd + uvd1
+        ds = pspecdata.PSpecData(dsets=[uvd2, uvd2], wgts=[None, None], beam=self.bm)
+        uvp = ds.pspec(bls, bls, (0, 1), [('xx','xx'), ('yy','yy')], spw_ranges=[(10, 24)], verbose=False)
 
+        uvd1 = copy.deepcopy(self.uvd)
+        uvd1.polarization_array = np.array([-6])
+        uvd2 = self.uvd + uvd1  
+        ds = pspecdata.PSpecData(dsets=[uvd2, uvd2], wgts=[None, None], beam=self.bm)
+        uvp = ds.pspec(bls, bls, (0, 1), [('xx','xx'), ('xy','xy')], spw_ranges=[(10, 24)], verbose=False)
+
     def test_normalization(self):
         # Test Normalization of pspec() compared to PAPER legacy techniques
         d1 = self.uvd.select(times=np.unique(self.uvd.time_array)[:-1:2], 
@@ -677,7 +681,7 @@ def test_normalization(self):
         legacy = np.fft.fftshift(np.fft.ifft(data1, axis=1) * np.conj(np.fft.ifft(data2, axis=1)) * scalar, axes=1)[0]
         # hera_pspec OQE
         ds = pspecdata.PSpecData(dsets=[d1, d2], wgts=[None, None], beam=beam)
-        uvp = ds.pspec(bls1, bls2, (0, 1), taper='none', input_data_weight='identity', norm='I')
+        uvp = ds.pspec(bls1, bls2, (0, 1), pols=('xx','xx'), taper='none', input_data_weight='identity', norm='I')
         oqe = uvp.get_data(0, ((24, 25), (37, 38)), 'xx')[0]
         # assert answers are same to within 3%
         nt.assert_true(np.isclose(np.real(oqe)/np.real(legacy), 1, atol=0.03, rtol=0.03).all())
@@ -691,7 +695,7 @@ def test_normalization(self):
         legacy = np.fft.fftshift(np.fft.ifft(data1*window[None, :], axis=1) * np.conj(np.fft.ifft(data2*window[None, :], axis=1)) * scalar, axes=1)[0]
         # hera_pspec OQE
         ds = pspecdata.PSpecData(dsets=[d1, d2], wgts=[None, None], beam=beam)
-        uvp = ds.pspec(bls1, bls2, (0, 1), taper='blackman-harris', input_data_weight='identity', norm='I')
+        uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), taper='blackman-harris', input_data_weight='identity', norm='I')
         oqe = uvp.get_data(0, ((24, 25), (37, 38)), 'xx')[0]
         # assert answers are same to within 3%
         nt.assert_true(np.isclose(np.real(oqe)/np.real(legacy), 1, atol=0.03, rtol=0.03).all())