# coding: utf-8 from __future__ import annotations import os import h5py import h5py._hl.selections as selection import numpy from h5py import h5s as h5py_h5s from silx.io.url import DataUrl from silx.io.utils import get_data, h5py_read_dataset from tomoscan.esrf.scan.utils import cwd_context from tomoscan.io import HDF5File from silx.io.utils import open as open_hdf5 from nxtomo.io import to_target_rel_path from nxtomomill.utils.h5pyutils import from_data_url_to_virtual_source from nxtomomill.utils.hdf5 import DatasetReader __all__ = [ "FrameAppender", ] class FrameAppender: """ Class to insert 2D frame(s) to an existing dataset """ def __init__( self, data: numpy.ndarray | DataUrl, file_path, data_path, where, logger=None, ): if where not in ("start", "end"): raise ValueError("`where` should be `start` or `end`") if not isinstance( data, (DataUrl, numpy.ndarray, list, tuple, h5py.VirtualSource) ): raise TypeError( f"data should be an instance of DataUrl or a numpy array not {type(data)}" ) self.data = data self.file_path = os.path.abspath(file_path) self.data_path = data_path self.where = where self.logger = logger def process(self) -> None: """ main function. Will start the insertion of frame(s) """ with HDF5File(self.file_path, mode="a") as h5s: if self.data_path in h5s: self._add_to_existing_dataset(h5s) else: self._create_new_dataset(h5s) if self.logger: self.logger.info(f"data added to {self.data_path}@{self.file_path}") def _add_to_existing_virtual_dataset(self, h5s): if ( h5py.version.hdf5_version_tuple[0] <= 1 and h5py.version.hdf5_version_tuple[1] < 12 ): if self.logger: self.logger.warning( "You are working on virtual dataset" "with a hdf5 version < 12. Frame " "you want to change might be " "modified depending on the working " "directory without notifying." "See https://github.com/silx-kit/silx/issues/3277" ) if isinstance(self.data, h5py.VirtualSource): self.__insert_virtual_source_in_vds(h5s=h5s, new_virtual_source=self.data) elif isinstance(self.data, DataUrl): if self.logger is not None: self.logger.debug( f"Update virtual dataset: {self.data_path}@{self.file_path}" ) # store DataUrl in the current virtual dataset url = self.data def check_dataset(dataset_frm_url): data_need_reshape = False """check if the dataset is valid or might need a reshape""" if dataset_frm_url.ndim not in (2, 3): raise ValueError(f"{url.path()} should point to 2D or 3D dataset ") if dataset_frm_url.ndim == 2: new_shape = 1, dataset_frm_url.shape[0], dataset_frm_url.shape[1] if self.logger is not None: self.logger.info( f"reshape provided data to 3D (from {dataset_frm_url.shape} to {new_shape})" ) data_need_reshape = True return data_need_reshape loaded_dataset = None if url.data_slice() is None: # case we can avoid to load the data in memory with DatasetReader(url) as data_frm_url: data_need_reshape = check_dataset(data_frm_url) # FIXME: avoid keeping some file open. not clear why this is needed data_frm_url = None else: data_frm_url = get_data(url) data_need_reshape = check_dataset(data_frm_url) loaded_dataset = data_frm_url if url.data_slice() is None and not data_need_reshape: # case we can avoid to load the data in memory with DatasetReader(self.data) as data_frm_url: self.__insert_url_in_vds(h5s, url, data_frm_url) # FIXME: avoid keeping some file open. not clear why this is needed data_frm_url = None else: if loaded_dataset is None: data_frm_url = get_data(url) else: data_frm_url = loaded_dataset self.__insert_url_in_vds(h5s, url, data_frm_url) else: raise TypeError( "Provided data is a numpy array when given" "dataset path is a virtual dataset. " "You must store the data somewhere else " "and provide a DataUrl" ) def __insert_url_in_vds(self, h5s, url, data_frm_url): if data_frm_url.ndim == 2: dim_2, dim_1 = data_frm_url.shape data_frm_url = data_frm_url.reshape(1, dim_2, dim_1) elif data_frm_url.ndim == 3: _, dim_2, dim_1 = data_frm_url.shape else: raise ValueError("data to had is expected to be 2 or 3 d") new_virtual_source = h5py.VirtualSource( path_or_dataset=url.file_path(), name=url.data_path(), shape=data_frm_url.shape, ) if url.data_slice() is not None: # in the case we have to process to a FancySelection with open_hdf5(os.path.abspath(url.file_path())) as h5sd: dst = h5sd[url.data_path()] sel = selection.select( h5sd[url.data_path()].shape, url.data_slice(), dst ) new_virtual_source.sel = sel self.__insert_virtual_source_in_vds( h5s=h5s, new_virtual_source=new_virtual_source, relative_path=True ) def __insert_virtual_source_in_vds( self, h5s, new_virtual_source: h5py.VirtualSource, relative_path=True ): if not isinstance(new_virtual_source, h5py.VirtualSource): raise TypeError( f"{new_virtual_source} is expected to be an instance of h5py.VirtualSource and not {type(new_virtual_source)}" ) if not len(new_virtual_source.shape) == 3: raise ValueError( f"virtual source shape is expected to be 3D and not {len(new_virtual_source.shape)}D." ) # preprocess virtualSource to insure having a relative path if relative_path: vds_file_path = to_target_rel_path(new_virtual_source.path, self.file_path) new_virtual_source_sel = new_virtual_source.sel new_virtual_source = h5py.VirtualSource( path_or_dataset=vds_file_path, name=new_virtual_source.name, shape=new_virtual_source.shape, dtype=new_virtual_source.dtype, ) new_virtual_source.sel = new_virtual_source_sel virtual_sources_len = [] virtual_sources = [] # we need to recreate the VirtualSource they are not # store or available from the API for vs_info in h5s[self.data_path].virtual_sources(): length, vs = self._recreate_vs(vs_info=vs_info, vds_file=self.file_path) virtual_sources.append(vs) virtual_sources_len.append(length) n_frames = h5s[self.data_path].shape[0] + new_virtual_source.shape[0] data_type = h5s[self.data_path].dtype if self.where == "start": virtual_sources.insert(0, new_virtual_source) virtual_sources_len.insert(0, new_virtual_source.shape[0]) else: virtual_sources.append(new_virtual_source) virtual_sources_len.append(new_virtual_source.shape[0]) # create the new virtual dataset layout = h5py.VirtualLayout( shape=( n_frames, new_virtual_source.shape[-2], new_virtual_source.shape[-1], ), dtype=data_type, ) last = 0 for v_source, vs_len in zip(virtual_sources, virtual_sources_len): layout[last : vs_len + last] = v_source last += vs_len if self.data_path in h5s: del h5s[self.data_path] h5s.create_virtual_dataset(self.data_path, layout) def _add_to_existing_none_virtual_dataset(self, h5s): """ for now when we want to add data *to a none virtual dataset* we always duplicate data if provided from a DataUrl. We could create a virtual dataset as well but seems to complicated for a use case that we don't really have at the moment. :param h5s: """ if self.logger is not None: self.logger.debug("Update dataset: {entry}@{file_path}") if isinstance(self.data, (numpy.ndarray, list, tuple)): new_data = self.data else: url = self.data new_data = get_data(url) if isinstance(new_data, numpy.ndarray): if not new_data.shape[1:] == h5s[self.data_path].shape[1:]: raise ValueError( f"Data shapes are incoherent: {new_data.shape} vs {h5s[self.data_path].shape}" ) new_shape = ( new_data.shape[0] + h5s[self.data_path].shape[0], new_data.shape[1], new_data.shape[2], ) data_to_store = numpy.empty(new_shape) if self.where == "start": data_to_store[: new_data.shape[0]] = new_data data_to_store[new_data.shape[0] :] = h5py_read_dataset( h5s[self.data_path] ) else: data_to_store[: h5s[self.data_path].shape[0]] = h5py_read_dataset( h5s[self.data_path] ) data_to_store[h5s[self.data_path].shape[0] :] = new_data else: assert isinstance( self.data, (list, tuple) ), f"Unmanaged data type {type(self.data)}" o_data = h5s[self.data_path] o_data = list(h5py_read_dataset(o_data)) if self.where == "start": new_data.extend(o_data) data_to_store = numpy.asarray(new_data) else: o_data.extend(new_data) data_to_store = numpy.asarray(o_data) del h5s[self.data_path] h5s[self.data_path] = data_to_store def _add_to_existing_dataset(self, h5s): """Add the frame to an existing dataset""" if h5s[self.data_path].is_virtual: self._add_to_existing_virtual_dataset(h5s=h5s) else: self._add_to_existing_none_virtual_dataset(h5s=h5s) def _create_new_dataset(self, h5s): """ needs to create a new dataset. In this case the policy is: - if a DataUrl is provided then we create a virtual dataset - if a numpy array is provided then we create a 'standard' dataset """ if isinstance(self.data, DataUrl): url = self.data url_file_path = to_target_rel_path(url.file_path(), self.file_path) url = DataUrl( file_path=url_file_path, data_path=url.data_path(), scheme=url.scheme(), data_slice=url.data_slice(), ) with cwd_context(os.path.dirname(self.file_path)): vs, vs_shape, data_type = from_data_url_to_virtual_source(url) layout = h5py.VirtualLayout(shape=vs_shape, dtype=data_type) layout[:] = vs h5s.create_virtual_dataset(self.data_path, layout) elif isinstance(self.data, h5py.VirtualSource): virtual_source = self.data layout = h5py.VirtualLayout( shape=virtual_source.shape, dtype=virtual_source.dtype, ) vds_file_path = to_target_rel_path(virtual_source.path, self.file_path) virtual_source_rel_path = h5py.VirtualSource( path_or_dataset=vds_file_path, name=virtual_source.name, shape=virtual_source.shape, dtype=virtual_source.dtype, ) virtual_source_rel_path.sel = virtual_source.sel layout[:] = virtual_source_rel_path # convert path to relative h5s.create_virtual_dataset(self.data_path, layout) elif not isinstance(self.data, numpy.ndarray): raise TypeError( f"self.data should be an instance of DataUrl, a numpy array or a VirtualSource. Not {type(self.data)}" ) else: h5s[self.data_path] = self.data @staticmethod def _recreate_vs(vs_info, vds_file): """Simple util to retrieve a h5py.VirtualSource from virtual source information. to understand clearly this function you might first have a look at the use case exposed in issue: https://gitlab.esrf.fr/tomotools/nxtomomill/-/issues/40 """ with cwd_context(os.path.dirname(vds_file)): dataset_file_path = vs_info.file_name # in case the virtual source is in the same file if dataset_file_path == ".": dataset_file_path = vds_file with open_hdf5(dataset_file_path) as vs_node: dataset = vs_node[vs_info.dset_name] select_bounds = vs_info.vspace.get_select_bounds() left_bound = select_bounds[0] right_bound = select_bounds[1] length = right_bound[0] - left_bound[0] + 1 # warning: for now step is not managed with virtual # dataset virtual_source = h5py.VirtualSource( vs_info.file_name, vs_info.dset_name, shape=dataset.shape, ) # here we could provide dataset but we won't to # insure file path will be relative. type_code = vs_info.src_space.get_select_type() # check for unlimited selections in case where selection is regular # hyperslab, which is the only allowed case for h5s.UNLIMITED to be # in the selection if ( type_code == h5py_h5s.SEL_HYPERSLABS and vs_info.src_space.is_regular_hyperslab() ): ( source_start, stride, count, block, ) = vs_info.src_space.get_regular_hyperslab() source_end = source_start[0] + length sel = selection.select( dataset.shape, slice(source_start[0], source_end), dataset=dataset, ) virtual_source.sel = sel return ( length, virtual_source, )