diff --git a/src/nexgen/beamlines/I19_2_nxs.py b/src/nexgen/beamlines/I19_2_nxs.py index 21868501..1c3be01b 100644 --- a/src/nexgen/beamlines/I19_2_nxs.py +++ b/src/nexgen/beamlines/I19_2_nxs.py @@ -15,7 +15,7 @@ from pydantic import field_validator from nexgen.nxs_utils.detector import EigerStreamFormat -from nexgen.tools.vds_w_tools import define_vds_dtype_from_bit_depth +from nexgen.tools.vds_tools import define_vds_dtype_from_bit_depth from nexgen.utils import get_iso_timestamp from .. import log diff --git a/src/nexgen/beamlines/SSX_Eiger_nxs.py b/src/nexgen/beamlines/SSX_Eiger_nxs.py index 745a8d92..0fd3d477 100644 --- a/src/nexgen/beamlines/SSX_Eiger_nxs.py +++ b/src/nexgen/beamlines/SSX_Eiger_nxs.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import Literal, get_args -from nexgen.tools.vds_w_tools import define_vds_dtype_from_bit_depth +from nexgen.tools.vds_tools import define_vds_dtype_from_bit_depth from .. import log from ..nxs_utils import ( diff --git a/src/nexgen/beamlines/i19_2/eiger.py b/src/nexgen/beamlines/i19_2/eiger.py index 82bdc332..c509be35 100644 --- a/src/nexgen/beamlines/i19_2/eiger.py +++ b/src/nexgen/beamlines/i19_2/eiger.py @@ -22,8 +22,8 @@ from nexgen.tools.vds_tools import ( VdsMapping, VdsSettings, + define_vds_dtype_from_bit_depth, ) -from nexgen.tools.vds_w_tools import define_vds_dtype_from_bit_depth logger = logging.getLogger("nexgen.beamlines.I19_2.eiger") @@ -291,7 +291,8 @@ def start_writer( vds_offset=vds_settings.vds_offset, vds_shape=vds_settings.vds_shape, vds_dtype=vds_settings.vds_dtype, - ) # TODO add mapping + vds_mapping=vds_settings.vds_mapping, + ) if parameters.timestamps[1]: NXmx_writer.update_timestamps(parameters.timestamps[1], "end_time") if notes: diff --git a/src/nexgen/command_line/nexus_generator.py b/src/nexgen/command_line/nexus_generator.py index 76b167f3..41bc725f 100644 --- a/src/nexgen/command_line/nexus_generator.py +++ b/src/nexgen/command_line/nexus_generator.py @@ -28,7 +28,7 @@ from nexgen.nxs_write.nxmx_writer import EventNXmxFileWriter, NXmxFileWriter from nexgen.nxs_write.write_utils import find_number_of_images from nexgen.tools.data_writer import generate_event_files, generate_image_files -from nexgen.tools.vds_w_tools import define_vds_dtype_from_bit_depth +from nexgen.tools.vds_tools import define_vds_dtype_from_bit_depth from nexgen.utils import ( get_filename_template, get_iso_timestamp, diff --git a/src/nexgen/nxs_write/ed_nxmx_writer.py b/src/nexgen/nxs_write/ed_nxmx_writer.py index 58fe2e2f..ba6d3eb8 100644 --- a/src/nexgen/nxs_write/ed_nxmx_writer.py +++ b/src/nexgen/nxs_write/ed_nxmx_writer.py @@ -14,7 +14,8 @@ from numpy.typing import DTypeLike from ..nxs_utils import Attenuator, Beam, Detector, Goniometer, Source -from ..tools.vds_w_tools import image_vds_writer, vds_file_writer +from ..tools.vds_tools.blocked_mapping import image_vds_writer +from ..tools.vds_tools.vds_file import vds_file_writer from ..utils import coord2mcstas from .nxclass_writers import ( write_NXcoordinate_system_set, diff --git a/src/nexgen/nxs_write/nxmx_writer.py b/src/nexgen/nxs_write/nxmx_writer.py index 7a22c2c0..e0d1ac55 100644 --- a/src/nexgen/nxs_write/nxmx_writer.py +++ b/src/nexgen/nxs_write/nxmx_writer.py @@ -2,8 +2,6 @@ A writer for NXmx format NeXus Files. """ -from __future__ import annotations - import logging import math from datetime import datetime @@ -13,15 +11,16 @@ import numpy as np from numpy.typing import DTypeLike +from nexgen.tools.vds_tools.strided_mapping import write_strided_vds +from nexgen.tools.vds_tools.utils import VdsMapping + from ..nxs_utils.detector import Detector from ..nxs_utils.goniometer import Goniometer from ..nxs_utils.sample import Sample from ..nxs_utils.source import Attenuator, Beam, Source -from ..tools.vds_w_tools import ( - clean_unused_links, - image_vds_writer, - jungfrau_vds_writer, -) +from ..tools.vds_tools import clean_unused_links +from ..tools.vds_tools.blocked_mapping import image_vds_writer +from ..tools.vds_tools.tiled_mapping import jungfrau_vds_writer from ..utils import ( MAX_FRAMES_PER_DATASET, MAX_SUFFIX_DIGITS, @@ -245,69 +244,98 @@ def write_vds( vds_offset: int = 0, vds_shape: tuple[int, int, int] = None, vds_dtype: DTypeLike = np.uint16, + vds_mapping: VdsMapping = VdsMapping.BLOCKED, + stride: int = 2, clean_up: bool = False, ): """Write a Virtual Dataset. - This method adds a VDS under /entry/data/data in the NeXus file, linking to either the full datasets or the subset defined by \ + This method adds a VDS under /entry/data/data in the NeXus file, linking to either the full datasets or the subset defined by vds_offset (used as start index) and vds_shape. WARNING. Only use clean up if the data collection is finished and all the files have already been written. Args: - vds_offset (int, optional): Start index for the vds writer. Defaults to 0. - vds_shape (tuple[int,int,int], optional): Shape of the data which will be linked in the VDS. If not passed, it will be defined as \ + vds_offset (int, optional): Start index for the vds writer. For the strided mapping this will be where the source starts. + Defaults to 0. + vds_shape (tuple[int,int,int], optional): Shape of the data which will be linked in the VDS. If not passed, it will be defined as (tot_num_imgs - start_idx, *image_size). Defaults to None. vds_dtype (DTypeLike, optional): The type of the input data. Defaults to np.uint16. + vds_mapping (VdsMapping, optional): Mapping to apply to build the vds. Defaults to "blocked". + stride (int, optional): Every how many images to map the strided vds. Defaults to 2 (map to every other image). clean_up(bool, optional): Clean up unused links in vds. Defaults to False. """ - if not vds_shape: - vds_shape = ( - self.tot_num_imgs - vds_offset, - *self.detector.detector_params.image_size, - ) - if self.goniometer.get_number_of_scan_points() != vds_shape[0]: - vds_shape = ( - self.goniometer.get_number_of_scan_points(), - *vds_shape[1:], - ) - nxmx_logger.warning( - "The number of scan points doesn't match the calculated vds_shape. \ - Resetting it to match the number of frames indicated by the scan." - ) - - nxmx_logger.debug(f"VDS shape set to {vds_shape}.") - - with h5py.File(self.filename, "r+") as nxs: - # For a coming ticket - for now write a separate file - # Here will be better to have a match-case for VDS mapping. Default is the same as blocked. - if "jungfrau" in self.detector.detector_params.description.lower(): + if "jungfrau" in self.detector.detector_params.description.lower(): + with h5py.File(self.filename, "r+") as nxs: jungfrau_vds_writer( nxs, (self.tot_num_imgs, *self.detector.detector_params.image_size), data_type=vds_dtype, ) - else: - image_vds_writer( - nxs, - (self.tot_num_imgs, *self.detector.detector_params.image_size), - start_index=vds_offset, - vds_shape=vds_shape, - data_type=vds_dtype, + nxmx_logger.info("VDS for JF1M written.") + return + + match vds_mapping: + case VdsMapping.BLOCKED: + nxmx_logger.debug("Writing 'standard' blocked VDS") + if not vds_shape: + vds_shape = ( + self.tot_num_imgs - vds_offset, + *self.detector.detector_params.image_size, + ) + if self.goniometer.get_number_of_scan_points() != vds_shape[0]: + vds_shape = ( + self.goniometer.get_number_of_scan_points(), + *vds_shape[1:], + ) + nxmx_logger.warning( + "The number of scan points doesn't match the calculated vds_shape. \ + Resetting it to match the number of frames indicated by the scan." + ) + + nxmx_logger.debug(f"VDS shape set to {vds_shape}.") + + with h5py.File(self.filename, "r+") as nxs: + image_vds_writer( + nxs, + (self.tot_num_imgs, *self.detector.detector_params.image_size), + start_index=vds_offset, + vds_shape=vds_shape, + data_type=vds_dtype, + ) + + if clean_up is True: + nxmx_logger.warning("Starting clean up of unused links.") + clean_unused_links( + nxs, + vds_shape=vds_shape, + start_index=vds_offset, + ) + + # If number of frames in the VDS is lower than the total, nimages in NXcollection should be overwritten to match this + if vds_shape[0] < self.tot_num_imgs: + del nxs["/entry/instrument/detector/detectorSpecific/nimages"] + nxs[ + "/entry/instrument/detector/detectorSpecific" + ].create_dataset("nimages", data=vds_shape[0]) + nxmx_logger.info(f"VDS correctly written to file {self.filename}") + case VdsMapping.STRIDED: + nxmx_logger.info( + f"Writing strided nexus file with start idx {vds_offset} and stride {stride}." ) - if clean_up is True: - nxmx_logger.warning("Starting clean up of unused links.") - clean_unused_links( - nxs, - vds_shape=vds_shape, - start_index=vds_offset, - ) - - # If number of frames in the VDS is lower than the total, nimages in NXcollection should be overwritten to match this - if vds_shape[0] < self.tot_num_imgs: - del nxs["/entry/instrument/detector/detectorSpecific/nimages"] - nxs["/entry/instrument/detector/detectorSpecific"].create_dataset( - "nimages", data=vds_shape[0] + with h5py.File(self.filename, "r+") as nxs: + write_strided_vds( + nxs, + (self.tot_num_imgs, *self.detector.detector_params.image_size), + vds_offset, + stride, + vds_dtype, + ) + nxmx_logger.info(f"VDS correctly written to file {self.filename}") + case VdsMapping.TILED: + nxmx_logger.error( + "Tiled VDS mapping not implemented yet, will not write vds." ) + return class EventNXmxFileWriter(NXmxFileWriter): diff --git a/src/nexgen/tools/vds_tools/__init__.py b/src/nexgen/tools/vds_tools/__init__.py index 9a0fd0fa..cb3ec0ae 100644 --- a/src/nexgen/tools/vds_tools/__init__.py +++ b/src/nexgen/tools/vds_tools/__init__.py @@ -1,3 +1,15 @@ -from nexgen.tools.vds_tools.utils import VdsMapping, VdsSettings, find_datasets_in_file +from nexgen.tools.vds_tools.utils import ( + VdsMapping, + VdsSettings, + clean_unused_links, + define_vds_dtype_from_bit_depth, + find_datasets_in_file, +) -__all__ = ["VdsMapping", "VdsSettings", "find_datasets_in_file"] +__all__ = [ + "VdsMapping", + "VdsSettings", + "find_datasets_in_file", + "define_vds_dtype_from_bit_depth", + "clean_unused_links", +] diff --git a/src/nexgen/tools/vds_w_tools.py b/src/nexgen/tools/vds_tools/blocked_mapping.py similarity index 56% rename from src/nexgen/tools/vds_w_tools.py rename to src/nexgen/tools/vds_tools/blocked_mapping.py index ae432212..d22be08b 100644 --- a/src/nexgen/tools/vds_w_tools.py +++ b/src/nexgen/tools/vds_tools/blocked_mapping.py @@ -1,13 +1,6 @@ -""" -Tools to write Virtual DataSets -""" - -from __future__ import annotations - import logging import operator from functools import reduce -from pathlib import Path from typing import Sequence import h5py @@ -15,22 +8,10 @@ from numpy.typing import DTypeLike from pydantic.dataclasses import dataclass -from nexgen.tools.vds_tools import find_datasets_in_file - -from ..utils import MAX_FRAMES_PER_DATASET -from .constants import jungfrau_fill_value, jungfrau_gap_size, jungfrau_mod_size +from nexgen.tools.vds_tools.utils import find_datasets_in_file +from nexgen.utils import MAX_FRAMES_PER_DATASET -vds_logger = logging.getLogger("nexgen.VDSWriter") - - -def define_vds_dtype_from_bit_depth(bit_depth: int) -> DTypeLike: - """Define dtype of VDS based on the passed bit depth.""" - if bit_depth == 32: - return np.uint32 - elif bit_depth == 8: - return np.uint8 - else: - return np.uint16 +block_vds_logger = logging.getLogger("nexgen.tools.vds_tools.block_mapping") @dataclass @@ -77,7 +58,7 @@ def __add__(self, x): def split_datasets( - dsets, + dsets: list[str], data_shape: tuple[int, int, int], start_idx: int = 0, vds_shape: tuple[int, int, int] = None, @@ -108,16 +89,20 @@ def split_datasets( raise ValueError("Start index must be positive") if not isinstance(data_shape[0], int): - vds_logger.warning("Datashape not passed as int, will attempt to cast") + block_vds_logger.warning("Datashape not passed as int, will attempt to cast") if not isinstance(start_idx, int): - vds_logger.warning("VDS start index not passed as int, will attempt to cast") + block_vds_logger.warning( + "VDS start index not passed as int, will attempt to cast" + ) if vds_shape and not isinstance(vds_shape[0], int): - vds_logger.warning("VDS start index not passed as int, will attempt to cast") + block_vds_logger.warning( + "VDS start index not passed as int, will attempt to cast" + ) if vds_shape is None: - vds_logger.debug( + block_vds_logger.debug( "VDS shape not chosen, it will be calculated from the full data shape and the chosen start index." ) vds_shape = (data_shape[0] - start_idx, *data_shape[1:]) @@ -198,7 +183,7 @@ def image_vds_writer( data_type (DTypeLike, optional): The type of the input data. Defaults to np.uint16. entry_key (str, optional): Entry key for the Virtual DataSet name. Defaults to data. """ - vds_logger.debug("Start creating VDS ...") + block_vds_logger.debug("Start creating VDS ...") # Where the vds will go nxdata = nxsfile["/entry/data"] dset_names = find_datasets_in_file(nxdata) @@ -228,123 +213,4 @@ def image_vds_writer( # Writea Virtual Dataset in NeXus file nxdata.create_virtual_dataset(entry_key, layout, fillvalue=-1) - vds_logger.debug("VDS correctly written to NeXus file.") - - -def jungfrau_vds_writer( - nxsfile: h5py.File, - vds_shape: tuple | list, - data_type: DTypeLike = np.uint16, - source_dsets: list[str] | None = None, -): - """Write VDS for Jungfrau 1M use case, with a tiled layout.""" - external_dsets = True - entry_key = "data" - frames = vds_shape[0] - - nxdata = nxsfile["/entry/data"] - if not source_dsets: - source_dsets = find_datasets_in_file(nxdata) - external_dsets = False - - sources = [] - for dset in source_dsets: - source_path = dset if external_dsets is True else "." - source_name = entry_key if external_dsets is True else f"/entry/data/{dset}" - source = h5py.VirtualSource( - source_path, source_name, shape=(frames, *jungfrau_mod_size) - ) - sources.append(source) - - layout = h5py.VirtualLayout(shape=vds_shape, dtype=data_type) - # The first one is the upper one - s0 = jungfrau_mod_size[0] + jungfrau_gap_size[0] - layout[:, : jungfrau_mod_size[0], :] = sources[1][:, :, :] - layout[:, s0:, :] = sources[0][:, :, :] - - nxdata.create_virtual_dataset(entry_key, layout, fillvalue=jungfrau_fill_value) - - -def vds_file_writer( - nxsfile: h5py.File, - datafiles: list[Path], - data_shape: tuple | list, - data_type: DTypeLike = np.uint16, - entry_key: str = "data", -): - """ - Write a Virtual DataSet _vds.h5 file for image data. - - Args: - nxsfile (h5py.File): NeXus file being written. - datafiles (list[Path]): list of paths to source files. - data_shape (tuple | list): Shape of the dataset, usually defined as (num_frames, *image_size). - data_type (DTypeLike, optional): Dtype. Defaults to np.uint16. - entry_key (str): Entry key for the Virtual DataSet name. Defaults to data. - """ - vds_logger.debug("Start creating VDS file ...") - # Where the vds will go - nxdata = nxsfile["/entry/data"] - # entry_key = "data" - - # For every source dataset define its shape and number of frames - # Once again, it is assumed that the maximum number of frames per dataset is 1000 - frames = (data_shape[0] // 1000) * [1000] + [data_shape[0] % 1000] - sshape = [(f, *data_shape[1:]) for f in frames] - - # Create virtual layout - layout = h5py.VirtualLayout(shape=data_shape, dtype=data_type) - start = 0 - for n, filename in enumerate(datafiles): - end = start + frames[n] - vsource = h5py.VirtualSource( - filename.name, entry_key, shape=sshape[n] - ) # Source definition - layout[start:end:1, :, :] = vsource - start = end - - # Create a _vds.h5 file and add link to nexus file - s = Path(nxsfile.filename).expanduser().resolve() - vds_filename = s.parent / f"{s.stem}_vds.h5" - del s - with h5py.File(vds_filename, "w") as vds: - vds.create_virtual_dataset("data", layout, fillvalue=-1) - nxdata["data"] = h5py.ExternalLink(vds_filename.name, "data") - vds_logger.debug(f"{vds_filename} written and link added to NeXus file.") - - -def clean_unused_links( - nxsfile: h5py.File, - vds_shape: tuple | list, - start_index: int = 0, -): - """ - Remove links to external data not used in VDS. - - Args: - nxsfile (h5py.File): Handle to NeXus file being written. - vds_shape (tuple | list): Actual shape of the VDS dataset, usually defined as (num_frames, *image_size). - start_index(int): The start point for the source data. Defaults to 0. - """ - vds_logger.debug("Cleaning links unused in VDS ...") - # Location of the VDS - nxdata = nxsfile["/entry/data"] - dataset_names = find_datasets_in_file(nxdata) - if len(dataset_names) == 1: - vds_logger.debug("Only one linked file, no need to remove it.") - return - datasets = [nxdata[name] for name in dataset_names] - dataset_lengths = [d.shape[0] for d in datasets] - if sum(dataset_lengths) == vds_shape[0]: - vds_logger.debug("All links are used in VDS, no need to remove any.") - return - for i, _ in enumerate(datasets): - # unlink datasets before the start of VDS - if sum(dataset_lengths[0 : i + 1]) < start_index: - vds_logger.debug(f"Removing {dataset_names[i]} link.") - del nxdata[dataset_names[i]] - # unlink datasets after the end of VDS - if sum(dataset_lengths[0:i]) > start_index + vds_shape[0]: - vds_logger.debug(f"Removing {dataset_names[i]} link.") - del nxdata[dataset_names[i]] - vds_logger.debug("Links unused in VDS removed from NeXus file.") + block_vds_logger.debug("VDS correctly written to NeXus file.") diff --git a/src/nexgen/tools/vds_tools/tiled_mapping.py b/src/nexgen/tools/vds_tools/tiled_mapping.py new file mode 100644 index 00000000..9a398a80 --- /dev/null +++ b/src/nexgen/tools/vds_tools/tiled_mapping.py @@ -0,0 +1,53 @@ +"""Create a Virtual DataSet with a strided mapping, ie mappaing every {n} frames from each file + +For now only the JF1M functionality has been added, the full vds capabilities will be added later. +""" + +import logging + +import h5py +import numpy as np +from numpy.typing import DTypeLike + +from nexgen.tools.constants import ( + jungfrau_fill_value, + jungfrau_gap_size, + jungfrau_mod_size, +) +from nexgen.tools.vds_tools.utils import find_datasets_in_file + +tiled_vds_logger = logging.getLogger("nexgen.tools.vds_tools.utils") + + +def jungfrau_vds_writer( + nxsfile: h5py.File, + vds_shape: tuple | list, + data_type: DTypeLike = np.uint16, + source_dsets: list[str] | None = None, +): + """Write VDS for Jungfrau 1M use case, with a tiled layout.""" + external_dsets = True + entry_key = "data" + frames = vds_shape[0] + + nxdata = nxsfile["/entry/data"] + if not source_dsets: + source_dsets = find_datasets_in_file(nxdata) + external_dsets = False + + sources = [] + for dset in source_dsets: + source_path = dset if external_dsets is True else "." + source_name = entry_key if external_dsets is True else f"/entry/data/{dset}" + source = h5py.VirtualSource( + source_path, source_name, shape=(frames, *jungfrau_mod_size) + ) + sources.append(source) + + layout = h5py.VirtualLayout(shape=vds_shape, dtype=data_type) + # The first one is the upper one + s0 = jungfrau_mod_size[0] + jungfrau_gap_size[0] + layout[:, : jungfrau_mod_size[0], :] = sources[1][:, :, :] + layout[:, s0:, :] = sources[0][:, :, :] + + nxdata.create_virtual_dataset(entry_key, layout, fillvalue=jungfrau_fill_value) diff --git a/src/nexgen/tools/vds_tools/utils.py b/src/nexgen/tools/vds_tools/utils.py index dbff00de..89991976 100644 --- a/src/nexgen/tools/vds_tools/utils.py +++ b/src/nexgen/tools/vds_tools/utils.py @@ -1,12 +1,16 @@ +import logging from enum import StrEnum from typing import Any, TypeAlias import h5py import numpy as np +from numpy.typing import DTypeLike from pydantic import BaseModel, ConfigDict PydanticDTypeLike: TypeAlias = np.dtype[Any] | type | str | None +vds_logger = logging.getLogger("nexgen.tools.vds_tools") + class VdsMapping(StrEnum): BLOCKED = "blocked" # default, usual one @@ -23,9 +27,10 @@ class VdsSettings(BaseModel): vds_mapping: VdsMapping = VdsMapping.BLOCKED -def find_datasets_in_file(nxdata: h5py.Group) -> list: +def find_datasets_in_file(nxdata: h5py.Group) -> list[str]: """ - Look for the source datasets in the NeXus file. Assumes that the source datasets are always h5py.ExternalLink. + Look for the source datasets in the NeXus file. + Assumes that the source datasets are always h5py.ExternalLink. Args: nxdata (h5py.Group): Group where the data should be linked. @@ -36,13 +41,60 @@ def find_datasets_in_file(nxdata: h5py.Group) -> list: Returns: dsets (list): The source datasets. """ - # FIXME for now this assumes that the source datasets are always links dsets = [] for k in nxdata.keys(): if isinstance(nxdata.get(k, getlink=True), h5py.ExternalLink): dsets.append(k) if not dsets: + vds_logger.error("No extrnale link datasets found.") raise KeyError( f"No External Link datasets found in NeXus file under {nxdata.name}" ) return dsets + + +def define_vds_dtype_from_bit_depth(bit_depth: int) -> DTypeLike: + """Define dtype of VDS based on the passed bit depth.""" + if bit_depth == 32: + return np.uint32 + elif bit_depth == 8: + return np.uint8 + else: + return np.uint16 + + +def clean_unused_links( + nxsfile: h5py.File, + vds_shape: tuple | list, + start_index: int = 0, +): + """ + Remove links to external data not used in VDS. + + Args: + nxsfile (h5py.File): Handle to NeXus file being written. + vds_shape (tuple | list): Actual shape of the VDS dataset, usually defined as (num_frames, *image_size). + start_index(int): The start point for the source data. Defaults to 0. + """ + vds_logger.debug("Cleaning links unused in VDS ...") + # Location of the VDS + nxdata = nxsfile["/entry/data"] + dataset_names = find_datasets_in_file(nxdata) + if len(dataset_names) == 1: + vds_logger.debug("Only one linked file, no need to remove it.") + return + datasets = [nxdata[name] for name in dataset_names] + dataset_lengths = [d.shape[0] for d in datasets] + if sum(dataset_lengths) == vds_shape[0]: + vds_logger.debug("All links are used in VDS, no need to remove any.") + return + for i, _ in enumerate(datasets): + # unlink datasets before the start of VDS + if sum(dataset_lengths[0 : i + 1]) < start_index: + vds_logger.debug(f"Removing {dataset_names[i]} link.") + del nxdata[dataset_names[i]] + # unlink datasets after the end of VDS + if sum(dataset_lengths[0:i]) > start_index + vds_shape[0]: + vds_logger.debug(f"Removing {dataset_names[i]} link.") + del nxdata[dataset_names[i]] + vds_logger.debug("Links unused in VDS removed from NeXus file.") diff --git a/src/nexgen/tools/vds_tools/vds_file.py b/src/nexgen/tools/vds_tools/vds_file.py new file mode 100644 index 00000000..6949952f --- /dev/null +++ b/src/nexgen/tools/vds_tools/vds_file.py @@ -0,0 +1,60 @@ +""" +Tools to write Virtual DataSets +""" + +import logging +from pathlib import Path + +import h5py +import numpy as np +from numpy.typing import DTypeLike + +vds_logger = logging.getLogger("nexgen.VDSWriter") + + +def vds_file_writer( + nxsfile: h5py.File, + datafiles: list[Path], + data_shape: tuple | list, + data_type: DTypeLike = np.uint16, + entry_key: str = "data", +): + """ + Write a Virtual DataSet _vds.h5 file for image data. + + Args: + nxsfile (h5py.File): NeXus file being written. + datafiles (list[Path]): list of paths to source files. + data_shape (tuple | list): Shape of the dataset, usually defined as (num_frames, *image_size). + data_type (DTypeLike, optional): Dtype. Defaults to np.uint16. + entry_key (str): Entry key for the Virtual DataSet name. Defaults to data. + """ + vds_logger.debug("Start creating VDS file ...") + # Where the vds will go + nxdata = nxsfile["/entry/data"] + # entry_key = "data" + + # For every source dataset define its shape and number of frames + # Once again, it is assumed that the maximum number of frames per dataset is 1000 + frames = (data_shape[0] // 1000) * [1000] + [data_shape[0] % 1000] + sshape = [(f, *data_shape[1:]) for f in frames] + + # Create virtual layout + layout = h5py.VirtualLayout(shape=data_shape, dtype=data_type) + start = 0 + for n, filename in enumerate(datafiles): + end = start + frames[n] + vsource = h5py.VirtualSource( + filename.name, entry_key, shape=sshape[n] + ) # Source definition + layout[start:end:1, :, :] = vsource + start = end + + # Create a _vds.h5 file and add link to nexus file + s = Path(nxsfile.filename).expanduser().resolve() + vds_filename = s.parent / f"{s.stem}_vds.h5" + del s + with h5py.File(vds_filename, "w") as vds: + vds.create_virtual_dataset("data", layout, fillvalue=-1) + nxdata["data"] = h5py.ExternalLink(vds_filename.name, "data") + vds_logger.debug(f"{vds_filename} written and link added to NeXus file.") diff --git a/tests/tools/test_VDS_tools.py b/tests/tools/vds/test_blocked_vds.py similarity index 70% rename from tests/tools/test_VDS_tools.py rename to tests/tools/vds/test_blocked_vds.py index 334933e3..d51fb3d5 100644 --- a/tests/tools/test_VDS_tools.py +++ b/tests/tools/vds/test_blocked_vds.py @@ -5,25 +5,14 @@ import numpy as np import pytest -from nexgen.tools.vds_w_tools import ( +from nexgen.tools.vds_tools.blocked_mapping import ( Dataset, create_virtual_layout, - define_vds_dtype_from_bit_depth, image_vds_writer, - jungfrau_vds_writer, split_datasets, ) -@pytest.mark.parametrize( - "bit_depth, expected_dtype", [(8, np.uint8), (16, np.uint16), (32, np.uint32)] -) -def test_vds_dtype_from_input(bit_depth, expected_dtype): - d = define_vds_dtype_from_bit_depth(bit_depth) - - assert d == expected_dtype - - def test_when_get_frames_and_shape_less_than_1000_then_correct(): sshape = split_datasets(["test1"], (500, 10, 10)) assert sshape == [Dataset("test1", (500, 10, 10), 0, 500)] @@ -110,28 +99,3 @@ def test_given_file_with_no_dataset_external_links_then_exception_is_sensible(): test_nexus_file["/entry/data/data_0001"] = MagicMock() with pytest.raises(KeyError): image_vds_writer(test_nexus_file, (1000, 10, 10)) - - -def test_jungfrau_vds_writer_with_external_dsets(): - test_hdf_file = tempfile.TemporaryFile() - test_nexus_file = h5py.File(test_hdf_file, "w") - test_nexus_file["/entry/data/data_0001"] = MagicMock() - source_dsets = ["path/to/file1", "path/to/file2"] - jungfrau_vds_writer( - test_nexus_file, - (100, 1066, 1030), - source_dsets=source_dsets, - ) - assert "data" in list(test_nexus_file["/entry/data"].keys()) - - -def test_jungfrau_vds_writer_not_failing_if_no_external_dsets(): - test_hdf_file = tempfile.TemporaryFile() - test_nexus_file = h5py.File(test_hdf_file, "w") - test_nexus_file["/entry/data/data_0001"] = h5py.ExternalLink("f1le1", "data") - test_nexus_file["/entry/data/data_0002"] = h5py.ExternalLink("f1le2", "data") - jungfrau_vds_writer( - test_nexus_file, - (100, 1066, 1030), - ) - assert "data" in list(test_nexus_file["/entry/data"].keys()) diff --git a/tests/tools/vds/test_tiled_mapping.py b/tests/tools/vds/test_tiled_mapping.py new file mode 100644 index 00000000..8250d8f4 --- /dev/null +++ b/tests/tools/vds/test_tiled_mapping.py @@ -0,0 +1,33 @@ +import tempfile +from unittest.mock import MagicMock + +import h5py + +from nexgen.tools.vds_tools.tiled_mapping import ( + jungfrau_vds_writer, +) + + +def test_jungfrau_vds_writer_with_external_dsets(): + test_hdf_file = tempfile.TemporaryFile() + test_nexus_file = h5py.File(test_hdf_file, "w") + test_nexus_file["/entry/data/data_0001"] = MagicMock() + source_dsets = ["path/to/file1", "path/to/file2"] + jungfrau_vds_writer( + test_nexus_file, + (100, 1066, 1030), + source_dsets=source_dsets, + ) + assert "data" in list(test_nexus_file["/entry/data"].keys()) + + +def test_jungfrau_vds_writer_not_failing_if_no_external_dsets(): + test_hdf_file = tempfile.TemporaryFile() + test_nexus_file = h5py.File(test_hdf_file, "w") + test_nexus_file["/entry/data/data_0001"] = h5py.ExternalLink("f1le1", "data") + test_nexus_file["/entry/data/data_0002"] = h5py.ExternalLink("f1le2", "data") + jungfrau_vds_writer( + test_nexus_file, + (100, 1066, 1030), + ) + assert "data" in list(test_nexus_file["/entry/data"].keys()) diff --git a/tests/tools/vds/test_utils.py b/tests/tools/vds/test_utils.py index 77c41a60..28c38f60 100644 --- a/tests/tools/vds/test_utils.py +++ b/tests/tools/vds/test_utils.py @@ -1,9 +1,13 @@ import tempfile import h5py +import numpy as np import pytest -from nexgen.tools.vds_tools.utils import find_datasets_in_file +from nexgen.tools.vds_tools.utils import ( + define_vds_dtype_from_bit_depth, + find_datasets_in_file, +) def test_find_datasets_in_file(nexus_file_with_single_dataset): @@ -26,3 +30,12 @@ def test_find_datasets_fails_if_no_links_in_file(): test_nexus_file["/entry/data/data_0001"] = [0, 0, 0] with pytest.raises(KeyError): find_datasets_in_file(test_nexus_file) + + +@pytest.mark.parametrize( + "bit_depth, expected_dtype", [(8, np.uint8), (16, np.uint16), (32, np.uint32)] +) +def test_vds_dtype_from_input(bit_depth, expected_dtype): + d = define_vds_dtype_from_bit_depth(bit_depth) + + assert d == expected_dtype