Merge 0ac0101 into 4a877f1

hera_pspec/tests/test_uvpspec_utils.py

@@ -98,7 +98,7 @@ def test_select_common():
 
 def test_get_blpairs_from_bls():
     """
-    Test conversion of 
+    Test conversion of bls to set of blpairs.
     """
     # setup uvp
     beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
@@ -109,7 +109,55 @@ def test_get_blpairs_from_bls():
     blps = uvputils._get_blpairs_from_bls(uvp, bls=101102)
     blps = uvputils._get_blpairs_from_bls(uvp, bls=(101,102))
     blps = uvputils._get_blpairs_from_bls(uvp, bls=[101102, 101103])
+
+
+def test_get_red_bls():
+    """
+    Test retrieval of redundant baseline groups.
+    """
+    # Setup uvp
+    beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
+    beam = pspecbeam.PSpecBeamUV(beamfile)
+    uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam)
+
+    # Get redundant baseline groups
+    bls, lens, angs = uvp.get_red_bls()
+
+    nt.assert_equal(len(bls), 3) # three red grps in this file
+    nt.assert_equal(len(bls), len(lens)) # Should be one length for each group
+    nt.assert_equal(len(bls), len(angs)) # Ditto, for angles
+
+    # Check that number of grouped baselines = total no. of baselines
+    num_bls = 0
+    for grp in bls:
+        for bl in grp:
+            num_bls += 1
+    nt.assert_equal(num_bls, np.unique(uvp.bl_array).size)
+
+
+def test_get_red_blpairs():
+    """
+    Test retrieval of redundant baseline groups for baseline-pairs.
+    """
+    # Setup uvp
+    beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')
+    beam = pspecbeam.PSpecBeamUV(beamfile)
+    uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam)
+
+    # Get redundant baseline groups
+    blps, lens, angs = uvp.get_red_blpairs()
 
+    nt.assert_equal(len(blps), 3) # three red grps in this file
+    nt.assert_equal(len(blps), len(lens)) # Should be one length for each group
+    nt.assert_equal(len(blps), len(angs)) # Ditto, for angles
+
+    # Check that number of grouped blps = total no. of blps
+    num_blps = 0
+    for grp in blps:
+        for blp in grp:
+            num_blps += 1
+    nt.assert_equal(num_blps, np.unique(uvp.blpair_array).size)
+
 
 def test_polpair_int2tuple():
     """

hera_pspec/utils.py

@@ -1022,3 +1022,5 @@ def get_reds(uvd, bl_error_tol=1.0, pick_data_ants=False, bl_len_range=(0, 1e4),
         reds, lens, angs = _reds, _lens, _angs
 
     return reds, lens, angs
+
+

hera_pspec/uvpspec.py

@@ -15,6 +15,7 @@ class UVPSpec(object):
     An object for storing power spectra generated by hera_pspec and a file-format
     for its data and meta-data.
     """
+
     def __init__(self):
         """
         An object for storing power spectra and associated metadata generated
@@ -148,6 +149,11 @@ def __init__(self):
 
         # Default parameter values
         self.folded = False
+
+        # Set function docstrings
+        self.get_red_bls.__func__.__doc__ = uvputils._get_red_bls.__doc__
+        self.get_red_blpairs.__func__.__doc__ = uvputils._get_red_blpairs.__doc__
+
 
     def get_cov(self, key, omit_flags=False):
         """
@@ -1679,8 +1685,11 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None,
         for i, blp in enumerate(blpairs):
             # get indices
             inds = self.blpair_to_indices(blp)
-            assert isinstance(Tsys[blp], (float, np.float, int, np.int)) or Tsys[blp].shape[0] == self.Ntimes, "Tsys must be a float or an ndarray with shape[0] == Ntimes"
+            assert isinstance(Tsys[blp], (float, np.float, int, np.int)) \
+                or Tsys[blp].shape[0] == self.Ntimes, \
+                "Tsys must be a float or an ndarray with shape[0] == Ntimes"
             P_blp = []
+
             # iterate over time axis
             for j, ind in enumerate(inds):
                 # get integration time and n_samp
@@ -1695,7 +1704,8 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None,
 
                 # Get noise power spectrum
                 pn = noise.calc_P_N(scalar, Tsys[blp][j], t_int, k=k,
-                                    Nincoherent=n_samp, form=form, component=component)
+                                    Nincoherent=n_samp, form=form, 
+                                    component=component)
 
                 # Put into appropriate form
                 if form == 'Pk':
@@ -1834,7 +1844,14 @@ def get_blpair_groups_from_bl_groups(self, blgroups, only_pairs_in_bls=False):
                 blpair_groups.append(blp)
 
         return blpair_groups
-
+
+    def get_red_bls(self, bl_len_tol=1., bl_ang_tol=1.):
+        return uvputils._get_red_bls(self, bl_len_tol=bl_len_tol, 
+                                     bl_ang_tol=bl_ang_tol)
+
+    def get_red_blpairs(self, bl_len_tol=1., bl_ang_tol=1.):
+        return uvputils._get_red_blpairs(self, bl_len_tol=bl_len_tol, 
+                                         bl_ang_tol=bl_ang_tol)
 
     def compute_scalar(self, spw, polpair, num_steps=1000, little_h=True,
                        noise_scalar=False):

hera_pspec/uvpspec_utils.py

@@ -1,5 +1,6 @@
 import numpy as np
 import copy, operator
+from . import utils
 from collections import OrderedDict as odict
 from pyuvdata.utils import polstr2num, polnum2str
 
@@ -945,3 +946,145 @@ def _fast_lookup_blpairts(src_blpts, query_blpts, time_prec=8):
     blpts_idxs = np.where(src_blpts == query_blpts[:,np.newaxis])[1]
 
     return blpts_idxs
+
+
+def _get_red_bls(uvp, bl_len_tol=1., bl_ang_tol=1.):
+    """
+    Get redundant baseline groups that are present in a UVPSpec object.
+    
+    Parameters
+    ----------
+    uvp : UVPSpec
+        UVPSpec object.
+    
+    bl_len_tol : float, optional
+        Maximum difference in length to use for grouping baselines. 
+        This does not guarantee that the maximum length difference 
+        between any two baselines in a group is less than bl_len_tol 
+        however. Default: 1.0.
+    
+    bl_ang_tol : float, optional
+        Maximum separation in angle to use for grouping baselines. 
+        This does not guarantee that the maximum angle between any 
+        two baselines in a group is less than bl_ang_tol however. 
+        Default: 1.0.
+    
+    Returns
+    -------
+    grp_bls : list of array_like
+        List of redundant baseline groups. Each list item contains 
+        an array of baseline integers corresponding to the members 
+        of the group.
+        
+    grp_lens : list of float
+        Average length of the baselines in each group.
+        
+    grp_angs : list of float
+        Average angle of the baselines in each group.
+    """
+    # Calculate length and angle of baseline vecs
+    bl_vecs = uvp.get_ENU_bl_vecs()
+
+    lens, angs = utils.get_bl_lens_angs(bl_vecs, bl_error_tol=1.0)
+
+    # Baseline indices
+    idxs = np.arange(len(lens)).astype(np.int)
+    grp_bls = []; grp_len = []; grp_ang = []
+
+    # Group baselines by length and angle
+    max_loops = idxs.size
+    nloops = 0
+    while len(idxs) > 0 and nloops < max_loops:
+        nloops += 1
+
+        # Match bls within some tolerance in length and angle
+        matches = np.where(np.logical_and(
+                            np.abs(lens - lens[0]) < bl_len_tol,
+                            np.abs(angs - angs[0]) < bl_ang_tol) )
+
+        # Save info about this group
+        grp_bls.append(uvp.bl_array[idxs[matches]])
+        grp_len.append(np.mean(lens[matches]))
+        grp_ang.append(np.mean(angs[matches]))
+
+        # Remove bls that were matched so we don't try to group them again
+        idxs = np.delete(idxs, matches)
+        lens = np.delete(lens, matches)
+        angs = np.delete(angs, matches)
+
+    return grp_bls, grp_len, grp_ang
+
+
+def _get_red_blpairs(uvp, bl_len_tol=1., bl_ang_tol=1.):
+    """
+    Group baseline-pairs from a UVPSpec object according to the 
+    redundant groups that their constituent baselines belong to.
+    
+    NOTE: Baseline-pairs made up of baselines from two different 
+    redundant groups are ignored.
+    
+    Parameters
+    ----------
+    uvp : UVPSpec
+        UVPSpec object.
+    
+    bl_len_tol : float, optional
+        Maximum difference in length to use for grouping baselines. 
+        This does not guarantee that the maximum length difference 
+        between any two baselines in a group is less than bl_len_tol 
+        however. Default: 1.0.
+    
+    bl_ang_tol : float, optional
+        Maximum separation in angle to use for grouping baselines. 
+        This does not guarantee that the maximum angle between any 
+        two baselines in a group is less than bl_ang_tol however. 
+        Default: 1.0.
+    
+    Returns
+    -------
+    grp_bls : list of array_like
+        List of redundant baseline groups. Each list item contains 
+        an array of baseline-pair integers corresponding to the 
+        members of the group.
+        
+    grp_lens : list of float
+        Average length of the baselines in each group.
+        
+    grp_angs : list of float
+        Average angle of the baselines in each group.
+    """
+    # Get redundant baseline groups
+    red_bls, red_lens, red_angs = _get_red_bls(uvp=uvp, 
+                                               bl_len_tol=bl_len_tol, 
+                                               bl_ang_tol=bl_ang_tol)
+
+    # Get all available blpairs and convert to pairs of integers
+    blps = [(uvp.antnums_to_bl(blp[0]), uvp.antnums_to_bl(blp[1]))
+            for blp in uvp.get_blpairs()]
+    bl1, bl2 = zip(*blps)
+
+    # Build bl -> group index dict
+    group_idx = {}
+    for i, grp in enumerate(red_bls):
+        for bl in grp:
+            group_idx[bl] = i
+
+    # Get red. group that each bl belongs to
+    bl1_grp = np.array([group_idx[bl] for bl in bl1])
+    bl2_grp = np.array([group_idx[bl] for bl in bl2])
+
+    # Convert to arrays for easier indexing
+    bl1 = np.array(bl1)
+    bl2 = np.array(bl2)
+
+    # Loop over redundant groups; assign blpairs to each group
+    red_grps = []
+    grp_ids = np.arange(len(red_bls))
+    for i in grp_ids:
+        # This line only keeps blpairs where both bls belong to the same red grp!
+        matches = np.where(np.logical_and(bl1_grp == i, bl2_grp == i))
+
+        # Unpack into tuple of bl integer pairs
+        red_grps.append( list(zip(bl1[matches], bl2[matches])) )
+
+    return red_grps, red_lens, red_angs