Revised CalCat API

This is based on our discussion last week, but it's still a WIP. In particular, I haven't implemented actually loading the constants yet.

I've tweaked the design slightly: rather than having CalibrationData subclasses for each detector type, you use CalibrationData.from_conditions() with a detector specific object like AGIPDConditions. Jump to the tests to see what this looks like.

Questions:

• How to deal with missing modules?
  • A module may be missing for one type of constant but not for another
  • How do we even tell if modules are missing - should we query the PDUs so we know how many modules to expect?
• How to deal with missing calibration types (e.g. SlopesPC)
  • Especially if one group (e.g. illuminated) not found - this currently raises an error, even if we have other constants we could return
• Naming:
  • 'Calibrations' vs 'calibration' types vs 'constant types' (e.g. Offset, SlopesPC
  • 'Modules' vs 'aggregators'
• How many of SingleConstantVersion fields do we actually want?
• Should .select_modules() and .select_calibrations() take *args or a list?

cc @schmidtp @ahmedk @hammerd

src/cal_tools/calcat_interface2.py

@@ -4,13 +4,14 @@ from datetime import datetime
 from pathlib import Path
 from typing import Dict, Optional, Sequence, Union
 
+import h5py
+import pasha as psh
 from calibration_client import CalibrationClient
 from calibration_client.modules import CalibrationConstantVersion
 
 from .calcat_interface import CalCatApi, CalCatError
 from .tools import module_index_to_qm
 
-global_client = None
 
 class ModuleNameError(KeyError):
     def __init__(self, name):
@@ -20,6 +21,9 @@ class ModuleNameError(KeyError):
         return f"No module named {self.name!r}"
 
 
+global_client = None
+
+
 def get_client():
     global global_client
     if global_client is None:
@@ -43,6 +47,24 @@ def setup_client(base_url, client_id, client_secret, user_email, **kwargs):
     )
 
 
+_default_caldb_root = ...
+
+
+def _get_default_caldb_root():
+    global _default_caldb_root
+    if _default_caldb_root is ...:
+        onc_path = Path("/common/cal/caldb_store")
+        maxwell_path = Path("/gpfs/exfel/d/cal/caldb_store")
+        if onc_path.is_dir():
+            _default_caldb_root = onc_path
+        elif maxwell_path.is_dir():
+            _default_caldb_root = maxwell_path
+        else:
+            _default_caldb_root = None
+
+    return _default_caldb_root
+
+
 @dataclass
 class SingleConstantVersion:
     """A Calibration Constant Version for 1 detector module"""
@@ -76,6 +98,18 @@ class SingleConstantVersion:
             physical_name=ccv["physical_detector_unit"]["physical_name"],
         )
 
+    def dataset_obj(self, caldb_root=None):
+        if caldb_root is not None:
+            caldb_root = Path(caldb_root)
+        else:
+            caldb_root = _get_default_caldb_root()
+
+        f = h5py.File(caldb_root / self.path, "r")
+        return f[self.dataset]["data"]
+
+    def ndarray(self, caldb_root=None):
+        return self.dataset_obj(caldb_root)[:]
+
 
 @dataclass
 class ModulesConstantVersions:
@@ -83,7 +117,7 @@ class ModulesConstantVersions:
 
     constants: Dict[str, SingleConstantVersion]  # Keys e.g. 'LPD00'
 
-    def select_modules(self, *aggregators) -> "ModulesConstantVersions":
+    def select_modules(self, aggregators) -> "ModulesConstantVersions":
         d = {aggr: scv for (aggr, scv) in self.constants.items() if aggr in aggregators}
         return ModulesConstantVersions(d)
 
@@ -133,7 +167,7 @@ class CalibrationData(Mapping):
 
         api = CalCatApi(client or get_client())
 
-        detector_id = api.detector(detector_name)['id']
+        detector_id = api.detector(detector_name)["id"]
         all_modules = api.physical_detector_units(detector_id, pdu_snapshot_at)
         if modules is None:
             modules = sorted(all_modules)
@@ -173,7 +207,7 @@ class CalibrationData(Mapping):
 
         res = cls(d, modules)
         if modules:
-            res = res.select_modules(*modules)
+            res = res.select_modules(modules)
         return res
 
     @classmethod
@@ -227,14 +261,14 @@ class CalibrationData(Mapping):
     def qm_names(self):
         return [module_index_to_qm(n) for n in self.module_nums]
 
-    def require_calibrations(self, *calibrations):
+    def require_calibrations(self, calibrations):
         """Drop any modules missing the specified constant types"""
         mods = set(self.aggregators)
         for cal_type in calibrations:
             mods.intersection_update(self[cal_type].constants)
         return self.select_modules(mods)
 
-    def select_calibrations(self, *calibrations, require_all=True):
+    def select_calibrations(self, calibrations, require_all=True):
         if require_all:
             missing = set(calibrations) - set(self.constant_groups)
             if missing:
@@ -248,11 +282,14 @@ class CalibrationData(Mapping):
         # TODO: missing for some modules?
         return type(self)(d, self.aggregators)
 
-    def select_modules(self, *aggregators):
-        return type(self)({
-                cal_type: mcv.select_modules(*aggregators).constants
+    def select_modules(self, aggregators):
+        return type(self)(
+            {
+                cal_type: mcv.select_modules(aggregators).constants
                 for (cal_type, mcv) in self.constant_groups.items()
-        }, sorted(aggregators))
+            },
+            sorted(aggregators),
+        )
 
     def merge(self, *others: "CalibrationData") -> "CalibrationData":
         d = {}
@@ -268,6 +305,27 @@ class CalibrationData(Mapping):
 
         return type(self)(d, sorted(aggregators))
 
+    def load_all(self, caldb_root=None):
+        res = {}
+
+        const_load_mp = psh.ProcessContext(num_workers=24)
+        keys = []
+        for cal_type, mcv in self.constant_groups.items():
+            res[cal_type] = {}
+            for module in mcv.aggregators:
+                dset = mcv.constants[module].dataset_obj(caldb_root)
+                res[cal_type][module] = const_load_mp.alloc(
+                    shape=dset.shape, dtype=dset.dtype
+                )
+                keys.append((cal_type, module))
+
+        def _load_constant_dataset(wid, index, key):
+            cal_type, mod = key
+            dset = self[cal_type].constants[mod].dataset_obj(caldb_root)
+            dset.read_direct(res[cal_type][mod])
+
+        const_load_mp.map(_load_constant_dataset, keys)
+
 
 class ConditionsBase:
     calibration_types = {}  # For subclasses: {calibration: [parameter names]}