import contextlib
import math
from collections.abc import Callable
from contextlib import nullcontext
from functools import partial
from pathlib import Path
from threading import Event
from typing import Optional
from warnings import warn
import h5netcdf
import h5py
import numpy as np
import xarray as xr
from dask.diagnostics import ProgressBar
from ..memory_allocator import empty_aligned
from ..util import make_timestamp
from ..util._tqdm import tqdm_base
from ..util.index_tricks import ndrange
from ..util.io_utils import check_disk_space
from ._compression import DataEncoder, compression_map, get_compression_type, round_to
from ._properties import get_stack_dims
# Mostly here we want to add the backends (if all the requirements are installed)
# to the list of available backends to xarray
from ._ramona_backend import RamonaH5NetCDFStore, has_direct_driver
from ._ramona_libnetcdf4_backend import RamonaNetCDF4DataStore # noqa
from ._update import update_dataset
[docs]
def save(
mcam_dataset,
filename: Path,
*,
mode='w',
engine='ramona',
include_timestamp: bool=True,
disk_space_tolerance=0.1E9,
tqdm=None,
_stacklevel_increment=0,
single_file=None
):
"""Save data as a hdf5 using netcdf4 API.
Default extension is nc if none is given
Returns the saved path name as a Path object.
Parameters
----------
mcam_dataset: xarray Dataset
The mcam_data you wish to save.
filename: Path-like
The filename where the data should be saved. If the filename has no
extension, then an ``'.nc'`` extension is added.
mode: str
The mode to open the file in. Valid options are ``'w'`` (write),
or ```x``` exclusive creation, failing if the file already exists.
engine:
Parameter passed to xarray.Dataset.to_netcdf to select the
backend used for writing data to disk.
include_timestamp: bool
If set to `True`, this will append a timestamp to the provided
path name. If set to false, no timestamp will be appended to the
path name possibly overwriting any files currently within the previous
path name.
disk_space_tolerance: float
The difference between the free disk space in the given directory and
the dataset to be save in bytes. If the directory has fewer free bytes
than this tolerance plus the dataset size an error will be raised.
single_file:
.. versionchanged :: 0.14.0
In version 0.14.0 this parameter is ignored, and the behavior is
always that of ``single_file=True``
.. versionchanged :: 0.18.9
In version 0.18.9 this parameter will emit a deprecation warning
indicating that it will be removed in version 0.20.0
Returns
-------
filename: Path-like
The filename where the the data has been stored. It includes the
optional timestamp, and the new suffix.
"""
if single_file is not None:
warn("The single_file parameter is deprecated and will be removed in "
"version 0.20.0. If you want to export the data, please see the "
"export function "
"https://docs.ramonaoptics.com/python_module.html#owl.mcam_data.export ",
stacklevel=2 + _stacklevel_increment)
for dim in mcam_dataset.dims:
if dim not in mcam_dataset.coords:
warn(
f'{dim} dimension has not been given coordinates within the dataset. '
'Please contact Ramona Optics at help@ramonaoptics.com to get this resolved.',
stacklevel=2 + _stacklevel_increment,
)
# 2025/09/22 Clay / Mark
# https://gitlab.com/ramonaoptics/mcam/python-owl/-/merge_requests/7810
# We are forwarding save into save_video which takes advantage
# of our auto-chunking and direct writing capabilities
# and provides a progress bar
return save_video(
mcam_dataset,
filename,
include_timestamp=include_timestamp,
disk_space_tolerance=disk_space_tolerance,
engine=engine,
mode=mode,
tqdm=tqdm,
_stacklevel_increment=_stacklevel_increment
)
def load(filename: Path, *,
delayed=True,
scheduler='threading',
progress=True,
update=True,
engine='ramona',
chunks=None,
_stacklevel_increment=0,
**_kwargs):
"""Load mcam_data from a hdf5 file using netcdf4 and convert
it to whatever we want.
Parameters
----------
filename: Path-like
The netcdf4 file where the data is stored.
delayed: bool, optional
If True, the computation will return a lazy object. As the user
of this library, you will have to explicitly force the computation
of the lazy object. This function will attempt to automatically
determine if the data should be loaded lazily or eagerly.
If a particular behavior is desired in your application, the delayed
parameter should be specified.
progress:
If True, then a progress bar is shown during loading operations. This
is only valid if delayed=True.
scheduler:
Parameter passed to ``dask.compute`` to select which schedule is used
to load the data in parallel.
update:
Set to False if you wish to load the raw un-updated data. Setting this
parameter to False keeps the raw metadata as is in the `.nc` file but
means that the loaded dataset is unsupported by the remainder of the
``owl`` analysis functions.
engine:
xarray engine used to open the dataset.
chunks:
A dictionary of the image axes to chunk and the size of the chunks
along the axes. Axes should be referenced by their coordinate name.
This is best used in conjunction with ``delayed=True``
Returns
-------
mcam_dataset: DataArray
Return the MCAM data in an xarray DataArray with all the metadata.
"""
legacy_kwargs_0_17_0 = [
'output_mode', 'dtype', 'make_square', 'preserve_range',
]
for key in legacy_kwargs_0_17_0:
if key in _kwargs:
raise RuntimeError(
f"The {key} parameter is no longer supported as of version 0.17.0. "
"Please contact Ramona Optics at help@ramonaoptics.com with a screenshot "
"of your screen including this error message to learn how to migrate your code.")
mcam_dataset = xr.open_dataset(filename, engine=engine, chunks=chunks)
if isinstance(mcam_dataset, xr.DataArray):
mcam_dataset.name = 'images'
mcam_dataset = mcam_dataset.to_dataset()
if update:
mcam_dataset = update_dataset(mcam_dataset, chunks=chunks)
if delayed is None:
# Automatically load datasets that are less than 1 GB
# otherwise, just delay them
delayed = mcam_dataset.nbytes > 1E9
if not delayed:
if progress:
ctx_manager = ProgressBar
else:
ctx_manager = nullcontext
with ctx_manager():
mcam_dataset.load()
return mcam_dataset
def _expand_variable(nc_variable, mcam_dataset, expanding_dim, index, added_length):
# For time deltas, we must ensure that we use the same encoding as
# what was previously stored.
# We likely need to do this as well for variables that had custom
# encodings too
# Save the old encoding in case we clobber it
old_encoding = mcam_dataset.encoding
if hasattr(nc_variable, 'calendar'):
mcam_dataset.encoding = {
'units': nc_variable.units,
'calendar': nc_variable.calendar,
}
data_encoded = xr.conventions.encode_cf_variable(mcam_dataset)
left_slices = mcam_dataset.dims.index(expanding_dim)
right_slices = mcam_dataset.ndim - left_slices - 1
nc_slice = (
(slice(None),) * left_slices +
(slice(index, index + added_length),) +
(slice(None),) * right_slices
)
nc_variable[nc_slice] = data_encoded.data
mcam_dataset.encoding = old_encoding
def append(filename, ds_to_append, unlimited_dims, *, engine=None):
"""Append dataset to netCDF4 file.
Append the provided dataset to the file along the unlimited_dim.
Parameters
----------
filename: Path-like or File-like object
The path to the file where the data should be written or alternatively,
a file object from netCDF4 or h5netcdf. If ``filename`` is a File-like
handle, the ``engine`` parameter is ignored.
.. versionchanged:: 0.18.15
The filename parameter can now accept a file object.
ds_to_append: xarray.Dataset
xarray dataset to append to the filename.
ulimited_dims: str or List[str]
Dimension over which to append the dataset to the file. Currently,
only one dimension is supported.
engine: str
The name of the backend to use to write the file. Should be one of
``"netcdf4"``, ``"h5netcdf"``, or ``"ramona"``.
.. versionadded:: 0.18.15
The ``engine`` parameter.
"""
if engine is None:
engine = 'netcdf4'
elif isinstance(engine, str):
engine = engine.lower()
if hasattr(filename, "close"):
@contextlib.contextmanager
def File():
yield filename
if isinstance(filename, h5netcdf.File):
engine = "h5netcdf"
else:
engine = "netcdf4"
elif engine == "netcdf4":
# 2026/03/23
# Lazy-load here to remove the hard dependency on netCDF4
# https://gitlab.com/ramonaoptics/mcam/python-owl/-/merge_requests/9398
import netCDF4
File = partial(
netCDF4.Dataset,
filename,
mode="a",
)
elif engine == "h5netcdf":
# Since we are opening the array for writing, the value that we
# set here doesn't actually change the data in the file.
# https://github.com/h5netcdf/h5netcdf/issues/132
# Note that this is contrary to what xarray does
# https://github.com/pydata/xarray/pull/4893/files
File = partial(
h5netcdf.File,
filename,
mode="a",
decode_vlen_strings=False,
alignment_threshold=15 * 4096, # Align pretty aggressively for speed
alignment_interval=4096, # align to one kernel page
)
elif engine == "ramona":
ramona_kwargs = dict(
decode_vlen_strings=False,
alignment_threshold=15 * 4096, # Align pretty aggressively for speed
alignment_interval=4096, # align to one kernel page
)
# The HDF5 'direct' VFD (direct I/O) only exists on Linux builds compiled
# with it (has_direct_driver); on macOS / unsupported builds it is absent
# and h5py raises "Unknown driver type 'direct'". Fall back to the default
# driver there so the file still opens (buffered I/O); driver-specific
# kwargs (cbuf_size) only apply to the direct VFD.
if has_direct_driver:
ramona_kwargs.update(
driver="direct",
cbuf_size=512 * 1024, # 0.5 MB buffer to copy unaligned data
)
File = partial(
h5netcdf.File,
filename,
mode="a",
**ramona_kwargs,
)
else:
raise ValueError(
f"Unknown engine. Got {engine} expecting 'h5netcdf', 'netcdf4', or 'ramona'."
)
if isinstance(unlimited_dims, str):
unlimited_dims = [unlimited_dims]
if len(unlimited_dims) != 1:
# TODO: change this so it can support multiple expanding dims
raise ValueError(
"We only support one unlimited dim for now, "
f"got {len(unlimited_dims)}.")
unlimited_dims = list(set(unlimited_dims))
expanding_dim = unlimited_dims[0]
with File() as nc:
nc_coord = nc[expanding_dim]
index = len(nc_coord)
added_length = len(ds_to_append[expanding_dim])
if engine in ["h5netcdf", "ramona"]:
# explicitly resize the dimension
#
nc.resize_dimension(expanding_dim, index + added_length)
variables, _attrs = xr.conventions.encode_dataset_coordinates(ds_to_append)
for name, mcam_dataset in variables.items():
if expanding_dim not in mcam_dataset.dims:
# Nothing to do, data assumed to the identical
continue
nc_variable = nc[name]
_expand_variable(nc_variable, mcam_dataset, expanding_dim, index, added_length)
def _write_metadata(mcam_dataset, filename, mode='w'):
# Open with a backend that "caches" to make metadata writing fast
# metadata typically consists of lots of small things, not single
# large arrays
# The HDF5 DIRECT backend is really slow at writing metadata
metadata = mcam_dataset.drop_vars('images', errors='ignore')
metadata.to_netcdf(filename, format='NETCDF4', engine="h5netcdf",
mode=mode)
def _largest_nontrivial_divisor_leq(n, limit):
if limit >= n:
return n
for d in range(limit, 1, -1):
if n % d == 0:
return d
return limit
def _snap_chunk_down_to_close_divisor(chunk_i, shape_i, *, minimum_fraction=0.5):
if chunk_i <= 1 or shape_i % chunk_i == 0:
return chunk_i
divisor = _largest_nontrivial_divisor_leq(shape_i, chunk_i)
if divisor / chunk_i >= minimum_fraction:
return divisor
return chunk_i
def _decrease_chunk_size(chunksize, maximum_chunk_bytes, minimum_chunksize):
# We aim to chunks that are smaller than a certain threshold that are
# "C-contiguous".
for i in range(len(chunksize)):
chunk_bytes = math.prod(chunksize)
if chunk_bytes < maximum_chunk_bytes:
break
chunk_bytes_no_i = math.prod(chunksize[:i] + chunksize[i + 1:])
chunk_i_preferred = maximum_chunk_bytes // chunk_bytes_no_i
chunk_i_preferred = max(chunk_i_preferred, minimum_chunksize[i])
chunksize = chunksize[:i] + (chunk_i_preferred,) + chunksize[i + 1:]
return chunksize
def _increase_chunk_size(chunksize, minimum_chunk_bytes, maximum_chunksize):
# We aim to create large chunks that are "C-contiguous".
for i in reversed(range(len(chunksize))):
chunk_bytes = math.prod(chunksize)
if chunk_bytes > minimum_chunk_bytes:
break
chunk_bytes_no_i = math.prod(chunksize[:i] + chunksize[i + 1:])
chunk_i_preferred = (minimum_chunk_bytes + chunk_bytes_no_i - 1) // chunk_bytes_no_i
chunk_i_preferred = min(chunk_i_preferred, maximum_chunksize[i])
chunksize = chunksize[:i] + (chunk_i_preferred,) + chunksize[i + 1:]
return chunksize
def _get_chunksizes(
dataset,
*,
minimum_chunk_bytes=None,
maximum_chunk_bytes=None,
preferred_chunksize=None
):
shape = dataset.images.shape
dtype = dataset.images.dtype
itemsize = dtype.itemsize
if maximum_chunk_bytes is None:
maximum_chunk_bytes = (3 * 1024 ** 3 + 512 * 1024 ** 2) # 3.5 GiB
if minimum_chunk_bytes is None:
minimum_chunk_bytes = 128 * 1024 ** 2
if minimum_chunk_bytes > maximum_chunk_bytes:
raise ValueError(
f"minimum_chunk_bytes ({minimum_chunk_bytes}) must "
f"be smaller than maximum_chunk_bytes ({maximum_chunk_bytes})"
)
if preferred_chunksize is None:
stack_dims = get_stack_dims(dataset)
preferred_chunksize = tuple(
dataset.sizes[s] if s not in stack_dims else 1
for s in dataset.images.dims
)
if len(preferred_chunksize) != len(shape):
raise ValueError(
f"Length of preferred_chunksize ({len(preferred_chunksize)}) "
f"must be the same as length of shape ({len(shape)})."
)
if minimum_chunk_bytes < itemsize:
raise ValueError(
f"minimum_chunk_bytes ({minimum_chunk_bytes}) is smaller "
f"than the itemsize ({itemsize})."
)
# Add itemsize to ensure we account for it in our products
chunksize = preferred_chunksize + (itemsize,)
chunksize = _increase_chunk_size(
chunksize, minimum_chunk_bytes,
maximum_chunksize=shape + (itemsize,)
)
chunksize = _decrease_chunk_size(
chunksize, maximum_chunk_bytes,
minimum_chunksize=(1,) * len(shape) + (itemsize,),
)
# Remove itemsize
chunksize = chunksize[:-1]
snapped = list(chunksize)
for i, (chunk_i, shape_i) in enumerate(zip(chunksize, shape)):
candidate = _snap_chunk_down_to_close_divisor(chunk_i, shape_i)
if candidate == chunk_i:
continue
trial = snapped.copy()
trial[i] = candidate
if math.prod(trial) * itemsize >= minimum_chunk_bytes:
snapped[i] = candidate
chunksize = tuple(snapped)
return chunksize
def _is_power_of_two(value):
return isinstance(value, int) and value > 0 and (value & (value - 1)) == 0
def _resolve_pyramid_config(var_name, variable):
explicit_levels = 'pyramid_levels' in variable.encoding
explicit_downsample = 'pyramid_downsample' in variable.encoding
compression = get_compression_type(variable.encoding)
has_compression = compression in compression_map
default_levels = 3 if has_compression else 1
pyramid_levels = variable.encoding.get('pyramid_levels', default_levels)
pyramid_downsample = variable.encoding.get('pyramid_downsample', 4)
if not isinstance(pyramid_levels, int) or pyramid_levels < 1:
raise ValueError(
f"Variable '{var_name}' has invalid pyramid_levels={pyramid_levels}. "
"pyramid_levels must be an integer >= 1."
)
if not _is_power_of_two(pyramid_downsample):
raise ValueError(
f"Variable '{var_name}' has invalid pyramid_downsample={pyramid_downsample}. "
"pyramid_downsample must be a positive power of two."
)
if pyramid_levels == 1:
return None
if 'x' not in variable.dims or 'y' not in variable.dims:
if explicit_levels or explicit_downsample:
raise ValueError(
f"Variable '{var_name}' requested pyramid levels but does not have "
"both 'x' and 'y' dimensions."
)
return None
return {
'levels': pyramid_levels,
'downsample': pyramid_downsample,
'x_index': variable.dims.index('x'),
'y_index': variable.dims.index('y'),
}
def _iter_pyramid_level_specs(var_name, variable, pyramid_config):
level_specs = []
levels = pyramid_config['levels']
downsample = pyramid_config['downsample']
x_size = variable.shape[pyramid_config['x_index']]
y_size = variable.shape[pyramid_config['y_index']]
for level_index in range(1, levels):
downsample_factor = downsample ** level_index
if (x_size % downsample_factor) != 0 or (y_size % downsample_factor) != 0:
warn(
f"Variable '{var_name}' requested pyramid_levels={levels}, but x/y "
f"shape {x_size}x{y_size} only supports downsampling by {downsample_factor}. "
"Higher levels will be skipped.",
stacklevel=3,
)
break
level_key = str(int(math.log2(downsample_factor)))
level_shape = list(variable.shape)
level_shape[pyramid_config['x_index']] = x_size // downsample_factor
level_shape[pyramid_config['y_index']] = y_size // downsample_factor
level_specs.append((level_key, downsample_factor, tuple(level_shape)))
return level_specs
def _get_level_chunks(base_chunks, level_shape, pyramid_config, downsample_factor):
if base_chunks is None:
return None
level_chunks = list(base_chunks)
for dim_index in (pyramid_config['x_index'], pyramid_config['y_index']):
level_chunks[dim_index] = max(1, level_chunks[dim_index] // downsample_factor)
return tuple(min(chunk, shape) for chunk, shape in zip(level_chunks, level_shape))
def _get_h5_compression_kwargs(h5ds):
kwargs = {}
if h5ds._filters:
filter_id, filter_options = next(iter(h5ds._filters.items()))
try:
kwargs['compression'] = int(filter_id)
except (TypeError, ValueError):
kwargs['compression'] = filter_id
if filter_options:
kwargs['compression_opts'] = filter_options
elif h5ds.compression is not None:
kwargs['compression'] = h5ds.compression
if h5ds.compression_opts is not None:
kwargs['compression_opts'] = h5ds.compression_opts
if h5ds.shuffle:
kwargs['shuffle'] = h5ds.shuffle
if h5ds.fletcher32:
kwargs['fletcher32'] = h5ds.fletcher32
return kwargs
def _write_pyramid_levels(store, var_name, variable, pyramid_config, base_h5ds, tqdm):
level_specs = _iter_pyramid_level_specs(var_name, variable, pyramid_config)
if not level_specs:
return
levels_group_name = f"{var_name}.levels"
h5file = store.ds._root._h5file
levels_group = h5file.require_group(levels_group_name)
compression_kwargs = _get_h5_compression_kwargs(base_h5ds)
for level_key, downsample_factor, level_shape in tqdm(level_specs):
if level_key in levels_group:
del levels_group[level_key]
level_chunks = _get_level_chunks(
base_h5ds.chunks, level_shape, pyramid_config, downsample_factor,
)
dataset_kwargs = dict(compression_kwargs)
if level_chunks is not None:
dataset_kwargs['chunks'] = level_chunks
level_dataset = levels_group.create_dataset(
level_key,
shape=level_shape,
dtype=base_h5ds.dtype,
**dataset_kwargs,
)
if level_dataset.chunks is not None:
chunk_iterable = ndrange(
(0,) * len(level_shape),
level_shape,
level_dataset.chunks,
)
else:
chunk_iterable = [tuple(0 for _ in level_shape)]
for chunk_index in chunk_iterable:
level_slice = tuple(
slice(start, min(start + chunk, shape))
for start, chunk, shape in zip(
chunk_index,
level_dataset.chunks or level_shape,
level_shape,
)
)
source_slice = list(level_slice)
for dim_index in (pyramid_config['x_index'], pyramid_config['y_index']):
source_slice[dim_index] = slice(
level_slice[dim_index].start * downsample_factor,
level_slice[dim_index].stop * downsample_factor,
downsample_factor,
)
level_dataset[level_slice] = np.ascontiguousarray(variable[tuple(source_slice)])
def write_chunked_variable(
store: RamonaH5NetCDFStore,
var_name: str,
variable: xr.Variable,
data_encoder: Optional[DataEncoder] = None,
tqdm: Optional[Callable] = None,
) -> None:
if tqdm is None:
tqdm = tqdm_base
target, _data = store.prepare_variable(var_name, variable)
# Get the underlying images dataset
images_ds = target.get_array()._h5ds
# 2022/02/16: Mark
# xarray may have chosen not to chunk the dataset in the case
# that the data was too small.
# We therefore test to see if the dataset has any chunks before
# using chunk specific operations
if images_ds.chunks is not None:
# Use the actual HDF5 chunk layout, not variable.encoding['chunksizes'].
# prepare_variable already passes encoding chunksizes to h5py, which
# clips them to fit the data shape. The encoding can be stale after
# isel() (see https://github.com/pydata/xarray/issues/11028), so we
# must iterate using the real chunk layout that h5py created.
chunksizes = images_ds.chunks
chunk_iterator = list(
ndrange(
(0,) * variable.ndim,
variable.shape,
chunksizes
)
)
iterable = chunk_iterator
else:
iterable = range(1)
chunksizes = variable.shape
n_items = int(np.prod(chunksizes))
# For compression, we need a larger buffer to handle worst-case expansion
compression = get_compression_type(variable.encoding)
compression_config = compression_map.get(compression, {})
max_size_func = compression_config.get('max_size_func')
if max_size_func and data_encoder is not None:
# For compressed data, allocate buffer based on worst-case size
chunk_nbytes = n_items * variable.dtype.itemsize
max_size = round_to(max_size_func(chunk_nbytes), 4096)
output_buffer = empty_aligned(max_size, dtype=np.uint8)
else:
output_buffer = empty_aligned(n_items, dtype=variable.dtype)
# Allocate the pad buffer lazily since many shapes can avoid edge
# chunks entirely after _get_chunksizes() snaps to a divisor.
pad_buffer = None
progress = tqdm(iterable)
use_encoder = max_size_func is not None
complevel = variable.encoding.get('complevel', None)
if images_ds.chunks is not None:
for chunk_index in progress:
data_slice = tuple(
slice(i, min(i + s, shape))
for i, s, shape in zip(
chunk_index,
chunksizes,
variable.shape,
)
)
data = np.ascontiguousarray(variable[data_slice])
images_id = images_ds.id
chunk_info = images_id.get_chunk_info_by_coord(chunk_index)
byte_offset = chunk_info.byte_offset
# Mark - 2026/01
# something really weird is happening.
# For some reason the dataset is created, but it already contains
# the metadata and the chunks of predetermined shape
# But for some reason, the spots seem pre-allocated, and too small
# Perhaps it is because at allocation time, the size of the chunks
# is very small (zero) for the encoding.
# Without this check,
# It would fail in really strange cases
# https://gitlab.com/ramonaoptics/mcam/python-owl/-/merge_requests/8791
# 1. Create dataset with compressed variable
# 2. Load it
# 3. Downsample it
# I'm pretty sure it is also due to some retention of metadata in the
# encoding
# variable
# Clearly under normal circumstances the byte_offset is a None value
# but under some edge cases it is predefined???
chunk_aligned = byte_offset is None or byte_offset % 4096 == 0
# Edge chunks (at dataset boundaries) may be smaller than the
# full HDF5 chunk size. write_direct_chunk requires the
# payload to represent exactly one full logical chunk. Copy
# edge chunks into the shared full-chunk buffer so they can
# use the fast path.
if data.shape != chunksizes:
if pad_buffer is None:
# Padding bytes sit beyond the dataset dimensions and
# are never read back by HDF5, so they can remain
# uninitialized.
pad_buffer = empty_aligned(chunksizes, dtype=variable.dtype)
# Note: data now aliases pad_buffer. The encoder call
# below consumes data immediately, so the next iteration's
# write into pad_buffer is safe.
pad_buffer[tuple(slice(0, s) for s in data.shape)] = data
data = pad_buffer
if not chunk_aligned:
images_ds[data_slice] = data
elif data_encoder is not None and use_encoder:
original_shape = data.shape
data = data.reshape(np.prod(original_shape))
encoded_data, nbytes = data_encoder(data, level=complevel, out=output_buffer)
# For compressed data, use only the actual compressed bytes
images_id.write_direct_chunk(chunk_index, encoded_data[:nbytes])
else:
encoded_data = data
images_ds.id.write_direct_chunk(chunk_index, encoded_data)
else:
# Use tqdm to ensure the progress bars "appears" then closes
for _ in progress:
# cast to numpy array to catch different types of lazy arrays (dask)
data = np.ascontiguousarray(variable.data)
offset = images_ds.id.get_offset()
chunk_aligned = offset is None or offset % 4096 == 0
if not chunk_aligned:
images_ds[...] = data
elif data_encoder is not None and use_encoder:
original_shape = data.shape
# reshape() may copy if the data is not contiguous, which is
# acceptable here (unlike the zero-copy views elsewhere).
data = data.reshape((np.prod(original_shape),))
encoded_data, nbytes = data_encoder(data, level=complevel, out=output_buffer)
# For compressed data, use only the actual compressed bytes
encoded_data = encoded_data[:nbytes]
images_ds.write_direct(encoded_data)
else:
images_ds.write_direct(data)
pyramid_config = _resolve_pyramid_config(var_name, variable)
if pyramid_config is not None:
_write_pyramid_levels(
store=store,
var_name=var_name,
variable=variable,
pyramid_config=pyramid_config,
base_h5ds=images_ds,
tqdm=tqdm,
)
[docs]
def write_netcdf(dataset, filename, mode='w', tqdm=None, vars_to_chunk=None,
virtual_dataset=None):
# I'm not really sure, but I feel like libnetcdf doesn't close files
# correctly and this results in bugs in HDF5 like:
# HDF5-DIAG: Error detected in HDF5 (1.14.0) thread 1:
# #000: H5A.c line 679 in H5Aopen_by_name(): unable to synchronously open attribute
# major: Attribute
# minor: Can't open object
# By using "h5netcdf" it seems to avoid this bug....
if tqdm is None:
tqdm = tqdm_base
if vars_to_chunk:
unchunk_dataset = dataset.drop_vars(vars_to_chunk, errors='ignore')
else:
unchunk_dataset = dataset
if not virtual_dataset:
# 20251031 Jed: the _write_netcdf function seems to be about 20% slower
# so we avoid it unless it is needed to write virtual datasets
unchunk_dataset.to_netcdf(filename, format='NETCDF4', engine='h5netcdf', mode=mode)
else:
_write_netcdf(unchunk_dataset, filename, mode=mode, virtual_dataset=virtual_dataset)
if not vars_to_chunk:
return
with RamonaH5NetCDFStore.open(filename, format='NETCDF4', mode='a') as store:
t = tqdm(vars_to_chunk)
for chunk_var in t:
# https://gitlab.com/ramonaoptics/mcam/python-owl/-/merge_requests/8185
# https://gitlab.com/ramonaoptics/mcam/python-owl/-/issues/813
dataset[chunk_var].encoding.pop('dtype', None)
compression_type = get_compression_type(dataset[chunk_var].encoding)
if compression_type and compression_type in compression_map:
data_encoder = compression_map[compression_type]['encoder']
else:
data_encoder = None
write_chunked_variable(
store,
chunk_var,
dataset[chunk_var].variable,
data_encoder=data_encoder,
tqdm=t if callable(t) else None
)
def _write_netcdf(
direct_dataset,
filename,
mode='w',
virtual_dataset=None,
_stacklevel_increment=0
):
if virtual_dataset is None:
virtual_dataset = {}
with h5py.File(filename, mode) as f:
# Track scalar coordinates to add to _Netcdf4Coordinates
# these seem to be:
# __owl_version__
# channel
# frame_number
scalar_coords = []
# Create coordinates
virtual_coords = {}
for key, value in virtual_dataset.items():
virtual_coords.update(value['coords'])
coords = {**direct_dataset.coords, **virtual_coords}
for name, data in coords.items():
coord_data = data.values if hasattr(data, 'values') else data
# Handle string dtypes for coordinates
# This seems to be mainly __owl_version__ and __owl_previous_versions__
if hasattr(coord_data, 'dtype') and coord_data.dtype.kind in ['U', 'S', 'O']:
coord_data = coord_data.astype('T')
# Add scalar coordinates as regular datasets as coords without dims are
# not supported as dimension scales
if coord_data.ndim == 0:
# Create as regular dataset, not a dimension scale
ds = f.create_dataset(name, data=coord_data)
# Add _Netcdf4Dimid = -1 to mark as a scalar coordinate
ds.attrs['_Netcdf4Dimid'] = np.int32(-1)
# Track this as a coordinate to add to data vars later
scalar_coords.append(name)
else:
# 1D coordinate - create as dimension scale
ds = f.create_dataset(name, data=coord_data)
ds.make_scale(name)
# Copy attributes
if hasattr(data, 'attrs'):
for attr_name, attr_value in data.attrs.items():
ds.attrs[attr_name] = attr_value
# Add dims - these are like coords, but are just numbered 0 - len(dim)
for dim_name in direct_dataset.dims:
if dim_name in f:
continue # already created as coord
dim_size = direct_dataset.sizes[dim_name]
dim_data = np.arange(dim_size)
dim_ds = f.create_dataset(dim_name, data=dim_data)
dim_ds.make_scale(dim_name)
warn(
f'{dim_name} dimension has not been given coordinates within the dataset. '
'Please contact Ramona Optics at help@ramonaoptics.com to get this resolved. '
'Please include the dataset and mention #1644 in your message.',
stacklevel=2 + _stacklevel_increment,
)
# Virtual dataset
aggregated_file_dir = Path(filename).parent
for key, value in virtual_dataset.items():
virtual_layout = h5py.VirtualLayout(shape=value['shape'], dtype=value['dtype'])
for t, source_filename in enumerate(value['filenames']):
shape = (
len(value['coords'][dim]) for dim in value['dims']
if dim not in value['virtual_dims']
)
source_path = Path(source_filename)
if source_path.is_relative_to(aggregated_file_dir):
source_path = source_path.relative_to(aggregated_file_dir)
else:
source_path = source_path.resolve()
vsource = h5py.VirtualSource(
str(source_path),
key,
shape=shape
)
virtual_layout[t, ...] = vsource
virtual_ds = f.create_virtual_dataset(key, virtual_layout, fillvalue=0)
for i, dim in enumerate(value['dims']):
virtual_ds.dims[i].attach_scale(f[dim])
virtual_ds.dims[i].label = dim
virtual_ds.attrs.update(value['attrs'])
# Copy all real variables from direct_dataset
for var_name in direct_dataset.data_vars:
var_data = np.asarray(direct_dataset[var_name])
# Handle special dtypes that h5py can't handle directly
if var_data.dtype.kind in ['U', 'S', 'O']: # Unicode, bytes, or object strings
var_data = var_data.astype('T')
var_ds = f.create_dataset(var_name, data=var_data)
elif var_data.dtype.kind == 'M': # Datetime
# Convert datetime to int64 nanoseconds since Unix epoch
int_data = var_data.astype('int64')
var_ds = f.create_dataset(var_name, data=int_data)
# Store with standard units since .astype('int64') gives us
# nanoseconds since 1970-01-01 (Unix epoch)
var_ds.attrs['units'] = 'nanoseconds since 1970-01-01T00:00:00'
encoding = direct_dataset[var_name].encoding
if 'calendar' in encoding:
var_ds.attrs['calendar'] = encoding['calendar']
elif var_data.dtype.kind == 'm': # Timedelta
# Convert timedelta to int64 (nanoseconds)
int_data = var_data.astype('int64')
var_ds = f.create_dataset(var_name, data=int_data)
var_ds.attrs['_timedelta_units'] = 'nanoseconds'
var_ds.attrs['_timedelta_dtype'] = str(var_data.dtype)
else:
# Non-special data - should work directly
var_ds = f.create_dataset(var_name, data=var_data)
# Copy attributes
for attr_name, attr_value in direct_dataset[var_name].attrs.items():
# Handle string attributes too
if isinstance(attr_value, str):
var_ds.attrs[attr_name] = attr_value
elif isinstance(attr_value, (np.ndarray, list)) and len(attr_value) > 0:
if isinstance(attr_value[0], str) or (
hasattr(attr_value, 'dtype') and attr_value.dtype.kind == 'U'):
var_ds.attrs[attr_name] = attr_value
else:
var_ds.attrs[attr_name] = attr_value
else:
var_ds.attrs[attr_name] = attr_value
# If the variable has dimensions that match coordinates, attach them
if hasattr(direct_dataset[var_name], 'dims'):
for dim_idx, dim_name in enumerate(direct_dataset[var_name].dims):
var_ds.dims[dim_idx].attach_scale(f[dim_name])
var_ds.dims[dim_idx].label = dim_name
# copy attributes from timelapse_metadata to root attrs
for attr_name, attr_value in direct_dataset.attrs.items():
f.attrs[attr_name] = attr_value
if scalar_coords:
# Add 'coordinates' attribute to ALL data variables pointing to scalar
# coordinates (CF convention)
# This is what makes xarray recognize __owl_version__ etc as coordinates
coords_str = ' '.join(scalar_coords)
for var_name in list(direct_dataset.data_vars) + list(virtual_dataset.keys()):
if var_name in f:
f[var_name].attrs['coordinates'] = coords_str
[docs]
def save_video(
mcam_dataset,
filename: Path,
*,
include_timestamp: bool = True,
engine='ramona',
disk_space_tolerance=0.1e9,
tqdm=None,
mode='w',
_stacklevel_increment=0,
):
"""Save datasets that contain stacks and provide progress information.
This method provides an optimized way to save datasets for speed of writing
and reading. It also provides the user with feedback during data writing
through the form of an optional progress bar.
Parameters
----------
mcam_dataset: xarray Dataset
The mcam_data you wish to save.
filename: Path-like
The filename where the data should be saved. If the filename has no
extension, then an ``'.nc'`` extension is added.
include_timestamp: bool
If set to `True`, this will append a timestamp to the provided
path name. If set to false, no timestamp will be appended to the
path name possibly overwriting any files currently within the previous
path name.
disk_space_tolerance: float
The difference between the free disk space in the given directory and
the dataset to be save in bytes. If the directory has fewer free bytes
than this tolerance plus the dataset size an error will be raised.
mode: str
The mode to open the file in. Valid options are ``'w'`` (write),
or ```x``` exclusive creation, failing if the file already exists.
Returns
-------
save_filepath:
The path where the data was saved.
"""
check_disk_space(filename, 1, mcam_dataset.nbytes, absolute_tolerance=disk_space_tolerance)
if mode not in ['x', 'w']:
raise ValueError(
f"Invalid file creation mode ({mode}). Valid options are 'w' or 'x'")
if tqdm is None:
tqdm = tqdm_base
filename = Path(filename)
ext = filename.suffix
if len(ext) == 0:
ext = '.nc'
stem = filename.stem
if include_timestamp:
filename = filename.parent / (stem + '_' + make_timestamp() + ext)
else:
filename = filename.parent / (stem + ext)
# 2022/02/05 Mark:
# Two step process:n
# Write the metadata using a driver that has support for caching
# This makes, writing small bits of data very fast!
chunked_variables = {}
for var_name in mcam_dataset.data_vars:
variable = mcam_dataset[var_name].variable
# Skip datetime variables as they need special encoding
if np.issubdtype(variable.dtype, np.datetime64):
continue
if np.issubdtype(variable.dtype, np.timedelta64):
continue
# Skip boolean variables as they are not supported by netCDF4
if np.issubdtype(variable.dtype, np.bool_):
continue
# Check if variable needs chunked writing:
# 1. Has explicit chunking in encoding
# 2. Has compression settings (zlib or zstd)
# 3. Is larger than 1GB
has_chunking = variable.encoding.get('chunksizes') is not None
compression = get_compression_type(variable.encoding)
has_compression = compression in compression_map
# Consider variables larger than 100MB as large
# We (Mark and Clay) chose 100MB to make testing simpler
# as default mcam_data.new_dataset is <1GB
is_large = variable.nbytes > 100E6
if has_chunking or has_compression or is_large:
chunked_variables[var_name] = variable
if not chunked_variables:
mcam_dataset.to_netcdf(filename, format='NETCDF4', engine=engine, mode=mode)
return filename
unlimited_dims = mcam_dataset.encoding.get('unlimited_dims', set())
if len(unlimited_dims) != 0:
raise ValueError("We still don't support unlimited_dims in saving videos.")
prepared_variables = {}
for var_name, variable in chunked_variables.items():
# Create a copy of the structures like the encoding so we can adjust it as
# needed, but do not copy the underlying data in the array since it can
# be huge
var_copy = variable.copy(deep=False)
# We are writing to a new file, this original shape isn't meaningful for that file
# Having this doesn't allow the creation of subsets with different chunksizes
var_copy.encoding.pop('original_shape', None)
# 2022/04/11: Clay, Mark
# We think there is abug in xarray
# https://gitlab.com/ramonaoptics/mcam/python-owl/-/issues/813
var_copy.encoding.pop('dtype', None)
chunksizes = var_copy.encoding.get('chunksizes', None)
# Recompute chunksizes if the dimensions have changed
# i.e. in projections / focus selection
if chunksizes is not None and len(chunksizes) != var_copy.ndim:
# https://github.com/pydata/xarray/issues/11028
warn(
f"Variable '{var_name}' has encoding chunksizes with "
f"{len(chunksizes)} dims but data has {var_copy.ndim} dims. "
f"This typically happens after isel() removes a dimension. "
f"Chunksizes will be recomputed. To silence this warning, "
f"clear the encoding before saving: "
f"dataset['{var_name}'].encoding.pop('chunksizes', None)",
stacklevel=2 + _stacklevel_increment,
)
chunksizes = None
if chunksizes is None:
use_image_chunking = (
var_name == 'images' or
var_copy.attrs.get('__owl_settings_type__') == 'analysis_output_mask'
)
chunksizes = _get_chunksizes(mcam_dataset) if use_image_chunking else var_copy.shape
else:
chunksizes = tuple(min(c, s) for c, s in zip(chunksizes, var_copy.shape))
var_copy.encoding["chunksizes"] = chunksizes
compression = get_compression_type(var_copy.encoding)
compression_config = compression_map.get(compression, {})
data_encoder = compression_config.get('encoder')
prepared_variables[var_name] = {
'variable': var_copy,
'data_encoder': data_encoder,
}
metadata = mcam_dataset.drop_vars(list(chunked_variables.keys()), errors='ignore')
metadata.to_netcdf(filename, format='NETCDF4', engine=engine, mode=mode)
def images_first_then_nbytes(x):
return (x != 'images', -prepared_variables[x]['variable'].nbytes)
sorted_variables = sorted(prepared_variables, key=images_first_then_nbytes)
with RamonaH5NetCDFStore.open(filename, mode='a') as store:
for var_name in (tqdm := tqdm(sorted_variables)):
var_info = prepared_variables[var_name]
write_chunked_variable(
store,
var_name,
var_info['variable'],
data_encoder=var_info['data_encoder'],
tqdm=tqdm if callable(tqdm) else None,
)
return filename
def _save_video_streamed(
*,
metadata,
images_variable,
filename: Path,
timeout=None,
frame_timeout=None,
data_read_ready_semaphore=None,
data_write_ready_semaphore=None,
metadata_ready=None,
include_timestamp: bool=True,
subindexed_data=None,
disk_space_tolerance=0.1E9,
ready_to_write_event=None,
mode='w',
tqdm=None,
stop_event=None,
# TODO: cleanup this API.
# We allow to absorb unknown keyword arguments to unify the api with the
# streamed video methods
**kwargs,
):
if frame_timeout is None:
frame_timeout = timeout
if stop_event is None:
stop_event = Event()
if tqdm is None:
tqdm = lambda x, **kwargs: x # noqa
if 'images' in metadata:
raise ValueError("Metadata cannot contain the 'images' variable.")
chunksizes = images_variable.encoding.get("chunksizes", None)
if chunksizes is None:
raise ValueError("Must specify chunksizes in the images variable")
write_step = chunksizes[0]
if images_variable.shape[0] % write_step != 0:
raise ValueError(
"The chunk size of the leading dimension has "
"to be an integer multiple of the data buffer's leading dimension."
f"Got a chunk size of {write_step} and a buffer shape of {images_variable.shape[0]}."
)
if chunksizes[1:] != images_variable.shape[1:]:
raise ValueError(
"The chunk size for all but the first dimension needs to be equal "
"to the shape of the data. "
f"Got a chunk size of {chunksizes} and a buffer shape of {images_variable.shape}."
)
N_frames = len(metadata['frame_number'])
check_disk_space(
filename, 1,
metadata.nbytes + images_variable.data[0].nbytes * N_frames,
absolute_tolerance=disk_space_tolerance
)
filename = Path(filename)
ext = filename.suffix
if len(ext) == 0:
ext = '.nc'
stem = filename.stem
if include_timestamp:
filename = (filename.parent / (stem + '_' + make_timestamp() + ext))
else:
filename = (filename.parent / (stem + ext))
# 2022/02/05 Mark:
# Three step process:
# 1. Write the coordinates that we need for the image_variable
# 2. Let xarray and netcdf manage creating the dataset dynamically
# 3. Write the rest of the metadata
# (likely overwriting the values of the previously written coordinates).
# Write the metadata using a driver that has support for caching
# This makes, writing small bits of data very fast!
# Write just the coordinates of the image array
# Not including any data variables that will be written
variables_to_drop = (
metadata.variables.keys() -
images_variable.dims -
set((
'__owl_version__',
'__sys_version__',
'__owl_sys_info__',
))
)
coords_metadata = metadata.drop_vars(variables_to_drop)
if stop_event.is_set():
return
# This is the first time
coords_metadata.to_netcdf(
filename, format='NETCDF4', engine="h5netcdf",
mode=mode,
)
# This function is created s othat we can have a single point of exit when
# the stop event is set so that we can bail quickly...
def write_worker():
if subindexed_data is not None:
from owl.memory_allocator import full_aligned
write_buffer = full_aligned(
chunksizes,
fill_value=subindexed_data._fill_value,
dtype=images_variable.dtype
)
with RamonaH5NetCDFStore.open(filename, mode='a') as store:
target, data = store.prepare_variable('images', images_variable)
# Get the underlying images dataset
images_h5ds = target.get_array()._h5ds
if ready_to_write_event is not None:
ready_to_write_event.set()
for i in tqdm(range(0, N_frames, write_step)):
# Wait for the data to be read y
if data_read_ready_semaphore is not None:
for j in range(i, min(N_frames, i + write_step)):
if stop_event.is_set():
return
if not data_read_ready_semaphore.acquire(
blocking=True,
timeout=frame_timeout
):
# It might have timedout because the user cancelled
if stop_event.is_set():
return
raise RuntimeError(
f"Timeout on obtained frames to write for frame {j}.")
i_data = i % data.shape[0]
if subindexed_data is None:
write_buffer = np.ascontiguousarray(
data[i_data:i_data + write_step]
)
else:
# Manually slicing because things aren't optimized to be super
# lazy yet
# Ideally, we would slice, then call `copyto`
# https://gitlab.com/ramonaoptics/mcam/python-owl/-/issues/1081
raw_data = subindexed_data._original_data[
i_data:i_data + write_step]
for i_well in np.ndindex(
subindexed_data._input_indexers.shape):
input_slice = subindexed_data._input_indexers[i_well]
output_slice = subindexed_data._output_indexers[i_well]
write_buffer[(slice(None),) + output_slice] = \
raw_data[(slice(None),) + input_slice]
images_h5ds.id.write_direct_chunk(
(i,) + (0,) * (data.ndim - 1),
write_buffer
)
if data_write_ready_semaphore is not None:
to_release = min(N_frames - i, write_step)
data_write_ready_semaphore.release(to_release)
write_worker()
if stop_event.is_set():
# Mark -- 2023/11
# what happens if mode == "a"?? the file is corrupt...fun...
if mode == "w":
filename.unlink()
return
if metadata_ready is not None:
if not metadata_ready.wait(
timeout=timeout
):
raise RuntimeError("Timeout obtained waiting for metadata")
_write_metadata(metadata, filename, mode='a')
return filename