DRAFT get_data method to Components

extra_data/components.py

@@ -8,19 +8,23 @@
 import numpy as np
 import pandas as pd
 import xarray
+from extra_geom import (AGIPD_500K2GGeometry, Epix100Geometry,
+                        JUNGFRAUGeometry, PNCCDGeometry)
 
-from .exceptions import SourceNameError
-from .reader import DataCollection, by_id, by_index
+from .exceptions import PropertyNameError, SourceNameError
 from .read_machinery import DataChunk, roi_shape, split_trains
+from .reader import DataCollection, by_id, by_index
+from .write_cxi import JUNGFRAUCXIWriter, XtdfCXIWriter
 from .writer import FileWriter
-from .write_cxi import XtdfCXIWriter, JUNGFRAUCXIWriter
 
 __all__ = [
     'AGIPD1M',
     'AGIPD500K',
     'DSSC1M',
     'LPD1M',
     'JUNGFRAU',
+    'PNCCD',
+    'Epix100',
     'identify_multimod_detectors',
 ]
 
@@ -442,6 +446,57 @@ def get_array(self, key, *, fill_value=None, roi=(), astype=None):
 
         return xarray.DataArray(out, dims=dims, coords=coords)
 
+    def _default_geometry(self):
+        """Return a default geometry for the data"""
+        raise Exception(
+            f'Component {self.__class__.__name__} with {self.n_modules} modules'
+            ' does not provide a default geometry')
+
+    def get_data(self, geometry=None, *, fill_value=None, roi=(), astype=None,
+                 mask=False, asic_seams=False):
+        """Get assembled data with geometry and mask applied.
+
+        geometry:
+            An extra_geom geometry object.
+        fill_value: int or float, optional
+            Value to use for missing values. If None (default) the fill value
+            is 0 for integers and np.nan for floats.
+        roi: tuple
+          Specify e.g. ``np.s_[10:60, 100:200]`` to select pixels within each
+          module when reading data. The selection is applied to each individual
+          module, so it may only be useful when working with a single module.
+        astype: Type
+          Data type of the output array. If None (default) the dtype matches the
+          input array dtype
+        mask: (int, bool)
+            if True, use the mask present in the data (if in use with CORR data). If an
+            int is provided, it will only mask the selected bits.
+        asic_seams: bool
+            mask asic edges if the geometry implements it.
+        """
+        geometry = geometry or self._default_geometry()
+        data = self.get_array(self._main_data_key, fill_value=fill_value, roi=roi, astype=astype)
+        data = data.transpose(..., 'module', *data.dims[-2:]).data
+
+        if hasattr(geometry, '_ensure_shape'):
+            data = geometry._ensure_shape(data)
+        if asic_seams and hasattr(geometry, 'asic_seams'):
+            data = data * ~geometry.asic_seams()
+
+        if mask:
+            try:
+                if isinstance(mask, int) and not isinstance(mask, bool):
+                    # mask only the selected bits if an int is passed to the mask arg
+                    mask = (self.get_array('image.mask', roi=roi) & mask).astype(np.bool_)
+                else:
+                    mask = self.get_array('image.mask', fill_value=False, roi=roi, astype=np.bool_)
+            except PropertyNameError:
+                pass  # no mask available in this data
+            else:
+                data = data * ~mask
+
+        assembled, centre = geometry.position_modules(data)
+        return assembled, centre
 
     def get_dask_array(self, key, fill_value=None, astype=None):
         """Get a labelled Dask array of detector data
@@ -1265,6 +1320,9 @@ class AGIPD500K(XtdfDetectorBase):
     module_shape = (512, 128)
     n_modules = 8
 
+    def _default_geometry(self):
+        return AGIPD_500K2GGeometry.from_origin()
+
 
 @multimod_detectors
 class DSSC1M(XtdfDetectorBase):
@@ -1447,6 +1505,11 @@ def __init__(self, data: DataCollection, detector_name=None, modules=None,
         src = next(iter(self.source_to_modno))
         self._frames_per_entry = self.data[src, self._main_data_key].entry_shape[0]
 
+    def _default_geometry(self):
+        if self.n_modules == 1:
+            return JUNGFRAUGeometry.from_module_positions()
+        super()._default_geometry()
+
     @staticmethod
     def _label_dims(arr):
         # Label dimensions to match the AGIPD/DSSC/LPD data access
@@ -1543,6 +1606,41 @@ def write_virtual_cxi(self, filename, fillvalues=None):
         """
         JUNGFRAUCXIWriter(self).write(filename, fillvalues=fillvalues)
 
+
+@multimod_detectors
+class PNCCD(MultimodDetectorBase):
+    """An interface to PNCCD data
+
+    Detector names are like ''
+    """
+    _source_re = re.compile(r'(?P<detname>.+_PNCCD1MP)/CAL/PNCCD_FMT-(?P<modno>\d+)')
+    _main_data_key = 'data.image'
+    module_shape = PNCCDGeometry.expected_data_shape[-2:]
+    n_modules = 2
+
+    def _default_geometry(self):
+        return PNCCDGeometry.from_relative_positions()
+
+
+@multimod_detectors
+class Epix100(MultimodDetectorBase):
+    """An interface to Epix100 data
+
+    Detector names are like ''
+    """
+    _source_re = re.compile(r'(?P<detname>.+_(EPX100|EPIX)-(?P<modno>\d+))/DET/RECEIVER')
+    _main_data_key = 'data.image.pixels'
+    module_shape = Epix100Geometry.expected_data_shape[-2:]
+    n_modules = 1
+
+    def _default_geometry(self):
+        if self.detector_name == 'HED_IA1_EPX100-2':
+            # module 2  at HED has a different geometry, for more details, see:
+            # https://extra-geom.readthedocs.io/en/latest/geometry.html#extra_geom.Epix100Geometry.from_origin
+            return Epix100Geometry.from_relative_positions(top=[386.5, 364.5, 0.], bottom=[386.5, -12.5, 0.])
+        return Epix100Geometry.from_origin()
+
+
 def identify_multimod_detectors(
         data, detector_name=None, *, single=False, clses=None
 ):

extra_data/tests/conftest.py

@@ -115,6 +115,20 @@ def mock_jungfrau_run():
         yield td
 
 
+@pytest.fixture(scope='session')
+def mock_epix100_run():
+    with TemporaryDirectory() as td:
+        make_examples.make_epix100_run(td)
+        yield td
+
+
+@pytest.fixture(scope='session')
+def mock_pnccd_run():
+    with TemporaryDirectory() as td:
+        make_examples.make_pnccd_run(td)
+        yield td
+
+
 @pytest.fixture(scope='session')
 def mock_scs_run():
     with TemporaryDirectory() as td:

extra_data/tests/make_examples.py

@@ -3,20 +3,21 @@
 
 import h5py
 import numpy as np
+from extra_data.components import PNCCD
 
 from .mockdata import write_file
 from .mockdata.adc import ADC
 from .mockdata.agipd import AGIPD1MFPGA, AGIPD1MPSC, AGIPD500KFPGA, AGIPDMDL
 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 (PNCCD, AGIPDModule, DSSCModule, Epix100,
+                                 LPDModule)
 from .mockdata.gauge import Gauge
 from .mockdata.gec_camera import GECCamera
 from .mockdata.imgfel import IMGFELCamera, IMGFELMotor
-from .mockdata.jungfrau import (
-    JUNGFRAUControl, JUNGFRAUModule, JUNGFRAUMonitor, JUNGFRAUPower
-)
+from .mockdata.jungfrau import (JUNGFRAUControl, JUNGFRAUModule,
+                                JUNGFRAUMonitor, JUNGFRAUPower)
 from .mockdata.motor import Motor
 from .mockdata.mpod import MPOD
 from .mockdata.proc import ReconstructedDLD6
@@ -373,6 +374,19 @@ def make_jungfrau_run(dir_path):
         JUNGFRAUPower('SPB_IRDA_JF4M/MDL/POWER'),
     ], ntrains=100, chunksize=1, format_version='1.0')
 
+
+def make_epix100_run(dir_path):
+    # Naming based on /gpfs/exfel/exp/MID/202201/p002834/raw/r0199
+    path = osp.join(dir_path, 'RAW-R0199-EPIX01-S00000.h5')
+    write_file(path, [Epix100('MID_EXP_EPIX-1/DET/RECEIVER')], ntrains=100, chunksize=1, format_version='1.0')
+
+
+def make_pnccd_run(dir_path):
+    # Naming based on /gpfs/exfel/exp/SQS/202201/p002857/raw/r0259/
+    path = osp.join(dir_path, 'RAW-R0259-PNCCD01-S00000.h5')
+    write_file(path, [PNCCD('SQS_NQS_PNCCD1MP/CAL/PNCCD_FMT-0')], ntrains=100, chunksize=1, format_version='1.0')
+
+
 def make_remi_run(dir_path):
     write_file(osp.join(dir_path, f'CORR-R0210-REMI01-S00000.h5'), [
         ReconstructedDLD6('SQS_REMI_DLD6/DET/TOP'),

extra_data/tests/mockdata/base.py

@@ -109,7 +109,7 @@ def write_instrument(self, f):
             if len(trainids) > 0:
                 tid[:self.nsamples] = trainids
             for (topic, datatype, dims) in self.instrument_keys:
-                f.create_dataset('INSTRUMENT/%s/%s' % (dev_chan, topic),
+                f.create_dataset('INSTRUMENT/%s/%s' % (dev_chan, topic.replace('.', '/')),
                                  (Npad,) + dims, datatype, maxshape=((None,) + dims))
 
     def datasource_ids(self):

extra_data/tests/mockdata/detectors.py

@@ -1,5 +1,8 @@
 import numpy as np
 
+from .base import DeviceBase
+
+
 class DetectorModule:
     # Overridden in subclasses:
     image_dims = ()
@@ -165,3 +168,35 @@ class LPDModule(DetectorModule):
 class DSSCModule(DetectorModule):
     image_dims = (1, 128, 512)
     detector_data_size = 416
+
+class PNCCD(DeviceBase):
+    output_channels = ('output/data',)
+
+    instrument_keys = [
+        ('image', 'u2', (1024, 1024)),
+    ]
+
+class Epix100(DeviceBase):
+    output_channels = ('daqOutput/data',)
+
+    instrument_keys = [
+        ('ambTemp', 'i4', ()),
+        ('analogCurr', 'u4', ()),
+        ('analogInputVolt', 'u4', ()),
+        ('backTemp', 'i4', ()),
+        ('digitalCurr', 'u4', ()),
+        ('digitalInputVolt', 'u4', ()),
+        ('guardCurr', 'u4', ()),
+        ('image.binning', 'u8', (2,)),
+        ('image.bitsPerPixel', 'i4', ()),
+        ('image.dimTypes', 'i4', (2,)),
+        ('image.dims', 'u8', (2,)),
+        ('image.encoding', 'i4', ()),
+        ('image.flipX', 'u1', ()),
+        ('image.flipY', 'u1', ()),
+        ('image.pixels', 'i2', (708, 768)),
+        ('image.roiOffsets', 'u8', (2,)),
+        ('image.rotation', 'i4', ()),
+        ('pulseId', 'u8', ()),
+        ('relHumidity', 'i4', ()),
+    ]

extra_data/tests/test_components.py

@@ -7,8 +7,10 @@
 
 from extra_data.reader import RunDirectory, H5File, by_id, by_index
 from extra_data.components import (
-    AGIPD1M, DSSC1M, LPD1M, JUNGFRAU, identify_multimod_detectors,
+    AGIPD1M, DSSC1M, LPD1M, JUNGFRAU, Epix100, PNCCD, identify_multimod_detectors,
 )
+from extra_geom import AGIPD_1MGeometry, LPD_1MGeometry, DSSC_1MGeometry
+from extra_geom.tests.test_jungfrau_geometry import jf4m_geometry
 
 
 def test_get_array(mock_fxe_raw_run):
@@ -581,3 +583,39 @@ def test_identify_multimod_detectors_multi(mock_fxe_raw_run, mock_spb_raw_run):
     name, cls = identify_multimod_detectors(combined, single=True, clses=[AGIPD1M])
     assert name == 'SPB_DET_AGIPD1M-1'
     assert cls is AGIPD1M
+
+
+def test_get_data(mock_fxe_raw_run, mock_jungfrau_run, mock_epix100_run, mock_pnccd_run):
+    # LPD
+    fxe = RunDirectory(mock_fxe_raw_run).select_trains(np.s_[:5])
+    lpd = LPD1M(fxe)
+    with pytest.raises(Exception):
+        lpd.get_data()
+    lpd_data, _ = lpd.get_data(geometry=LPD_1MGeometry.from_quad_positions([(11.4, 299), (-11.5, 8), (254.5, -16), (278.5, 275)]))
+    assert lpd_data.shape == (5, 128, 1202, 1104)
+
+    # JUNGFRAU 4M
+    jf_run = RunDirectory(mock_jungfrau_run).select_trains(np.s_[:5])
+    jf = JUNGFRAU(jf_run)
+    with pytest.raises(Exception):
+        jf_data, _ = jf.get_data()
+    jf_data, _ = jf.get_data(geometry=jf4m_geometry())
+    assert jf_data.shape == (5, 16, 2156, 2250), jf_data.shape
+
+    # JUNGFRAU single module
+    jf_run = jf_run.select('SPB_IRDA_JF4M/DET/JNGFR01:daqOutput', '*')
+    jf = JUNGFRAU(jf_run, n_modules=1)
+    jf_data, _ = jf.get_data()
+    assert jf_data.shape == (5, 16, 514, 1030), jf_data.shape
+
+    # Epix100
+    epix_run = RunDirectory(mock_epix100_run).select_trains(np.s_[:5])
+    epix100 = Epix100(epix_run)
+    epix_data, _ = epix100.get_data()
+    assert epix_data.shape == (5, 709, 773), epix_data.shape
+
+    # PNCCD
+    pnccd_run = RunDirectory(mock_pnccd_run).select_trains(np.s_[:10])
+    pnccd = PNCCD(pnccd_run)
+    pnccd_data, _ = pnccd.get_data()
+    assert pnccd_data.shape == (10, 1, 1078, 1024)

setup.py

@@ -50,6 +50,7 @@ def find_version(*parts):
           ],
       },
       install_requires=[
+          'extra_geom',
           'h5py>=2.10',
           'numpy',
           'pandas',