fix: try/except for _add_new_group

hera_cal/red_groups.py

@@ -345,7 +345,11 @@ def _add_new_group(self, bls: Iterable[BlLike], ubl: BlLike | None = None):
         that do the appropriate checking).
         """
         if ubl is None:
-            ubl = self.key_chooser(bls)
+            try:
+                ubl = self.key_chooser(bls)
+            except NonExistentBaselineError:
+                ubl = bls[0]
+
         rubl = reverse_bl(ubl)
         self._red_key_to_bls_map[ubl] = bls
         self._red_key_to_bls_map[rubl] = [reverse_bl(bl) for bl in bls]
@@ -590,16 +594,67 @@ def get_reds_in_bl_set(
         bl: BLLike,
         bl_set: Sequence[BlLike],
         include_conj: bool = False,
+        match_conj_to_set: bool = False,
         include_conj_only_if_missing: bool = False,
     ) -> set[BlLike]:
-        """Return a list of redundant groups that contain baselines in the given set."""
+        """Return a list of baselines redundant with bl that are in bl_set.
+
+        Parameters
+        ----------
+        bl
+            Only return baselines redundant with this bl. Can be an antpair or baseline
+            triple (with pol).
+        bl_set
+            Only return baselines that are in this set.
+        include_conj
+            Whether to include conjugates of the baselines that are redundant with bl,
+            if they are in bl_set. This will only include the conjugate if *the conjugate*
+            is in bl_set (i.e. it will not return both the original and the conjugate
+            if only the original is listed in bl_set).
+        match_conj_to_set
+            If include_conj is True, the default behaviour is to include a baseline if
+            it or its conjugate is in the bl_set, but the returned baseline will match
+            the ordering of the original bl. If match_conj_to_set is True, the returned
+            baseline will match the ordering of the baseline in bl_set.
+        include_conj_only_if_missing
+            If include_conj is True *and* match_conj_to_set is True, then if both
+            a baseline and its conjugate are in bl_set, then both will be in the
+            returned set. To include only the original ordering in this case, set
+            include_conj_only_if_missing to True.
+
+        Returns
+        -------
+        set
+            A set of baselines redundant with bl that are in bl_set.
+
+        Examples
+        --------
+        >>> rg = RedundantGroups.from_antpos([[(0, 1), (1, 2)], [(0, 2)]])
+        >>> rg.get_reds_in_bl_set((0, 1), {(0, 1), (1, 2)})
+        {(0, 1), (1, 2)}
+        >>> rg.get_reds_in_bl_set((0, 1), {(0, 1)})
+        {(0, 1)}
+        >>> rg.get_reds_in_bl_set((0, 1), {(1, 0)}, include_conj=True)
+        {(0, 1)}
+        >>> rg.get_reds_in_bl_set((0, 1), {(1, 0)}, include_conj=True, match_conj_to_set=True)
+        {(1, 0)}
+        >>> rg.get_reds_in_bl_set((0, 1), {(0, 1), (1, 0)}, include_conj=True, match_conj_to_set=True)
+        {(0, 1), (1, 0)}
+        >>> rg.get_reds_in_bl_set((0, 1), {(0, 1), (1, 0)}, include_conj=True, match_conj_to_set=True, include_conj_only_if_missing=True)
+        {(0, 1)}
+
+        """
         all_red_bls = self[bl]
 
         out = {key for key in all_red_bls if key in bl_set}
 
         if include_conj:
             all_red_bls = self[reverse_bl(bl)]
-            newout = {reverse_bl(key) for key in all_red_bls if (key in bl_set)}
+            if match_conj_to_set:
+                newout = {key for key in all_red_bls if (key in bl_set)}
+            else:
+                newout = {reverse_bl(key) for key in all_red_bls if (key in bl_set)}
+
             if include_conj_only_if_missing:
                 newout = {key for key in newout if reverse_bl(key) not in out}
             out = out | newout

hera_cal/tests/test_red_groups.py

@@ -235,16 +235,12 @@ def test_keyed_on_bls(self):
         rg = RedundantGroups(
             [[(0, 1)], [(0, 2), (0, 3)]],
         )
-        print(rg._red_list)
         new = rg.keyed_on_bls(bls=[(0, 3)])
-        print(new._red_key_to_bls_map)
         assert (0, 2) in new
         assert new.get_ubl_key((0, 2)) == (0, 3)
         assert new.get_ubl_key((0, 1)) == (0, 1)
 
         new.keyed_on_bls(bls=[(0, 2)], inplace=True)
-        print(new._red_key_to_bls_map)
-        print(new._bl_to_red_map)
         assert new.get_ubl_key((0, 2)) == (0, 2)
 
     def test_delitem(self):
@@ -269,3 +265,58 @@ def test_data_bls(self):
             [[(0, 1)], [(0, 2), (0, 3)]],
         )
         assert len(rg.data_bls) == 3
+
+
+class TestRedundantGroupsGetRedsInBlSet:
+    def setup_method(self, method):
+        self.rg = RedundantGroups(
+            [[(0, 1)], [(0, 2), (0, 3)]],
+        )
+
+    def test_default(self):
+        assert self.rg.get_reds_in_bl_set((0, 1), {(0, 1), (0, 2)}) == {(0, 1)}
+        assert self.rg.get_reds_in_bl_set((0, 1), {(0, 2)}) == set()
+        assert self.rg.get_reds_in_bl_set((0, 2), {(0, 2)}) == {(0, 2)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(0, 2), (0, 3)}) == {(0, 2), (0, 3)}
+
+    def test_conj_in_set(self):
+        """Test that passing a baseline that exists and only having the conj in the set returns nothing."""
+        assert self.rg.get_reds_in_bl_set((0, 1), {(1, 0), (0, 2)}) == set()
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0)}) == set()
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0), (0, 3)}) == {(0, 3)}
+
+        # However, if we *give* it the conjugate, it should work
+        assert self.rg.get_reds_in_bl_set((1, 0), {(1, 0), (0, 2)}) == {(1, 0)}
+
+    def test_bad_key(self):
+        with pytest.raises(KeyError):
+            # The baseline (0, 4) doesn't exist
+            self.rg.get_reds_in_bl_set((0, 4), {})
+
+        with pytest.raises(KeyError):
+            self.rg.get_reds_in_bl_set((0, 1, 'ee'), {})
+
+    def test_conj_in_set_include_conj_match_bl(self):
+        # The following test that at least the conjugates are returned
+        assert self.rg.get_reds_in_bl_set((0, 1), {(1, 0), (0, 2)}, include_conj=True) == {(0, 1)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0)}, include_conj=True) == {(0, 2)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0), (0, 3)}, include_conj=True) == {(0, 2), (0, 3)}
+
+        # If both the original and conjugate are in the set, only the original is returned, unless match_conj_to_set is True
+        assert self.rg.get_reds_in_bl_set((0, 1), {(1, 0), (0, 2), (0, 1)}, include_conj=True) == {(0, 1)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0), (0, 2)}, include_conj=True) == {(0, 2)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0), (0, 3), (0, 2)}, include_conj=True) == {(0, 2), (0, 3)}
+
+    def test_conj_in_set_include_conj_match_set(self):
+        # If both the original and conjugate are in the set, both returned if match_conj_to_set is True
+        kw = dict(include_conj=True, match_conj_to_set=True)
+        assert self.rg.get_reds_in_bl_set((0, 1), {(1, 0), (0, 2), (0, 1)}, **kw) == {(0, 1), (1, 0)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0), (0, 2)}, **kw) == {(0, 2), (2, 0)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0), (0, 3), (0, 2)}, **kw) == {(0, 2), (2, 0), (0, 3)}
+
+    def test_conj_in_set_only_if_missing(self):
+        # If the original is in the set, don't return the conjugate
+        kw = dict(include_conj=True, include_conj_only_if_missing=True, match_conj_to_set=True)
+        assert self.rg.get_reds_in_bl_set((0, 1), {(0, 1), (1, 0)}, **kw) == {(0, 1)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(2, 0), (3, 0)}, **kw) == {(2, 0), (3, 0)}
+        assert self.rg.get_reds_in_bl_set((0, 2), {(0, 2), (3, 0)}, **kw) == {(0, 2), (3, 0)}