"""Writing CXI files from AGIPD/LPD data""" import h5py import logging import numpy as np log = logging.getLogger(__name__) class VirtualCXIWriterBase: """ Base class for machinery to write a CXI file containing virtual datasets. You don't normally need to use this class directly. Instead, use the write_virtual_cxi() method on a multi-module detector data interface object. CXI specifies a particular layout of data in the HDF5 file format. It is documented here: http://www.cxidb.org/cxi.html This code writes version 1.5 CXI files. Parameters ---------- detdata: extra_data.components.MultimodDetectorBase The detector data interface for the data to gather in this file. """ # 1 entry is an index along the first (time) dimension in the source files. # XTDF detectors (AGIPD etc.) arrange pulses along this dimension, so each # entry is one frame & one memory cell. JUNGFRAU in burst mode makes one # entry with a separate dimension for several pulses, so overrides this. cells_per_entry = 1 def __init__(self, detdata): self.detdata = detdata self.group_label, self.image_label = detdata._main_data_key.split('.') frame_counts = detdata.frame_counts * self.cells_per_entry self.nframes = frame_counts.sum() log.info("Up to %d frames per train, %d frames in total", frame_counts.max(), self.nframes) self.train_ids_perframe = np.repeat( frame_counts.index.values, frame_counts.values.astype(np.intp) ) # For AGIPD, DSSC & LPD detectors modules are numbered from 0. # Overridden for JUNGFRAU to number from 1. self.modulenos = list(range(self.nmodules)) @property def nmodules(self): """Number of detector modules.""" return self.detdata.n_modules @property def data(self): """DataCollection with detector data from a run.""" return self.detdata.data def _get_module_index(self, module): """Returns an index for the specified module.""" return self.modulenos.index(module) def collect_pulse_ids(self): """ Gather pulse/cell ID labels for all modules and check consistency. Raises ------ Exception: Some data has no pulse ID values for any module. Exception: Inconsistent pulse IDs between detector modules. Returns ------- pulse_ids_min: np.array Array of pulse IDs per frame common for all detector modules. """ # Gather pulse IDs NO_PULSE_ID = 9999 pulse_ids = np.full((self.nframes, self.nmodules), NO_PULSE_ID, dtype=np.uint64) pulse_key = self.group_label + '.' + self.pulse_id_label for source, modno in self.detdata.source_to_modno.items(): module_ix = self._get_module_index(modno) for chunk in self.data._find_data_chunks(source, pulse_key): chunk_data = chunk.dataset self._map_chunk(chunk, chunk_data, pulse_ids, module_ix) # Sanity checks on pulse IDs pulse_ids_min = pulse_ids.min(axis=1) if (pulse_ids_min == NO_PULSE_ID).any(): raise Exception("Failed to find pulse IDs for some data") pulse_ids[pulse_ids == NO_PULSE_ID] = 0 if (pulse_ids_min != pulse_ids.max(axis=1)).any(): raise Exception("Inconsistent pulse IDs for different modules") # Pulse IDs make sense. Drop the modules dimension, giving one # pulse ID for each frame. return pulse_ids_min def _map_chunk(self, chunk, chunk_data, target, tgt_ax1, have_data=None): """ Map data from chunk into target. Chunk points to contiguous source data, but if this misses a train, it might not correspond to a contiguous region in the output. So this may perform multiple mappings. Parameters ---------- chunk: read_machinery::DataChunk Reference to a contiguous chunk of data to be mapped. chunk_data: h5py.Dataset / h5py.VirtualSource Dataset / VirtualSource to map data from. target: np.array / h5py.VirtualLayout Target to map data to. tgt_ax1: int Value for the target axis 1 - index corresponding to the detector module. have_data: np.array(dtype=bool), optional An array to monitor which part of the target have been mapped with data. Defaults to None. """ # Expand the list of train IDs to one per frame for tgt_slice, chunk_slice in self.detdata._split_align_chunk( chunk, self.detdata.train_ids_perframe ): tgt_start = tgt_slice.start * self.cells_per_entry tgt_end = tgt_slice.stop * self.cells_per_entry if self.cells_per_entry == 1: # In some cases, there's an extra dimension of length 1. # E.g. JUNGFRAU data with 1 memory cell per train or # DSSC/LPD raw data. if (len(chunk_data.shape) > 1 and chunk_data.shape[1] == 1): matched = chunk_data[chunk_slice, 0] else: matched = chunk_data[chunk_slice] target[tgt_start:tgt_end, tgt_ax1] = matched else: matched = chunk_data[chunk_slice] if isinstance(chunk_data, h5py.VirtualSource): # Use broadcasting of h5py.VirtualSource target[tgt_start:tgt_end, tgt_ax1] = matched else: target[tgt_start:tgt_end, tgt_ax1] = matched.reshape( (-1,) + matched.shape[2:]) # Fill in the map of what data we have if have_data is not None: have_data[tgt_start:tgt_end, tgt_ax1] = True def _map_layouts(self, layouts): """ Map virtual sources into virtual layouts. Parameters ---------- layouts: dict A dictionary of unmapped virtual layouts. Returns ------- layouts: dict A dictionary of virtual layouts mapped to the virtual sources. """ for name, layout in layouts.items(): key = '{}.{}'.format(self.group_label, name) have_data = np.zeros((self.nframes, self.nmodules), dtype=bool) for source, modno in self.detdata.source_to_modno.items(): print(f" ### Source: {source}, ModNo: {modno}, Key: {key}") module_ix = self._get_module_index(modno) for chunk in self.data._find_data_chunks(source, key): vsrc = h5py.VirtualSource(chunk.dataset) self._map_chunk(chunk, vsrc, layout, module_ix, have_data) filled_pct = 100 * have_data.sum() / have_data.size if hasattr(layout, 'sources'): n_mappings = len(layout.sources) # h5py < 3.3 else: n_mappings = layout.dcpl.get_virtual_count() # h5py >= 3.3 log.info(f"Assembled {n_mappings:d} chunks for {key:s}, " f"filling {filled_pct:.2f}% of the hyperslab") return layouts def write(self, filename, fillvalues=None): """ Write the file on disc to filename. Parameters ---------- filename: str Path of the file to be written. fillvalues: dict, optional Keys are datasets names (one of: data, gain, mask) and associated fill value for missing data. defaults are: - data: nan (proc, float32) or 0 (raw, uint16) - gain: 0 (uint8) - mask: 0xffffffff (uint32) """ pulse_ids = self.collect_pulse_ids() experiment_ids = np.char.add(np.char.add( self.train_ids_perframe.astype(str), ':'), pulse_ids.astype(str)) layouts = self.collect_data() data_label = self.image_label _fillvalues = { # Data can be uint16 (raw) or float32 (proc) data_label: np.nan if layouts[data_label].dtype.kind == 'f' else 0, 'gain': 0, 'mask': 0xffffffff } if fillvalues: _fillvalues.update(fillvalues) # Enforce that fill values are compatible with array dtype _fillvalues[data_label] = layouts[data_label].dtype.type( _fillvalues[data_label]) if 'gain' in layouts: _fillvalues['gain'] = layouts['gain'].dtype.type( _fillvalues['gain']) if 'mask' in layouts: _fillvalues['mask'] = layouts['mask'].dtype.type( _fillvalues['mask']) log.info("Writing to %s", filename) # Virtual datasets require HDF5 >= 1.10. # Specifying this up front should mean it fails before touching # the file if run on an older version. We also specify this as # the maximum version, to ensure we're creating files that can # be read by HDF5 1.10. with h5py.File(filename, 'w', libver=('v110', 'v110')) as f: f.create_dataset('cxi_version', data=[150]) d = f.create_dataset('entry_1/experiment_identifier', shape=experiment_ids.shape, dtype=h5py.special_dtype(vlen=str)) d[:] = experiment_ids # pulseId, trainId, cellId are not part of the CXI standard, # but it allows extra data. f.create_dataset(f'entry_1/{self.pulse_id_label}', data=pulse_ids) f.create_dataset('entry_1/trainId', data=self.train_ids_perframe) cellids = f.create_virtual_dataset('entry_1/cellId', layouts[self.cell_id_label]) cellids.attrs['axes'] = 'experiment_identifier:module_identifier' dgrp = f.create_group('entry_1/instrument_1/detector_1') if len(layouts[data_label].shape) == 4: axes_s = 'experiment_identifier:module_identifier:y:x' else: # 5D dataset, with extra axis for axes_s = 'experiment_identifier:module_identifier:data_gain:y:x' ndg = layouts[data_label].shape[2] d = f.create_dataset('entry_1/data_gain', shape=(ndg,), dtype=h5py.special_dtype(vlen=str)) d[:] = ([data_label, 'gain'] if ndg == 2 else [data_label]) dgrp['data_gain'] = h5py.SoftLink('/entry_1/data_gain') data = dgrp.create_virtual_dataset( 'data', layouts[data_label], fillvalue=_fillvalues[data_label] ) data.attrs['axes'] = axes_s if 'gain' in layouts: gain = dgrp.create_virtual_dataset( 'gain', layouts['gain'], fillvalue=_fillvalues['gain'] ) gain.attrs['axes'] = axes_s if 'mask' in layouts: mask = dgrp.create_virtual_dataset( 'mask', layouts['mask'], fillvalue=_fillvalues['mask'] ) mask.attrs['axes'] = axes_s dgrp['experiment_identifier'] = h5py.SoftLink( '/entry_1/experiment_identifier') f['entry_1/data_1'] = h5py.SoftLink( '/entry_1/instrument_1/detector_1') dgrp.create_dataset('module_identifier', data=self.modulenos) log.info("Finished writing virtual CXI file") class XtdfCXIWriter(VirtualCXIWriterBase): """ Machinery to write VDS files for a group of detectors with similar data format - AGIPD, DSSC & LPD. You don't normally need to use this class directly. Instead, use the write_virtual_cxi() method on a multi-module detector data interface object. CXI specifies a particular layout of data in the HDF5 file format. It is documented here: http://www.cxidb.org/cxi.html This code writes version 1.5 CXI files. Parameters ---------- detdata: extra_data.components.XtdfDetectorBase The detector data interface for the data to gather in this file. """ def __init__(self, detdata) -> None: self.cells_per_entry = 1 self.pulse_id_label = 'pulseId' self.cell_id_label = 'cellId' super().__init__(detdata) def collect_data(self): """ Prepare virtual layouts and map them to the virtual sources in the data chunks. Returns ------- layouts: dict A dictionary mapping virtual datasets names (e.g. ``data``) to h5py virtual layouts. """ src = next(iter(self.detdata.source_to_modno)) h5file = self.data[src].files[0].file image_grp = h5file['INSTRUMENT'][src][self.group_label] VLayout = h5py.VirtualLayout det_name = type(self.detdata).__name__ if 'gain' in image_grp: log.info(f"Identified {det_name} calibrated data") shape = (self.nframes, self.nmodules) + self.detdata.module_shape log.info("Virtual data shape: %r", shape) layouts = { self.image_label: VLayout( shape, dtype=image_grp[self.image_label].dtype), 'gain': VLayout(shape, dtype=image_grp['gain'].dtype), } if 'mask' in image_grp: layouts['mask'] = VLayout(shape, dtype=image_grp['mask'].dtype) else: log.info(f"Identified {det_name} raw data") shape = (self.nframes, self.nmodules) + image_grp['data'].shape[1:] log.info("Virtual data shape: %r", shape) layouts = { self.image_label: VLayout( shape, dtype=image_grp[self.image_label].dtype), } layouts[self.cell_id_label] = VLayout( (self.nframes, self.nmodules), dtype=image_grp[self.cell_id_label].dtype ) return self._map_layouts(layouts) class JUNGFRAUCXIWriter(VirtualCXIWriterBase): """ Machinery to write VDS files for JUNGFRAU data in the same format as AGIPD/LPD virtual datasets. You don't normally need to use this class directly. Instead, use the write_virtual_cxi() method on a multi-module detector data interface object. CXI specifies a particular layout of data in the HDF5 file format. It is documented here: http://www.cxidb.org/cxi.html This code writes version 1.5 CXI files. Parameters ---------- detdata: extra_data.components.JUNGFRAU The detector data interface for the data to gather in this file. """ def __init__(self, detdata) -> None: # Check number of cells src = next(iter(detdata.source_to_modno)) keydata = detdata.data[src, 'data.adc'] self.cells_per_entry = keydata.entry_shape[0] self.pulse_id_label = 'memoryCell' self.cell_id_label = 'memoryCell' super().__init__(detdata) # For JUNGFRAU detectors modules are numbered from 1 self.modulenos = list(range(1, self.nmodules + 1)) def collect_data(self): """ Prepare virtual layouts and map them to the virtual sources in the data chunks. Returns ------- layouts: dict A dictionary mapping virtual datasets names (e.g. ``data``) to h5py virtual layouts. """ src = next(iter(self.detdata.source_to_modno)) h5file = self.data[src].files[0].file image_grp = h5file['INSTRUMENT'][src][self.group_label] VLayout = h5py.VirtualLayout det_name = type(self.detdata).__name__ log.info(f"Identified {det_name} data") shape = (self.nframes, self.nmodules) + self.detdata.module_shape log.info("Virtual data shape: %r", shape) layouts = { self.image_label: VLayout( shape, dtype=image_grp[self.image_label].dtype), 'gain': VLayout(shape, dtype=image_grp['gain'].dtype), self.cell_id_label: VLayout( (self.nframes, self.nmodules), dtype=image_grp[self.cell_id_label].dtype ), } if 'mask' in image_grp: layouts['mask'] = VLayout(shape, dtype=image_grp['mask'].dtype) return self._map_layouts(layouts)