diff --git a/src/parcels/_chunk_cached_array/core.py b/src/parcels/_chunk_cached_array/core.py index d9b938aa2..e155989e4 100644 --- a/src/parcels/_chunk_cached_array/core.py +++ b/src/parcels/_chunk_cached_array/core.py @@ -42,7 +42,7 @@ def wrap_dataset(ds: xr.Dataset, max_cache_bytes: int) -> xr.Dataset: for name in list(ds.coords): ds[name].load() for name in ds.data_vars: - var = ds[name].variable + var = ds.variables[name] if is_duck_array(var._data) and is_dask_collection(var._data): var._data = ChunkCachedArray(var._data, max_cache_bytes) # type: ignore[assignment, arg-type] return ds diff --git a/tests/test_backends.py b/tests/test_backends.py index bf8880ae8..136837dc9 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -10,6 +10,7 @@ import parcels import parcels.tutorial from parcels._datasets.structured.generated import simple_UV_dataset +from parcels._datasets.unstructured.generated import simple_small_delaunay from parcels.kernels import AdvectionRK2, AdvectionRK4 BackendT = Literal["WindowedArray", "Dask", "NumPy", "ChunkCachedArray"] @@ -170,3 +171,10 @@ def test_fieldset_describe_backend(nemo_dataset, backend: BackendT): for other_backend in BACKENDS - {backend}: assert not fieldset_uses_backend(fieldset, other_backend) + + +def test_chunk_cached_arrays_on_unstructured_fieldset(): + ds = simple_small_delaunay().chunk({"time": 1}) + fset = parcels.FieldSet.from_ugrid_conventions(ds) + fset.to_chunk_cached_arrays() + assert fieldset_uses_backend(fset, "ChunkCachedArray")