calibration.fringefit now works with single-spw splits across a ps

src/astroviper/calibration/fringefit.py

@@ -7,15 +7,17 @@
 from typing import Dict, Union
 
 def _fringe_node_task(input_params: Dict):
-    ps = input_params['ps']
+    ps = input_params['ps'] 
     data_selection = input_params['data_selection']
+    # FIXME: for now we do single band
+    if len(data_selection.keys())>1:
+        print(f'{data_selection.keys()=}')
+        raise RuntimeError("We only do single xdses so far")
     name = list(data_selection.keys())[0]
     xds = ps[name]
-    # FIXME: for now we do single band
     data_sub_selection = input_params['data_sub_selection']
-    # breakpoint()
     pol = data_sub_selection['polarization']
-    xds2 = xds.isel(**input_params['data_selection'][name])
+    xds2 = xds.isel(**data_selection[name])
     xds2 = xds2.sel(polarization=pol)
     vis = xds2.VISIBILITY
     ang = np.angle(vis)
@@ -31,14 +33,23 @@ def _fringe_node_task(input_params: Dict):
     )
     ref_ant = 1
     res = {}
-    bl_slice = input_params['data_selection'][name]["baseline_id"]
-    baselines = list(xds.baseline_id[bl_slice].values)
-    for i, bl in enumerate(baselines):
-        ant1 = int(xds2.baseline_antenna1_id[bl].values)
-        ant2 = int(xds2.baseline_antenna2_id[bl].values)
-        a = np.abs(fftvis[:,i,:])
-        ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
-        res.setdefault(xds2.time.values[0], {})[bl] = (ind, a[ind], a.shape)
+    bl_slice = data_selection[name]["baseline_id"]
+    baselines = xds.baseline_id[bl_slice].values
+    ant1s = xds2.baseline_antenna1_id.values
+    ant2s = xds2.baseline_antenna1_id.values
+    try:
+        # FIXME:
+        #
+        # In the case of subcubes we *don't* get all the antenna1_id stuff!
+        # antenna1_id.values: [ 6  6  6  7  7  7  7  7  8  8  8  8  9  9  9 10 10 11]
+        # baselines :         [60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77]
+        for i, (bl, ant1, ant2) in enumerate(zip(baselines, ant1s, ant2s)):
+            a = np.abs(fftvis[:,i,:])
+            ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
+            res.setdefault(xds2.time.values[0], {})[bl] = (ind, a[ind], a.shape)
+    except IndexError as e: 
+        print(f'{xds2.baseline_antenna1_id.values}\n{baselines=}')
+        raise e
     return res
 
 def _fringefit_reduce(graph_inputs: xr.Dataset, input_params: Dict):