Add loader and det cal functionality to load in bandpass info using bg

sotodlib/io/metadata.py

@@ -220,6 +220,46 @@ def batch_from_loadspec(self, load_params, **kwargs):
                                                   row_order=mask.nonzero()[0])
         return results
 
+class ModifyFieldsResultSetHdfLoader(ResultSetHdfLoader):
+    def __init__(self):
+        self.load_fields = None
+
+    def _check_key_map(self, key_map, data_in):
+        for key in key_map.keys():
+            if key not in data_in.dtype.names:
+                raise KeyError(f"{key} not included inside dataset with keys"
+                               f"{data_in.dtype.names}")
+
+
+    def _prefilter_data(self, data_in):
+        """Extend to allow us to specify which data fields to load."""
+        key_map = {}
+        if self.load_fields is not None:
+            for field in self.load_fields:
+                if isinstance(field, dict):
+                    key_map.update(field)
+                else:
+                    key_map[field] = field
+
+            for k in data_in.dtype.names:
+                if k not in key_map.keys():
+                    key_map[k] = None
+
+            self._check_key_map(key_map, data_in)
+
+        return _decode_array(data_in, key_map=key_map)
+
+    def batch_from_loadspec(self, load_params, **kwargs):
+        """Extend to allow us to input 'load_fields' which will modify our
+        class variables as used in _prefilter_data(...).
+
+        """
+        if 'load_fields' in kwargs.keys():
+            self.load_fields = kwargs['load_fields']
+        results = super().batch_from_loadspec(load_params)
+        return results
+
+
 
 def _decode_array(data_in, key_map={}):
     """Converts a structured numpy array to a structured numpy array,
@@ -278,6 +318,7 @@ def __exit__(self, *excinfo):
 SuperLoader.register_metadata('DefaultHdf', DefaultHdfLoader)
 SuperLoader.register_metadata('AxisManagerHdf', AxisManagerHdfLoader)
 SuperLoader.register_metadata('ResultSetHdf', ResultSetHdfLoader)
+SuperLoader.register_metadata('ModifyFieldsResultSetHdf', ModifyFieldsResultSetHdfLoader)
 
 # The old name... remove some day.
 SuperLoader.register_metadata('PerDetectorHdf5', ResultSetHdfLoader)

sotodlib/preprocess/processes.py

@@ -1324,6 +1324,7 @@ class PCARelCal(_Preprocess):
     def __init__(self, step_cfgs):
         self.signal = step_cfgs.get('signal', 'signal')
         self.run = step_cfgs.get('pca_run', 'run1')
+        self.bandpass = step_cfgs.get('bandpass', 'wafer.bandpass')
         self.run_name = f'{self.signal}_{self.run}'
 
         super().__init__(step_cfgs)
@@ -1346,15 +1347,15 @@ def calc_and_save(self, aman, proc_aman):
             if self.plot_cfgs:
                 self.plot_signal = filt_aman[self.signal]
 
-        bands = np.unique(aman.det_info.wafer.bandpass)
+        bands = np.unique(aman.det_info[self.bandpass])
         bands = bands[bands != 'NC']
         # align samps w/ proc_aman to include samps restriction when loading back from db.
         rc_aman = core.AxisManager(proc_aman.dets, proc_aman.samps)
         pca_det_mask = np.full(aman.dets.count, False, dtype=bool)
         relcal = np.zeros(aman.dets.count)
         pca_weight0 = np.zeros(aman.dets.count)
         for band in bands:
-            m0 = aman.det_info.wafer.bandpass == band
+            m0 = aman.det_info[self.bandpass] == band
             rc_aman.wrap(f'{band}_idx', m0, [(0, 'dets')])
             band_aman = aman.restrict('dets', aman.dets.vals[m0], in_place=False)
 
@@ -1408,7 +1409,7 @@ def plot(self, aman, proc_aman, filename):
             det = aman.dets.vals[0]
             ufm = det.split('_')[2]
 
-            bands = np.unique(aman.det_info.wafer.bandpass)
+            bands = np.unique(aman.det_info[self.bandpass])
             bands = bands[bands != 'NC']
             for band in bands:
                 pca_aman = aman.restrict('dets', aman.dets.vals[proc_aman[self.run_name][f'{band}_idx']], in_place=False)
@@ -1900,9 +1901,9 @@ def process(self, aman, proc_aman):
 _Preprocess.register(HWPAngleModel)
 _Preprocess.register(GetStats)
 _Preprocess.register(UnionFlags)
-_Preprocess.register(RotateQU) 
-_Preprocess.register(SubtractQUCommonMode) 
-_Preprocess.register(FocalplaneNanFlags) 
-_Preprocess.register(PointingModel)  
+_Preprocess.register(RotateQU)
+_Preprocess.register(SubtractQUCommonMode)
+_Preprocess.register(FocalplaneNanFlags)
+_Preprocess.register(PointingModel)
 _Preprocess.register(BadSubscanFlags)
 _Preprocess.register(CorrectIIRParams)

sotodlib/site_pipeline/update_det_cal.py

@@ -33,6 +33,12 @@
 DEFAULT_pA_per_phi0 = 9e6
 TES_BIAS_COUNT = 12  # per detset / primary file group
 
+# For converting bias group to bandpass.
+BGS = {'lb': [0, 1, 4, 5, 8, 9], 'hb': [2, 3, 6, 7, 10, 11]}
+BAND_STR = {'mf': {'lb': 'f090', 'hb': 'f150'},
+            'uhf': {'lb': 'f220', 'hb': 'f280'},
+            'lf': {'lb': 'f030', 'hb': 'f040'}}
+
 logger = logging.getLogger("det_cal")
 if not logger.hasHandlers():
     sp_util.init_logger("det_cal")
@@ -248,6 +254,8 @@ class CalInfo:
         Current responsivity of the TES [1/V] computed using bias steps at the
         bias point. This is based on the naive bias step estimation without
         using any additional corrections.
+    bandpass: str
+        Detector bandpass, computed from bias group information.
     """
 
     readout_id: str = ""
@@ -269,6 +277,7 @@ class CalInfo:
     naive_r_frac: float = np.nan
     naive_p_bias: float = np.nan
     naive_s_i: float = np.nan
+    bandpass: str = "NC"
 
     @classmethod
     def dtype(cls) -> List[Tuple[str, Any]]:
@@ -617,6 +626,12 @@ def find_correction_results(band, chan, dset):
             else:
                 cal.phase_to_pW = pA_per_phi0 / (2 * np.pi) / cal.s_i * cal.polarity
 
+            # Add bandpass informaton from bias group
+            if cal.bg in bgs['lb']:
+                cal.bandpass = band_str[tube_flavor]['lb']
+            elif cal.bg in bgs['hb']:
+                cal.bandpass = band_str[tube_flavor]['hb']
+
         res.result_set = np.array([astuple(c) for c in cals], dtype=CalInfo.dtype())
         res.success = True
     except Exception as e: