Merge 42ab8d7 into e3b3b15

hera_pspec/pspecdata.py

@@ -4,7 +4,7 @@
 import copy, operator, itertools
 from collections import OrderedDict as odict
 import hera_cal as hc
-from hera_pspec import uvpspec, utils, version
+from hera_pspec import uvpspec, utils, version, pspecbeam
 from hera_pspec import uvpspec_utils as uvputils
 from pyuvdata import utils as uvutils
 
@@ -1144,9 +1144,9 @@ def pspec(self, bls1, bls2, dsets, pols, input_data_weight='identity', norm='I',
             A list of spectral window channel ranges to select within the total 
             bandwidth of the datasets, each of which forms an independent power 
             spectrum estimate. Example: [(220, 320), (650, 775)].
-            Each tuple should contain a start and stop channel used to index 
-            the `freq_array` of each dataset. The default (None) is to use the 
-            entire band provided in each dataset.
+            Each tuple should contain a start (inclusive) and stop (exclusive) 
+            channel used to index the `freq_array` of each dataset. The default 
+            (None) is to use the entire band provided in each dataset.
         
         verbose : bool, optional
             If True, print progress, warnings and debugging info to stdout.
@@ -1215,6 +1215,8 @@ def pspec(self, bls1, bls2, dsets, pols, input_data_weight='identity', norm='I',
         dset2 = self.dsets[self.dset_idx(dsets[1])]
 
         # assert form of bls1 and bls2
+        assert isinstance(bls1, list), "bls1 and bls2 must be fed as a list of antpair tuples"
+        assert isinstance(bls2, list), "bls1 and bls2 must be fed as a list of antpair tuples"
         assert len(bls1) == len(bls2), "length of bls1 must equal length of bls2"
         for i in range(len(bls1)):
             if isinstance(bls1[i], tuple):
@@ -1399,8 +1401,8 @@ def pspec(self, bls1, bls2, dsets, pols, input_data_weight='identity', norm='I',
                     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)
+                    # take inverse average of nsamp1 and nsamp2 and multiply by integration time [seconds] to get total integration
+                    pol_ints.extend(1./np.mean([1./nsamp1.clip(1e-10, np.inf), 1./nsamp2.clip(1e-10, np.inf)], axis=0) * dset1.integration_time)
 
                     # combined weight is geometric mean
                     pol_wgts.extend(np.concatenate([wgts1[:, :, None], wgts2[:, :, None]], axis=2))
@@ -1505,12 +1507,11 @@ def pspec(self, bls1, bls2, dsets, pols, input_data_weight='identity', norm='I',
     def rephase_to_dset(self, dset_index=0, inplace=True):
         """
         Rephase visibility data in self.dsets to the LST grid of dset[dset_index] 
-        using hera_cal.utils.lst_rephase. 
+        using hera_cal.utils.lst_rephase. Each integration in all other dsets are 
+        phased to the center of the corresponding LST bin (by index) in dset[dset_index].
 
-        Each integration in all other dsets are phased to the center of the 
-        corresponding LST bin (by index) in dset[dset_index].
-
-        Will only phase if the dataset's phase type is 'drift'.
+        Will only phase if the dataset's phase type is 'drift'. Note that if you intend
+        to use Jy_to_mK(), it must be run after rephase_to_dset().
 
         Parameters
         ----------
@@ -1597,6 +1598,45 @@ def rephase_to_dset(self, dset_index=0, inplace=True):
         if inplace is False:
             return dsets
 
+    def Jy_to_mK(self, beam=None):
+        """
+        Convert internal datasets from a Jy-scale to mK scale using a primary beam
+        model if available. Note that if you intend to rephase_to_dset(), Jy to mK conversion
+        must be done after that step.
+
+        Parameters
+        ----------
+        beam : 
+        """
+        # get all unique polarizations of all the datasets
+        pols = set(np.ravel([dset.polarization_array for dset in self.dsets]))
+
+        # assign beam
+        if beam is None:
+            beam = self.primary_beam
+        else:
+            if self.primary_beam is not None:
+                print "Warning: feeding a beam model when self.primary_beam already exists..."
+
+        # assert type of beam
+        assert isinstance(beam, pspecbeam.PSpecBeamBase), "beam model must be a subclass of pspecbeam.PSpecBeamBase"
+
+        # iterate over all pols and get conversion factors
+        factors = {}
+        for p in pols:
+            factors[p] = beam.Jy_to_mK(self.freqs, pol=p)
+
+        # iterate over datasets and apply factor
+        for i, dset in enumerate(self.dsets):
+            # check dset vis units
+            if dset.vis_units != 'Jy':
+                print "Cannot convert dset {} Jy -> mK because vis_units = {}".format(i, dset.vis_units)
+                continue
+            for j, p in enumerate(dset.polarization_array):
+                dset.data_array[:, :, :, j] *= factors[p][None, None, :]
+            dset.vis_units = 'mK'
+
+
 def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False, group=False, Nblps_per_group=1):
     """
     Construct a list of baseline-pairs from a baseline-group. This function can be used to easily convert a 

hera_pspec/tests/test_pspecdata.py

@@ -522,6 +522,34 @@ def test_rephase_to_dst(self):
         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())
 
+    def test_Jy_to_mK(self):
+        # test basic execution
+        uvd = self.uvd
+        uvd.vis_units = 'Jy'
+        ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], 
+                                 wgts=[None, None], beam=self.bm)
+        ds.Jy_to_mK()
+        nt.assert_true(ds.dsets[0].vis_units, 'mK')
+        nt.assert_true(ds.dsets[1].vis_units, 'mK')
+        nt.assert_true(uvd.get_data(24, 25, 'xx')[30, 30] / ds.dsets[0].get_data(24, 25, 'xx')[30, 30] < 1.0)
+
+        # test feeding beam
+        ds2 = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], 
+                                 wgts=[None, None], beam=self.bm)
+        ds2.Jy_to_mK(beam=self.bm)
+        nt.assert_equal(ds.dsets[0], ds2.dsets[0])
+
+        # test vis_units no Jansky
+        uvd2 = copy.deepcopy(uvd)
+        uvd2.polarization_array[0] = 1
+        uvd2.vis_units = 'UNCALIB'
+        ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd2)], 
+                                 wgts=[None, None], beam=self.bm)
+        ds.Jy_to_mK()
+        nt.assert_equal(ds.dsets[0].vis_units, "mK")
+        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_units(self):
         ds = pspecdata.PSpecData()
         # test exception
@@ -654,6 +682,13 @@ def test_pspec(self):
         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)
 
+        # test with nsamp set to zero
+        uvd = copy.deepcopy(self.uvd)
+        uvd.nsample_array[uvd.antpair2ind(24, 25)] = 0.0
+        ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm)
+        uvp = ds.pspec([(24, 25)], [(37, 38)], (0, 1), [('xx', 'xx')])
+        nt.assert_true(np.all(np.isclose(uvp.integration_array[0], 0.0)))
+
     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], 

hera_pspec/tests/test_uvpspec.py

@@ -160,6 +160,13 @@ def test_get_funcs(self):
         key = {'spw':0, 'blpair':((1, 2), (1, 2)), 'pol': 'xx'}
         d = self.uvp.get_data(key)
         nt.assert_equal(d.shape, (10, 50))
+        # test get_blpairs
+        blps = self.uvp.get_blpairs()
+        nt.assert_equal(blps, [((1, 2), (1, 2)), ((1, 3), (1, 3)), ((2, 3), (2, 3))])
+        # test get all keys
+        keys = self.uvp.get_all_keys()
+        nt.assert_equal(keys, [(0, ((1, 2), (1, 2)), 'XX'), (0, ((1, 3), (1, 3)), 'XX'),
+                               (0, ((2, 3), (2, 3)), 'XX')])
 
     def test_convert_deltasq(self):
         uvp = copy.deepcopy(self.uvp)

hera_pspec/uvpspec.py

@@ -335,6 +335,29 @@ def get_kparas(self, spw, little_h=True):
 
         return k_para
 
+    def get_blpairs(self):
+        """
+        Returns a list of all blpair tuples in the data_array.
+        """
+        return [self.blpair_to_antnums(blp) for blp in np.unique(self.blpair_array)]
+
+    def get_all_keys(self):
+        """
+        Returns a list of all possible tuple keys in the data_array.
+        """
+        # get unique blpair tuples
+        blps = self.get_blpairs()
+
+        all_keys = []
+
+        # loop over spw and pols and add to keys
+        for spw in range(self.Nspws):
+            for p in self.pol_array:
+                pstr = uvutils.polnum2str(p)
+                all_keys.extend([(spw, blp, pstr) for blp in blps])
+
+        return all_keys
+
     def convert_to_deltasq(self, little_h=True, inplace=True):
         """
         Convert from P(k) to Delta^2(k) by multiplying by k^3 / (2pi^2).