Rework LPD Mini components class

extra_data/components.py

@@ -178,14 +178,16 @@ def __getitem__(self, item):
         return MultimodKeyData(self, item)
 
     @classmethod
-    def _find_detector_name(cls, data):
+    def _find_detector_name(cls, data, source_re=None):
         detector_names = set()
+        if source_re is None:
+            source_re = cls._source_re
         for source in data.instrument_sources:
-            m = cls._source_re.match(source)
+            m = source_re.match(source)
             if m:
                 detector_names.add(m.group('detname'))
         if not detector_names:
-            raise SourceNameError(cls._source_re.pattern)
+            raise SourceNameError(source_re.pattern)
         elif len(detector_names) > 1:
             raise ValueError(
                 "Multiple detectors found in the data: {}. "
@@ -1561,8 +1563,114 @@ class LPDMini(LPDBase, XtdfDetectorBase):
       repeat the pulse & cell IDs from the first 1/3 of each train, and add gain
       stage labels from 0 (high-gain) to 2 (low-gain).
     """
-    _source_re = re.compile(r'(?P<detname>.+_LPD_MINI.*)/DET/(?P<modno>\d+)CH')
-    module_shape = (256, 256)
+    # Some code uses cls._source_re, but when creating an instance we replace
+    # this with either the raw or corrected variant.
+    _source_re = re.compile(r'(?P<detname>.+_LPD_MINI.*)/(DET|CORR)/(?P<modno>\d+)CH')
+    _source_re_raw = re.compile(r'(?P<detname>.+_LPD_MINI.*)/DET/(?P<modno>\d+)CH')
+    _source_re_corr = re.compile(r'(?P<detname>.+_LPD_MINI.*)/CORR/(?P<modno>\d+)CH')
+    module_shape = (32, 256)
+
+    def __init__(self, data: DataCollection, detector_name=None, modules=None,
+                 *, corrected=True, parallel_gain=False):
+        self.corrected = corrected
+        self._source_re = self._source_re_corr if corrected else self._source_re_raw
+        if detector_name is None:
+            detector_name = self._find_detector_name(data, self._source_re)
+        super().__init__(data, detector_name, modules=[0], parallel_gain=parallel_gain)
+
+    def __getitem__(self, item):
+        if item.startswith('image.'):
+            return LPDMiniImageKey(self, item, corrected=self.corrected)
+        return super().__getitem__(item)
+
+
+class LPDMiniImageKey(XtdfImageMultimodKeyData):
+    def __init__(self, det: XtdfDetectorBase, key, pulse_sel=by_index[0:MAX_PULSES:1], modules=None, corrected=True):
+        super().__init__(det, key, pulse_sel)
+        if modules is None and self._has_modules:
+            eshape = self._eg_keydata.entry_shape
+            nmod = eshape[-3] if corrected else (eshape[-2] // 32)
+            modules = list(range(nmod))
+        self._modules = modules
+        self.corrected = corrected
+
+    @property
+    def _has_modules(self):
+        return (self._eg_keydata.ndim - self._extraneous_dim) > 1
+
+    @property
+    def ndim(self):
+        return super().ndim if self._has_modules else (super().ndim - 1)
+
+    @property
+    def dimensions(self):
+        ndim_inner = self.ndim - 1 - self._has_modules
+        if ndim_inner == 2:
+            # 2D pixel data
+            entry_dims = ['slow_scan', 'fast_scan']
+        else:
+            # Everything else seems to be 1D, but just in case
+            entry_dims = [f'dim_{i}' for i in range(ndim_inner)]
+
+        if self._has_modules:
+            return ['train_pulse', 'module'] + entry_dims
+        return ['train_pulse'] + entry_dims
+
+    @property
+    def modules(self):
+        return self._modules
+
+    def buffer_shape(self, module_gaps=False, roi=()):
+        """Get the array shape for this data
+
+        If *module_gaps* is True, include space for modules which are missing
+        from the data. *roi* may be a tuple of slices defining a region of
+        interest on the inner dimensions of the data.
+        """
+        if self._has_modules:
+            nframes_sel = len(self.train_id_coordinates())
+            module_dim = 8 if module_gaps else len(self.modules)
+
+            return (nframes_sel, module_dim) + roi_shape(self.det.module_shape, roi)
+        else:
+            return super().buffer_shape(module_gaps, roi)[1:]
+
+    # Used for .select_trains() and .split_trains()
+    def _with_selected_det(self, det_selected):
+        return LPDMiniImageKey(det_selected, self.key, self._pulse_sel, modules=self.modules, corrected=self.corrected)
+
+    def select_pulses(self, pulses):
+        pulses = _check_pulse_selection(pulses)
+
+        return LPDMiniImageKey(self.det, self.key, pulses, modules=self.modules, corrected=self.corrected)
+
+    def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, module_gaps=False):
+        """Get an array of per-pulse data (image.*) for xtdf detector"""
+        if roi or module_gaps:
+            raise NotImplementedError
+
+        # trains, modules, 32, 256
+        out_shape = self.buffer_shape(module_gaps=module_gaps, roi=roi)
+
+        if out is None:
+            dtype = self._eg_keydata.dtype if astype is None else np.dtype(astype)
+            out = _out_array(out_shape, dtype, fill_value=fill_value)
+        elif out.shape != out_shape:
+            raise ValueError(f'requires output array of shape {out_shape}')
+
+        reading_view = out.view()
+        if not self.corrected:
+            reading_view.shape = (
+                out_shape[:1]
+                + ((1,) if self._extraneous_dim else ())
+                + ((out_shape[1] * out_shape[2], out_shape[3]) if self._has_modules else ())
+            )
+
+        for _modno, kd in self.modno_to_keydata.items():
+            for chunk in kd._data_chunks:
+                self._read_chunk(chunk, reading_view, roi)
+
+        return out
 
 
 @multimod_detectors

extra_data/tests/make_examples.py

@@ -10,7 +10,7 @@
 from .mockdata.base import write_base_index
 from .mockdata.basler_camera import BaslerCamera as BaslerCam
 from .mockdata.dctrl import DCtrl
-from .mockdata.detectors import AGIPDModule, DSSCModule, LPDModule
+from .mockdata.detectors import AGIPDModule, DSSCModule, LPDModule, LPDMini
 from .mockdata.gauge import Gauge
 from .mockdata.gec_camera import GECCamera
 from .mockdata.imgfel import IMGFELCamera, IMGFELMotor
@@ -243,8 +243,8 @@ def make_fxe_run(dir_path, raw=True, format_version='0.5'):
 
     path = osp.join(dir_path, f'{prefix}-R0450-LPDMINI00-S00000.h5')
     write_file(path, [
-        LPDModule('FXE_DET_LPD_MINI/DET/0CH0', raw=raw, frames_per_train=128)
-        ], ntrains=480, chunksize=32, format_version=format_version)
+        LPDMini('FXE_DET_LPD_MINI/DET/0CH0', raw=raw, frames_per_train=128, modules=2)
+    ], ntrains=480, chunksize=32, format_version=format_version)
 
     if not raw:
         return
@@ -274,7 +274,7 @@ def make_lpd_parallelgain_run(dir_path, raw=True, format_version='0.5'):
 
     path = osp.join(dir_path, f'{prefix}-R0450-LPDMINI00-S00000.h5')
     write_file(path, [
-        LPDModule('FXE_DET_LPD_MINI/DET/0CH0', raw=raw, frames_per_train=300)
+        LPDMini('FXE_DET_LPD_MINI/DET/0CH0', raw=raw, frames_per_train=300, modules=2)
     ], ntrains=100, chunksize=32, format_version=format_version)
 
 def make_lpd_run_mini_missed_train(dir_path):

extra_data/tests/mockdata/detectors.py

@@ -162,6 +162,15 @@ class LPDModule(DetectorModule):
     image_dims = (1, 256, 256)
     detector_data_size = 416
 
+class LPDMini(DetectorModule):
+    def __init__(self, device_id, frames_per_train=64, raw=True, modules=1):
+        if raw:
+            self.image_dims = (1, modules * 32, 256)
+        else:
+            # The 1 is removed in parent class
+            self.image_dims = (1, modules, 32, 256)
+        super().__init__(device_id, frames_per_train, raw)
+
 class DSSCModule(DetectorModule):
     image_dims = (1, 128, 512)
     detector_data_size = 416

extra_data/tests/test_components.py

@@ -179,12 +179,11 @@ def test_get_array_lpd_parallelgain(mock_lpd_parallelgain_run):
     np.testing.assert_array_equal(arr.coords['gain'], np.arange(3))
     np.testing.assert_array_equal(arr.coords['pulse'], np.arange(100))
 
-    run = RunDirectory(mock_lpd_parallelgain_run)
-    det = LPDMini(run.select_trains(by_index[:2]), parallel_gain=True)
+    det = LPDMini(run.select_trains(by_index[:2]), parallel_gain=True, corrected=False)
     assert det.detector_name == 'FXE_DET_LPD_MINI'
 
     arr = det.get_array('image.data')
-    assert arr.shape == (1, 2, 3, 100, 256, 256)
+    assert arr.shape == (2, 2, 3, 100, 32, 256)
     assert arr.dims == ('module', 'train', 'gain', 'pulse', 'slow_scan', 'fast_scan')
     np.testing.assert_array_equal(arr.coords['gain'], np.arange(3))
     np.testing.assert_array_equal(arr.coords['pulse'], np.arange(100))
@@ -206,18 +205,17 @@ def test_get_array_lpd_parallelgain_select_pulses(mock_lpd_parallelgain_run):
     assert arr.shape == (16, 2, 3, 5, 256, 256)
     np.testing.assert_array_equal(arr.coords['pulse'], np.arange(5))
 
-    run = RunDirectory(mock_lpd_parallelgain_run)
-    det = LPDMini(run.select_trains(by_index[:2]), parallel_gain=True)
+    det = LPDMini(run.select_trains(by_index[:2]), parallel_gain=True, corrected=False)
     assert det.detector_name == 'FXE_DET_LPD_MINI'
 
     arr = det.get_array('image.data', pulses=np.s_[:5])
-    assert arr.shape == (1, 2, 3, 5, 256, 256)
+    assert arr.shape == (2, 2, 3, 5, 32, 256)
     assert arr.dims == ('module', 'train', 'gain', 'pulse', 'slow_scan', 'fast_scan')
     np.testing.assert_array_equal(arr.coords['gain'], np.arange(3))
     np.testing.assert_array_equal(arr.coords['pulse'], np.arange(5))
 
     arr = det.get_array('image.data', pulses=by_id[:5])
-    assert arr.shape == (1, 2, 3, 5, 256, 256)
+    assert arr.shape == (2, 2, 3, 5, 32, 256)
     np.testing.assert_array_equal(arr.coords['pulse'], np.arange(5))
 
 

extra_data/tests/test_stacking.py

@@ -22,7 +22,7 @@ def test_stack_detector_data(mock_fxe_raw_run):
 
 def test_stack_detector_data_missing(mock_fxe_raw_run):
     test_run = RunDirectory(mock_fxe_raw_run)
-    tid, data = test_run.train_from_id(10000, devices=[('*/DET/*', 'image.data')])
+    tid, data = test_run.train_from_id(10000, devices=[('*_LPD1M-1/DET/*', 'image.data')])
 
     # Three variants of missing data:
     # 1. Source missing
@@ -52,7 +52,7 @@ def test_stack_detector_data_missing(mock_fxe_raw_run):
 
 def test_stack_detector_data_stackview(mock_fxe_raw_run):
     test_run = RunDirectory(mock_fxe_raw_run)
-    tid, data = test_run.train_from_id(10000, devices=[('*/DET/*', 'image.data')])
+    tid, data = test_run.train_from_id(10000, devices=[('*_LPD1M-1/DET/*', 'image.data')])
 
     # Three variants of missing data:
     # 1. Source missing
@@ -126,7 +126,7 @@ def test_stack_detector_data_type_error(mock_fxe_raw_run):
 
 def test_stack_detector_data_extra_mods(mock_fxe_raw_run):
     test_run = RunDirectory(mock_fxe_raw_run)
-    tid, data = test_run.train_from_id(10000, devices=[('*/DET/*', 'image.data')])
+    tid, data = test_run.train_from_id(10000, devices=[('*_LPD1M-1/DET/*', 'image.data')])
 
     data.setdefault(
         'FXE_DET_LPD1M-1/DET/16CH0:xtdf',