diff --git a/extra_data/components.py b/extra_data/components.py index 4d3d56e9..ddbb5156 100644 --- a/extra_data/components.py +++ b/extra_data/components.py @@ -10,10 +10,12 @@ import numpy as np import pandas as pd +from . import direct_read from .exceptions import SourceNameError from .reader import DataCollection, by_id, by_index -from .read_machinery import DataChunk, roi_shape, split_trains -from .utils import default_num_threads +from .read_machinery import ( + ReadOp, contiguous_regions, roi_shape, split_trains, +) from .writer import FileWriter from .write_cxi import XtdfCXIWriter, JUNGFRAUCXIWriter @@ -957,10 +959,43 @@ def split_trains(self, parts=None, trains_per_part=None, frames_per_part=None): for det_split in self.det.split_trains(parts, trains_per_part, frames_per_part): yield self._with_selected_det(det_split) - def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, module_gaps=False): - """Get data as a plain NumPy array with no labels""" + def _module_indices(self, module_gaps): + """(index in the output array, KeyData) for each module, in order.""" + for i, (modno, kd) in enumerate(sorted(self.modno_to_keydata.items())): + yield ((modno - self.det._modnos_start_at) if module_gaps else i), kd + + def _read_ops(self, module_gaps, entries): + """The reads needed to fill an array with *entries* rows per module. + + The destination indices count through the modules dimension and the + entries dimension as if they were one, so that :meth:`_read` can fill + every module in a single pass. + """ train_ids = np.asarray(self.det.train_ids) + ops = [] + for mod_ix, kd in self._module_indices(module_gaps): + for chunk in kd._data_chunks: + for tgt_slice, chunk_slice in self.det._split_align_chunk( + chunk, train_ids): + ops.append(ReadOp( + chunk.file, chunk.dataset_path, chunk_slice.start, + (mod_ix * entries) + tgt_slice.start, + chunk_slice.stop - chunk_slice.start, + )) + + return ops + + def _read(self, out, ops, roi, parallel): + """Fill *out*, shaped (modules, entries, ...), from *ops*.""" + # The ops index the modules and entries dimensions as one, which is only + # the same array if it's contiguous. + assert out.flags.c_contiguous, "output array must be C-contiguous" + direct_read.read(out.reshape((-1,) + out.shape[2:]), ops, roi, parallel) + + def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, + module_gaps=False, parallel=-1): + """Get data as a plain NumPy array with no labels""" out_shape = self.buffer_shape(module_gaps, roi) if out is None: @@ -969,13 +1004,7 @@ def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, module_gaps elif out.shape != out_shape: raise ValueError(f'requires output array of shape {out_shape}') - for i, (modno, kd) in enumerate(sorted(self.modno_to_keydata.items())): - mod_ix = (modno - self.det._modnos_start_at) if module_gaps else i - for chunk in kd._data_chunks: - for tgt_slice, chunk_slice in self.det._split_align_chunk(chunk, train_ids): - chunk.dataset.read_direct( - out[mod_ix, tgt_slice], source_sel=(chunk_slice,) + roi - ) + self._read(out, self._read_ops(module_gaps, out_shape[1]), roi, parallel) return out def _wrap_xarray(self, arr): @@ -984,8 +1013,9 @@ def _wrap_xarray(self, arr): coords = {'module': self.modules, 'trainId': self.train_id_coordinates()} return DataArray(arr, dims=self.dimensions, coords=coords) - def xarray(self, *, fill_value=None, roi=(), astype=None): - arr = self.ndarray(fill_value=fill_value, roi=roi, astype=astype) + def xarray(self, *, fill_value=None, roi=(), astype=None, parallel=-1): + arr = self.ndarray(fill_value=fill_value, roi=roi, astype=astype, + parallel=parallel) return self._wrap_xarray(arr) def dask_array(self, *, labelled=False, fill_value=None, astype=None): @@ -1181,115 +1211,47 @@ def _sel_frames(self): self._sel_frames_cached = s return self._sel_frames_cached - def _read_chunk(self, chunk: DataChunk, mod_out, roi): - """Read per-pulse data from file into an output array (of 1 module)""" - # Limit to 5 GB sections of the dataset at once, so the temporary - # arrays used in the workaround below are not too large. - nbytes_frame = chunk.dataset.dtype.itemsize - for dim in chunk.dataset.shape[1:]: - nbytes_frame *= dim - frame_limit = 5 * (1024 ** 3) // nbytes_frame - - for tgt_slice, chunk_slice in self.det._split_align_chunk( - chunk, self.det.train_ids_perframe, length_limit=frame_limit - ): - inc_pulses_chunk = self._sel_frames[tgt_slice] - if inc_pulses_chunk.sum() == 0: # No data from this chunk selected - continue - elif inc_pulses_chunk.all(): # All pulses in chunk - chunk.dataset.read_direct( - mod_out[tgt_slice], source_sel=(chunk_slice,) + roi - ) - continue + def _read_ops(self, module_gaps, entries): + """The reads needed to fill an array with *entries* frames per module. - # Read a subset of pulses from the chunk: - - # Reading a non-contiguous selection in HDF5 seems to be slow: - # https://forum.hdfgroup.org/t/performance-reading-data-with-non-contiguous-selection/8979 - # Except it's fast if you read the data to a matching selection in - # memory (one weird trick). - # So as a workaround, this allocates a temporary array of the same - # shape as the full chunk, reads into it, and then copies the selected - # data to the output array. The extra memory copy is not optimal, - # but it's better than the HDF5 performance issue, at least in some - # realistic cases. - # N.B. tmp should only use memory for the data it contains - - # zeros() uses calloc, so the OS can do virtual memory tricks. - # Don't change this to zeros_like() ! - tmp = np.zeros( - shape=inc_pulses_chunk.shape + chunk.dataset.shape[1:], - dtype=chunk.dataset.dtype - ) - tmp_sel = np.nonzero(inc_pulses_chunk)[0] - dataset_sel = tmp_sel + chunk_slice.start - chunk.dataset.read_direct( - tmp, source_sel=(dataset_sel,) + roi, dest_sel=(tmp_sel,) + roi, - ) - # Where does this data go in the target array? - tgt_start_ix = self._sel_frames[:tgt_slice.start].sum() - tgt_pulse_sel = slice( - tgt_start_ix, tgt_start_ix + inc_pulses_chunk.sum() - ) - # Copy data from temp array to output array - np.compress( - inc_pulses_chunk, tmp[np.index_exp[:] + roi], - axis=0, out=mod_out[tgt_pulse_sel] - ) - - def _read_parallel_decompress(self, out, module_gaps, threads=16): - try: - from .compression import multi_dataset_decompressor, parallel_decompress_chunks - except ImportError: - return False - - modno_to_keydata_no_virtual = {} - all_datasets = [] - for (m, vkd) in self.modno_to_keydata.items(): - modno_to_keydata_no_virtual[m] = kd = vkd._without_virtual_overview() - all_datasets.extend([f.file[kd.hdf5_data_path] for f in kd.files]) - - if any(d.chunks != (1,) + d.shape[1:] for d in all_datasets): - return False # Chunking not as we expect - - decomp_proto = multi_dataset_decompressor(all_datasets) - if decomp_proto is None: - return False # No suitable fast decompression path + A pulse selection makes holes in what we want from each chunk, so this + makes one read per contiguous run of selected frames. + """ + sel_frames = self._sel_frames + ops = [] - load_tasks = [] - for i, (modno, kd) in enumerate(sorted(modno_to_keydata_no_virtual.items())): - mod_ix = (modno - self.det._modnos_start_at) if module_gaps else i - # 'chunk' in the lines below means a range of consecutive indices - # in one HDF5 dataset, as elsewhere in EXtra-data. - # We use this to build a list of HDF5 chunks (1 frame per chunk) - # to be loaded & decompressed. Sorry about that. + for mod_ix, kd in self._module_indices(module_gaps): for chunk in kd._data_chunks: - dset = chunk.dataset - for tgt_slice, chunk_slice in self.det._split_align_chunk( - chunk, self.det.train_ids_perframe, - ): - inc_pulses_chunk = self._sel_frames[tgt_slice] - if inc_pulses_chunk.sum() == 0: # No data from this chunk selected + chunk, self.det.train_ids_perframe): + inc_frames = sel_frames[tgt_slice] + if not inc_frames.any(): # Nothing selected from this chunk continue - dataset_ixs = np.nonzero(inc_pulses_chunk)[0] + chunk_slice.start - - # Where does this data go in the target array? - tgt_start_ix = self._sel_frames[:tgt_slice.start].sum() + # Where the frames selected so far end up in the output + dest = ((mod_ix * entries) + + int(sel_frames[:tgt_slice.start].sum())) - # Each task is a h5py.h5d.DatasetID, coordinates & array destination - load_tasks.extend([ - (dset.id, (ds_ix, 0, 0), out[mod_ix, tgt_start_ix + i]) - for i, ds_ix in enumerate(dataset_ixs)] - ) + for first, stop in contiguous_regions(inc_frames): + ops.append(ReadOp( + chunk.file, chunk.dataset_path, + chunk_slice.start + first, dest, stop - first, + )) + dest += stop - first - parallel_decompress_chunks(load_tasks, decomp_proto, threads=threads) - - return True + return ops def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, - module_gaps=False, decompress_threads=None): + module_gaps=False, parallel=-1, decompress_threads=None): """Get an array of per-pulse data (image.*) for xtdf detector""" + if decompress_threads is not None: + warn("decompress_threads is deprecated, use parallel= instead", + DeprecationWarning, stacklevel=2) + if parallel == -1: # Don't override an explicit parallel= + # 1 thread meant letting HDF5 decompress the data, and more + # than that meant using threads where the data allowed it. + parallel = 0 if decompress_threads == 1 else -1 + out_shape = self.buffer_shape(module_gaps=module_gaps, roi=roi) if out is None: @@ -1298,14 +1260,6 @@ def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, elif out.shape != out_shape: raise ValueError(f'requires output array of shape {out_shape}') - if roi == () and astype is None: - if decompress_threads is None: - decompress_threads = default_num_threads(fixed_limit=16) - - if decompress_threads > 1: - if self._read_parallel_decompress(out, module_gaps, decompress_threads): - return out - reading_view = out.view() if self._extraneous_dim: reading_view.shape = out.shape[:2] + (1,) + out.shape[2:] @@ -1313,10 +1267,8 @@ def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, # dim in raw data (except AGIPD, where it is data/gain) roi = np.index_exp[:] + roi - for i, (modno, kd) in enumerate(sorted(self.modno_to_keydata.items())): - mod_ix = (modno - self.det._modnos_start_at) if module_gaps else i - for chunk in kd._data_chunks: - self._read_chunk(chunk, reading_view[mod_ix], roi) + self._read(reading_view, self._read_ops(module_gaps, out_shape[1]), + roi, parallel) return out @@ -1332,11 +1284,13 @@ def _wrap_xarray(self, arr, subtrain_index='pulseId'): }) def xarray(self, *, pulses=None, fill_value=None, roi=(), astype=None, - subtrain_index='pulseId', unstack_pulses=False, decompress_threads=None): + subtrain_index='pulseId', unstack_pulses=False, parallel=-1, + decompress_threads=None): arr = self.ndarray( fill_value=fill_value, roi=roi, astype=astype, + parallel=parallel, decompress_threads=decompress_threads, ) out = self._wrap_xarray(arr, subtrain_index) diff --git a/extra_data/compression.py b/extra_data/compression.py index d6723b8a..ba2faeb3 100644 --- a/extra_data/compression.py +++ b/extra_data/compression.py @@ -1,6 +1,4 @@ -import threading from copy import copy -from multiprocessing.pool import ThreadPool import h5py import numpy as np @@ -82,36 +80,3 @@ def dataset_decompressor(dset): return inst return None - - -def multi_dataset_decompressor(dsets): - if not dsets: - return None - - chunk = dsets[0].chunks - dtype = dsets[0] - filters = filter_ids(dsets[0]) - for d in dsets[1:]: - if d.chunks != chunk or d.dtype != dtype or filter_ids(d) != filters: - return None # Datasets are not consistent - - return dataset_decompressor(dsets[0]) - - -def parallel_decompress_chunks(tasks, decompressor_proto, threads=16): - tlocal = threading.local() - - def load_one(dset_id, coord, dest): - try: - decomp = tlocal.decompressor - except AttributeError: - tlocal.decompressor = decomp = decompressor_proto.clone() - - if dset_id.get_chunk_info_by_coord(coord).byte_offset is None: - return # Chunk not allocated in file - - filter_mask, compdata = dset_id.read_direct_chunk(coord) - decomp.apply_filters(compdata, filter_mask, dest) - - with ThreadPool(threads) as pool: - pool.starmap(load_one, tasks) diff --git a/extra_data/direct_read.py b/extra_data/direct_read.py new file mode 100644 index 00000000..0ab6ad17 --- /dev/null +++ b/extra_data/direct_read.py @@ -0,0 +1,638 @@ +import os +import os.path as osp +import sys +from concurrent.futures import ThreadPoolExecutor + +import h5py +import numpy as np + +from .utils import default_num_threads + +# Size to aim for per read of uncompressed data. Anything from 4 to 44 MiB is +# the same speed, but on GPFS a request of 48 MiB or more stops overlapping with +# the other threads, dropping to what a single thread gets whatever the thread +# count (~40 -> ~5 GB/s at 16 threads). 32MiB seems like a nice big number. +SPLIT_BYTES = 32 * 1024 ** 2 + +# Compressed reads are decompression-bound and get slower past ~16 +# threads. Unfiltered reads are flat from 16 when cold and gain ~10% warm up to +# ~48. Capped by core count. +THREADS_FILTERED = 16 +THREADS_UNFILTERED = 48 + + +class UnsupportedDataset(Exception): + """Raised while planning, for a dataset we can't read ourselves.""" + + +def read_debug(msg, *args): + """Say what a read did, if EXTRA_DATA_DIRECT_READ_DEBUG is set. + + Whether a read falls back to h5py is otherwise invisible. The environment + is checked each time, so this can be switched on part way through a session. + """ + if os.environ.get('EXTRA_DATA_DIRECT_READ_DEBUG'): + print('extra_data.direct_read:', msg % args, file=sys.stderr) + + +# What's in the file: --------------------------------------------------------- + +def sel_block(space): + """The solid rectangle a dataspace selects: (starts, ends, shape). + + `ends` are exclusive. Raises UnsupportedDataset if the selection isn't one solid + rectangle. The number of blocks says nothing about that: h5py builds these + selections out of one unit block per element. + """ + shape = space.get_simple_extent_dims() + + if space.get_select_type() == h5py.h5s.SEL_ALL: + return (0,) * len(shape), tuple(shape), tuple(shape) + + if not space.is_regular_hyperslab(): + raise UnsupportedDataset("virtual mapping is not a regular hyperslab") + + starts, strides, counts, blocks = space.get_regular_hyperslab() + ends = [] + for start, stride, count, block in zip(starts, strides, counts, blocks): + if count > 1 and stride != block: + raise UnsupportedDataset("virtual mapping has gaps in it") + ends.append(int(start) + int(count) * int(block)) + + return tuple(int(i) for i in starts), tuple(ends), tuple(shape) + + +def virtual_maps(dset): + """Where a virtual dataset gets its data from. + + Returns a list of ``(start, stop, filename, dataset_path, src_start)``, + sorted by `start`, saying that entries ``[start:stop]`` of this dataset are + entries ``[src_start:src_start + stop - start]`` of another one. + + Only mappings that shift a contiguous run of entries along the first + dimension are understood, which is all EXtra-data itself makes. + """ + dirname = osp.dirname(dset.file.filename) + maps = [] + + for vspace, filename, ds_path, src_space in dset.virtual_sources(): + v_start, v_end, v_shape = sel_block(vspace) + s_start, s_end, s_shape = sel_block(src_space) + + # Dimensions after the first must map straight through, in full. + if v_shape[1:] != s_shape[1:]: + raise UnsupportedDataset("virtual mapping changes the entry shape") + for start, end, length in zip(v_start[1:] + s_start[1:], + v_end[1:] + s_end[1:], + v_shape[1:] + s_shape[1:]): + if start != 0 or end != length: + raise UnsupportedDataset("virtual mapping covers part of an entry") + + if (v_end[0] - v_start[0]) != (s_end[0] - s_start[0]): + raise UnsupportedDataset("virtual mapping is not one-to-one") + + maps.append((v_start[0], v_end[0], osp.join(dirname, filename), + ds_path, s_start[0])) + + maps.sort() + return maps + + +def get_decompressor(dset): + """A prototype decompressor for this dataset, or None if unfiltered.""" + # Counting filters needs only h5py, so uncompressed data can be read + # without the optional decompressors installed. + if dset.id.get_create_plist().get_nfilters() == 0: + return None + + try: + from .compression import dataset_decompressor + except ImportError as e: + raise UnsupportedDataset(f"can't load the decompressors ({e})") + + decompressor = dataset_decompressor(dset) + if decompressor is None: + raise UnsupportedDataset("no fast decompressor for these filters") + + return decompressor + + +class DatasetInfo: + """What we need to know about one real (non-virtual) dataset. + + The chunk table isn't cheap to read, and never changes for a given file, so + this is cached on the FileAccess. + """ + def __init__(self, dset, ds_path): + if dset.dtype.hasobject or dset.dtype.kind not in 'biufc': + raise UnsupportedDataset(f"dtype {dset.dtype} is not a plain number") + if dset.chunks is not None and dset.chunks[1:] != dset.shape[1:]: + # We rely on each chunk holding whole entries, contiguously. + raise UnsupportedDataset("chunks don't cover whole entries") + + self.filename = dset.file.filename + self.key = (self.filename, ds_path) # Identifies it in dicts + self.dtype = dset.dtype + self.shape = dset.shape + self.entry_shape = dset.shape[1:] + self.frame_bytes = self.dtype.itemsize * int( + np.prod(self.entry_shape, dtype=np.intp)) + self.decompressor = get_decompressor(dset) + + # {first entry of chunk: (byte offset, stored size, filter mask)} + if dset.chunks is None: + # Unchunked datasets are treated as a single unfiltered chunk + self.chunk_frames = max(1, dset.shape[0]) + offset = dset.id.get_offset() + if offset is None: + raise UnsupportedDataset("dataset has no data in the file") + self.chunks = {0: (offset, self.chunk_frames * self.frame_bytes, 0)} + else: + self.chunk_frames = dset.chunks[0] + self.chunks = {} + dset.id.chunk_iter(lambda info: self.chunks.__setitem__( + info.chunk_offset[0], + (info.byte_offset, info.size, info.filter_mask) + )) + + self.chunk_nbytes = self.chunk_frames * self.frame_bytes + + def chunk_first(self, entry): + """Where the chunk holding `entry` starts.""" + return (entry // self.chunk_frames) * self.chunk_frames + + def chunk_of(self, entry): + """Where the chunk holding `entry` starts, and what's in the table.""" + first = self.chunk_first(entry) + try: + return first, self.chunks[first] + except KeyError: + # HDF5 would give the fill value here, we don't support that + raise UnsupportedDataset("chunk is not allocated in the file") + + +def band_index(roi_dim0, nrows): + """Which rows of an entry to read for `roi_dim0`, and how to index them. + + Returns ``(first, stop, index)``: rows ``[first:stop]`` of each entry are + read, and indexing those rows with `index` gives what the ROI asked for. + """ + if isinstance(roi_dim0, slice): + start, stop, step = roi_dim0.indices(nrows) + if step > 0 and stop > start: + last = start + ((stop - start - 1) // step) * step + return start, last + 1, slice(0, last - start + 1, step) + + elif isinstance(roi_dim0, (int, np.integer)): + row = int(roi_dim0) + if row < 0: + row += nrows + if 0 <= row < nrows: + return row, row + 1, 0 + + elif isinstance(roi_dim0, (list, np.ndarray)): + index = np.asarray(roi_dim0) + if index.dtype.kind in 'iu' and index.size and index.min() >= 0: + first = int(index.min()) + return first, int(index.max()) + 1, index - first + + # Anything else (boolean mask, empty selection, negative step): read whole + # entries and let numpy index them. + return 0, nrows, roi_dim0 + + +def noop_roi(roi, entry_shape): + """Whether `roi` selects every element of an entry, unchanged.""" + return all(isinstance(r, slice) and r.indices(n) == (0, n, 1) + for r, n in zip(roi, entry_shape)) + + +def buffer(scratch, key, nbytes): + """A byte buffer of exactly `nbytes`, re-used between jobs.""" + buf = scratch.get(key) + if buf is None or buf.size < nbytes: + buf = scratch[key] = np.empty(nbytes, dtype=np.uint8) + + return buf[:nbytes] + + +class DatasetReader: + """A dataset to read from, with the details of one particular read. + + Dataset facts come from a cached :class:`DatasetInfo`; this adds the file + descriptor, whatever depends on the ROI and the output array, and the reads. + + The descriptor is deliberately not cached: files get closed and reopened + elsewhere, and a recycled fd number would read the wrong file rather than + fail. + """ + def __init__(self, info, fd, roi, out_dtype): + self.info = info + self.fd = fd + + if len(roi) > len(info.entry_shape): + raise UnsupportedDataset("the ROI has more dimensions than an entry") + + # Padded to the entry's dimensions, so there's no separate no-ROI case. + # Decompressing gives back whole entries, so that path applies this as + # written; band reads use the band-relative version made below. + self.roi = roi + (np.s_[:],) * (len(info.entry_shape) - len(roi)) + + self._plan_band() + + # With whole entries and no conversion, the file's bytes are exactly + # what belongs in the output array, so we can read straight into it. + self.whole_frame = (self.band_bytes == info.frame_bytes) + self.direct = (self.whole_frame and out_dtype == info.dtype + and noop_roi(self.roi, info.entry_shape)) + + self.read_run = self._read_direct if self.direct else self._read_band + + def _plan_band(self): + """Which rows of an entry the ROI needs, and how to index them. + + The rows it touches are read in full: rows are contiguous within an + entry, columns are not. + """ + info = self.info + + if info.entry_shape: + first, stop, index = band_index(self.roi[0], info.entry_shape[0]) + self.band_shape = (stop - first,) + info.entry_shape[1:] + self.band_index = (index,) + self.roi[1:] + else: + first = 0 + self.band_shape = () + self.band_index = () + + self.band_offset = first * info.dtype.itemsize * int( + np.prod(info.entry_shape[1:], dtype=np.intp)) + self.band_bytes = info.dtype.itemsize * int( + np.prod(self.band_shape, dtype=np.intp)) + + # Doing the reading: ------------------------------------------------------ + + # A job is a (bound method, arguments) pair, called as + # method(scratch, out, *args), and is always a single read. `scratch` holds + # one worker's buffers and decompressors. + + def _pread(self, buf, offset): + got = os.preadv(self.fd, [buf], offset) + if got != buf.nbytes: + raise EOFError(f"read {got} of {buf.nbytes} bytes at {offset} from " + f"{self.info.filename}") + + def _read_direct(self, scratch, out, offset, dest_first, count): + """Copy a run of entries from the file straight into the output array.""" + dest = out[dest_first:dest_first + count] + self._pread(dest.reshape(-1).view(np.uint8), offset) + + def _read_band(self, scratch, out, offset, dest_first, count): + """Read part of each entry, or convert it, on the way into `out`. + + `count` is 1 unless whole entries are being read: only whole entries are + contiguous with each other in the file. + """ + buf = buffer(scratch, 'band', count * self.band_bytes) + self._pread(buf, offset) + + data = buf.view(self.info.dtype).reshape((count,) + self.band_shape) + out[dest_first:dest_first + count] = data[(np.s_[:],) + self.band_index] + + def _decompress(self, scratch, offset, size, filter_mask, dest): + """Read one compressed chunk from the file and unpack it into `dest`.""" + compressed = buffer(scratch, 'compressed', size) + self._pread(compressed, offset) + + # Decompressors hold a buffer, so each worker needs its own copy of the + # prototype made while planning. + key = ('decompressor', self.info.key) + decompressor = scratch.get(key) + if decompressor is None: + decompressor = scratch[key] = self.info.decompressor.clone() + + decompressor.apply_filters(compressed, filter_mask, dest) + + def read_chunk(self, scratch, out, offset, size, filter_mask, segments): + """Read one compressed chunk and unpack the wanted entries out of it. + + `segments` are ``(first entry in the chunk, count, destination index)``. + A chunk is decompressed whole however little of it is wanted, so + everything needed from one chunk is a single job. + """ + chunk = buffer(scratch, 'chunk', self.info.chunk_nbytes) + self._decompress(scratch, offset, size, filter_mask, chunk) + + data = chunk.view(self.info.dtype).reshape( + (self.info.chunk_frames,) + self.info.entry_shape) + + for first, count, dest_first in segments: + # These are whole entries, so the ROI applies as the caller wrote it + out[dest_first:dest_first + count] = \ + data[first:first + count][(np.s_[:],) + self.roi] + + def read_chunk_inplace(self, scratch, out, offset, size, filter_mask, + dest_first): + """Decompress a whole chunk into the output array, with no copy. + + Used where the chunk's entries all belong in one contiguous piece of + `out`, exactly as they are in the file. That is the common case of + reading everything. + """ + dest = out[dest_first:dest_first + self.info.chunk_frames] + self._decompress(scratch, offset, size, filter_mask, + dest.reshape(-1).view(np.uint8)) + + +# Planning: ------------------------------------------------------------------- + +class Planner: + """Turns a list of :class:`ReadOp` into a list of jobs. + + All the HDF5 work (resolving virtual datasets, reading chunk tables) happens + here so that running the jobs doesn't need HDF5. + """ + def __init__(self, out_dtype, roi): + self.out_dtype = out_dtype + self.roi = roi + self.jobs = [] + self.fds = {} # filename -> file descriptor + self._files = {} # filename -> h5py.File we opened and must close + self._borrowed = {} # filename -> FileAccess whose file we may use + self._readers = {} # (filename, dataset path) -> DatasetReader + + # Which entries are wanted from each compressed chunk, turned into jobs + # at the end: separate runs of entries can land in the same chunk, and + # that chunk is only worth reading once. + self._chunk_segments = {} # (reader, first entry of chunk) -> segments + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + self._close_files() + for fd in self.fds.values(): + os.close(fd) + self.fds.clear() + + def _close_files(self): + """Close the HDF5 files, keeping the descriptors we read through.""" + for file in self._files.values(): + file.close() + self._files.clear() + self._borrowed.clear() + + def _dataset(self, filename, ds_path): + """An h5py dataset, through a FileAccess where we have one. + + A borrowed file is fetched each time rather than held onto: the open + file limiter may close it between ops. + """ + file_access = self._borrowed.get(filename) + if file_access is not None: + file = file_access.file + elif filename in self._files: + file = self._files[filename] + else: + file = self._files[filename] = h5py.File(filename, 'r') + + try: + return file[ds_path] + except KeyError: + raise UnsupportedDataset(f"{ds_path} not found in {filename}") + + def _fd(self, filename): + fd = self.fds.get(filename) + if fd is None: + fd = self.fds[filename] = os.open(filename, os.O_RDONLY) + + return fd + + def _reader(self, cache, filename, ds_path): + """The dataset at `ds_path`, which must not be a virtual dataset.""" + reader = self._readers.get((filename, ds_path)) + if reader is not None: + return reader + + info = cache.get(('info', filename, ds_path)) + if info is None: + dset = self._dataset(filename, ds_path) + if dset.is_virtual: + # One level only: overview files point at real files, and we + # never nest them. + raise UnsupportedDataset("virtual dataset points to another one") + + info = DatasetInfo(dset, ds_path) + cache[('info', filename, ds_path)] = info + + reader = DatasetReader(info, self._fd(filename), self.roi, self.out_dtype) + self._readers[(filename, ds_path)] = reader + return reader + + def plan(self, ops): + try: + for op in ops: + self._add_op(op) + + for (reader, chunk_first), segments in self._chunk_segments.items(): + self._add_chunk_job(reader, chunk_first, segments) + finally: + # Nothing from here on needs HDF5, only the file descriptors. + self._close_files() + + return self.jobs + + def all_unfiltered(self): + """True if none of the datasets we planned are compressed.""" + return all(r.info.decompressor is None for r in self._readers.values()) + + def _add_op(self, op): + """Plan one contiguous run of entries from one dataset.""" + # Chunk tables and virtual mappings are cached on the FileAccess, so + # reading the same data again doesn't pay for them again. + cache = op.file._direct_read_cache + filename, ds_path = op.file.filename, op.dataset_path + self._borrowed.setdefault(filename, op.file) + + dset = self._dataset(filename, ds_path) + if not dset.is_virtual: + reader = self._reader(cache, filename, ds_path) + self._add(reader, op.src_first, op.dest_first, op.count) + return + + maps = cache.get(('virtual', ds_path)) + if maps is None: + maps = cache[('virtual', ds_path)] = virtual_maps(dset) + + # Sorted, so walking them in order fills the run from its start; a gap + # means part of it maps to no real data. + cursor, end = op.src_first, op.src_first + op.count + for map_start, map_stop, map_file, map_path, map_src in maps: + if map_stop <= cursor: + continue + if map_start > cursor or map_start >= end: + break + + reader = self._reader(cache, map_file, map_path) + count = min(map_stop, end) - cursor + self._add(reader, map_src + (cursor - map_start), + op.dest_first + (cursor - op.src_first), count) + + cursor += count + if cursor >= end: + return + + raise UnsupportedDataset("virtual dataset doesn't map all the data") + + def _add(self, reader, src_first, dest_first, count): + if count == 0: + return + + if src_first + count > reader.info.shape[0]: + raise UnsupportedDataset("read runs past the end of the dataset") + + if reader.info.decompressor is None: + self._add_unfiltered(reader, src_first, dest_first, count) + else: + self._add_filtered(reader, src_first, dest_first, count) + + def _add_unfiltered(self, reader, src_first, dest_first, count): + """Plan reads of stored-as-is data, which can be split anywhere. + + Entries are contiguous within a chunk, so a chunk needs at most one + read. Long runs are split further, so that a few big chunks can still be + read by many threads at once. + """ + if reader.whole_frame: + # A job covers at most SPLIT_BYTES, which bounds a worker's scratch. + max_per_read = max(1, SPLIT_BYTES // max(1, reader.band_bytes)) + else: + # Only whole entries are contiguous with each other in the file, so + # a job reading part of an entry covers exactly one. + max_per_read = 1 + + cursor, end = src_first, src_first + count + while cursor < end: + chunk_first, (byte_offset, _, filter_mask) = \ + reader.info.chunk_of(cursor) + if filter_mask: + raise UnsupportedDataset("filters skipped on an unfiltered dataset") + + n = min(end, chunk_first + reader.info.chunk_frames) - cursor + offset = (byte_offset + + ((cursor - chunk_first) * reader.info.frame_bytes) + + reader.band_offset) + + for start in range(0, n, max_per_read): + self.jobs.append((reader.read_run, ( + offset + (start * reader.info.frame_bytes), + dest_first + (cursor - src_first) + start, + min(max_per_read, n - start), + ))) + + cursor += n + + def _add_filtered(self, reader, src_first, dest_first, count): + """Note which entries are wanted from each chunk of compressed data.""" + cursor, end = src_first, src_first + count + while cursor < end: + chunk_first = reader.info.chunk_first(cursor) + n = min(end, chunk_first + reader.info.chunk_frames) - cursor + + self._chunk_segments.setdefault((reader, chunk_first), []).append( + (cursor - chunk_first, n, dest_first + (cursor - src_first)) + ) + + cursor += n + + def _add_chunk_job(self, reader, chunk_first, segments): + _, (byte_offset, size, filter_mask) = reader.info.chunk_of(chunk_first) + args = (byte_offset, size, filter_mask) + + # A whole chunk landing unchanged in one contiguous piece of the output + # can be decompressed straight into it. The last chunk may hang over the + # end of the dataset, where not all of it belongs there. + first, count, dest_first = segments[0] + if (len(segments) == 1 and reader.direct + and count == reader.info.chunk_frames + and chunk_first + count <= reader.info.shape[0]): + self.jobs.append((reader.read_chunk_inplace, args + (dest_first,))) + else: + self.jobs.append((reader.read_chunk, args + (segments,))) + + +# Running the plan: ----------------------------------------------------------- + +def run_batch(jobs, out): + scratch = {} + for func, args in jobs: + func(scratch, out, *args) + + +def run(jobs, out, threads): + """Run the jobs, spread over `threads` workers. + """ + n_batches = min(threads, len(jobs)) + if n_batches <= 1: + run_batch(jobs, out) + return + + # Submit batches of jobs to the pool + size = -(-len(jobs) // n_batches) + batches = [jobs[i * size:(i + 1) * size] for i in range(n_batches)] + with ThreadPoolExecutor(n_batches) as pool: + results = [pool.submit(run_batch, batch, out) for batch in batches] + for result in results: + result.result() + + +def read_directly(out, ops, roi, threads): + """Fill `out` from `ops` by reading chunks, or raise UnsupportedDataset.""" + # We write into slices of the output array as raw bytes. + if not out.flags.c_contiguous: + raise UnsupportedDataset("the output array is not C-contiguous") + if out.dtype.hasobject: + raise UnsupportedDataset(f"dtype {out.dtype} is not a plain number") + + with Planner(out.dtype, roi) as planner: + jobs = planner.plan(ops) + + if threads is None: + threads = default_num_threads( + THREADS_UNFILTERED if planner.all_unfiltered() + else THREADS_FILTERED + ) + + read_debug("reading %.2f MB as %d jobs on %d threads", + out.nbytes / 1e6, len(jobs), min(threads, len(jobs))) + run(jobs, out, threads) + + +def read(out, ops, roi=(), parallel=-1): + """Fill `out` from `ops`, taking `roi` of each entry. + + `parallel` is how many threads to read the chunks with. 0 goes through h5py + on this thread instead. -1 reads the chunks directly if these datasets allow + it, and through h5py if they don't. A positive number insists on reading the + chunks directly, raising :exc:`UnsupportedDataset` if that isn't possible. + + Returns True if we read the chunks ourselves, False if h5py did. Set + ``EXTRA_DATA_DIRECT_READ_DEBUG=1`` to have each read report which it was. + """ + if len(ops) == 0: + return True + + if parallel != 0: + try: + read_directly(out, ops, roi, parallel if parallel > 0 else None) + return True + except UnsupportedDataset as e: + if parallel > 0: + raise + + read_debug("h5py: %s", e) + + for op in ops: + op.read_into(out[op.dest_slice], roi) + + return False diff --git a/extra_data/file_access.py b/extra_data/file_access.py index 806a1bc8..f23cab69 100644 --- a/extra_data/file_access.py +++ b/extra_data/file_access.py @@ -202,6 +202,8 @@ def __init__(self, filename, _cache_info=None): self._run_keys_cache = {} # {source: set(keys)} - including incomplete sets self._known_keys = defaultdict(set) + # Chunk tables & virtual dataset mappings, see direct_read.py + self._direct_read_cache = {} @property def file(self): diff --git a/extra_data/keydata.py b/extra_data/keydata.py index 94bcf87c..c9de8c21 100644 --- a/extra_data/keydata.py +++ b/extra_data/keydata.py @@ -3,11 +3,12 @@ import h5py import numpy as np +from . import direct_read from .exceptions import TrainIDError, NoDataError from .file_access import FileAccess from .read_machinery import ( - contiguous_regions, DataChunk, select_train_ids, split_trains, roi_shape, - trains_files_index, + contiguous_regions, DataChunk, ReadOp, select_train_ids, split_trains, + roi_shape, trains_files_index, ) @@ -473,33 +474,56 @@ def as_single_value(self, rtol=1e-5, atol=0.0, reduce_by=None): # Getting data as different kinds of array: ------------------------------- - def ndarray(self, roi=(), out=None): + def buffer_shape(self, roi=()): + """Get the final shape of this data.""" + return self.shape[:1] + roi_shape(self.entry_shape, roi) + + def _read_ops(self): + """The reads needed to load this data, one per contiguous chunk.""" + ops = [] + dest_cursor = 0 + for chunk in self._data_chunks_nonempty: + ops.append(ReadOp(chunk.file, chunk.dataset_path, chunk.first, + dest_cursor, chunk.total_count)) + dest_cursor += chunk.total_count + + return ops + + def ndarray(self, roi=(), out=None, parallel=-1): """Load this data as a numpy array - *roi* may be a ``numpy.s_[]`` expression to load e.g. only part of each - image from a camera. If *out* is not given, a suitable array will be - allocated. + Parameters + ---------- + + roi: numpy.s_[], slice, or tuple of slices + The region of interest. This expression selects data in all + dimensions apart from the first (trains) dimension. If the data + holds a 1D array for each entry, roi=np.s_[:8] would get the first 8 + values from every train. If the data is 2D or more at each entry, + selection looks like roi=np.s_[:8, 5:10] . + out: numpy.ndarray + An array to read the data into, of the shape given by + :meth:`buffer_shape`. If not given, a suitable array is allocated. + An array that isn't C-contiguous is read through HDF5. + parallel: int + How many threads to read the data with, which is much faster for + large amounts of data. The default (-1) uses several threads where + the data allows it, and reads through HDF5 where it doesn't. 0 + always reads through HDF5, on this thread. A positive number + requires the faster path, raising an exception if this data can't be + read that way. """ if not isinstance(roi, tuple): roi = (roi,) - req_shape = self.shape[:1] + roi_shape(self.entry_shape, roi) + req_shape = self.buffer_shape(roi) if out is None: out = np.empty(req_shape, dtype=self.dtype) elif out is not None and out.shape != req_shape: raise ValueError(f'requires output array of shape {req_shape}') - # Read the data from each chunk into the result array - dest_cursor = 0 - for chunk in self._data_chunks_nonempty: - dest_chunk_end = dest_cursor + chunk.total_count - - slices = (chunk.slice,) + roi - chunk.dataset.read_direct( - out[dest_cursor:dest_chunk_end], source_sel=slices - ) - dest_cursor = dest_chunk_end + direct_read.read(out, self._read_ops(), roi, parallel) if out.dtype.hasobject: # Can current only occur for string properties, convert from @@ -558,7 +582,8 @@ def train_index_bounds(self, labelled=False): start[1:] = counts.cumsum()[:-1] return start, start + counts - def xarray(self, extra_dims=None, roi=(), name=None, extra_coords=None): + def xarray(self, extra_dims=None, roi=(), name=None, extra_coords=None, + out=None, parallel=-1): """Load this data as a labelled xarray array or dataset. The first dimension is labelled with train IDs. Other dimensions may be @@ -596,10 +621,14 @@ def xarray(self, extra_dims=None, roi=(), name=None, extra_coords=None): coordinates will match the selected region of interest. If a dict is given, it should map dimension names to coordinate arrays. If True, default coordinate arrays will be generated. + out: numpy.ndarray + An array to read the data into, as for :meth:`ndarray`. + parallel: int + How many threads to read the data with, as for :meth:`ndarray`. """ import xarray - ndarr = self.ndarray(roi=roi) + ndarr = self.ndarray(roi=roi, out=out, parallel=parallel) # Train ID index coords = {'trainId': self.train_id_coordinates()} diff --git a/extra_data/read_machinery.py b/extra_data/read_machinery.py index 625883f5..68590a38 100644 --- a/extra_data/read_machinery.py +++ b/extra_data/read_machinery.py @@ -201,6 +201,32 @@ def trains_files_index(train_ids, files, inc_suspect_trains=True) -> list: tids_files[ix] = file return tids_files +class ReadOp: + """One contiguous run of entries to copy from a dataset into an array. + + This is the unit of work for reading: `src_first` and `count` say where the + entries are in the dataset, `dest_first` where they belong in the + destination array. + """ + def __init__(self, file, dataset_path, src_first, dest_first, count): + self.file = file + self.dataset_path = dataset_path + # Converted from numpy integers: on numpy < 2, uint64 + int gives a + # float, which h5py won't accept as a slice bound. + self.src_first = int(src_first) + self.dest_first = int(dest_first) + self.count = int(count) + + @property + def dest_slice(self): + return slice(self.dest_first, self.dest_first + self.count) + + def read_into(self, out, roi=()): + """Read this run of entries into `out`, which must hold `count` of them.""" + source_sel = (np.s_[self.src_first:self.src_first + self.count],) + roi + self.file.file[self.dataset_path].read_direct(out, source_sel=source_sel) + + class DataChunk: """Reference to a contiguous chunk of data for one or more trains.""" def __init__(self, file, dataset_path, first, train_ids, counts): diff --git a/extra_data/tests/conftest.py b/extra_data/tests/conftest.py index 86fcd379..f6bda192 100644 --- a/extra_data/tests/conftest.py +++ b/extra_data/tests/conftest.py @@ -200,6 +200,12 @@ def mock_modern_spb_proc_run(): make_examples.make_modern_spb_proc_run(td) yield td +@pytest.fixture(scope='session') +def mock_small_agipd_proc_run(): + with TemporaryDirectory() as td: + make_examples.make_small_agipd_proc_run(td) + yield td + @pytest.fixture() def mock_spb_raw_and_modern_proc_run(monkeypatch): with TemporaryDirectory() as td: diff --git a/extra_data/tests/make_examples.py b/extra_data/tests/make_examples.py index d9b9c66d..b3f641bb 100644 --- a/extra_data/tests/make_examples.py +++ b/extra_data/tests/make_examples.py @@ -224,7 +224,7 @@ def make_data_file_bad_device_name(path, format_version='0.5'): """Not all devices have the Karabo standard A/B/C naming convention""" write_file(path, [ BaslerCam('SPB_IRU_SIDEMIC_CAM', sensor_size=(1000, 1000)) - ], ntrains=500, chunksize=50, format_version=format_version) + ], ntrains=10, chunksize=50, format_version=format_version) def make_agipd_file(path, format_version='0.5'): write_file(path, [ @@ -275,11 +275,13 @@ def make_lpd_parallelgain_run(dir_path, raw=True, format_version='0.5'): def make_lpd_run_mini_missed_train(dir_path): write_file(osp.join(dir_path, 'RAW-R0450-LPD00-S00000.h5'), [ - LPDModule('FXE_DET_LPD1M-1/DET/0CH0', frames_per_train=10), + LPDModule('FXE_DET_LPD1M-1/DET/0CH0', frames_per_train=10, + fill_image=True), ], ntrains=5, chunksize=5, format_version='1.0') mod1_f = osp.join(dir_path, 'RAW-R0450-LPD01-S00000.h5') write_file(mod1_f, [ - LPDModule('FXE_DET_LPD1M-1/DET/1CH0', frames_per_train=10), + LPDModule('FXE_DET_LPD1M-1/DET/1CH0', frames_per_train=10, + fill_image=True), ], ntrains=4, chunksize=5, format_version='1.0') # Modify the file for module 1, as if it missed train 10002 @@ -369,6 +371,35 @@ def make_modern_spb_proc_run(dir_path, format_version='1.2'): ds[0, 0, 5] = 1 +def make_small_agipd_proc_run(dir_path, format_version='1.2', nmodules=2, + ntrains=20, frames=8, dims=(32, 16)): + # Holds all three dataset layouts we can read. image.data is chunked and + # uncompressed, image.mask is chunked and gzipped, and image.gain is + # rewritten as a contiguous dataset. + for modno in range(nmodules): + module = AGIPDModule(f'SPB_DET_AGIPD1M-1/DET/{modno}CH0', raw=False, + frames_per_train=frames) + module.image_dims = dims + path = osp.join(dir_path, f'CORR-R0142-AGIPD{modno:0>2}-S00000.h5') + write_file(path, [module], ntrains=ntrains, chunksize=4, + format_version=format_version) + + with h5py.File(path, 'r+') as f: + group = f[f'INSTRUMENT/SPB_DET_AGIPD1M-1/DET/{modno}CH0:xtdf/image'] + # Random values so that reading the wrong bytes doesn't match reading + # the right ones + for key in ['data', 'mask', 'gain']: + dset = group[key] + dset[:] = np.random.uniform(0, 100, dset.shape).astype(dset.dtype) + + gain = group['gain'][:] + del group['gain'] + # Contiguous + group.create_dataset('gain', data=gain) + + return dir_path + + def make_agipd1m_run( dir_path, rep_rate=True, diff --git a/extra_data/tests/mockdata/base.py b/extra_data/tests/mockdata/base.py index bd9f9461..bf0c415e 100644 --- a/extra_data/tests/mockdata/base.py +++ b/extra_data/tests/mockdata/base.py @@ -6,6 +6,28 @@ import numpy as np from packaging import version +def make_dataset(f, path, shape, dtype, maxshape=None, fill=True, nrows=None, + **kwargs): + """Create a dataset and write zeros into the first `nrows` rows. + + A dataset which is never written has no chunks in the file, so reading it + gives fill values without touching the disk and never exercises the direct + reader. The rows past `nrows` are padding so we don't care about them. + """ + if shape[1:] and 'chunks' not in kwargs: + # Choose a chunk size of at most 8MiB to save on disk space + max_chunk_bytes = 8 * 1024 ** 2 + entry_bytes = int(np.prod(shape[1:])) * np.dtype(dtype).itemsize + rows = max(1, min(nrows or shape[0], max_chunk_bytes // entry_bytes)) + kwargs['chunks'] = (rows,) + shape[1:] + + ds = f.create_dataset(path, shape, dtype, maxshape=maxshape, **kwargs) + if fill and ds.size and ds.dtype.kind not in 'SUO': + ds[:nrows] = 0 + + return ds + + class DeviceBase: # Override these in subclasses control_keys = [] @@ -13,6 +35,8 @@ class DeviceBase: output_channels = () instrument_keys = [] + fill_instrument = True + # These are set by write_file ntrains = 400 firsttrain = 10000 @@ -47,8 +71,6 @@ def write_control(self, f): i_count[:] = 0 if self.no_ctrl_data else 1 # CONTROL & RUN - # Creating empty datasets for now. - if self.no_ctrl_data: N = 0 @@ -57,16 +79,16 @@ def write_control(self, f): f.create_group(f'CONTROL/{self.device_id}') for (topic, datatype, dims) in self.control_keys: - f.create_dataset('CONTROL/%s/%s/timestamp' % (self.device_id, topic), - (N,), 'u8', maxshape=(None,)) - f.create_dataset('CONTROL/%s/%s/value' % (self.device_id, topic), - (N,)+dims, datatype, maxshape=((None,)+dims)) + make_dataset(f, 'CONTROL/%s/%s/timestamp' % (self.device_id, topic), + (N,), 'u8', maxshape=(None,)) + make_dataset(f, 'CONTROL/%s/%s/value' % (self.device_id, topic), + (N,)+dims, datatype, maxshape=((None,)+dims)) # RUN is the value at the start of the run - f.create_dataset('RUN/%s/%s/timestamp' % (self.device_id, topic), - (1,), 'u8', maxshape=(None,)) - f.create_dataset('RUN/%s/%s/value' % (self.device_id, topic), - (1,)+dims, datatype, maxshape=((None,)+dims)) + make_dataset(f, 'RUN/%s/%s/timestamp' % (self.device_id, topic), + (1,), 'u8', maxshape=(None,)) + make_dataset(f, 'RUN/%s/%s/value' % (self.device_id, topic), + (1,)+dims, datatype, maxshape=((None,)+dims)) for row in self.extra_run_values: if len(row) == 3: @@ -135,8 +157,9 @@ 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), - (Npad,) + dims, datatype, maxshape=((None,) + dims)) + make_dataset(f, 'INSTRUMENT/%s/%s' % (dev_chan, topic), + (Npad,) + dims, datatype, maxshape=((None,) + dims), + fill=self.fill_instrument, nrows=self.nsamples) def datasource_ids(self): if self.control_keys or self.extra_run_values: diff --git a/extra_data/tests/mockdata/basler_camera.py b/extra_data/tests/mockdata/basler_camera.py index 2d0d2ee5..8082c2e8 100644 --- a/extra_data/tests/mockdata/basler_camera.py +++ b/extra_data/tests/mockdata/basler_camera.py @@ -7,6 +7,8 @@ class BaslerCamera(DeviceBase): Based on example /gpfs/exfel/exp/SPB/201930/p900061/raw/r0055/RAW-R0055-DA01-S00000.h5 """ + fill_instrument = False + def __init__(self, device_id, nsamples=None, sensor_size=None): """Create a dummy basler device that inherits from Device Base""" self.sensor_size = sensor_size or (2058, 2456) diff --git a/extra_data/tests/mockdata/detectors.py b/extra_data/tests/mockdata/detectors.py index b7a881a4..30bc9160 100644 --- a/extra_data/tests/mockdata/detectors.py +++ b/extra_data/tests/mockdata/detectors.py @@ -1,6 +1,8 @@ import numpy as np import h5py +from .base import make_dataset + class DetectorModule: # Overridden in subclasses: image_dims = () @@ -19,7 +21,8 @@ class DetectorModule: ] def __init__(self, device_id, frames_per_train=64, raw=True, - channel_name='xtdf', legacy_name=None): + channel_name='xtdf', legacy_name=None, fill_image=False): + self.fill_image = fill_image self.device_id = device_id self._frames_per_train = frames_per_train if not raw: @@ -143,26 +146,27 @@ def write_instrument(self, f): max_len = None if self.raw else nframes for (key, datatype, dims, kw) in self.image_keys: - if dims == self.image_dims and 'chunks' not in kw: - kw['chunks'] = (1,) + dims - f.create_dataset( - f'INSTRUMENT/{inst_source}/image/{key}', + make_dataset( + f, f'INSTRUMENT/{inst_source}/image/{key}', shape=(nframes,) + dims, dtype=datatype, maxshape=((max_len,) + dims), + fill=self.fill_image, **kw ) # INSTRUMENT (other parts) for part in ['detector', 'header', 'trailer']: - ds = f.create_dataset(f'INSTRUMENT/{inst_source}/{part}/trainId', - (ntrains_pad,), 'u8', maxshape=(None,)) + ds = make_dataset(f, f'INSTRUMENT/{inst_source}/{part}/trainId', + (ntrains_pad,), 'u8', maxshape=(None,), + nrows=self.ntrains) ds[:self.ntrains] = trainids for (key, datatype, dims) in self.other_keys: - f.create_dataset(f'INSTRUMENT/{inst_source}/{key}', - (ntrains_pad,) + dims, datatype, maxshape=((None,) + dims)) + make_dataset(f, f'INSTRUMENT/{inst_source}/{key}', + (ntrains_pad,) + dims, datatype, + maxshape=((None,) + dims), nrows=self.ntrains) if self.legacy_name is not None: # The legacy source name for corrected data is the same as for diff --git a/extra_data/tests/mockdata/gec_camera.py b/extra_data/tests/mockdata/gec_camera.py index 72deb8a9..4ec106af 100644 --- a/extra_data/tests/mockdata/gec_camera.py +++ b/extra_data/tests/mockdata/gec_camera.py @@ -1,6 +1,8 @@ from .base import DeviceBase class GECCamera(DeviceBase): + fill_instrument = False + control_keys = [ ('acquisitionTime', 'f4', ()), ('binningX', 'i4', ()), diff --git a/extra_data/tests/mockdata/imgfel.py b/extra_data/tests/mockdata/imgfel.py index 1e5faa46..6676e42f 100644 --- a/extra_data/tests/mockdata/imgfel.py +++ b/extra_data/tests/mockdata/imgfel.py @@ -3,6 +3,8 @@ class IMGFELCamera(DeviceBase): + fill_instrument = False + control_keys = [ ('Logger/file/maxBackupIndex', 'u4', ()), ('Logger/file/maxFileSize', 'u4', ()), diff --git a/extra_data/tests/mockdata/jungfrau.py b/extra_data/tests/mockdata/jungfrau.py index aebeaa93..00f638db 100644 --- a/extra_data/tests/mockdata/jungfrau.py +++ b/extra_data/tests/mockdata/jungfrau.py @@ -2,6 +2,7 @@ class JUNGFRAUModule(DeviceBase): output_channels = ('daqOutput/data',) + fill_instrument = False def __init__(self, device_id, nsamples=None, raw=False): super().__init__(device_id, nsamples) diff --git a/extra_data/tests/mockdata/sidemic_camera.py b/extra_data/tests/mockdata/sidemic_camera.py index 35864dfc..b9b5685c 100644 --- a/extra_data/tests/mockdata/sidemic_camera.py +++ b/extra_data/tests/mockdata/sidemic_camera.py @@ -1,6 +1,8 @@ from .base import DeviceBase class SidemicCamera(DeviceBase): + fill_instrument = False + # Based on example in /gpfs/exfel/d/raw/SPB/201701/p002012/r0309/RAW-R0309-DA01-S00000.h5 # Technically, only the part before the / is the output channel. diff --git a/extra_data/tests/test_compression.py b/extra_data/tests/test_compression.py index 32e3f7d1..f31071eb 100644 --- a/extra_data/tests/test_compression.py +++ b/extra_data/tests/test_compression.py @@ -5,7 +5,7 @@ from extra_data.compression import ( DeflateDecompressor, ShuffleDeflateDecompressor, - dataset_decompressor, multi_dataset_decompressor, filter_ids, + dataset_decompressor, filter_ids, ) def test_deflate(tmp_path): @@ -39,20 +39,3 @@ def test_shuffle_deflate(tmp_path): decomp.apply_filters(data, filter_mask, out) np.testing.assert_array_equal(out, arr[:, :10]) - -def test_multi_dataset_decompressor(tmp_path): - f = h5py.File(tmp_path / 'test.h5', 'w') - ds1 = f.create_dataset('a', shape=(10, 50), chunks=(1, 50), - compression='gzip', dtype=np.uint32) - ds2 = f.create_dataset('b', shape=(20, 50), chunks=(1, 50), - compression='gzip', dtype=np.uint32) - ds3 = f.create_dataset('c', shape=(10, 50), chunks=(1, 50), - compression='gzip', dtype=np.uint8) - - # Differing shape is OK - assert isinstance(multi_dataset_decompressor([ds1, ds2]), DeflateDecompressor) - - # But dtype needs to match - assert multi_dataset_decompressor([ds1, ds3]) is None - - assert multi_dataset_decompressor([]) is None diff --git a/extra_data/tests/test_direct_read.py b/extra_data/tests/test_direct_read.py new file mode 100644 index 00000000..6db63090 --- /dev/null +++ b/extra_data/tests/test_direct_read.py @@ -0,0 +1,114 @@ +import numpy as np +import pytest + +from extra_data import H5File, RunDirectory, direct_read, voview + +from . import make_examples + +SRC = 'SPB_DET_AGIPD1M-1/DET/0CH0:xtdf' + + +def assert_direct(kd, roi=(), dtype=None, threads=4): + """Read *kd* with the direct reader, and check it against h5py.""" + if not isinstance(roi, tuple): + roi = (roi,) + + expected = kd.ndarray(roi=roi, parallel=0) + out = np.zeros(expected.shape, dtype or expected.dtype) + + assert direct_read.read(out, kd._read_ops(), roi, parallel=threads) + np.testing.assert_array_equal(out, expected.astype(out.dtype)) + + +@pytest.mark.parametrize('key', ['image.data', 'image.mask', 'image.gain']) +def test_layouts(mock_small_agipd_proc_run, key): + # Uncompressed chunked, gzip chunked, and contiguous, respectively + assert_direct(RunDirectory(mock_small_agipd_proc_run)[SRC, key]) + + +@pytest.mark.parametrize('threads', [1, 4, 32]) +def test_thread_counts(mock_small_agipd_proc_run, threads): + assert_direct(RunDirectory(mock_small_agipd_proc_run)[SRC, 'image.data'], threads=threads) + + +@pytest.mark.parametrize('roi', [ + np.s_[8:24], # A band of rows + np.s_[:16, :8], # Part of each row as well + np.s_[::4], # Strided + np.s_[5], # A single row, dropping that dimension + np.s_[[1, 3, 7]], # Arbitrary rows + np.s_[:, 2:6], # Columns only, so whole rows are read +]) +@pytest.mark.parametrize('key', ['image.data', 'image.mask']) +def test_roi(mock_small_agipd_proc_run, key, roi): + assert_direct(RunDirectory(mock_small_agipd_proc_run)[SRC, key], roi=roi) + + +@pytest.mark.parametrize('key', ['image.data', 'image.mask']) +def test_dtype_conversion(mock_small_agipd_proc_run, key): + assert_direct(RunDirectory(mock_small_agipd_proc_run)[SRC, key], dtype=np.float64) + + +@pytest.mark.parametrize('key', ['image.data', 'image.mask']) +def test_train_selection(mock_small_agipd_proc_run, key): + # Scattered trains give many separate reads; neighbouring ones can share an + # HDF5 chunk, which should only be decompressed once. + kd = RunDirectory(mock_small_agipd_proc_run)[SRC, key].select_trains(np.s_[::3]) + assert_direct(kd) + assert_direct(kd, roi=np.s_[:16]) + + +def test_ndarray_uses_it(mock_small_agipd_proc_run): + kd = RunDirectory(mock_small_agipd_proc_run)[SRC, 'image.mask'] + np.testing.assert_array_equal(kd.ndarray(parallel=4), + kd.ndarray(parallel=0)) + + +def test_virtual_overview(tmp_path): + # Its own run, because this test closes the files + make_examples.make_small_agipd_proc_run(tmp_path) + overview = tmp_path / 'overview.h5' + run = RunDirectory(tmp_path) + voview.VirtualOverviewFileWriter(overview, run).write() + + # HDF5 won't resolve a VDS mapping onto a file this process has open, so + # the parallel=0 reads in assert_direct need these closed. The direct + # reader resolves mappings itself. + for file in run.files: + file.close() + del run + + kd = H5File(overview)[SRC, 'image.data'] + assert kd.files[0].file[kd.hdf5_data_path].is_virtual + assert_direct(kd) + assert_direct(kd, roi=np.s_[:16, :8]) + assert_direct(kd.select_trains(np.s_[::3])) + + +def test_falls_back_without_data(mock_small_agipd_proc_run): + # A dataset with no data in it has no chunks to read, so h5py handles it + kd = RunDirectory(mock_small_agipd_proc_run)[SRC, 'image.length'] + out = np.zeros(kd.shape, kd.dtype) + + assert not direct_read.read(out, kd._read_ops()) + np.testing.assert_array_equal(out, kd.ndarray(parallel=0)) + + # Requiring the direct reader throws + with pytest.raises(direct_read.UnsupportedDataset, match='not allocated'): + kd.ndarray(parallel=4) + + +@pytest.mark.parametrize('kwargs', [ + dict(), + dict(roi=np.s_[:16, :8]), + dict(module_gaps=True), + dict(astype=np.float64), +]) +def test_multimod_detector(mock_small_agipd_proc_run, kwargs): + # Detector data is read for all of the modules in one pass + from extra_data.components import AGIPD1M + + det = AGIPD1M(RunDirectory(mock_small_agipd_proc_run), raw=False) + for kd in [det['image.data'], det['image.data'].select_pulses(np.s_[::2])]: + np.testing.assert_array_equal(kd.ndarray(parallel=4, **kwargs), + kd.ndarray(parallel=0, **kwargs)) diff --git a/extra_data/tests/test_keydata.py b/extra_data/tests/test_keydata.py index 52727727..72dc8797 100644 --- a/extra_data/tests/test_keydata.py +++ b/extra_data/tests/test_keydata.py @@ -6,7 +6,9 @@ import h5py -from extra_data import RunDirectory, H5File, open_run +from extra_data import ( + RunDirectory, H5File, open_run, by_index, direct_read, +) from extra_data.keydata import expand_indexing from extra_data.exceptions import TrainIDError, NoDataError from . import make_examples @@ -375,6 +377,44 @@ def test_ndarray_out(mock_spb_raw_run): assert buf_in is buf_out +def test_ndarray_parallel(mock_spb_raw_run): + run = RunDirectory(mock_spb_raw_run).select_trains(by_index[:8]) + am0 = run['SPB_DET_AGIPD1M-1/DET/0CH0:xtdf', 'image.data'] + serial = am0.ndarray(parallel=0) + + np.testing.assert_array_equal(am0.ndarray(), serial) + + # Reading into an array the caller supplied + buf = np.zeros(am0.shape, dtype=am0.dtype) + assert am0.ndarray(out=buf) is buf + np.testing.assert_array_equal(buf, serial) + + roi = np.s_[0, :16, :8] + np.testing.assert_array_equal(am0.ndarray(roi=roi), + am0.ndarray(roi=roi, parallel=0)) + + +def test_xarray_parallel(mock_spb_raw_run): + run = RunDirectory(mock_spb_raw_run) + xgm = run['SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos.value'] + + xr.testing.assert_identical(xgm.xarray(), xgm.xarray(parallel=0)) + + +def test_parallel_fallback(mock_spb_raw_run): + run = RunDirectory(mock_spb_raw_run) + state = run['SPB_XTD9_XGM/DOOCS/MAIN', 'state'] + + # Strings fall back to h5py + assert state.dtype.hasobject + data = state.ndarray() + assert (data[3:8] == ['OFF', 'OFF', 'ON', 'ON', 'ON']).all() + + # But requiring the direct reader throws an exception with the reason + with pytest.raises(direct_read.UnsupportedDataset, match='not a plain number'): + state.ndarray(parallel=4) + + def test_string_arrays(mock_spb_raw_run): f = RunDirectory(mock_spb_raw_run) state = f['SPB_XTD9_XGM/DOOCS/MAIN', 'state']