From 590ea48712cf6a94f5305e8dab889da7e9240be7 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 27 May 2026 12:41:40 -0700 Subject: [PATCH 1/8] adios2: Add write support `save` can now write datasets to Adios2 .bp format. Includes a round-trip unit test. --- pyproject.toml | 3 + xbout/adioswriter.py | 222 +++++++++++++++++++++++++++ xbout/boutdataset.py | 52 +++++-- xbout/tests/test_adios2_roundtrip.py | 47 ++++++ xbout/xarraybackend.py | 59 ++++++- 5 files changed, 362 insertions(+), 21 deletions(-) create mode 100644 xbout/adioswriter.py create mode 100644 xbout/tests/test_adios2_roundtrip.py diff --git a/pyproject.toml b/pyproject.toml index 350b5a90..b5d63d96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ dependencies = [ ] [project.optional-dependencies] +adios2 = [ + "adios2" +] calc = [ "numpy >= 1.18.0", "scipy >= 1.3.0", diff --git a/xbout/adioswriter.py b/xbout/adioswriter.py new file mode 100644 index 00000000..e31b5cbd --- /dev/null +++ b/xbout/adioswriter.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +import json +from collections.abc import Iterable, Mapping +from typing import Any + +import numpy as np + +DATASET_ATTR_PREFIX = "__xarray_dataset_attrs__/" +VAR_ATTR_PREFIX_DIMENSIONS = "__xarray_dimensions__" +VAR_ATTR_PREFIX_ORIGINAL_DTYPE = "__xarray_original_dtype__" + + +class Adios2NotInstalledError(ImportError): + pass + + +def _normalize_attr_value(value: Any) -> Any: + if value is None: + return "null" + if isinstance(value, (str, bytes, int, float, bool, np.number)): + return value + if isinstance(value, np.ndarray): + return value + if isinstance(value, (list, tuple)): + if all(isinstance(v, (str, bytes, int, float, bool, np.number)) for v in value): + return list(value) + return json.dumps(value, default=str) + if isinstance(value, Mapping): + return json.dumps(value, default=str) + return json.dumps(value, default=str) + + +def _numpy_for_write(data: Any) -> np.ndarray: + arr = np.asarray(data) + + if arr.dtype == np.dtype("bool"): + return arr.astype(np.uint8) + + if arr.dtype.kind in {"M", "m"}: + raise TypeError( + "datetime64/timedelta64 not supported for ADIOS2 writer yet; " + "encode to int64 with units attrs first" + ) + + if arr.dtype == np.dtype("O"): + # Conservative: only allow scalar string-like objects. + if arr.ndim == 0 and isinstance(arr.item(), (str, bytes)): + return np.asarray(arr.item()) + raise TypeError( + "object dtype not supported for ADIOS2 writer (except scalar strings)" + ) + + return arr + + +def _attr_to_numpy(value: Any) -> Any: + norm = _normalize_attr_value(value) + if isinstance(norm, np.ndarray): + if norm.ndim == 0: + return norm.item() + return norm + if isinstance(norm, (str, bytes, int, float, bool, np.number)): + return norm + return np.asarray(norm) + + +def write_dataset_bp( + ds, + path: str, + *, + time_dim: str = "t", + engine: str = "BP4", + parameters: Mapping[str, str] | None = None, + overwrite: bool = True, + variables: Iterable[str] | None = None, +) -> None: + """ + Write an xarray Dataset to an ADIOS2 .bp output. + + Design choices: + - Map ``time_dim`` to ADIOS2 steps (variables containing time_dim are written + one slice per step). + - Store per-variable dimension names as attribute ``{var}/__xarray_dimensions__`` + excluding ``time_dim`` when writing steps. + - Store dataset attributes under ``__xarray_dataset_attrs__/{key}``. + """ + try: + import adios2 # type: ignore + except ImportError as e: # pragma: no cover + raise Adios2NotInstalledError( + "adios2 is required to write .bp files; install adios2-python" + ) from e + if engine != "BP4": + raise NotImplementedError( + "Only BP4 is supported by the current writer implementation" + ) + if parameters: + raise NotImplementedError( + "ADIOS2 engine parameters are not supported by the current writer implementation" + ) + + if variables is None: + names_to_write = list(ds.variables) + else: + names_to_write = list(variables) + + # Determine step count across all time-dependent variables. + step_count: int | None = None + for name in names_to_write: + var = ds.variables[name] + if time_dim in var.dims: + if var.dims[0] != time_dim: + raise ValueError( + f"Only supports {time_dim!r} as the first dimension for {name!r}; " + f"got dims={var.dims!r}" + ) + nsteps = int(var.sizes[time_dim]) + step_count = nsteps if step_count is None else max(step_count, nsteps) + if step_count is None: + step_count = 1 + + mode = "w" if overwrite else "a" + stream = adios2.Stream(path, mode) + try: + # Dataset attrs. + for k, v in ds.attrs.items(): + stream.write_attribute(DATASET_ATTR_PREFIX + str(k), _attr_to_numpy(v)) + + # Define all variables up-front to avoid Stream.write(name, ndarray) inference, + # which is unreliable on some adios2-python builds. + io = stream.io + adios_vars: dict[str, Any] = {} + sample_buffers: dict[str, np.ndarray] = {} + + for name in names_to_write: + var = ds.variables[name] + write_dtype = _numpy_for_write(var.data).dtype + + if time_dim in var.dims: + shape = [int(var.sizes[d]) for d in var.dims if d != time_dim] + else: + shape = [int(s) for s in np.asarray(var.data).shape] + + start = [0] * len(shape) + count = shape[:] + + if write_dtype.kind in {"U", "S"}: + sample = np.asarray([""], dtype=write_dtype) + else: + sample = np.array([0], dtype=write_dtype) + + # Keep a reference to the sample buffer alive for the lifetime of the + # stream. Some adios2-python builds appear to keep a view/pointer to the + # provided numpy object when defining variables. + sample_buffers[name] = sample + + adios_var = io.define_variable(name, sample, shape, start, count) + if shape: + adios_var.set_shape(shape) + adios_var.set_selection([start, count]) + adios_vars[name] = adios_var + + # Per-variable attrs/dims (now that variables exist). + for name in names_to_write: + var = ds.variables[name] + dims_wo_time = [d for d in var.dims if d != time_dim] + stream.write_attribute( + VAR_ATTR_PREFIX_DIMENSIONS, + _attr_to_numpy(dims_wo_time), + variable_name=name, + separator="/", + ) + for ak, av in var.attrs.items(): + stream.write_attribute( + str(ak), + _attr_to_numpy(av), + variable_name=name, + separator="/", + ) + + original_dtype = np.asarray(var.data).dtype + write_arr = _numpy_for_write(var.data) + if write_arr.dtype != original_dtype: + stream.write_attribute( + VAR_ATTR_PREFIX_ORIGINAL_DTYPE, + _attr_to_numpy(str(original_dtype)), + variable_name=name, + separator="/", + ) + + for step in range(step_count): + stream.begin_step() + try: + for name in names_to_write: + var = ds.variables[name] + adios_var = adios_vars[name] + if time_dim in var.dims: + if step >= int(var.sizes[time_dim]): + continue + data = _numpy_for_write(var.data)[step] + elif step == 0: + data = _numpy_for_write(var.data) + else: + continue + + data_arr = np.asarray(data) + if data_arr.ndim: + # Use an owned, contiguous buffer for ADIOS2 Put(). + # Some adios2-python builds appear to segfault when passed + # non-owned views. + data_arr = np.array(data_arr, copy=True, order="C") + # Ensure selection matches the provided buffer. + adios_var.set_selection([[0] * data_arr.ndim, list(data_arr.shape)]) + else: + # Ensure a stable 0-d buffer with the intended dtype. + data_arr = np.asarray(data_arr.item(), dtype=data_arr.dtype).reshape(()) + stream.write(adios_var, data_arr) + finally: + stream.end_step() + finally: + stream.close() diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index c316884c..077e904f 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -875,6 +875,10 @@ def dict_to_attrs(obj, section): else: encoding = None + savepath_path = Path(savepath) + is_adios_bp = savepath_path.suffix == ".bp" + time_dim = "t" + if separate_vars: # Save each major variable to a different netCDF file @@ -910,13 +914,23 @@ def dict_to_attrs(obj, section): var_encoding = None print("Saving " + major_var + " data...") with ProgressBar(): - single_var_ds.to_netcdf( - path=str(var_savepath), - format=filetype, - engine=_check_filetype(Path(var_savepath)), - compute=True, - encoding=var_encoding, - ) + if Path(var_savepath).suffix == ".bp": + from xbout.adioswriter import write_dataset_bp + + write_dataset_bp( + single_var_ds, + str(var_savepath), + time_dim=time_dim, + overwrite=True, + ) + else: + single_var_ds.to_netcdf( + path=str(var_savepath), + format=filetype, + engine=_check_filetype(Path(var_savepath)), + compute=True, + encoding=var_encoding, + ) # Force memory deallocation to limit RAM usage single_var_ds.close() @@ -926,13 +940,23 @@ def dict_to_attrs(obj, section): # Save data to a single file print("Saving data...") with ProgressBar(): - to_save.to_netcdf( - path=savepath, - engine=_check_filetype(Path(savepath)), - format=filetype, - compute=True, - encoding=encoding, - ) + if is_adios_bp: + from xbout.adioswriter import write_dataset_bp + + write_dataset_bp( + to_save, + str(savepath), + time_dim=time_dim, + overwrite=True, + ) + else: + to_save.to_netcdf( + path=savepath, + engine=_check_filetype(Path(savepath)), + format=filetype, + compute=True, + encoding=encoding, + ) return diff --git a/xbout/tests/test_adios2_roundtrip.py b/xbout/tests/test_adios2_roundtrip.py new file mode 100644 index 00000000..38b67e06 --- /dev/null +++ b/xbout/tests/test_adios2_roundtrip.py @@ -0,0 +1,47 @@ +import numpy as np +import pytest + + +adios2 = pytest.importorskip("adios2") +xr = pytest.importorskip("xarray") + + +def test_adios2_roundtrip_dataset_attrs_and_vars(tmp_path): + from xbout.adioswriter import write_dataset_bp + from xbout.xarraybackend import BoutAdiosBackendEntrypoint + + ds = xr.Dataset( + data_vars={ + "a": (("t", "x"), np.arange(6, dtype=np.float32).reshape(3, 2)), + "flag": (("t",), np.array([True, False, True], dtype=bool)), + "scalar": ((), np.int64(7)), + }, + coords={ + "t": ("t", np.array([0.0, 0.5, 1.0], dtype=np.float64)), + "x": ("x", np.array([10, 20], dtype=np.int32)), + }, + attrs={"title": "roundtrip", "answer": 42}, + ) + ds["a"].attrs["units"] = "arb" + + path = tmp_path / "roundtrip.bp" + write_dataset_bp(ds, str(path), time_dim="t", overwrite=True) + + ds2 = BoutAdiosBackendEntrypoint().open_dataset(str(path)) + try: + assert ds2.attrs["title"] == "roundtrip" + assert int(ds2.attrs["answer"]) == 42 + + np.testing.assert_allclose(ds2["t"].values, ds["t"].values) + np.testing.assert_allclose(ds2["x"].values, ds["x"].values) + np.testing.assert_allclose(ds2["a"].values, ds["a"].values) + assert ds2["a"].attrs["units"] == "arb" + + assert ds2["flag"].dtype == bool + assert ds2["flag"].dims == ("t",) + np.testing.assert_array_equal(ds2["flag"].values, ds["flag"].values) + + assert int(ds2["scalar"].values) == 7 + finally: + ds2.close() + diff --git a/xbout/xarraybackend.py b/xbout/xarraybackend.py index 8f1b5e7d..346a27a4 100644 --- a/xbout/xarraybackend.py +++ b/xbout/xarraybackend.py @@ -31,6 +31,9 @@ # need some special secret attributes to tell us the dimensions DIMENSION_KEY = "time_dimension" +DATASET_ATTR_PREFIX = "__xarray_dataset_attrs__/" +XARRAY_DIMS_ATTR = "__xarray_dimensions__" +XARRAY_ORIGINAL_DTYPE_ATTR = "__xarray_original_dtype__" adios_to_numpy_type = { "char": np.char, @@ -55,13 +58,21 @@ class BoutADIOSBackendArray(BackendArray): """ADIOS2 backend for lazily indexed arrays""" def __init__( - self, shape: list, dtype: np.dtype, lock, adiosfile: FileReader, varname: str + self, + shape: list, + dtype: np.dtype, + lock, + adiosfile: FileReader, + varname: str, + *, + cast_dtype: np.dtype | None = None, ): self.shape = shape self.dtype = dtype self.lock = lock self.fh = adiosfile self.varname = varname + self.cast_dtype = cast_dtype self.adiosvar = self.fh.inquire_variable(varname) self.steps = self.adiosvar.steps() # print(f"BoutADIOSBackendArray.__init__: {dtype} {varname} {shape} {dtype.itemsize}") @@ -113,6 +124,8 @@ def _raw_indexing_method(self, key: tuple) -> np.typing.ArrayLike: if self.steps > 1 and first_sl: # key[0] is the step selection # print(f" data step selection start = {st} count = {ct}") self.adiosvar.set_step_selection([st, ct]) + # Advance past the implicit steps dimension in self.shape + dimid += 1 else: start.append(st) count.append(ct) @@ -133,6 +146,8 @@ def _raw_indexing_method(self, key: tuple) -> np.typing.ArrayLike: f"shape={data.shape}, steps={self.steps}" ) data = data.reshape((self.steps, dim0) + data.shape[1:]) + if self.cast_dtype is not None: + data = np.asarray(data).astype(self.cast_dtype, copy=False) return data @@ -239,12 +254,18 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti dims = None vlen = len(varname) + 1 # include / xattrs = {} + original_dtype: np.dtype | None = None for aname, ainfo in varattrs: # print(f"\t{ainfo['Type']} {aname}\t = {ainfo['Value']}") attr_value = self._fh.read_attribute(aname) - if aname == varname + "/__xarray_dimensions__": + if aname == varname + "/" + XARRAY_DIMS_ATTR: dims = attr_value # print(f"\t\tDIMENSIONS = {dims}") + elif aname == varname + "/" + XARRAY_ORIGINAL_DTYPE_ATTR: + try: + original_dtype = np.dtype(str(attr_value)) + except TypeError: + original_dtype = None else: xattrs[aname[vlen:]] = attr_value attrs.pop(aname) @@ -261,23 +282,42 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti dims.insert(0, "t") # print(f"\tAdd time to shape {shape_list} {dims}") nptype = np.dtype(adios_to_numpy_type[varinfo["Type"]]) + cast_dtype = ( + original_dtype if original_dtype is not None and original_dtype != nptype else None + ) xdata = indexing.LazilyIndexedArray( - BoutADIOSBackendArray(shape_list, nptype, None, self._fh, varname) + BoutADIOSBackendArray( + shape_list, + nptype, + None, + self._fh, + varname, + cast_dtype=cast_dtype, + ) ) # print(f"\tDefine VARIABLE {varname} with dims {dims}") - xvar = Variable(dims, xdata, attrs=xattrs, encoding={"dtype": nptype}) + xvar = Variable( + dims, + xdata, + attrs=xattrs, + encoding={"dtype": (original_dtype or nptype)}, + ) # print(f"{xvar.dtype} {xvar.attrs["name"]} {xvar.dims} {xvar.encoding}") else: if steps > 1: avar = self._fh.inquire_variable(varname) avar.set_step_selection([0, avar.steps()]) data = self._fh.read(avar) + if original_dtype is not None and data.dtype != original_dtype: + data = np.asarray(data).astype(original_dtype, copy=False) # print(f"\tCreate timed scalar variable {varname}") xvar = Variable( "t", data, attrs=xattrs, encoding={"dtype": data.dtype} ) else: data = self._fh.read(varname) + if original_dtype is not None and np.asarray(data).dtype != original_dtype: + data = np.asarray(data).astype(original_dtype, copy=False) if varinfo["Type"] == "string": # print(f"\tCreate string scalar variable {varname}") xvar = Variable([], data, attrs=xattrs, encoding=None) @@ -287,9 +327,14 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti xvars[varname] = xvar # print(f"--- {xvar}") - for attname, attinfo in attrs.items(): - print(f"{attinfo['Type']} {attname}\t = {attinfo['Value']}") + ds_attrs = {} + for attname in list(attrs.keys()): + attr_value = self._fh.read_attribute(attname) + if isinstance(attname, str) and attname.startswith(DATASET_ATTR_PREFIX): + ds_attrs[attname[len(DATASET_ATTR_PREFIX) :]] = attr_value + else: + ds_attrs[attname] = attr_value - ds = Dataset(xvars, None, None) + ds = Dataset(xvars, None, ds_attrs) ds.set_close(BoutAdiosBackendEntrypoint.close) return ds From 11a16acabf5eb5ef1f6a29cca00ee17e3f7d935b Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 27 May 2026 12:49:11 -0700 Subject: [PATCH 2/8] adios2: Tidy xarray backend Removing commented-out code, adding docstrings. --- xbout/xarraybackend.py | 146 +++++++++++++++++++++-------------------- 1 file changed, 74 insertions(+), 72 deletions(-) diff --git a/xbout/xarraybackend.py b/xbout/xarraybackend.py index 346a27a4..c7d466fe 100644 --- a/xbout/xarraybackend.py +++ b/xbout/xarraybackend.py @@ -1,14 +1,20 @@ -"""License: -Distributed under the OSI-approved Apache License, Version 2.0. See -accompanying file Copyright.txt for details. +""" +xarray backend for reading ADIOS2 ``.bp`` files. + +This backend provides an xarray ``BackendEntrypoint`` that can open ADIOS2 +datasets via the ``adios2`` Python package. Variables are represented as +``LazilyIndexedArray`` objects backed by ADIOS2 selections. + +On-disk conventions used by this backend: +- Per-variable dimension names are stored as attributes + ``{varname}/__xarray_dimensions__``. +- Dataset-level attributes are stored under ``__xarray_dataset_attrs__/{key}``. """ from __future__ import annotations import os -# import warnings - from collections.abc import Iterable from typing import TYPE_CHECKING, Any, ItemsView @@ -23,14 +29,12 @@ from xarray.core import indexing -if TYPE_CHECKING: +if TYPE_CHECKING: # pragma: no cover from adios2 import FileReader from io import BufferedIOBase from xarray.backends.common import AbstractDataStore -# need some special secret attributes to tell us the dimensions -DIMENSION_KEY = "time_dimension" DATASET_ATTR_PREFIX = "__xarray_dataset_attrs__/" XARRAY_DIMS_ATTR = "__xarray_dimensions__" XARRAY_ORIGINAL_DTYPE_ATTR = "__xarray_original_dtype__" @@ -55,7 +59,13 @@ class BoutADIOSBackendArray(BackendArray): - """ADIOS2 backend for lazily indexed arrays""" + """ + Lazily indexed array backed by an ADIOS2 Variable. + + xarray calls ``__getitem__`` with an ``ExplicitIndexer``; this class maps the + indexer into ADIOS2 ``set_step_selection`` (for time/steps) and + ``set_selection`` (for spatial dimensions), then reads the selection. + """ def __init__( self, @@ -67,6 +77,24 @@ def __init__( *, cast_dtype: np.dtype | None = None, ): + """ + Parameters + ---------- + shape + Full xarray-visible shape. If the ADIOS2 variable has steps, the first + dimension is the synthetic xarray time dimension. + dtype + Numpy dtype used by ADIOS2 for the stored variable. + lock + Optional lock for thread-safety. ADIOS2 reads are not thread-safe. + adiosfile + Open ADIOS2 ``FileReader`` handle. + varname + Name of the ADIOS2 variable to read. + cast_dtype + Optional dtype to cast to after reading (used to round-trip types that + are stored differently on disk, e.g. ``bool`` stored as ``uint8``). + """ self.shape = shape self.dtype = dtype self.lock = lock @@ -75,11 +103,9 @@ def __init__( self.cast_dtype = cast_dtype self.adiosvar = self.fh.inquire_variable(varname) self.steps = self.adiosvar.steps() - # print(f"BoutADIOSBackendArray.__init__: {dtype} {varname} {shape} {dtype.itemsize}") def __getitem__(self, key: indexing.ExplicitIndexer) -> np.typing.ArrayLike: - # print(f"**** BoutADIOSBackendArray.__getitem__: {self.varname} key = {key}") - + """Read a selection defined by an xarray ``ExplicitIndexer``.""" return indexing.explicit_indexing_adapter( key, self.shape, @@ -88,13 +114,16 @@ def __getitem__(self, key: indexing.ExplicitIndexer) -> np.typing.ArrayLike: ) def _raw_indexing_method(self, key: tuple) -> np.typing.ArrayLike: - # print(f"****BoutADIOSBackendArray._raw_indexing_method: {self.varname} " - # f"key = {key} steps = {self.steps}") - # print(f" data shape {data.shape}") - - # thread safe method that access to data on disk needed because - # adios is not thread safe even for reading - # with self.lock: + """ + Convert xarray basic indexing into ADIOS2 selection calls and read. + + Notes + ----- + - ADIOS2 does not support stepped slicing (``slice.step != 1``). + - ADIOS2 time is represented as "steps". If an ADIOS2 variable has steps, + xarray's first dimension is treated as time and mapped via + ``set_step_selection``. + """ start = [] count = [] dimid = 0 @@ -131,7 +160,6 @@ def _raw_indexing_method(self, key: tuple) -> np.typing.ArrayLike: count.append(ct) dimid += 1 first_sl = False - # print(f" data selection start = {start} count = {count}") self.adiosvar.set_selection([start, count]) data = self.fh.read(self.adiosvar) @@ -152,7 +180,7 @@ def _raw_indexing_method(self, key: tuple) -> np.typing.ArrayLike: def attrs_of_var(varname: str, items: ItemsView, separator: str = "/"): - """Return attributes whose name starts with a variable's name""" + """Return (name, info) pairs for attributes scoped to a variable.""" return [(key, value) for key, value in items if key.startswith(varname + separator)] @@ -161,7 +189,7 @@ def attrs_of_var(varname: str, items: ItemsView, separator: str = "/"): # pylint: disable=E1121 # too-many-function-args class BoutAdiosBackendEntrypoint(BackendEntrypoint): """ - Backend for ".bp" folders based on the adios2 package. + xarray backend entrypoint for ADIOS2 ``.bp`` datasets. For more information about the underlying library, visit: https://adios2.readthedocs.io/en/stable @@ -177,18 +205,17 @@ class BoutAdiosBackendEntrypoint(BackendEntrypoint): def __init__(self): self._fh = None - def close(): - """Close the ADIOS file""" - # print("BoutAdiosBackendEntrypoint.close() called") - # Note that this is a strange method without 'self', so we cannot close the file because - # we don't have any handle to it - # if self._fh is not None: - # self._fh.close() + def close(self) -> None: + """Close the underlying ADIOS2 file handle, if one is open.""" + if self._fh is not None: + self._fh.close() + self._fh = None def guess_can_open( self, filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore, ) -> bool: + """Return True if this backend can open the provided filename.""" if isinstance(filename_or_obj, (str, os.PathLike)): _, ext = os.path.splitext(filename_or_obj) return ext in {".bp"} @@ -199,44 +226,26 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti self, filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore, *, - # mask_and_scale=True, - # decode_times=True, - # concat_characters=True, - # decode_coords=True, drop_variables: str | Iterable[str] | None = None, - # use_cftime=None, - # decode_timedelta=None, - # group=None, - # mode="r", - # synchronizer=None, - # consolidated=None, - # chunk_store=None, - # storage_options=None, - # stacklevel=3, - # adios_version=None, ) -> Dataset: + """ + Open an ADIOS2 ``.bp`` file/folder as an xarray Dataset. + + Parameters + ---------- + filename_or_obj + Path to a ``.bp`` dataset (directory or file, depending on ADIOS2 engine). + drop_variables + Optional variable name or iterable of variable names to exclude. + """ from adios2 import FileReader filename_or_obj = _normalize_path(filename_or_obj) - # print(f"BoutAdiosBackendEntrypoint: path = {filename_or_obj} type = {type(filename_or_obj)}") - - # if isinstance(filename_or_obj, os.PathLike): - # print(f" os.PathLike: {os.fspath(filename_or_obj)}") - # - # if isinstance(filename_or_obj, str): - # print(f" str: {os.path.abspath(os.path.expanduser(filename_or_obj))}") - - # if isinstance(filename_or_obj, BufferedIOBase): - # raise ValueError("ADIOS2 does not support BufferedIOBase input") - # - # if isinstance(filename_or_obj, AbstractDataStore): - # raise ValueError("ADIOS2 does not support AbstractDataStore input") self._fh = FileReader(filename_or_obj) vars = self._fh.available_variables() attrs = self._fh.available_attributes() attr_items = attrs.items() - # print(f"BoutAdiosBackendEntrypoint: {len(vars)} variables, {len(attrs)} attributes") xvars = {} for varname, varinfo in vars.items(): @@ -249,18 +258,15 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti shape_list = [] shape_str = [] steps = int(varinfo["AvailableStepsCount"]) - # print(f"{varinfo['Type']} {varname}\t {shape_list}") varattrs = attrs_of_var(varname, attr_items, "/") dims = None vlen = len(varname) + 1 # include / xattrs = {} original_dtype: np.dtype | None = None for aname, ainfo in varattrs: - # print(f"\t{ainfo['Type']} {aname}\t = {ainfo['Value']}") attr_value = self._fh.read_attribute(aname) if aname == varname + "/" + XARRAY_DIMS_ATTR: dims = attr_value - # print(f"\t\tDIMENSIONS = {dims}") elif aname == varname + "/" + XARRAY_ORIGINAL_DTYPE_ATTR: try: original_dtype = np.dtype(str(attr_value)) @@ -269,21 +275,19 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti else: xattrs[aname[vlen:]] = attr_value attrs.pop(aname) - # print(f"\txattrs = {xattrs}") # Create the xarray variable if dims is None: dims = shape_str if shape_list != []: - # for i in range(len(shape_str)): - # shape_str[i] = "d" + shape_str[i] if steps > 1: shape_list.insert(0, steps) dims.insert(0, "t") - # print(f"\tAdd time to shape {shape_list} {dims}") nptype = np.dtype(adios_to_numpy_type[varinfo["Type"]]) cast_dtype = ( - original_dtype if original_dtype is not None and original_dtype != nptype else None + original_dtype + if original_dtype is not None and original_dtype != nptype + else None ) xdata = indexing.LazilyIndexedArray( BoutADIOSBackendArray( @@ -302,7 +306,6 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti attrs=xattrs, encoding={"dtype": (original_dtype or nptype)}, ) - # print(f"{xvar.dtype} {xvar.attrs["name"]} {xvar.dims} {xvar.encoding}") else: if steps > 1: avar = self._fh.inquire_variable(varname) @@ -310,22 +313,21 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti data = self._fh.read(avar) if original_dtype is not None and data.dtype != original_dtype: data = np.asarray(data).astype(original_dtype, copy=False) - # print(f"\tCreate timed scalar variable {varname}") xvar = Variable( "t", data, attrs=xattrs, encoding={"dtype": data.dtype} ) else: data = self._fh.read(varname) - if original_dtype is not None and np.asarray(data).dtype != original_dtype: + if ( + original_dtype is not None + and np.asarray(data).dtype != original_dtype + ): data = np.asarray(data).astype(original_dtype, copy=False) if varinfo["Type"] == "string": - # print(f"\tCreate string scalar variable {varname}") xvar = Variable([], data, attrs=xattrs, encoding=None) else: - # print(f"\tCreate scalar variable {varname}") xvar = Variable([], data, attrs=xattrs, encoding=None) xvars[varname] = xvar - # print(f"--- {xvar}") ds_attrs = {} for attname in list(attrs.keys()): @@ -336,5 +338,5 @@ def open_dataset( # type: ignore[override] # allow LSP violation, not supporti ds_attrs[attname] = attr_value ds = Dataset(xvars, None, ds_attrs) - ds.set_close(BoutAdiosBackendEntrypoint.close) + ds.set_close(self.close) return ds From a4b14f74da2e256b9944f3afd77c0a902ed96df7 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 27 May 2026 13:40:01 -0700 Subject: [PATCH 3/8] boutdataset: Formatting, don't use ProgressBar for adios2 ProgressBar works when a Dask task graph is being executed, but that is not the case for write_dataset_bp. --- xbout/boutdataset.py | 42 ++++++++++++++++++++---------------------- 1 file changed, 20 insertions(+), 22 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 077e904f..1c4a8db0 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -1,30 +1,28 @@ import collections +import gc +import warnings from copy import copy -from pprint import pformat as prettyformat from functools import partial from itertools import chain from pathlib import Path -import warnings -import gc +from pprint import pformat as prettyformat -import xarray as xr import animatplot as amp +import numpy as np +import xarray as xr +from dask.diagnostics import ProgressBar from matplotlib import pyplot as plt from matplotlib.animation import PillowWriter - from mpl_toolkits.axes_grid1 import make_axes_locatable -import numpy as np -from dask.diagnostics import ProgressBar - from .geometries import apply_geometry from .plotting.animate import ( - animate_poloidal, - animate_pcolormesh, - animate_line, _add_controls, _normalise_time_coord, _parse_coord_option, + animate_line, + animate_pcolormesh, + animate_poloidal, ) from .region import _from_region from .utils import ( @@ -939,17 +937,17 @@ def dict_to_attrs(obj, section): else: # Save data to a single file print("Saving data...") - with ProgressBar(): - if is_adios_bp: - from xbout.adioswriter import write_dataset_bp - - write_dataset_bp( - to_save, - str(savepath), - time_dim=time_dim, - overwrite=True, - ) - else: + if is_adios_bp: + from xbout.adioswriter import write_dataset_bp + + write_dataset_bp( + to_save, + str(savepath), + time_dim=time_dim, + overwrite=True, + ) + else: + with ProgressBar(): to_save.to_netcdf( path=savepath, engine=_check_filetype(Path(savepath)), From c81efe46842052e1085a409579ae9da092fd05cf Mon Sep 17 00:00:00 2001 From: bendudson Date: Wed, 27 May 2026 20:43:48 +0000 Subject: [PATCH 4/8] Apply black formatting --- xbout/adioswriter.py | 8 ++++++-- xbout/tests/test_adios2_roundtrip.py | 2 -- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/xbout/adioswriter.py b/xbout/adioswriter.py index e31b5cbd..c54fddd4 100644 --- a/xbout/adioswriter.py +++ b/xbout/adioswriter.py @@ -211,10 +211,14 @@ def write_dataset_bp( # non-owned views. data_arr = np.array(data_arr, copy=True, order="C") # Ensure selection matches the provided buffer. - adios_var.set_selection([[0] * data_arr.ndim, list(data_arr.shape)]) + adios_var.set_selection( + [[0] * data_arr.ndim, list(data_arr.shape)] + ) else: # Ensure a stable 0-d buffer with the intended dtype. - data_arr = np.asarray(data_arr.item(), dtype=data_arr.dtype).reshape(()) + data_arr = np.asarray( + data_arr.item(), dtype=data_arr.dtype + ).reshape(()) stream.write(adios_var, data_arr) finally: stream.end_step() diff --git a/xbout/tests/test_adios2_roundtrip.py b/xbout/tests/test_adios2_roundtrip.py index 38b67e06..2d7c5a7c 100644 --- a/xbout/tests/test_adios2_roundtrip.py +++ b/xbout/tests/test_adios2_roundtrip.py @@ -1,7 +1,6 @@ import numpy as np import pytest - adios2 = pytest.importorskip("adios2") xr = pytest.importorskip("xarray") @@ -44,4 +43,3 @@ def test_adios2_roundtrip_dataset_attrs_and_vars(tmp_path): assert int(ds2["scalar"].values) == 7 finally: ds2.close() - From f11274615e3388a29a1b5dd013e26a66e453008c Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 28 May 2026 09:44:05 -0700 Subject: [PATCH 5/8] adioswriter: write_ints_as_int32 keyword The ADIOS2 reader in BOUT++ only handles int32 integers, so provide an option to convert all output integers to this size. --- xbout/__init__.py | 120 ++++++++++++++++++++++++--- xbout/adioswriter.py | 49 +++++++++-- xbout/tests/conftest.py | 28 +++++-- xbout/tests/test_adios2_roundtrip.py | 45 ++++++++++ 4 files changed, 219 insertions(+), 23 deletions(-) diff --git a/xbout/__init__.py b/xbout/__init__.py index b94242c6..4fdf42bf 100644 --- a/xbout/__init__.py +++ b/xbout/__init__.py @@ -1,25 +1,38 @@ -from .load import open_boutdataset, collect -from .lazyload import lazy_open_boutdataset +""" +xbout package public API. -from . import geometries -from .geometries import register_geometry, REGISTERED_GEOMETRIES +This module is intentionally lightweight: many xbout features depend on optional +third-party packages (e.g. matplotlib, dask, boutdata). Import those lazily so +that submodules like ``xbout.adioswriter`` can be used without pulling in the +full dependency set. +""" -from .boutdataset import BoutDatasetAccessor -from .boutdataarray import BoutDataArrayAccessor +from __future__ import annotations -from .plotting.animate import animate_pcolormesh, animate_poloidal -from .plotting.utils import plot_separatrix +from typing import TYPE_CHECKING -from .fastoutput import open_fastoutput +if TYPE_CHECKING: # pragma: no cover + from . import geometries + from .boutdataarray import BoutDataArrayAccessor + from .boutdataset import BoutDatasetAccessor + from .fastoutput import open_fastoutput + from .geometries import REGISTERED_GEOMETRIES, register_geometry + from .lazyload import lazy_open_boutdataset + from .load import collect, open_boutdataset + from .plotting.animate import animate_pcolormesh, animate_poloidal + from .plotting.utils import plot_separatrix -from importlib.metadata import version, PackageNotFoundError +from importlib.metadata import PackageNotFoundError, version try: __version__ = version(__name__) except PackageNotFoundError: - from setuptools_scm import get_version + try: + from setuptools_scm import get_version - __version__ = get_version(root="..", relative_to=__file__) + __version__ = get_version(root="..", relative_to=__file__) + except Exception: # pragma: no cover + __version__ = "0+unknown" __all__ = [ "open_boutdataset", @@ -35,3 +48,86 @@ "plot_separatrix", "open_fastoutput", ] + + +def __getattr__(name: str): + if name in {"geometries", "register_geometry", "REGISTERED_GEOMETRIES"}: + try: + from . import geometries + from .geometries import REGISTERED_GEOMETRIES, register_geometry + except ImportError as e: # pragma: no cover + raise ImportError( + "xbout geometries require optional dependencies (e.g. xarray)." + ) from e + return { + "geometries": geometries, + "register_geometry": register_geometry, + "REGISTERED_GEOMETRIES": REGISTERED_GEOMETRIES, + }[name] + + if name in {"open_boutdataset", "collect"}: + try: + from .load import collect, open_boutdataset + except ImportError as e: # pragma: no cover + raise ImportError( + "xbout.load requires optional dependencies (e.g. 'boutdata'). " + "Install the full xbout extras, or import submodules that do not " + "require boutdata." + ) from e + return {"open_boutdataset": open_boutdataset, "collect": collect}[name] + + if name == "lazy_open_boutdataset": + try: + from .lazyload import lazy_open_boutdataset + except ImportError as e: # pragma: no cover + raise ImportError( + "xbout.lazyload requires optional dependencies (e.g. dask, h5py, xarray)." + ) from e + return lazy_open_boutdataset + + if name in {"BoutDatasetAccessor", "BoutDataArrayAccessor"}: + try: + from .boutdataarray import BoutDataArrayAccessor + from .boutdataset import BoutDatasetAccessor + except ImportError as e: # pragma: no cover + raise ImportError( + "xbout accessors require optional dependencies (e.g. xarray)." + ) from e + return { + "BoutDatasetAccessor": BoutDatasetAccessor, + "BoutDataArrayAccessor": BoutDataArrayAccessor, + }[name] + + if name in {"animate_pcolormesh", "animate_poloidal"}: + try: + from .plotting.animate import animate_pcolormesh, animate_poloidal + except ImportError as e: # pragma: no cover + raise ImportError( + "xbout plotting requires optional dependencies (e.g. matplotlib)." + ) from e + return { + "animate_pcolormesh": animate_pcolormesh, + "animate_poloidal": animate_poloidal, + }[name] + + if name == "plot_separatrix": + try: + from .plotting.utils import plot_separatrix + except ImportError as e: # pragma: no cover + raise ImportError( + "xbout plotting requires optional dependencies (e.g. matplotlib)." + ) from e + return plot_separatrix + + if name == "open_fastoutput": + try: + from .fastoutput import open_fastoutput + except ImportError as e: # pragma: no cover + raise ImportError("xbout.fastoutput requires optional dependencies.") from e + return open_fastoutput + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__(): + return sorted(list(globals().keys()) + __all__) diff --git a/xbout/adioswriter.py b/xbout/adioswriter.py index c54fddd4..f8e867a3 100644 --- a/xbout/adioswriter.py +++ b/xbout/adioswriter.py @@ -15,6 +15,30 @@ class Adios2NotInstalledError(ImportError): pass +def _safe_int32_cast(arr: np.ndarray) -> np.ndarray: + if arr.size == 0: + return arr.astype(np.int32, copy=False) + + info = np.iinfo(np.int32) + + if arr.dtype.kind == "u": + max_value = int(np.asarray(arr).max()) + if max_value > info.max: + raise ValueError( + f"Cannot safely cast unsigned integer data to int32: max={max_value} > {info.max}" + ) + return arr.astype(np.int32, copy=False) + + min_value = int(np.asarray(arr).min()) + max_value = int(np.asarray(arr).max()) + if min_value < info.min or max_value > info.max: + raise ValueError( + "Cannot safely cast integer data to int32: " + f"min={min_value}, max={max_value}, int32=[{info.min}, {info.max}]" + ) + return arr.astype(np.int32, copy=False) + + def _normalize_attr_value(value: Any) -> Any: if value is None: return "null" @@ -31,12 +55,15 @@ def _normalize_attr_value(value: Any) -> Any: return json.dumps(value, default=str) -def _numpy_for_write(data: Any) -> np.ndarray: +def _numpy_for_write(data: Any, *, write_ints_as_int32: bool = False) -> np.ndarray: arr = np.asarray(data) if arr.dtype == np.dtype("bool"): return arr.astype(np.uint8) + if write_ints_as_int32 and arr.dtype.kind in {"i", "u"} and arr.dtype != np.int32: + arr = _safe_int32_cast(arr) + if arr.dtype.kind in {"M", "m"}: raise TypeError( "datetime64/timedelta64 not supported for ADIOS2 writer yet; " @@ -74,6 +101,7 @@ def write_dataset_bp( parameters: Mapping[str, str] | None = None, overwrite: bool = True, variables: Iterable[str] | None = None, + write_ints_as_int32: bool = False, ) -> None: """ Write an xarray Dataset to an ADIOS2 .bp output. @@ -84,6 +112,9 @@ def write_dataset_bp( - Store per-variable dimension names as attribute ``{var}/__xarray_dimensions__`` excluding ``time_dim`` when writing steps. - Store dataset attributes under ``__xarray_dataset_attrs__/{key}``. + - Optionally store all integer variables as int32 on disk (see + ``write_ints_as_int32``), while preserving the original dtype via the + ``__xarray_original_dtype__`` attribute for round-tripping. """ try: import adios2 # type: ignore @@ -135,7 +166,9 @@ def write_dataset_bp( for name in names_to_write: var = ds.variables[name] - write_dtype = _numpy_for_write(var.data).dtype + write_dtype = _numpy_for_write( + var.data, write_ints_as_int32=write_ints_as_int32 + ).dtype if time_dim in var.dims: shape = [int(var.sizes[d]) for d in var.dims if d != time_dim] @@ -180,7 +213,9 @@ def write_dataset_bp( ) original_dtype = np.asarray(var.data).dtype - write_arr = _numpy_for_write(var.data) + write_arr = _numpy_for_write( + var.data, write_ints_as_int32=write_ints_as_int32 + ) if write_arr.dtype != original_dtype: stream.write_attribute( VAR_ATTR_PREFIX_ORIGINAL_DTYPE, @@ -198,9 +233,13 @@ def write_dataset_bp( if time_dim in var.dims: if step >= int(var.sizes[time_dim]): continue - data = _numpy_for_write(var.data)[step] + data = _numpy_for_write( + var.data, write_ints_as_int32=write_ints_as_int32 + )[step] elif step == 0: - data = _numpy_for_write(var.data) + data = _numpy_for_write( + var.data, write_ints_as_int32=write_ints_as_int32 + ) else: continue diff --git a/xbout/tests/conftest.py b/xbout/tests/conftest.py index 1867f312..ce1f660c 100644 --- a/xbout/tests/conftest.py +++ b/xbout/tests/conftest.py @@ -6,13 +6,24 @@ import numpy as np import pytest -from xarray import DataArray -from xbout.tests.utils_for_tests import ( - _get_kwargs, - create_bout_ds_list, - create_bout_grid_ds, -) +try: # optional test dependency + import xarray as xr + from xarray import DataArray +except ModuleNotFoundError: # pragma: no cover + xr = None + DataArray = None + +try: # optional test dependency chain (boutdata, xarray, netCDF tooling, ...) + from xbout.tests.utils_for_tests import ( + _get_kwargs, + create_bout_ds_list, + create_bout_grid_ds, + ) +except (ModuleNotFoundError, ImportError): # pragma: no cover + _get_kwargs = None + create_bout_ds_list = None + create_bout_grid_ds = None @pytest.fixture(scope="session") @@ -47,6 +58,11 @@ def _bout_xyt_example_files( containing them, deleting the temporary directory once that test is done (if write_to_disk=True). """ + if xr is None: + pytest.skip("xarray is required for xbout tests") + if _get_kwargs is None or create_bout_ds_list is None: + pytest.skip("boutdata is required for xbout tests") + call_args = _get_kwargs(ignore="tmp_path_factory") try: diff --git a/xbout/tests/test_adios2_roundtrip.py b/xbout/tests/test_adios2_roundtrip.py index 2d7c5a7c..288aa273 100644 --- a/xbout/tests/test_adios2_roundtrip.py +++ b/xbout/tests/test_adios2_roundtrip.py @@ -43,3 +43,48 @@ def test_adios2_roundtrip_dataset_attrs_and_vars(tmp_path): assert int(ds2["scalar"].values) == 7 finally: ds2.close() + + +def test_adios2_write_ints_as_int32_on_disk(tmp_path): + from xbout.adioswriter import write_dataset_bp + + ds = xr.Dataset( + data_vars={ + "i64": (("t",), np.array([1, 2, 3], dtype=np.int64)), + "u64_small": (("t",), np.array([0, 7, 42], dtype=np.uint64)), + "f32": (("t",), np.array([1.0, 2.0, 3.0], dtype=np.float32)), + }, + coords={"t": ("t", np.array([0, 1, 2], dtype=np.int32))}, + ) + + path = tmp_path / "ints_as_int32.bp" + write_dataset_bp( + ds, str(path), time_dim="t", overwrite=True, write_ints_as_int32=True + ) + + fh = adios2.FileReader(str(path)) + try: + vars = fh.available_variables() + assert vars["i64"]["Type"] == "int32_t" + assert vars["u64_small"]["Type"] == "int32_t" + assert vars["t"]["Type"] == "int32_t" + assert vars["f32"]["Type"] == "float" + finally: + fh.close() + + +def test_adios2_write_ints_as_int32_overflow_raises(tmp_path): + from xbout.adioswriter import write_dataset_bp + + ds = xr.Dataset( + data_vars={ + "too_big": (("t",), np.array([np.iinfo(np.int32).max + 1], dtype=np.int64)), + }, + coords={"t": ("t", np.array([0], dtype=np.int32))}, + ) + + path = tmp_path / "overflow.bp" + with pytest.raises(ValueError, match=r"Cannot safely cast"): + write_dataset_bp( + ds, str(path), time_dim="t", overwrite=True, write_ints_as_int32=True + ) From e2453d7b73f05b3eba8f894b3f33567a49a8edef Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 28 May 2026 10:05:37 -0700 Subject: [PATCH 6/8] Revert __init__.py changes --- xbout/__init__.py | 120 +++++----------------------------------------- 1 file changed, 12 insertions(+), 108 deletions(-) diff --git a/xbout/__init__.py b/xbout/__init__.py index 4fdf42bf..b94242c6 100644 --- a/xbout/__init__.py +++ b/xbout/__init__.py @@ -1,38 +1,25 @@ -""" -xbout package public API. +from .load import open_boutdataset, collect +from .lazyload import lazy_open_boutdataset -This module is intentionally lightweight: many xbout features depend on optional -third-party packages (e.g. matplotlib, dask, boutdata). Import those lazily so -that submodules like ``xbout.adioswriter`` can be used without pulling in the -full dependency set. -""" +from . import geometries +from .geometries import register_geometry, REGISTERED_GEOMETRIES -from __future__ import annotations +from .boutdataset import BoutDatasetAccessor +from .boutdataarray import BoutDataArrayAccessor -from typing import TYPE_CHECKING +from .plotting.animate import animate_pcolormesh, animate_poloidal +from .plotting.utils import plot_separatrix -if TYPE_CHECKING: # pragma: no cover - from . import geometries - from .boutdataarray import BoutDataArrayAccessor - from .boutdataset import BoutDatasetAccessor - from .fastoutput import open_fastoutput - from .geometries import REGISTERED_GEOMETRIES, register_geometry - from .lazyload import lazy_open_boutdataset - from .load import collect, open_boutdataset - from .plotting.animate import animate_pcolormesh, animate_poloidal - from .plotting.utils import plot_separatrix +from .fastoutput import open_fastoutput -from importlib.metadata import PackageNotFoundError, version +from importlib.metadata import version, PackageNotFoundError try: __version__ = version(__name__) except PackageNotFoundError: - try: - from setuptools_scm import get_version + from setuptools_scm import get_version - __version__ = get_version(root="..", relative_to=__file__) - except Exception: # pragma: no cover - __version__ = "0+unknown" + __version__ = get_version(root="..", relative_to=__file__) __all__ = [ "open_boutdataset", @@ -48,86 +35,3 @@ "plot_separatrix", "open_fastoutput", ] - - -def __getattr__(name: str): - if name in {"geometries", "register_geometry", "REGISTERED_GEOMETRIES"}: - try: - from . import geometries - from .geometries import REGISTERED_GEOMETRIES, register_geometry - except ImportError as e: # pragma: no cover - raise ImportError( - "xbout geometries require optional dependencies (e.g. xarray)." - ) from e - return { - "geometries": geometries, - "register_geometry": register_geometry, - "REGISTERED_GEOMETRIES": REGISTERED_GEOMETRIES, - }[name] - - if name in {"open_boutdataset", "collect"}: - try: - from .load import collect, open_boutdataset - except ImportError as e: # pragma: no cover - raise ImportError( - "xbout.load requires optional dependencies (e.g. 'boutdata'). " - "Install the full xbout extras, or import submodules that do not " - "require boutdata." - ) from e - return {"open_boutdataset": open_boutdataset, "collect": collect}[name] - - if name == "lazy_open_boutdataset": - try: - from .lazyload import lazy_open_boutdataset - except ImportError as e: # pragma: no cover - raise ImportError( - "xbout.lazyload requires optional dependencies (e.g. dask, h5py, xarray)." - ) from e - return lazy_open_boutdataset - - if name in {"BoutDatasetAccessor", "BoutDataArrayAccessor"}: - try: - from .boutdataarray import BoutDataArrayAccessor - from .boutdataset import BoutDatasetAccessor - except ImportError as e: # pragma: no cover - raise ImportError( - "xbout accessors require optional dependencies (e.g. xarray)." - ) from e - return { - "BoutDatasetAccessor": BoutDatasetAccessor, - "BoutDataArrayAccessor": BoutDataArrayAccessor, - }[name] - - if name in {"animate_pcolormesh", "animate_poloidal"}: - try: - from .plotting.animate import animate_pcolormesh, animate_poloidal - except ImportError as e: # pragma: no cover - raise ImportError( - "xbout plotting requires optional dependencies (e.g. matplotlib)." - ) from e - return { - "animate_pcolormesh": animate_pcolormesh, - "animate_poloidal": animate_poloidal, - }[name] - - if name == "plot_separatrix": - try: - from .plotting.utils import plot_separatrix - except ImportError as e: # pragma: no cover - raise ImportError( - "xbout plotting requires optional dependencies (e.g. matplotlib)." - ) from e - return plot_separatrix - - if name == "open_fastoutput": - try: - from .fastoutput import open_fastoutput - except ImportError as e: # pragma: no cover - raise ImportError("xbout.fastoutput requires optional dependencies.") from e - return open_fastoutput - - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -def __dir__(): - return sorted(list(globals().keys()) + __all__) From 2e0419b6b1a0f381cc1ff3e4c8778c4f8aff427c Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 28 May 2026 10:39:58 -0700 Subject: [PATCH 7/8] dataset.bout.save: write_ints_as_int32 keyword Passes through to ADIOS write_dataset_bp --- xbout/boutdataset.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 1c4a8db0..19b346d4 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -778,6 +778,7 @@ def save( save_dtype=None, separate_vars=False, pre_load=False, + write_ints_as_int32: bool = False, ): """ Save data variables to a netCDF file. @@ -920,6 +921,7 @@ def dict_to_attrs(obj, section): str(var_savepath), time_dim=time_dim, overwrite=True, + write_ints_as_int32=write_ints_as_int32, ) else: single_var_ds.to_netcdf( @@ -945,6 +947,7 @@ def dict_to_attrs(obj, section): str(savepath), time_dim=time_dim, overwrite=True, + write_ints_as_int32=write_ints_as_int32, ) else: with ProgressBar(): From 5ddcd18ea978f82db76870706f58d233916210fb Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 11 Jun 2026 13:35:23 -0700 Subject: [PATCH 8/8] Remove fallbacks for missing xarray and boutdata Tests require these packages, so fail early if they're missing. --- xbout/tests/conftest.py | 27 ++++++--------------------- 1 file changed, 6 insertions(+), 21 deletions(-) diff --git a/xbout/tests/conftest.py b/xbout/tests/conftest.py index ce1f660c..b73c8ea8 100644 --- a/xbout/tests/conftest.py +++ b/xbout/tests/conftest.py @@ -6,24 +6,13 @@ import numpy as np import pytest +from xarray import DataArray -try: # optional test dependency - import xarray as xr - from xarray import DataArray -except ModuleNotFoundError: # pragma: no cover - xr = None - DataArray = None - -try: # optional test dependency chain (boutdata, xarray, netCDF tooling, ...) - from xbout.tests.utils_for_tests import ( - _get_kwargs, - create_bout_ds_list, - create_bout_grid_ds, - ) -except (ModuleNotFoundError, ImportError): # pragma: no cover - _get_kwargs = None - create_bout_ds_list = None - create_bout_grid_ds = None +from xbout.tests.utils_for_tests import ( + _get_kwargs, + create_bout_ds_list, + create_bout_grid_ds, +) @pytest.fixture(scope="session") @@ -58,10 +47,6 @@ def _bout_xyt_example_files( containing them, deleting the temporary directory once that test is done (if write_to_disk=True). """ - if xr is None: - pytest.skip("xarray is required for xbout tests") - if _get_kwargs is None or create_bout_ds_list is None: - pytest.skip("boutdata is required for xbout tests") call_args = _get_kwargs(ignore="tmp_path_factory")