added exclude_cross_bls to utils.construct_blpairs

hera_pspec/pspecdata.py

@@ -2893,7 +2893,7 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True,
               groupname=None, dset_labels=None, dset_pairs=None, psname_ext=None,
               spw_ranges=None, n_dlys=None, pol_pairs=None, blpairs=None,
               input_data_weight='identity', norm='I', taper='none',
-              exclude_auto_bls=False, exclude_permutations=True,
+              exclude_auto_bls=False, exclude_cross_bls=False, exclude_permutations=True,
               Nblps_per_group=None, bl_len_range=(0, 1e10),
               bl_deg_range=(0, 180), bl_error_tol=1.0,
               beam=None, cosmo=None, interleave_times=False, rephase_to_dset=None,
@@ -2986,6 +2986,11 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True,
         exclude_auto_bls is True, eliminate all instances of a bl crossed
         with itself. Default: False
 
+    exclude_cross_bls : boolean
+        If True and if blpairs is None, exclude all bls crossed with a
+        different baseline. Note if this and exclude_auto_bls are True
+        then no blpairs will exist. Default: False
+
     exclude_permutations : boolean
         If blpairs is None, redundant baseline groups will be formed and
         all cross-multiplies will be constructed. In doing so, if
@@ -3277,6 +3282,7 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True,
                                       dsets[dsetp[0]], dsets[dsetp[1]],
                                       filter_blpairs=True,
                                       exclude_auto_bls=exclude_auto_bls,
+                                      exclude_cross_bls=exclude_cross_bls,
                                       exclude_permutations=exclude_permutations,
                                       Nblps_per_group=Nblps_per_group,
                                       bl_len_range=bl_len_range,
@@ -3368,6 +3374,7 @@ def list_of_tuple_tuples(v):
     a.add_argument("--time_thresh", default=0.2, type=float, help="Fractional flagging threshold across time to trigger flag broadcast if broadcast_dset_flags is True")
     a.add_argument("--Jy2mK", default=False, action='store_true', help="Convert datasets from Jy to mK if a beam model is provided.")
     a.add_argument("--exclude_auto_bls", default=False, action='store_true', help='If blpairs is not provided, exclude all baselines paired with itself.')
+    a.add_argument("--exclude_cross_bls", default=False, action='store_true', help='If blpairs is not provided, exclude all baselines paired with a different baseline.')
     a.add_argument("--exclude_permutations", default=False, action='store_true', help='If blpairs is not provided, exclude a basline-pair permutations. Ex: if (A, B) exists, exclude (B, A).')
     a.add_argument("--Nblps_per_group", default=None, type=int, help="If blpairs is not provided and group == True, set the number of blpairs in each group.")
     a.add_argument("--bl_len_range", default=(0, 1e10), nargs='+', type=float, help="If blpairs is not provided, limit the baselines used based on their minimum and maximum length in meters.")

hera_pspec/tests/test_utils.py

@@ -147,7 +147,6 @@ def test_spw_range_from_redshifts(self):
         nt.ok_( isinstance(spw5, tuple) )
         nt.ok_( spw5[0] is not None )
 
-
     def test_calc_blpair_reds(self):
         fname = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA')
         uvd = UVData()
@@ -194,10 +193,17 @@ def test_calc_blpair_reds(self):
                                    bl_len_range=(10.0, 20.0))
         nt.assert_equal(blps, [((24, 25), (37, 38))])
 
+        # test exclude_cross_bls
+        (bls1, bls2, blps, xants1,
+         xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_cross_bls=True)
+        for bl1, bl2 in blps:
+            assert bl1 == bl2
+
         # test exceptions
         uvd2 = copy.deepcopy(uvd)
         uvd2.antenna_positions[0] += 2
         nt.assert_raises(AssertionError, utils.calc_blpair_reds, uvd, uvd2)
+        nt.assert_raises(AssertionError, utils.calc_blpair_reds, uvd, uvd, exclude_auto_bls=True, exclude_cross_bls=True)
 
     def test_get_delays(self):
         utils.get_delays(np.linspace(100., 200., 50)*1e6)

hera_pspec/utils.py

@@ -75,8 +75,8 @@ 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_cross_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
@@ -93,6 +93,10 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False,
         If True, exclude all baselines crossed with themselves from the final
         blpairs list. Default: False.
 
+    exclude_cross_bls : bool, optional
+        If True, exclude all bls crossed with a different baseline. Note if
+        this and exclude_auto_bls are True then no blpairs will exist.
+
     exclude_permutations : bool, optional
         If True, exclude permutations and only form combinations of the bls
         list.
@@ -127,6 +131,7 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False,
     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."
+    assert (not exclude_auto_bls) or (not exclude_cross_bls), "Can't exclude both auto and cross blpairs"
 
     # form blpairs w/o explicitly forming auto blpairs
     # however, if there are repeated bl in bls, there will be auto bls in blpairs
@@ -146,6 +151,14 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False,
                 new_blpairs.append(blp)
         blpairs = new_blpairs
 
+    # same for cross
+    if exclude_cross_bls:
+        new_blpairs = []
+        for blp in blpairs:
+            if blp[0] == blp[1]:
+                new_blpairs.append(blp)
+        blpairs = new_blpairs
+
     # create bls1 and bls2 list
     bls1 = [blp[0] for blp in blpairs]
     bls2 = [blp[1] for blp in blpairs]
@@ -171,6 +184,7 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_permutations=False,
 
 def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True,
                      xant_flag_thresh=0.95, exclude_auto_bls=False,
+                     exclude_cross_bls=False,
                      exclude_permutations=True, Nblps_per_group=None,
                      bl_len_range=(0, 1e10), bl_deg_range=(0, 180)):
     """
@@ -198,6 +212,10 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True,
     exclude_auto_bls: boolean, optional
         If True, exclude all bls crossed with itself from the blpairs list
 
+    exclude_cross_bls : boolean, optional
+        If True, exclude all bls crossed with a different baseline. Note if
+        this and exclude_auto_bls are True then no blpairs will exist.
+
     exclude_permutations : boolean, optional
         If True, exclude permutations and only form combinations of the bls list.
 
@@ -291,7 +309,7 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True,
     for r in reds:
         (bls1, bls2,
          blps) = construct_blpairs(r, exclude_auto_bls=exclude_auto_bls,
-                                   group=False,
+                                   exclude_cross_bls=exclude_cross_bls, group=False,
                                    exclude_permutations=exclude_permutations)
         if len(bls1) < 1:
             continue