From b56453e2a6d9b8528cd2b2ed6c9c214bf738e9ec Mon Sep 17 00:00:00 2001 From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com> Date: Tue, 16 Jun 2026 10:24:35 +0200 Subject: [PATCH 1/4] Add open_raw_zarr helper --- src/parcels/__init__.py | 2 ++ src/parcels/_xarray.py | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) create mode 100644 src/parcels/_xarray.py diff --git a/src/parcels/__init__.py b/src/parcels/__init__.py index 1378340fd..78227247b 100644 --- a/src/parcels/__init__.py +++ b/src/parcels/__init__.py @@ -10,6 +10,7 @@ import warnings as _stdlib_warnings from parcels._core.fieldset import FieldSet +from parcels._xarray import open_raw_zarr from parcels._core.particleset import ParticleSet from parcels._core.particlefile import ParticleFile, read_particlefile from parcels._core.particle import ( @@ -42,6 +43,7 @@ __all__ = [ # noqa: RUF022 # Core classes "FieldSet", + "open_raw_zarr", "ParticleSet", "ParticleFile", "Variable", diff --git a/src/parcels/_xarray.py b/src/parcels/_xarray.py new file mode 100644 index 000000000..7765ea929 --- /dev/null +++ b/src/parcels/_xarray.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import xarray as xr +import zarr +import zarr.storage +from zarr.abc.store import Store + + +def _not_implemented(*args, **kwargs): + raise NotImplementedError("This function it not implemented") + + +def open_raw_zarr(store: Store): + """Open a Zarr dataset in an Xarray dataset, bypassing Dask.""" + with xr.open_zarr(store) as ds: + var_to_dims = {name: var.dims for name, var in ds.variables.items()} + coord_names = list(ds.coords) + + group = zarr.open(store, mode="r") + assert isinstance(group, zarr.Group) + + data_vars = {} + coords = {} + for name, array in group.members(): + if not isinstance(array, zarr.Array): + raise ValueError("Discovered a zarr.Group in the root group. open_raw_zarr doesn't work with nested groups") + is_coord = name in coord_names + + if not is_coord: + array.__array_function__ = _not_implemented # trick xarray to prevent coersion to a numpy array + + var = xr.Variable(var_to_dims[name], array) + + if is_coord: + coords[name] = var + else: # name is a data var + data_vars[name] = var + + return xr.Dataset(data_vars, coords) From 42609f16ae14ec78f04ed880270cbcb67b7aa897 Mon Sep 17 00:00:00 2001 From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com> Date: Tue, 16 Jun 2026 10:32:36 +0200 Subject: [PATCH 2/4] Make open_raw_zarr read metadata And add test --- src/parcels/_xarray.py | 21 ++++++++------------- tests/test_xarray.py | 15 +++++++++++++++ 2 files changed, 23 insertions(+), 13 deletions(-) create mode 100644 tests/test_xarray.py diff --git a/src/parcels/_xarray.py b/src/parcels/_xarray.py index 7765ea929..88f37bbd1 100644 --- a/src/parcels/_xarray.py +++ b/src/parcels/_xarray.py @@ -14,26 +14,21 @@ def open_raw_zarr(store: Store): """Open a Zarr dataset in an Xarray dataset, bypassing Dask.""" with xr.open_zarr(store) as ds: var_to_dims = {name: var.dims for name, var in ds.variables.items()} - coord_names = list(ds.coords) + var_to_attrs = {name: var.attrs for name, var in ds.variables.items()} + coords = {name: ds[name].variable.load() for name in ds.coords} + ds_attrs = ds.attrs group = zarr.open(store, mode="r") assert isinstance(group, zarr.Group) data_vars = {} - coords = {} for name, array in group.members(): if not isinstance(array, zarr.Array): raise ValueError("Discovered a zarr.Group in the root group. open_raw_zarr doesn't work with nested groups") - is_coord = name in coord_names + if name in coords: + continue - if not is_coord: - array.__array_function__ = _not_implemented # trick xarray to prevent coersion to a numpy array + array.__array_function__ = _not_implemented # trick xarray to prevent coersion to a numpy array + data_vars[name] = xr.Variable(var_to_dims[name], array, attrs=var_to_attrs[name]) - var = xr.Variable(var_to_dims[name], array) - - if is_coord: - coords[name] = var - else: # name is a data var - data_vars[name] = var - - return xr.Dataset(data_vars, coords) + return xr.Dataset(data_vars, coords, attrs=ds_attrs) diff --git a/tests/test_xarray.py b/tests/test_xarray.py new file mode 100644 index 000000000..a4d8f0349 --- /dev/null +++ b/tests/test_xarray.py @@ -0,0 +1,15 @@ +import pytest +import xarray as xr + +from parcels import open_raw_zarr +from parcels._datasets.structured.generic import datasets + + +@pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets.items()]) +def test_open_raw_zarr_roundtrip(ds, tmp_path): + path = tmp_path / "ds.zarr" + ds.to_zarr(path) + + result = open_raw_zarr(path) + + xr.testing.assert_identical(result.load(), ds) From bb4bd7b7a55d810bf34a4ba16928dfad776c608b Mon Sep 17 00:00:00 2001 From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:31:22 +0200 Subject: [PATCH 3/4] Add check that data_vars are Zarr arrays --- tests/test_xarray.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_xarray.py b/tests/test_xarray.py index a4d8f0349..3f509e1ab 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -1,15 +1,20 @@ import pytest import xarray as xr +import zarr from parcels import open_raw_zarr from parcels._datasets.structured.generic import datasets @pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets.items()]) -def test_open_raw_zarr_roundtrip(ds, tmp_path): +def test_open_raw_zarr(ds, tmp_path): path = tmp_path / "ds.zarr" ds.to_zarr(path) result = open_raw_zarr(path) + for k in result.data_vars: + # tests that the internal representation within Xarray isn't coerced into a numpy array + assert isinstance(result[k]._variable._data, zarr.Array) + xr.testing.assert_identical(result.load(), ds) From 93c4edf8ff608a14a243271339a802d9b5ca0d2a Mon Sep 17 00:00:00 2001 From: Vecko <36369090+VeckoTheGecko@users.noreply.github.com> Date: Thu, 2 Jul 2026 11:51:42 +0200 Subject: [PATCH 4/4] Fix mypy --- src/parcels/_xarray.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/parcels/_xarray.py b/src/parcels/_xarray.py index 88f37bbd1..a39761c4e 100644 --- a/src/parcels/_xarray.py +++ b/src/parcels/_xarray.py @@ -28,7 +28,8 @@ def open_raw_zarr(store: Store): if name in coords: continue - array.__array_function__ = _not_implemented # trick xarray to prevent coersion to a numpy array + # trick xarray to prevent coersion to a numpy array + array.__array_function__ = _not_implemented # type: ignore[attr-defined] data_vars[name] = xr.Variable(var_to_dims[name], array, attrs=var_to_attrs[name]) return xr.Dataset(data_vars, coords, attrs=ds_attrs)