Merge pull request #191 from HERA-Team/fix_pspec_input_args

Minor tweaks to improve error handling if arguments are in the wrong format

hera_pspec/pspecdata.py

@@ -16,7 +16,8 @@
 
 class PSpecData(object):
 
-    def __init__(self, dsets=[], wgts=[], dsets_std=None, labels=None, beam=None):
+    def __init__(self, dsets=[], wgts=None, dsets_std=None, labels=None, 
+                 beam=None):
         """
         Object to store multiple sets of UVData visibilities and perform
         operations such as power spectrum estimation on them.
@@ -30,12 +31,12 @@ def __init__(self, dsets=[], wgts=[], dsets_std=None, labels=None, beam=None):
 
         dsets_std: list or dict of UVData objects, optional
             Set of UVData objects containing the standard deviations of each
-            data point in UVData objects in dsets. If specified as a dict, the key names
-            will be used to tag each dataset. Default: Empty list.
+            data point in UVData objects in dsets. If specified as a dict, 
+            the key names will be used to tag each dataset. Default: [].
 
         wgts : list or dict of UVData objects, optional
             Set of UVData objects containing weights for the input data.
-            Default: Empty list.
+            Default: None (will use the flags of each input UVData object).
 
         labels : list of str, optional
             An ordered list of names/labels for each dataset, if dsets was
@@ -59,7 +60,11 @@ def __init__(self, dsets=[], wgts=[], dsets_std=None, labels=None, beam=None):
         # and taper to none by default
         self.data_weighting = 'identity'
         self.taper = 'none'
-
+
+        # Set all weights to None if wgts=None
+        if wgts is None:
+            wgts = [None for dset in dsets]
+
         # set dsets_std to None if any are None.
         if not dsets_std is None and None in dsets_std:
             dsets_std = None
@@ -84,7 +89,7 @@ def add(self, dsets, wgts, labels=None, dsets_std=None):
         wgts : UVData or list or dict
             UVData object or list of UVData objects containing weights to add
             to the collection. Must be the same length as dsets. If a weight is
-            set to None, the flags of the corresponding
+            set to None, the flags of the corresponding dset are used. 
 
         labels : list of str
             An ordered list of names/labels for each dataset, if dsets was
@@ -143,9 +148,14 @@ def add(self, dsets, wgts, labels=None, dsets_std=None):
                             "or lists of UVData")
 
         # Make sure enough weights were specified
-        assert(len(dsets) == len(wgts))
-        assert(len(dsets_std) == len(dsets))
-        if labels is not None: assert(len(dsets) == len(labels))
+        assert len(dsets) == len(wgts), \
+            "The dsets and wgts lists must have equal length"
+        assert len(dsets_std) == len(dsets), \
+            "The dsets and dsets_std lists must have equal length"
+        if labels is not None:
+            assert len(dsets) == len(labels), \
+                "If labels are specified, the dsets and labels lists " \
+                "must have equal length"
 
         # Check that everything is a UVData object
         for d, w, s in zip(dsets, wgts, dsets_std):
@@ -162,7 +172,8 @@ def add(self, dsets, wgts, labels=None, dsets_std=None):
         if self.labels is None:
             self.labels = []
         if labels is None:
-            labels = ["dset{:d}".format(i) for i in range(len(self.dsets), len(dsets)+len(self.dsets))]
+            labels = ["dset{:d}".format(i) 
+                    for i in range(len(self.dsets), len(dsets)+len(self.dsets))]
         self.labels += labels
 
         # Append to list
@@ -1807,8 +1818,8 @@ def validate_pol(self, dsets, pol_pair):
 
     def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, 
               input_data_weight='identity', norm='I', taper='none', 
-              sampling=False, little_h=True, spw_ranges=None,
-              store_cov=False, verbose=True, history=''):
+              sampling=False, little_h=True, spw_ranges=None, 
+              baseline_tol=1.0, store_cov=False, verbose=True, history=''):
         """
         Estimate the delay power spectrum from a pair of datasets contained in
         this object, using the optimal quadratic estimator of arXiv:1502.06016.
@@ -1841,11 +1852,12 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None,
             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 of strings or integers, or list of length-2 
-            tuples of strings or integers
+        pols : tuple or list of tuple
             Contains polarization pairs to use in forming power spectra
-            e.g. ('XX','XX') or [('XX','XX'),('pI','pI')] or list of 
-            polarization pairs.
+            e.g. ('XX','XX') or [('XX','XX'),('pI','pI')] or a list of 
+            polarization pairs. Individual strings are also supported, and will 
+            be expanded into a matching pair of polarizations, e.g. 'xx' 
+            becomes ('xx', 'xx'). 
             
             If a primary_beam is specified, only equal-polarization pairs can 
             be cross-correlated, as the beam scalar normalization is only 
@@ -1885,7 +1897,11 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None,
             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.
-
+        
+        baseline_tol : float, optional
+            Distance tolerance for notion of baseline "redundancy" in meters. 
+            Default: 1.0.
+        
         store_cov : boolean, optional
             If True, calculate an analytic covariance between bandpowers
             given an input visibility noise model, and store the output
@@ -1994,13 +2010,18 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None,
                     [ (bls1[i][j], bls2[i][j]) for j in range(len(bls1[i])) ] )
 
         # validate bl-pair redundancy
-        validate_blpairs(bl_pairs, dset1, dset2, baseline_tol=1.0)
+        validate_blpairs(bl_pairs, dset1, dset2, baseline_tol=baseline_tol)
 
         # configure spectral window selections
         if spw_ranges is None:
             spw_ranges = [(0, self.Nfreqs)]
-        else:
-            assert np.isclose([len(t) for t in spw_ranges], 2).all(), \
+
+        # convert to list if only a tuple was given
+        if isinstance(spw_ranges, tuple):
+            spw_ranges = [spw_ranges,]
+
+        # Check that spw_ranges is list of len-2 tuples    
+        assert np.isclose([len(t) for t in spw_ranges], 2).all(), \
                 "spw_ranges must be fed as a list of length-2 tuples"
 
         # if using default setting of number of delay bins equal to number 
@@ -2022,11 +2043,14 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None,
             "Need to specify number of delay bins for each spw"
 
         # setup polarization selection
-        if isinstance(pols, tuple): pols = [pols]
+        if isinstance(pols, (tuple, str)): pols = [pols]
 
         # convert all polarizations to integers if fed as strings
         _pols = []
         for p in pols:
+            if isinstance(p, str):
+                # Convert string to pol-integer pair
+                p = (uvutils.polstr2num(p), uvutils.polstr2num(p))
             if isinstance(p[0], (str, np.str)):
                 p = (uvutils.polstr2num(p[0]), p[1])
             if isinstance(p[1], (str, np.str)):
@@ -3017,7 +3041,8 @@ def validate_blpairs(blpairs, uvd1, uvd2, baseline_tol=1.0, verbose=True):
     shared = sorted(set(ap1.keys()) & set(ap2.keys()))
     for k in shared:
         assert np.linalg.norm(ap1[k] - ap2[k]) <= baseline_tol, \
-            "uvd1 and uvd2 don't agree on antenna positions within tolerance of {} m".format(baseline_tol)
+            "uvd1 and uvd2 don't agree on antenna positions within " \
+            "tolerance of {} m".format(baseline_tol)
     ap = ap1
     ap.update(ap2)
 
@@ -3030,7 +3055,8 @@ def validate_blpairs(blpairs, uvd1, uvd2, baseline_tol=1.0, verbose=True):
             bl1_vec = ap[blp[0][0]] - ap[blp[0][1]]
             bl2_vec = ap[blp[1][0]] - ap[blp[1][1]]
             if np.linalg.norm(bl1_vec - bl2_vec) >= baseline_tol:
-                raise_warning("blpair {} exceeds redundancy tolerance of {} m".format(blp, baseline_tol), verbose=verbose)
+                raise_warning("blpair {} exceeds redundancy tolerance of "
+                              "{} m".format(blp, baseline_tol), verbose=verbose)
 
 
 def raise_warning(warning, verbose=True):

hera_pspec/utils.py

@@ -75,39 +75,59 @@ def cov(d1, w1, d2=None, w2=None, conj_1=False, conj_2=True):
     return C
 
 
-def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False, group=False, Nblps_per_group=1):
+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
-    single list of baselines into the input needed by PSpecData.pspec(bls1, bls2, ...).
+    Construct a list of baseline-pairs from a baseline-group. This function 
+    can be used to easily convert a single list of baselines into the input 
+    needed by PSpecData.pspec(bls1, bls2, ...).
 
     Parameters
     ----------
-    bls : list of baseline tuples, Ex. [(1, 2), (2, 3), (3, 4)]
-
-    exclude_auto_bls: boolean, if True, exclude all baselines crossed with itself from the final blpairs list
-
-    exclude_permutations : boolean, if True, exclude permutations and only form combinations of the bls list.
-        For example, if bls = [1, 2, 3] (note this isn't the proper form of bls, but makes this example clearer)
-        and exclude_permutations = False, then blpairs = [11, 12, 13, 21, 22, 23,, 31, 32, 33].
-        If however exclude_permutations = True, then blpairs = [11, 12, 13, 22, 23, 33].
-        Furthermore, if exclude_auto_bls = True then 11, 22, and 33 would additionally be excluded.
+    bls : list of tuple
+        List of baseline tuples, Ex. [(1, 2), (2, 3), (3, 4)]. Baseline 
+        integers are not supported, and must first be converted to tuples 
+        using UVData.baseline_to_antnums().
+
+    exclude_auto_bls: bool, optional
+        If True, exclude all baselines crossed with themselves from the final 
+        blpairs list. Default: False.
+
+    exclude_permutations : bool, optional
+        If True, exclude permutations and only form combinations of the bls 
+        list.
+        
+        For example, if bls = [1, 2, 3] (note this isn't the proper form of 
+        bls, but makes the example clearer) and exclude_permutations = False, 
+        then blpairs = [11, 12, 13, 21, 22, 23,, 31, 32, 33]. If however 
+        exclude_permutations = True, then blpairs = [11, 12, 13, 22, 23, 33].
+        
+        Furthermore, if exclude_auto_bls = True then 11, 22, and 33 would 
+        also be excluded.
+        
+        Default: False.
 
-    group : boolean, optional
-        if True, group each consecutive Nblps_per_group blpairs into sub-lists
+    group : bool, optional
+        If True, group each consecutive Nblps_per_group blpairs into sub-lists. 
+        Default: False.
 
-    Nblps_per_group : integer, number of baseline-pairs to put into each sub-group
+    Nblps_per_group : int, optional
+        Number of baseline-pairs to put into each sub-group if group = True. 
+        Default: 1.
 
     Returns (bls1, bls2, blpairs)
     -------
-    bls1 : list of baseline tuples from the zeroth index of the blpair
+    bls1, bls2 : list of tuples
+        List of baseline tuples from the zeroth/first index of the blpair.
 
-    bls2 : list of baseline tuples from the first index of the blpair
-
-    blpairs : list of blpair tuples
+    blpairs : list of tuple
+        List of blpair tuples.
     """
     # assert form
-    assert isinstance(bls, list) and isinstance(bls[0], tuple), "bls must be fed as list of baseline tuples"
-
+    assert isinstance(bls, (list, np.ndarray)) and isinstance(bls[0], tuple), \
+        "bls must be fed as list or ndarray of baseline antnum tuples. Use " \
+        "UVData.baseline_to_antnums() to convert baseline integers to tuples."
+
     # form blpairs w/o explicitly forming auto blpairs
     # however, if there are repeated bl in bls, there will be auto bls in blpairs
     if exclude_permutations:
@@ -160,9 +180,9 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True,
 
     Parameters
     ----------
-    uvd1 : UVData instance with visibility data
-
-    uvd2 : UVData instance with visibility data
+    uvd1, uvd2 : UVData
+        UVData instances with visibility data for the first/second visibilities 
+        in the cross-spectra that will be formed.
 
     bl_tol : float, optional
         Baseline-vector redundancy tolerance in meters
@@ -179,7 +199,8 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True,
         If True, exclude all bls crossed with itself from the blpairs list
 
     exclude_permutations : boolean, optional
-        if True, exclude permutations and only form combinations of the bls list.
+        If True, exclude permutations and only form combinations of the bls list.
+        
         For example, if bls = [1, 2, 3] (note this isn't the proper form of bls,
         but makes this example clearer) and exclude_permutations = False,
         then blpairs = [11, 12, 13, 21, 22, 23, 31, 32, 33]. If however