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',