Repository: simonsobs/sotodlib

Issue:
Allow `dets` in `context.get_obs` to take a numpy array
I tried to run `aman = context.get_obs(obs, dets=meta.dets.vals)` and it threw an error because that list of detectors is a numpy array and not a list. It would be nice to not have to remember cast it.


Discussion:

Patch:
diff --git a/sotodlib/core/context.py b/sotodlib/core/context.py
index 8f4ce188..a4370b68 100644
--- a/sotodlib/core/context.py
+++ b/sotodlib/core/context.py
@@ -188,8 +188,8 @@ class Context(odict):
 
         Args:
           obs_id (multiple): The observation to load (see Notes).
-          dets (list, dict or ResultSet): The detectors to read.  If
-            None, all dets will be read.
+          dets (list, array, dict or ResultSet): The detectors to
+            read.  If None, all dets will be read.
           samples (tuple of ints): The start and stop sample indices.
             If None, read all samples.  (Note that some loader
             functions might not support this argument.)
@@ -198,8 +198,8 @@ class Context(odict):
             obsfiledb, but this shortcut will automatically determine
             the obs_id and the detector and sample range selections
             that correspond to this single file.
-          detsets (list): The detsets to read (with None equivalent to
-            requesting all detsets).
+          detsets (list, array): The detsets to read (with None
+            equivalent to requesting all detsets).
           meta (AxisManager): An AxisManager returned by get_meta
             (though possibly with additional axis restrictions
             applied) to use as a starting point for detector selection
@@ -299,6 +299,20 @@ class Context(odict):
         if aman is None:
             return meta
         if meta is not None:
+            if 'det_info' in aman:
+                # If the loader added det_info, then perform a special
+                # merge.  Duplicate keys should be avoided, because
+                # checking the values are the same is annoying.
+                _det_info = aman['det_info']
+                del aman['det_info']
+                _det_info.restrict_axes([meta.dets])
+                for k in _det_info._fields:
+                    if k in meta['det_info']:
+                        logger.warning(f'Key "{k}" is present in det_info returned by '
+                                       f'observation loader as well as in metadata '
+                                       f'databases; dropping the loader version.')
+                    else:
+                        meta.wrap(k, _det_info[k])
             aman.merge(meta)
         return aman
 
@@ -428,17 +442,23 @@ class Context(odict):
         request = {'obs:obs_id': obs_id}
         if detsets is not None:
             request['dets:detset'] = detsets
-        if isinstance(dets, list):
-            request['dets:readout_id'] = dets
-        elif isinstance(dets, metadata.ResultSet):
-            request['dets:readout_id'] = dets['readout_id']
-        elif dets is not None:
+
+        # Convert dets argument to request entry(s)
+        if isinstance(dets, dict):
             for k, v in dets.items():
                 if not k.startswith('dets:'):
                     k = 'dets:' + k
                 if k in request:
                     raise ValueError(f'Duplicate specification of dets field "{k}"')
                 request[k] = v
+        elif isinstance(dets, metadata.ResultSet):
+            request['dets:readout_id'] = dets['readout_id']
+        elif hasattr(dets, '__getitem__'):
+            # lists, tuples, arrays ...
+            request['dets:readout_id'] = dets
+        elif dets is not None:
+            # Try a cast ...
+            request['dets:readout_id'] = list(dets)
 
         metadata_list = self._get_warn_missing('metadata', [])
         meta = self.loader.load(metadata_list, request, det_info=det_info, check=check,
diff --git a/sotodlib/io/load_book.py b/sotodlib/io/load_book.py
index f77345f8..09f775c1 100644
--- a/sotodlib/io/load_book.py
+++ b/sotodlib/io/load_book.py
@@ -254,7 +254,7 @@ def load_obs_book(db, obs_id, dets=None, prefix=None, samples=None,
               [(0, 'bias_lines'), (1, 'samps')])
 
     det_info = core.metadata.ResultSet(
-        ['detset', 'readout_id', 'stream_id'],
+        ['detset_book', 'readout_id_book', 'stream_id'],
         src=list(det_info.values()))
-    aman.wrap('det_info', det_info.to_axismanager(axis_key='readout_id'))
+    aman.wrap('det_info', det_info.to_axismanager(axis_key='readout_id_book'))
     return aman


Test patch:
diff --git a/tests/test_context.py b/tests/test_context.py
index aebf82b7..c5f61733 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -79,6 +79,9 @@ class ContextTest(unittest.TestCase):
         n = len(dataset_sim.dets)
         obs_id = dataset_sim.obss['obs_id'][1]
         for selection, count in [
+                (['read05'], 1),
+                (np.array(['read05']), 1),
+                (metadata.ResultSet(['readout_id'], [['read05']]), 1),
                 ({'dets:readout_id': ['read05']}, 1),
                 ({'dets:readout_id': np.array(['read05'])}, 1),
                 ({'dets:detset': 'neard'}, 4),
@@ -94,6 +97,17 @@ class ContextTest(unittest.TestCase):
             self.assertTrue('cal' in meta)
             self.assertTrue('flags' in meta)
 
+        # And tolerance of the detsets argument ...
+        for selection, count in [
+                ('neard', 4),
+                (['neard'], 4),
+                (np.array(['neard', 'fard']), 8),
+        ]:
+            meta = ctx.get_meta(obs_id, detsets=selection)
+            self.assertEqual(meta.dets.count, count, msg=f"{selection}")
+            self.assertTrue('cal' in meta)
+            self.assertTrue('flags' in meta)
+
         # And without detdb nor metadata
         ctx = dataset_sim.get_context(with_detdb=False)
         for selection, count in [