diff --git a/python/omfiles/_chunk_utils.py b/python/omfiles/_chunk_utils.py new file mode 100644 index 0000000..6ed02c6 --- /dev/null +++ b/python/omfiles/_chunk_utils.py @@ -0,0 +1,46 @@ +import math + + +def _validate_chunk_alignment( + data_chunks: tuple, + om_chunks: list[int], + array_shape: tuple, +) -> None: + """ + Validate chunked array blocks for compatibility with OM chunk streaming. + + Every non-last data chunk along each dimension must be an exact multiple + of the corresponding OM chunk size (the last chunk may be smaller). + Additionally, for the leftmost dimension where a data block contains more + than one OM chunk, every trailing dimension must be fully covered by each + data block. This ensures the local chunk traversal inside a block matches + the global file order. + """ + ndim = len(om_chunks) + + for d in range(ndim): + dim_chunks = data_chunks[d] + for i, c in enumerate(dim_chunks[:-1]): + if c % om_chunks[d] != 0: + raise ValueError( + f"Dask chunk size {c} along dimension {d} (block {i}) " + f"is not a multiple of the OM chunk size {om_chunks[d]}." + ) + + first_multi = None + for d in range(ndim): + local_n = max(math.ceil(c / om_chunks[d]) for c in data_chunks[d]) + if local_n > 1: + first_multi = d + break + + if first_multi is not None: + for d in range(first_multi + 1, ndim): + dim_chunks = data_chunks[d] + if not (len(dim_chunks) == 1 and dim_chunks[0] == array_shape[d]): + raise ValueError( + f"Dask blocks have multiple OM chunks in dimension {first_multi}, " + f"but dimension {d} is not fully covered by each dask block " + f"(dask chunks {dim_chunks} vs array size {array_shape[d]}). " + f"Rechunk so trailing dimensions are fully covered." + ) diff --git a/python/omfiles/dask.py b/python/omfiles/dask.py index 27b9b89..f6cef56 100644 --- a/python/omfiles/dask.py +++ b/python/omfiles/dask.py @@ -1,10 +1,10 @@ """Dask array integration for writing to OM files.""" -import math from typing import Iterator, Optional, Sequence import numpy as np +from omfiles._chunk_utils import _validate_chunk_alignment from omfiles._rust import OmFileWriter, OmWriterVariable try: @@ -13,51 +13,6 @@ raise ImportError("omfiles[dask] is required for dask functionality") -def _validate_chunk_alignment( - data_chunks: tuple, - om_chunks: list[int], - array_shape: tuple, -) -> None: - """ - Validate dask chunks are compatible with OM chunks for block-level streaming. - - Every non-last dask chunk along each dimension must be an exact multiple - of the corresponding OM chunk size (the last chunk may be smaller). - Additionally, for the leftmost dimension where a dask block contains more - than one OM chunk, every trailing dimension must be fully covered by each - dask block. This ensures the local chunk traversal inside a block matches - the global file order. - """ - ndim = len(om_chunks) - - for d in range(ndim): - dim_chunks = data_chunks[d] - for i, c in enumerate(dim_chunks[:-1]): - if c % om_chunks[d] != 0: - raise ValueError( - f"Dask chunk size {c} along dimension {d} (block {i}) " - f"is not a multiple of the OM chunk size {om_chunks[d]}." - ) - - first_multi = None - for d in range(ndim): - local_n = max(math.ceil(c / om_chunks[d]) for c in data_chunks[d]) - if local_n > 1: - first_multi = d - break - - if first_multi is not None: - for d in range(first_multi + 1, ndim): - dim_chunks = data_chunks[d] - if not (len(dim_chunks) == 1 and dim_chunks[0] == array_shape[d]): - raise ValueError( - f"Dask blocks have multiple OM chunks in dimension {first_multi}, " - f"but dimension {d} is not fully covered by each dask block " - f"(dask chunks {dim_chunks} vs array size {array_shape[d]}). " - f"Rechunk so trailing dimensions are fully covered." - ) - - def _dask_block_iterator(dask_array: da.Array) -> Iterator[np.ndarray]: """ Yield computed numpy arrays from a dask array in C-order block traversal. diff --git a/python/omfiles/xarray.py b/python/omfiles/xarray.py index 0b1fb40..892f0fd 100644 --- a/python/omfiles/xarray.py +++ b/python/omfiles/xarray.py @@ -1,10 +1,15 @@ """OmFileReader backend for Xarray.""" # ruff: noqa: D101, D102, D105, D107 -from __future__ import annotations +import itertools +import os +import warnings +from typing import Any, Generator import numpy as np +from omfiles._chunk_utils import _validate_chunk_alignment + try: from xarray.core import indexing except ImportError: @@ -21,7 +26,7 @@ from xarray.core.utils import FrozenDict from xarray.core.variable import Variable -from ._rust import OmFileReader, OmVariable +from ._rust import OmFileReader, OmFileWriter, OmVariable # Special metadata child used to declare array dimension names. DIMENSION_KEY = "coordinates" @@ -55,17 +60,21 @@ def open_dataset( class OmDataStore(AbstractDataStore): root_variable: OmFileReader variables_store: dict[str, OmVariable] + _direct_children_store: dict[str, dict[str, OmVariable]] + _attributes_store: dict[str, dict[str, Any]] + _known_arrays_store: dict[str, OmVariable] | None + _dimension_declarations_store: dict[str, tuple[str, ...] | None] def __init__(self, root_variable: OmFileReader): self.root_variable = root_variable self.variables_store = self.root_variable._get_flat_variable_metadata() - self._children_by_parent: dict[str, dict[str, OmVariable]] = {} - self._attributes_by_path: dict[str, dict] = {} - self._known_arrays: dict[str, OmVariable] | None = None - - for path, variable in self.variables_store.items(): - parent_path, _, child_name = path.rpartition("/") - self._children_by_parent.setdefault(parent_path, {})[child_name] = variable + self._direct_children_store = {} + for var_key, variable in self.variables_store.items(): + parent_path, _, child_name = var_key.rpartition("/") + self._direct_children_store.setdefault(parent_path, {})[child_name] = variable + self._attributes_store = {} + self._known_arrays_store = None + self._dimension_declarations_store = {} def get_variables(self): datasets = self._get_datasets(self.root_variable) @@ -78,87 +87,149 @@ def get_attrs(self): return FrozenDict(self._get_attributes_for_variable(self.root_variable, f"/{self.root_variable.name}")) def _get_attributes_for_variable(self, reader: OmFileReader, path: str): - cached = self._attributes_by_path.get(path) - if cached is not None: - return cached + if path in self._attributes_store: + return self._attributes_store[path] attrs = {} direct_children = self._find_direct_children_in_store(path) for k, variable in direct_children.items(): + if k == DIMENSION_KEY: + continue child_reader = reader._init_from_variable(variable) if child_reader.is_scalar: + # Skip scalars that have dimension metadata — they are 0-d coordinate variables, + # not plain attributes. + dim_key = path + "/" + k + "/" + DIMENSION_KEY + if dim_key in self.variables_store: + continue attrs[k] = child_reader.read_scalar() - - self._attributes_by_path[path] = attrs + self._attributes_store[path] = attrs return attrs def _find_direct_children_in_store(self, path: str): - return self._children_by_parent.get(path, {}) + return self._direct_children_store.get(path, {}) def _get_known_arrays(self): - if self._known_arrays is not None: - return self._known_arrays + if self._known_arrays_store is not None: + return self._known_arrays_store arrays = {} - for var_key, variable in self.variables_store.items(): - if self.root_variable._init_from_variable(variable).is_array: - arrays[var_key] = variable - - self._known_arrays = arrays + for var_key, var in self.variables_store.items(): + reader = self.root_variable._init_from_variable(var) + if reader.is_array: + arrays[var_key] = var + self._known_arrays_store = arrays return arrays - def _get_known_dimensions(self, arrays: dict[str, OmVariable]): - """ - Get a set of all dimension names used in the dataset. - - This scans all array variables for their dimension metadata. - """ - dimensions = set() - - # Scan all array variables for dimension names. - for var_key, variable in arrays.items(): - reader = self.root_variable._init_from_variable(variable) - attrs = self._get_attributes_for_variable(reader, var_key) - if DIMENSION_KEY in attrs: - dim_names = attrs[DIMENSION_KEY] - if isinstance(dim_names, str): - dimensions.update(dim_names.split()) - elif isinstance(dim_names, list): - dimensions.update(dim_names) - return dimensions + @staticmethod + def _display_path(path: str) -> str: + return "/" + path.lstrip("/") + + def _get_dimension_declaration(self, path: str) -> tuple[str, ...] | None: + if path in self._dimension_declarations_store: + return self._dimension_declarations_store[path] + + dimension_variable = self._find_direct_children_in_store(path).get(DIMENSION_KEY) + if dimension_variable is None: + self._dimension_declarations_store[path] = None + return None + + dimension_path = f"{self._display_path(path).rstrip('/')}/{DIMENSION_KEY}" + dimension_reader = self.root_variable._init_from_variable(dimension_variable) + if not dimension_reader.is_scalar: + raise ValueError(f"Invalid dimension metadata at '{dimension_path}': expected a scalar string.") + + value = dimension_reader.read_scalar() + if not isinstance(value, str): + raise ValueError( + f"Invalid dimension metadata at '{dimension_path}': expected a string, got {type(value).__name__}." + ) + + declaration = tuple(value.split()) + self._dimension_declarations_store[path] = declaration + return declaration + + def _get_parent_reader(self, path: str) -> OmFileReader | None: + if path == f"/{self.root_variable.name}": + return self.root_variable + + variable = self.variables_store.get(path) + if variable is None: + return None + return self.root_variable._init_from_variable(variable) + + def _validate_array_dimensions(self, path: str, declaration: tuple[str, ...], ndim: int) -> None: + if len(declaration) == ndim: + return + + dimension_path = f"{self._display_path(path).rstrip('/')}/{DIMENSION_KEY}" + raise ValueError( + f"Invalid dimension metadata at '{dimension_path}' for array '{self._display_path(path)}': " + f"declared {len(declaration)} dimension(s) {declaration}, but the array has {ndim}." + ) def _get_datasets(self, reader: OmFileReader): datasets = {} - arrays = self._get_known_arrays() - known_dimensions = self._get_known_dimensions(arrays) + known_arrays = self._get_known_arrays() + arrays = [] + sibling_dimensions: dict[str, set[str]] = {} - for var_key, variable in arrays.items(): + for var_key, variable in known_arrays.items(): child_reader = reader._init_from_variable(variable) backend_array = OmBackendArray(reader=child_reader) shape = backend_array.reader.shape - - # Get attributes to check for dimension information. attrs = self._get_attributes_for_variable(child_reader, var_key) - attrs_for_var = {attr_k: attr_v for attr_k, attr_v in attrs.items() if attr_k != DIMENSION_KEY} - - # Look for dimension names in the dimension metadata. - if DIMENSION_KEY in attrs: - dim_names = attrs[DIMENSION_KEY] - if isinstance(dim_names, str): - # With no explicit separator, split() treats consecutive whitespace as one separator and - # does not produce empty names for leading or trailing whitespace. - dim_names = dim_names.split() + declaration = self._get_dimension_declaration(var_key) + if declaration is not None: + self._validate_array_dimensions(var_key, declaration, len(shape)) + parent_path = var_key.rpartition("/")[0] + sibling_dimensions.setdefault(parent_path, set()).update(declaration) + arrays.append((var_key, backend_array, shape, attrs, declaration)) + + parent_declarations: dict[str, tuple[str, ...] | None] = {} + for var_key, backend_array, shape, attrs, declaration in arrays: + if declaration is not None: + dim_names = declaration else: - # Default to generic dimension names if not specified. - dim_names = [f"dim{i}" for i in range(len(shape))] - - # Check if this variable is itself a dimension variable. - variable_name = var_key.split("/")[-1] - if len(shape) == 1 and variable_name in known_dimensions: - dim_names = [variable_name] + parent_path, _, variable_name = var_key.rpartition("/") + if parent_path not in parent_declarations: + parent_declarations[parent_path] = self._get_dimension_declaration(parent_path) + parent_declaration = parent_declarations[parent_path] + known_parent_dimensions = sibling_dimensions.get(parent_path, set()) + if parent_declaration is not None: + known_parent_dimensions = known_parent_dimensions.union(parent_declaration) + + parent_reader = self._get_parent_reader(parent_path) + if len(shape) == 1 and variable_name in known_parent_dimensions: + dim_names = (variable_name,) + elif ( + parent_reader is not None + and parent_reader.is_group + and parent_declaration is not None + and len(parent_declaration) == len(shape) + ): + dim_names = parent_declaration + else: + dim_names = tuple(f"dim{i}" for i in range(len(shape))) data = indexing.LazilyIndexedArray(backend_array) - datasets[var_key] = Variable(dims=dim_names, data=data, attrs=attrs_for_var, encoding=None, fastpath=True) + datasets[var_key] = Variable(dims=dim_names, data=data, attrs=attrs, encoding=None, fastpath=True) + + # Handle 0-d (scalar) variables that have dimension metadata. + for var_key, var in self.variables_store.items(): + if var_key in datasets: + continue + child_reader = reader._init_from_variable(var) + if not child_reader.is_scalar: + continue + declaration = self._get_dimension_declaration(var_key) + if declaration is None: + continue + self._validate_array_dimensions(var_key, declaration, 0) + scalar_value = child_reader.read_scalar() + attrs = self._get_attributes_for_variable(child_reader, var_key) + datasets[var_key] = Variable(dims=(), data=np.array(scalar_value), attrs=attrs) + return datasets def close(self): @@ -187,3 +258,234 @@ def __getitem__(self, key: indexing.ExplicitIndexer) -> np.typing.ArrayLike: indexing.IndexingSupport.BASIC, self.reader.__getitem__, ) + + +def _write_scalar_safe(writer: OmFileWriter, value: Any, name: str) -> OmVariable | None: + """Write a scalar, returning None and warning if the type is unsupported.""" + try: + return writer.write_scalar(value, name=name) + except (ValueError, TypeError) as e: + warnings.warn( + f"Skipping attribute '{name}' with value {value!r}: {e}", + UserWarning, + stacklevel=3, + ) + return None + + +def _chunked_block_iterator(data: Any) -> Generator[np.ndarray, None, None]: + """ + Yield numpy arrays from a chunked array in C-order block traversal. + + Works with any array that exposes ``.numblocks``, ``.blocks[idx]``, + and ``.compute()`` (e.g. dask arrays). No dask import required. + """ + block_index_ranges = [range(n) for n in data.numblocks] + for block_indices in itertools.product(*block_index_ranges): + block = data.blocks[block_indices] + if hasattr(block, "compute"): + yield block.compute() + else: + yield np.asarray(block) + + +def _resolve_chunks_for_variable( + var_name: str, + var: Variable, + encoding: dict[str, dict[str, Any]] | None, + global_chunks: dict[str, int] | None, + data_chunks: tuple | None = None, +) -> list[int]: + """Resolve chunk sizes for a variable using the priority chain.""" + if encoding and var_name in encoding and "chunks" in encoding[var_name]: + return list(encoding[var_name]["chunks"]) + + if global_chunks is not None: + return [global_chunks.get(dim, min(size, 512)) for dim, size in zip(var.dims, var.shape)] + + if data_chunks is not None: + return [int(c[0]) for c in data_chunks] + + return [min(size, 512) for size in var.shape] + + +def _resolve_encoding_for_variable( + var_name: str, + encoding: dict[str, dict[str, Any]] | None, + global_scale_factor: float, + global_add_offset: float, + global_compression: str, +) -> tuple[float, float, str]: + """Resolve compression parameters for a variable.""" + var_enc = (encoding or {}).get(var_name, {}) + sf = var_enc.get("scale_factor", global_scale_factor) + ao = var_enc.get("add_offset", global_add_offset) + comp = var_enc.get("compression", global_compression) + return sf, ao, comp + + +def _validate_om_name(name: Any, description: str) -> None: + """Validate a name that will become an OM hierarchy child or dimension token.""" + if not isinstance(name, str): + raise ValueError(f"{description} must be a string, got {type(name).__name__}.") + if not name: + raise ValueError(f"{description} must not be empty.") + if "/" in name: + raise ValueError(f"{description} '{name}' must not contain '/'.") + if any(character.isspace() for character in name): + raise ValueError(f"{description} '{name}' must not contain whitespace.") + + +def _validate_dataset_for_writing(ds: Dataset) -> None: + """Validate that a dataset can be represented without extending the Open-Meteo convention.""" + variable_names = set(ds.variables) + + for dimension_name in ds.dims: + _validate_om_name(dimension_name, "Dimension name") + + for variable_name, variable in ds.variables.items(): + _validate_om_name(variable_name, "Variable name") + if DIMENSION_KEY in variable.attrs: + raise ValueError( + f"Variable '{variable_name}' attribute '{DIMENSION_KEY}' conflicts with OM dimension metadata." + ) + for attribute_name in variable.attrs: + _validate_om_name(attribute_name, f"Attribute name on variable '{variable_name}'") + + if np.issubdtype(variable.dtype, np.datetime64) or np.issubdtype(variable.dtype, np.timedelta64): + raise TypeError( + f"Variable '{variable_name}' has dtype {variable.dtype}. " + "OM files do not support datetime64/timedelta64 natively. " + "Convert to a numeric type before writing." + ) + + for coordinate_name, coordinate in ds.coords.items(): + if coordinate.ndim != 1 or coordinate.dims != (coordinate_name,): + raise ValueError( + f"Coordinate '{coordinate_name}' with dimensions {coordinate.dims} is not supported. " + "OM dataset writing supports only one-dimensional dimension coordinates." + ) + + for attribute_name in ds.attrs: + _validate_om_name(attribute_name, "Global attribute name") + if attribute_name == DIMENSION_KEY: + raise ValueError(f"Global attribute '{DIMENSION_KEY}' conflicts with OM dimension metadata.") + if attribute_name in variable_names: + raise ValueError(f"Global attribute '{attribute_name}' conflicts with a dataset variable of the same name.") + + +def write_dataset( + ds: Dataset, + path: str | os.PathLike, + *, + fs: Any | None = None, + encoding: dict[str, dict[str, Any]] | None = None, + chunks: dict[str, int] | None = None, + scale_factor: float = 1.0, + add_offset: float = 0.0, + compression: str = "pfor_delta_2d", +) -> None: + """ + Write an xarray Dataset to an OM file. + + The resulting file can be read back with ``xr.open_dataset(path, engine="om")``. + + Only one-dimensional dimension coordinates are supported. Auxiliary and + scalar coordinates cannot be represented by the Open-Meteo coordinate + convention and are rejected before the output file is created. + + Args: + ds: The xarray Dataset to write. + path: Output file path (local path or path within the fsspec filesystem). + fs: Optional fsspec filesystem object. When provided, the file is written + via ``OmFileWriter.from_fsspec(fs, path)`` instead of the default + local-file writer. + encoding: Per-variable overrides. Keys per variable: ``"chunks"``, + ``"scale_factor"``, ``"add_offset"``, ``"compression"``. + chunks: Global default chunk sizes as ``{dim_name: chunk_size}``. + scale_factor: Global default scale factor for float compression. + add_offset: Global default offset for float compression. + compression: Global default compression algorithm. + """ + _validate_dataset_for_writing(ds) + path = str(path) + if fs is not None: + writer = OmFileWriter.from_fsspec(fs, path) + else: + writer = OmFileWriter(path) + all_children: list[OmVariable] = [] + + def _write_variable(name: str, var: Variable, is_dim_coord: bool) -> None: + """Write a data variable or dimension coordinate.""" + dim_var = writer.write_scalar(" ".join(var.dims), name=DIMENSION_KEY) + var_children: list[OmVariable] = [dim_var] + + for attr_name, attr_value in var.attrs.items(): + scalar = _write_scalar_safe(writer, attr_value, attr_name) + if scalar is not None: + var_children.append(scalar) + + if var.ndim == 0: + om_var = writer.write_scalar( + var.values[()], + name=name, + children=var_children if var_children else None, + ) + all_children.append(om_var) + return + + data = var.data + is_chunked = not is_dim_coord and hasattr(data, "chunks") and data.chunks is not None + + if is_dim_coord: + resolved_chunks = [var.shape[0]] + else: + resolved_chunks = _resolve_chunks_for_variable( + name, + var, + encoding, + chunks, + data_chunks=data.chunks if is_chunked else None, + ) + + sf, ao, comp = _resolve_encoding_for_variable(name, encoding, scale_factor, add_offset, compression) + + if is_chunked: + _validate_chunk_alignment(data.chunks, resolved_chunks, var.shape) + om_var = writer.write_array_streaming( + dimensions=[int(d) for d in var.shape], + chunks=[int(c) for c in resolved_chunks], + chunk_iterator=_chunked_block_iterator(data), + dtype=var.dtype, + scale_factor=sf, + add_offset=ao, + compression=comp, + name=name, + children=var_children if var_children else None, + ) + else: + om_var = writer.write_array( + var.values, + chunks=resolved_chunks, + scale_factor=sf, + add_offset=ao, + compression=comp, + name=name, + children=var_children if var_children else None, + ) + all_children.append(om_var) + + for var_name in ds.data_vars: + _write_variable(var_name, ds[var_name].variable, is_dim_coord=False) + + for coord_name in ds.coords: + coord = ds.coords[coord_name] + _write_variable(coord_name, coord.variable, is_dim_coord=True) + + for attr_name, attr_value in ds.attrs.items(): + scalar = _write_scalar_safe(writer, attr_value, attr_name) + if scalar is not None: + all_children.append(scalar) + + root_var = writer.write_group(name="", children=all_children) + writer.close(root_var) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 31edebf..4afa6bc 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -11,6 +11,7 @@ import xarray as xr from fsspec.implementations.local import LocalFileSystem from fsspec.implementations.memory import MemoryFileSystem +from omfiles.xarray import write_dataset from s3fs import S3FileSystem from .test_utils import filter_numpy_size_warning, find_chunk_for_timestamp @@ -224,10 +225,10 @@ def test_write_hierarchical_fsspec(memory_fs): humid_var = writer.write_array(humidity, chunks=[5, 5, 5], name="humidity", scale_factor=100.0) temp_units = writer.write_scalar("celsius", name="units") temp_desc = writer.write_scalar("Surface temperature", name="description") - temp_dims = writer.write_scalar("lat,lon,time", name="_ARRAY_DIMENSIONS") + temp_dims = writer.write_scalar("lat lon time", name="coordinates") humid_units = writer.write_scalar("percent", name="units") humid_desc = writer.write_scalar("Relative humidity", name="description") - humid_dims = writer.write_scalar("lat,lon,time", name="_ARRAY_DIMENSIONS") + humid_dims = writer.write_scalar("lat lon time", name="coordinates") temp_metadata = writer.write_group("temp_metadata", [temp_units, temp_desc, temp_dims]) humid_metadata = writer.write_group("humid_metadata", [humid_units, humid_desc, humid_dims]) root_group = writer.write_group("weather_data", [temp_var, humid_var, temp_metadata, humid_metadata]) @@ -277,3 +278,117 @@ def read_slice(idx, start): assert len(results) == num_threads for i, arr in enumerate(results): np.testing.assert_array_equal(arr, data[i * slice_size : (i + 1) * slice_size, :]) + + +# --- write_dataset fsspec tests --- + + +@filter_numpy_size_warning +def test_write_dataset_memory_fsspec(memory_fs): + """write_dataset with fs= writes to a memory filesystem and reads back.""" + ds = xr.Dataset( + {"temperature": (["lat", "lon"], np.random.rand(5, 5).astype(np.float32))}, + coords={ + "lat": np.arange(5, dtype=np.float32), + "lon": np.arange(5, dtype=np.float32), + }, + attrs={"description": "Test dataset"}, + ) + write_dataset(ds, "dataset_test.om", fs=memory_fs, scale_factor=100000.0) + assert_file_exists(memory_fs, "dataset_test.om") + + reader = omfiles.OmFileReader.from_fsspec(memory_fs, "dataset_test.om") + assert reader.num_children > 0 + reader.close() + + +@filter_numpy_size_warning +def test_write_dataset_memory_fsspec_roundtrip(memory_fs): + """Full roundtrip: write_dataset via memory fs, read back with xarray.""" + temperature_data = np.random.rand(5, 5).astype(np.float32) + ds = xr.Dataset( + {"temperature": (["lat", "lon"], temperature_data)}, + coords={ + "lat": np.arange(5, dtype=np.float32), + "lon": np.arange(5, dtype=np.float32), + }, + attrs={"description": "fsspec roundtrip test"}, + ) + path = "roundtrip_dataset.om" + write_dataset(ds, path, fs=memory_fs, scale_factor=100000.0) + + # Dump from memory fs to a temp file so xarray can read it back + with tempfile.NamedTemporaryFile(suffix=".om", delete=False) as tmp: + tmp.write(memory_fs.cat(path)) + tmp_path = tmp.name + try: + ds2 = xr.open_dataset(tmp_path, engine="om") + np.testing.assert_array_almost_equal(ds2["temperature"].values, temperature_data, decimal=4) + np.testing.assert_array_equal(ds2.coords["lat"].values, ds.coords["lat"].values) + np.testing.assert_array_equal(ds2.coords["lon"].values, ds.coords["lon"].values) + assert ds2.attrs["description"] == "fsspec roundtrip test" + finally: + os.unlink(tmp_path) + + +@filter_numpy_size_warning +def test_write_dataset_local_fsspec(local_fs): + """write_dataset with a local fsspec filesystem produces a valid file.""" + ds = xr.Dataset( + {"temperature": (["lat", "lon"], np.random.rand(8, 8).astype(np.float32))}, + coords={ + "lat": np.arange(8, dtype=np.float32), + "lon": np.arange(8, dtype=np.float32), + }, + ) + with tempfile.NamedTemporaryFile(suffix=".om", delete=False) as tmp: + tmp_path = tmp.name + try: + write_dataset(ds, tmp_path, fs=local_fs, scale_factor=100000.0) + assert os.path.exists(tmp_path) + assert os.path.getsize(tmp_path) > 0 + + ds2 = xr.open_dataset(tmp_path, engine="om") + np.testing.assert_array_almost_equal(ds2["temperature"].values, ds["temperature"].values, decimal=4) + finally: + os.unlink(tmp_path) + + +@filter_numpy_size_warning +def test_write_dataset_fs_none_backward_compatible(): + """Passing fs=None behaves identically to the default (local path).""" + ds = xr.Dataset( + {"data": (["x"], np.arange(5, dtype=np.float32))}, + ) + with tempfile.NamedTemporaryFile(suffix=".om", delete=False) as tmp: + tmp_path = tmp.name + try: + write_dataset(ds, tmp_path, fs=None) + ds2 = xr.open_dataset(tmp_path, engine="om") + np.testing.assert_array_equal(ds2["data"].values, ds["data"].values) + finally: + os.unlink(tmp_path) + + +@filter_numpy_size_warning +def test_write_and_read_dataset_fsspec_roundtrip(memory_fs): + """Full fsspec roundtrip: write_dataset via fs, read back via fsspec file object.""" + temperature_data = np.random.rand(5, 5).astype(np.float32) + ds = xr.Dataset( + {"temperature": (["lat", "lon"], temperature_data)}, + coords={ + "lat": np.arange(5, dtype=np.float32), + "lon": np.arange(5, dtype=np.float32), + }, + attrs={"description": "full fsspec roundtrip"}, + ) + path = "fsspec_full_roundtrip.om" + write_dataset(ds, path, fs=memory_fs, scale_factor=100000.0) + + # Read back via fsspec.core.OpenFile, which xr.open_dataset supports + backend = fsspec.core.OpenFile(memory_fs, path, mode="rb") + ds2 = xr.open_dataset(backend, engine="om") + np.testing.assert_array_almost_equal(ds2["temperature"].values, temperature_data, decimal=4) + np.testing.assert_array_equal(ds2.coords["lat"].values, ds.coords["lat"].values) + np.testing.assert_array_equal(ds2.coords["lon"].values, ds.coords["lon"].values) + assert ds2.attrs["description"] == "full fsspec roundtrip" diff --git a/tests/test_xarray.py b/tests/test_xarray.py index c98f71b..ca1b4a8 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -3,6 +3,7 @@ import pytest import xarray as xr from omfiles import OmFileReader, OmFileWriter +from omfiles.xarray import DIMENSION_KEY, write_dataset from xarray.core import indexing from .test_utils import create_test_om_file, filter_numpy_size_warning @@ -77,10 +78,70 @@ def test_xarray_metadata_discovery_is_cached(empty_temp_om_file): variable_reader = reader._init_from_variable(variable) attrs = store._get_attributes_for_variable(variable_reader, path) - assert attrs == {"coordinates": "x", "units": "m"} + assert attrs == {"units": "m"} assert store._get_attributes_for_variable(variable_reader, path) is attrs +@filter_numpy_size_warning +def test_xarray_open_meteo_run_coordinates(empty_temp_om_file): + writer = OmFileWriter(empty_temp_om_file) + coordinates = writer.write_scalar("lat lon time", name="coordinates") + time = writer.write_array(np.arange(4, dtype=np.int64), chunks=[4], name="time") + root = writer.write_array( + np.arange(24, dtype=np.float32).reshape(2, 3, 4), + chunks=[2, 3, 4], + name="", + children=[time, coordinates], + ) + writer.close(root) + + ds = xr.open_dataset(empty_temp_om_file, engine="om") + + assert ds[""].dims == ("lat", "lon", "time") + assert ds["time"].dims == ("time",) + assert "time" in ds.coords + assert DIMENSION_KEY not in ds.attrs + + +@filter_numpy_size_warning +def test_xarray_open_meteo_group_coordinates(empty_temp_om_file): + writer = OmFileWriter(empty_temp_om_file) + coordinates = writer.write_scalar("lat lon", name="coordinates") + lat = writer.write_array(np.arange(2, dtype=np.float32), chunks=[2], name="lat") + lon = writer.write_array(np.arange(3, dtype=np.float32), chunks=[3], name="lon") + temperature = writer.write_array( + np.zeros((2, 3), dtype=np.float32), + chunks=[2, 3], + name="temperature", + ) + root = writer.write_group("", children=[lat, lon, temperature, coordinates]) + writer.close(root) + + ds = xr.open_dataset(empty_temp_om_file, engine="om") + + assert ds["temperature"].dims == ("lat", "lon") + assert ds["lat"].dims == ("lat",) + assert ds["lon"].dims == ("lon",) + assert "lat" in ds.coords + assert "lon" in ds.coords + + +@filter_numpy_size_warning +def test_xarray_rejects_direct_dimension_rank_mismatch(empty_temp_om_file): + writer = OmFileWriter(empty_temp_om_file) + coordinates = writer.write_scalar("lat lon", name="coordinates") + root = writer.write_array( + np.zeros(4, dtype=np.float32), + chunks=[4], + name="data", + children=[coordinates], + ) + writer.close(root) + + with pytest.raises(ValueError, match=r"declared 2 dimension\(s\).*array has 1"): + xr.open_dataset(empty_temp_om_file, engine="om") + + @filter_numpy_size_warning def test_xarray_hierarchical_file(empty_temp_om_file): # Create test data @@ -183,3 +244,364 @@ def test_xarray_hierarchical_file(empty_temp_om_file): mean_temp = ds["temperature"].mean(dim="TIME") assert mean_temp.shape == (5, 5, 5) assert mean_temp.dims == ("LATITUDE", "LONGITUDE", "ALTITUDE") + + +@filter_numpy_size_warning +def test_write_dataset_basic_roundtrip(empty_temp_om_file): + ds = xr.Dataset( + {"temperature": (["lat", "lon"], np.random.rand(5, 5).astype(np.float32))}, + coords={ + "lat": xr.DataArray(np.arange(5, dtype=np.float32), dims="lat", attrs={"units": "degrees_north"}), + "lon": np.arange(5, dtype=np.float32), + }, + attrs={"description": "Test dataset"}, + ) + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_almost_equal(ds2["temperature"].values, ds["temperature"].values, decimal=4) + np.testing.assert_array_equal(ds2.coords["lat"].values, ds.coords["lat"].values) + np.testing.assert_array_equal(ds2.coords["lon"].values, ds.coords["lon"].values) + assert ds2.coords["lat"].attrs == {"units": "degrees_north"} + assert ds2.attrs["description"] == "Test dataset" + + +@filter_numpy_size_warning +def test_write_dataset_standalone_dimension_coordinate(empty_temp_om_file): + ds = xr.Dataset( + coords={ + "x": xr.DataArray( + np.arange(4, dtype=np.float32), + dims="x", + attrs={"units": "m"}, + ) + } + ) + + write_dataset(ds, empty_temp_om_file) + loaded = xr.open_dataset(empty_temp_om_file, engine="om") + + assert set(loaded.coords) == {"x"} + assert loaded["x"].dims == ("x",) + assert loaded["x"].attrs == {"units": "m"} + + +@filter_numpy_size_warning +def test_write_dataset_hierarchical_roundtrip(empty_temp_om_file): + """Mirrors test_xarray_hierarchical_file but uses write_dataset.""" + temperature_data = np.random.rand(5, 5, 5, 10).astype(np.float32) + precipitation_data = np.random.rand(5, 5, 10).astype(np.float32) + + ds = xr.Dataset( + { + "temperature": ( + ["LATITUDE", "LONGITUDE", "ALTITUDE", "TIME"], + temperature_data, + {"units": "celsius", "description": "Surface temperature"}, + ), + "precipitation": ( + ["LATITUDE", "LONGITUDE", "TIME"], + precipitation_data, + {"units": "mm", "description": "Precipitation"}, + ), + }, + coords={ + "LATITUDE": np.arange(5, dtype=np.float32), + "LONGITUDE": np.arange(5, dtype=np.float32), + "ALTITUDE": np.arange(5, dtype=np.float32), + "TIME": np.arange(10, dtype=np.float32), + }, + attrs={"description": "This is a hierarchical OM File"}, + ) + + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + assert ds2.attrs["description"] == "This is a hierarchical OM File" + assert set(ds2.data_vars) == {"temperature", "precipitation"} + + np.testing.assert_array_almost_equal(ds2["temperature"].values, temperature_data, decimal=4) + assert ds2["temperature"].dims == ("LATITUDE", "LONGITUDE", "ALTITUDE", "TIME") + assert ds2["temperature"].attrs["units"] == "celsius" + assert ds2["temperature"].attrs["description"] == "Surface temperature" + + np.testing.assert_array_almost_equal(ds2["precipitation"].values, precipitation_data, decimal=4) + assert ds2["precipitation"].dims == ("LATITUDE", "LONGITUDE", "TIME") + assert ds2["precipitation"].attrs["units"] == "mm" + + assert ds2["LATITUDE"].dims == ("LATITUDE",) + assert ds2["LONGITUDE"].dims == ("LONGITUDE",) + assert ds2["ALTITUDE"].dims == ("ALTITUDE",) + assert ds2["TIME"].dims == ("TIME",) + + +@filter_numpy_size_warning +def test_write_dataset_per_variable_encoding(empty_temp_om_file): + ds = xr.Dataset( + { + "high_res": (["x", "y"], np.random.rand(10, 10).astype(np.float32)), + "low_res": (["x", "y"], np.random.rand(10, 10).astype(np.float32)), + }, + coords={ + "x": np.arange(10, dtype=np.float32), + "y": np.arange(10, dtype=np.float32), + }, + ) + + write_dataset( + ds, + empty_temp_om_file, + scale_factor=1000.0, + encoding={ + "high_res": {"scale_factor": 100000.0, "chunks": [5, 5]}, + "low_res": {"chunks": [10, 10]}, + }, + ) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_almost_equal(ds2["high_res"].values, ds["high_res"].values, decimal=4) + np.testing.assert_array_almost_equal(ds2["low_res"].values, ds["low_res"].values, decimal=2) + + +@filter_numpy_size_warning +@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.uint32, np.uint64]) +def test_write_dataset_integer_dtypes(dtype, empty_temp_om_file): + data = np.arange(25, dtype=dtype).reshape(5, 5) + ds = xr.Dataset({"values": (["x", "y"], data)}) + + write_dataset(ds, empty_temp_om_file) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_equal(ds2["values"].values, data) + assert ds2["values"].dtype == dtype + + +@filter_numpy_size_warning +def test_write_dataset_unsupported_attrs_warning(empty_temp_om_file): + ds = xr.Dataset( + {"data": (["x"], np.arange(5, dtype=np.float32))}, + attrs={"valid": "hello", "invalid": [1, 2, 3]}, + ) + + with pytest.warns(UserWarning, match="Skipping attribute"): + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + assert ds2.attrs["valid"] == "hello" + assert "invalid" not in ds2.attrs + + +def test_write_dataset_datetime_raises(empty_temp_om_file): + time_values = np.array( + ["2020-01-01", "2020-01-02", "2020-01-03", "2020-01-04", "2020-01-05"], dtype="datetime64[ns]" + ) + ds = xr.Dataset( + {"data": (["time"], np.arange(5, dtype=np.float32))}, + coords={"time": time_values}, + ) + + with pytest.raises(TypeError, match="datetime64"): + write_dataset(ds, empty_temp_om_file) + + +def test_write_dataset_scalar_data_variable_attrs(empty_temp_om_file): + ds = xr.Dataset( + { + "height": xr.DataArray( + np.float32(12.5), + attrs={"units": "m", "long_name": "station height"}, + ) + } + ) + + write_dataset(ds, empty_temp_om_file) + loaded = xr.open_dataset(empty_temp_om_file, engine="om") + + assert "height" in loaded.data_vars + assert loaded["height"].ndim == 0 + np.testing.assert_almost_equal(float(loaded["height"]), 12.5) + assert loaded["height"].attrs == {"units": "m", "long_name": "station height"} + + +@filter_numpy_size_warning +def test_write_dataset_scalar_coordinate(empty_temp_om_file): + """Scalar coordinates are outside the Open-Meteo coordinate convention.""" + temperature_data = np.random.rand(5, 5).astype(np.float32) + ds = xr.Dataset( + {"temperature": (["lat", "lon"], temperature_data)}, + coords={ + "lat": np.arange(5, dtype=np.float32), + "lon": np.arange(5, dtype=np.float32), + "time": xr.DataArray(np.float32(42.0), attrs={"long_name": "forecast step"}), + }, + ) + with pytest.raises(ValueError, match=r"Coordinate 'time'.*only one-dimensional dimension coordinates"): + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + + +@filter_numpy_size_warning +def test_write_dataset_non_dimension_coordinate(empty_temp_om_file): + """Auxiliary coordinates are outside the Open-Meteo coordinate convention.""" + valid_time_data = np.arange(6, dtype=np.float32) + ds = xr.Dataset( + {"t2m": (("step", "lat"), np.zeros((6, 10), dtype=np.float32))}, + coords={"valid_time": ("step", valid_time_data)}, + ) + + with pytest.raises(ValueError, match=r"Coordinate 'valid_time'.*only one-dimensional dimension coordinates"): + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + + +@pytest.mark.parametrize( + ("attrs", "variable_attrs", "match"), + [ + ({"coordinates": "custom"}, {}, "Global attribute 'coordinates'"), + ({"data": "custom"}, {}, "conflicts with a dataset variable"), + ({}, {"coordinates": "custom"}, "Variable 'data' attribute 'coordinates'"), + ], +) +def test_write_dataset_rejects_metadata_name_collisions(empty_temp_om_file, attrs, variable_attrs, match): + ds = xr.Dataset( + {"data": (["x"], np.arange(4, dtype=np.float32), variable_attrs)}, + attrs=attrs, + ) + + with pytest.raises(ValueError, match=match): + write_dataset(ds, empty_temp_om_file) + + +@pytest.mark.parametrize("variable_name", ["", "bad/name", "bad name", 1]) +def test_write_dataset_rejects_invalid_variable_names(empty_temp_om_file, variable_name): + ds = xr.Dataset({variable_name: (["x"], np.arange(4, dtype=np.float32))}) + + with pytest.raises(ValueError, match="Variable name"): + write_dataset(ds, empty_temp_om_file) + + +@filter_numpy_size_warning +def test_write_dataset_dask_roundtrip(empty_temp_om_file): + da = pytest.importorskip("dask.array") + + np_data = np.random.rand(10, 20).astype(np.float32) + dask_data = da.from_array(np_data, chunks=(5, 10)) + + ds = xr.Dataset( + {"temperature": (["lat", "lon"], dask_data)}, + coords={ + "lat": np.arange(10, dtype=np.float32), + "lon": np.arange(20, dtype=np.float32), + }, + ) + + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_almost_equal(ds2["temperature"].values, np_data, decimal=4) + np.testing.assert_array_equal(ds2.coords["lat"].values, ds.coords["lat"].values) + np.testing.assert_array_equal(ds2.coords["lon"].values, ds.coords["lon"].values) + + +@filter_numpy_size_warning +def test_write_dataset_dask_mixed_variables(empty_temp_om_file): + da = pytest.importorskip("dask.array") + + np_temp = np.random.rand(10, 20).astype(np.float32) + dask_temp = da.from_array(np_temp, chunks=(5, 10)) + np_precip = np.random.rand(10, 20).astype(np.float32) + + ds = xr.Dataset( + { + "temperature": (["lat", "lon"], dask_temp), + "precipitation": (["lat", "lon"], np_precip), + }, + coords={ + "lat": np.arange(10, dtype=np.float32), + "lon": np.arange(20, dtype=np.float32), + }, + ) + + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_almost_equal(ds2["temperature"].values, np_temp, decimal=4) + np.testing.assert_array_almost_equal(ds2["precipitation"].values, np_precip, decimal=4) + + +@filter_numpy_size_warning +def test_write_dataset_dask_boundary_chunks(empty_temp_om_file): + da = pytest.importorskip("dask.array") + + np_data = np.arange(91, dtype=np.float32).reshape(7, 13) + dask_data = da.from_array(np_data, chunks=(4, 5)) + + ds = xr.Dataset({"data": (["x", "y"], dask_data)}) + + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_almost_equal(ds2["data"].values, np_data, decimal=4) + + +@filter_numpy_size_warning +def test_write_dataset_dask_with_attributes(empty_temp_om_file): + da = pytest.importorskip("dask.array") + + np_data = np.random.rand(5, 5).astype(np.float32) + dask_data = da.from_array(np_data, chunks=(5, 5)) + + ds = xr.Dataset( + {"temp": (["x", "y"], dask_data, {"units": "K", "long_name": "temperature"})}, + attrs={"source": "test"}, + ) + + write_dataset(ds, empty_temp_om_file, scale_factor=100000.0) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_almost_equal(ds2["temp"].values, np_data, decimal=4) + assert ds2["temp"].attrs["units"] == "K" + assert ds2["temp"].attrs["long_name"] == "temperature" + assert ds2.attrs["source"] == "test" + + +@filter_numpy_size_warning +@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.uint32]) +def test_write_dataset_dask_integer_dtypes(dtype, empty_temp_om_file): + da = pytest.importorskip("dask.array") + + np_data = np.arange(25, dtype=dtype).reshape(5, 5) + dask_data = da.from_array(np_data, chunks=(5, 5)) + + ds = xr.Dataset({"values": (["x", "y"], dask_data)}) + + write_dataset(ds, empty_temp_om_file) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_equal(ds2["values"].values, np_data) + assert ds2["values"].dtype == dtype + + +@filter_numpy_size_warning +def test_write_dataset_dask_larger_chunks_than_om(empty_temp_om_file): + """Dask blocks larger than OM chunks with explicit smaller OM chunk sizes.""" + da = pytest.importorskip("dask.array") + + np_data = np.random.rand(10, 20).astype(np.float32) + dask_data = da.from_array(np_data, chunks=(10, 20)) + + ds = xr.Dataset( + {"temperature": (["lat", "lon"], dask_data)}, + coords={ + "lat": np.arange(10, dtype=np.float32), + "lon": np.arange(20, dtype=np.float32), + }, + ) + + write_dataset( + ds, + empty_temp_om_file, + chunks={"lat": 5, "lon": 10}, + scale_factor=100000.0, + ) + ds2 = xr.open_dataset(empty_temp_om_file, engine="om") + + np.testing.assert_array_almost_equal(ds2["temperature"].values, np_data, decimal=4)