From 37ef04c2b451e334fb4ee0758e7fa2924867bddc Mon Sep 17 00:00:00 2001
From: Peijin Zhang <peijin@ustc.edu>
Date: Mon, 17 Aug 2026 21:19:43 +0000
Subject: [PATCH] Fix beam pointing antenna-health fallback

---
 mnc/anthealth.py      |  9 ++++++--
 mnc/control.py        | 24 +++++++++++++++++++--
 tests/test_control.py | 50 ++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 78 insertions(+), 5 deletions(-)

diff --git a/mnc/anthealth.py b/mnc/anthealth.py
index ce70e78..11c7492 100644
--- a/mnc/anthealth.py
+++ b/mnc/anthealth.py
@@ -115,8 +115,14 @@ def get_badants(method, time=None, naming='antname'):
             dd = ls.get_dict(f'/mon/anthealth/{method}')
         else:
             dd = ls.get_dict(f'/mon/anthealth/{method}/{float(mjd0)}')
+        if dd is None:
+            raise RuntimeError(f"No antenna-health data are available for {method}")
+        if not all(key in dd for key in ('flagged', 'antname', 'time')):
+            raise RuntimeError(f"Antenna-health data for {method} are incomplete")
         antstatus = dd['flagged']
         antnames = dd['antname']
+        if len(antstatus) != len(antnames):
+            raise RuntimeError(f"Antenna-health data for {method} have inconsistent lengths")
         if mjd0 is None:
             mjd0 = float(dd['time'])
     elif method == 'union_and':
@@ -132,7 +138,7 @@ def get_badants(method, time=None, naming='antname'):
 
     badants = np.array(antnames)[np.where(antstatus)].tolist()
 
-    if naming is "corr_num":
+    if naming == "corr_num":
         logger.debug("mapping antname to corr_num")
         badants2 = []
         for antname in badants:
@@ -158,4 +164,3 @@ def caltable_flags(caltable):
                       for (corrnum, pol) in zip(*np.where(allflg))])
 
     return badants
-
diff --git a/mnc/control.py b/mnc/control.py
index c5eb30b..9c49249 100644
--- a/mnc/control.py
+++ b/mnc/control.py
@@ -26,6 +26,18 @@
 
 
 CONFIG_FILE = '/home/pipeline/proj/lwa-shell/mnc_python/config/lwa_config_calim.yaml'
+
+
+def _flag_antname_to_correlator(antname):
+    """Convert a polarized LWA antenna name to its correlator number."""
+    antname = str(antname)
+    if antname.endswith(('A', 'B')):
+        antname = antname[:-1]
+    if not antname.startswith('LWA-'):
+        antname = 'LWA-' + antname
+    return mapping.antname_to_correlator(antname)
+
+
 FPG_FILE = '/home/pipeline/proj/lwa-shell/caltech-lwa/snap2_f_200msps_64i_4096c/outputs/snap2_f_200msps_64i_4096c.fpg'
 
 CORE_RADIUS_M = 200.0
@@ -371,9 +383,17 @@ def control_bf(self, num=1, coord=None, coordtype='celestial', targetname=None,
 
         # we convert antnames into corr_nums and ignore pol info
         if isinstance(flag_ants, str):
-            mjd, flag_ants = anthealth.get_badants(flag_ants)
+            flag_method = flag_ants
+            try:
+                mjd, flag_ants = anthealth.get_badants(flag_method)
+            except RuntimeError:
+                if flag_method != 'caltable':
+                    raise
+                logger.warning("No caltable antenna-health list is available; "
+                               "falling back to selfcorr")
+                mjd, flag_ants = anthealth.get_badants('selfcorr')
         assert isinstance(flag_ants, list)
-        flag_ants = list({f"{mapping.antname_to_correlator('LWA-'+antname.rstrip('A').rstrip('B')):03}" for antname in flag_ants})
+        flag_ants = list({f"{_flag_antname_to_correlator(antname):03}" for antname in flag_ants})
 
         if (callable(uvweight)):
             assert uvweight.__code__.co_argcount == 1, "uvweight function must only take one argument"
diff --git a/tests/test_control.py b/tests/test_control.py
index 8d748a7..354982c 100644
--- a/tests/test_control.py
+++ b/tests/test_control.py
@@ -13,7 +13,55 @@
 ):
     sys.modules.setdefault(_mod, MagicMock())
 
-from mnc.control import Controller, _recording_path_from_response
+from mnc.control import (
+    Controller,
+    _flag_antname_to_correlator,
+    _recording_path_from_response,
+)
+
+
+class TestFlagAntnameToCorrelator(unittest.TestCase):
+    @patch("mnc.control.mapping.antname_to_correlator")
+    def test_accepts_full_polarized_name(self, convert):
+        convert.return_value = 17
+
+        self.assertEqual(_flag_antname_to_correlator("LWA-002A"), 17)
+        convert.assert_called_once_with("LWA-002")
+
+    @patch("mnc.control.mapping.antname_to_correlator")
+    def test_accepts_numeric_polarized_name(self, convert):
+        convert.return_value = 18
+
+        self.assertEqual(_flag_antname_to_correlator("003B"), 18)
+        convert.assert_called_once_with("LWA-003")
+
+
+class TestControlBfAntennaHealth(unittest.TestCase):
+    @patch("mnc.control._flag_antname_to_correlator", return_value=2)
+    @patch("mnc.control.anthealth.get_badants")
+    def test_caltable_falls_back_to_selfcorr(self, get_badants, convert):
+        get_badants.side_effect = [
+            RuntimeError("missing caltable state"),
+            (61269.0, ["LWA-002A"]),
+        ]
+        controller = Controller.__new__(Controller)
+        beam = MagicMock()
+        beam.cal_set = True
+        controller.bfc = {5: beam}
+
+        controller.control_bf(
+            num=5,
+            targetname="sun",
+            track=False,
+            flag_ants="caltable",
+        )
+
+        self.assertEqual(
+            get_badants.call_args_list,
+            [unittest.mock.call("caltable"), unittest.mock.call("selfcorr")],
+        )
+        convert.assert_called_once_with("LWA-002A")
+        beam.set_beam_target.assert_called_once_with("sun")
 
 
 class TestRecordingPathFromResponse(unittest.TestCase):