diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index 9f31711d3..305fb0e8d 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -1,5 +1,6 @@ from __future__ import annotations +import warnings from abc import ABC, abstractmethod from collections.abc import Hashable, Sequence from typing import Any, Self @@ -7,6 +8,7 @@ import cf_xarray # noqa: F401 import uxarray as ux import xarray as xr +import zarr from dask import is_dask_collection import parcels._sgrid as sgrid @@ -131,6 +133,17 @@ def preprocess_sgrid_model_data(ds: xr.Dataset) -> xr.Dataset: return ds +def validate_field_data(ds: xr.Dataset) -> xr.Dataset: + if any(isinstance(da.variable._data, zarr.Array) for da in ds.data_vars.values()): + warnings.warn( + "Changing a Zarr-backed dataset. This may convert the Parcels backend to NumPy. " + "If you want to keep the Zarr backend, please use `skip_field_data_validation=True` when creating the FieldSet.", + UserWarning, + stacklevel=2, + ) + return ds.fillna(0) + + class StructuredModelData(ModelData): def __init__( self, @@ -144,7 +157,7 @@ def __init__( data = preprocess_sgrid_model_data(data) if not skip_field_data_validation: - data = data.fillna(0) + data = validate_field_data(data) grid = XGrid(data, mesh) self.data = data diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index 26310f703..f72ec7d17 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -397,6 +397,21 @@ def test_fieldset_add_error_on_duplicate_context_values(): fset1 + fset2 +@pytest.mark.parametrize("skip", [True, False]) +def test_zarr_warning_on_fieldset_creation(skip, tmp_path): + """Test that creating a FieldSet from a Zarr-backed dataset raises a warning about potential backend changes.""" + ds = parcels.tutorial.open_dataset("CopernicusMarine_data_for_Argo_tutorial/data") + ds = convert.copernicusmarine_to_sgrid(fields={"U": ds["uo"], "V": ds["vo"]}) + path = tmp_path / "ds.zarr" + ds.to_zarr(path) + ds_zarr = open_raw_zarr(path) + if not skip: + with pytest.warns(UserWarning, match="Changing a Zarr-backed dataset"): + FieldSet.from_sgrid_conventions(ds_zarr, skip_field_data_validation=skip) + else: + FieldSet.from_sgrid_conventions(ds_zarr, skip_field_data_validation=skip) + + def test_fieldset_add_context_values(): """Test that context values from both FieldSets are present in the combined FieldSet.""" ds1 = datasets_structured["ds_2d_left"][["U_A_grid", "grid"]].rename({"U_A_grid": "U1"}) @@ -494,15 +509,15 @@ def test_fieldset_describe_backends(tmp_path): path = tmp_path / "ds.zarr" ds_fset.to_zarr(path) ds_zarr = open_raw_zarr(path) - fieldset = FieldSet.from_sgrid_conventions(ds_zarr) + fieldset = FieldSet.from_sgrid_conventions(ds_zarr, skip_field_data_validation=True) io = StringIO() expected = """\ | Name | Type | Grid number | Interp method / value | Parcels backend | |:-------|:------------|--------------:|:------------------------|:------------------| -| U | Field | 0 | XLinear(...) | NumPy | -| V | Field | 0 | XLinear(...) | NumPy | -| W | Field | 0 | XLinear(...) | NumPy | +| U | Field | 0 | XLinear(...) | Zarr | +| V | Field | 0 | XLinear(...) | Zarr | +| W | Field | 0 | XLinear(...) | Zarr | | UV | VectorField | 0 | CGrid_Velocity(...) | - | | UVW | VectorField | 0 | CGrid_Velocity(...) | - |