diff --git a/src/parcels/_core/_windowed_array.py b/src/parcels/_core/_windowed_array.py index 77cb3cf1c..49526e687 100644 --- a/src/parcels/_core/_windowed_array.py +++ b/src/parcels/_core/_windowed_array.py @@ -108,4 +108,6 @@ def maybe_windowed(data: xr.DataArray, max_levels: int | None = None): return data if data.dims and data.dims[0] == "time" and is_dask_collection(data.data): return WindowedArray(data, max_levels=max_levels) + elif data.dims and data.dims[0] == "mockT" and is_dask_collection(data.data): + return data.compute() return data diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index 93f403adb..7a0945573 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -7,6 +7,7 @@ import cf_xarray # noqa: F401 import uxarray as ux import xarray as xr +from dask import is_dask_collection import parcels._sgrid as sgrid import parcels._typing as ptyping @@ -102,6 +103,12 @@ def to_windowed_arrays(self, *, max_levels: int | None = None) -> Self: and ``max_levels`` then bounds its size. """ windowed = self.__dict__.setdefault("_windowed", {}) + for dim in ["lon", "lat"]: + # ensure lon and lat are loaded into memory for dask-backed datasets + if is_dask_collection(self.data[dim]): + self.data[dim].load() + if "depth" in self.data.dims and is_dask_collection(self.data["depth"]): + self.data["depth"].load() for name in self.scalar_field_names: current = windowed.get(name, self.data[name]) windowed[name] = maybe_windowed(current, max_levels=max_levels) diff --git a/src/parcels/_reprs.py b/src/parcels/_reprs.py index b0a24ce17..5a6ae2944 100644 --- a/src/parcels/_reprs.py +++ b/src/parcels/_reprs.py @@ -8,7 +8,10 @@ import numpy as np import xarray as xr +import zarr +from dask.base import is_dask_collection +from parcels._core._windowed_array import WindowedArray from parcels._python import isinstance_noimport if TYPE_CHECKING: @@ -190,6 +193,7 @@ class _FieldSetDescriptionRow: model_id: int | None name: str interp_method_or_value: str + backend: str | None = None def to_dict(self) -> dict[str, str]: return { @@ -197,6 +201,7 @@ def to_dict(self) -> dict[str, str]: "Type": self.type_, "Grid number": str(self.model_id) if self.model_id is not None else "-", "Interp method / value": self.interp_method_or_value, + "Parcels backend": self.backend if self.backend is not None else "-", } @@ -213,6 +218,22 @@ def _print_time_interval(time_interval: TimeInterval | None) -> str: return repr((time_interval.left, time_interval.right)) +def _field_backend(field: Field | VectorField) -> str | None: + if hasattr(field, "data"): + if isinstance(field.data, WindowedArray): + return "WindowedArray" + elif is_dask_collection(field.data.data): + return "Dask" + elif isinstance(field.data.variable._data, zarr.Array): + return "Zarr" + elif isinstance(field.data.data, np.ndarray): + return "NumPy" + else: + return type(field.data).__name__ + else: + return None + + def fieldset_describe(fieldset: FieldSet) -> str: rows: list[_FieldSetDescriptionRow] = [] models: dict[int, int] = {} # mapping of memory ID to a human readable ID @@ -235,6 +256,7 @@ def fieldset_describe(fieldset: FieldSet) -> str: model_id=model_id, name=field.name, interp_method_or_value=repr(field.interp_method), + backend=_field_backend(field), ) ) for k, v in fieldset.context.items(): @@ -244,6 +266,7 @@ def fieldset_describe(fieldset: FieldSet) -> str: model_id=None, name=k, interp_method_or_value=repr(v), + backend=None, ) ) return ( diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index 1c288b301..5ec45ac0d 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -7,7 +7,8 @@ import pandas as pd import pytest -from parcels import Field, ParticleFile, ParticleSet, XGrid, convert +import parcels.tutorial +from parcels import Field, ParticleFile, ParticleSet, XGrid, convert, open_raw_zarr from parcels._core.fieldset import FieldSet, _datetime_to_msg from parcels._core.model import _default_vector_field_components from parcels._datasets.structured.generic import datasets as datasets_structured @@ -398,17 +399,18 @@ def test_fieldset_describe(fieldset_two_models: FieldSet): fieldset = fieldset_two_models io = StringIO() expected = """\ -| Name | Type | Grid number | Interp method / value | -|:---------------|:------------|:--------------|:------------------------| -| my_list | Context | - | [1, 2, 'hello'] | -| my_value | Context | - | 2.0 | -| U | Field | 0 | XLinear(...) | -| V | Field | 0 | XLinear(...) | -| UV | VectorField | 0 | XLinear_Velocity(...) | -| U_wind | Field | 1 | XLinear(...) | -| V_wind | Field | 1 | XLinear(...) | -| UV_wind | VectorField | 1 | XLinear_Velocity(...) | -| constant_field | Field | 2 | XConstantField(...) | +| Name | Type | Grid number | Interp method / value | Backend | +| Name | Type | Grid number | Interp method / value | Parcels backend | +|:---------------|:------------|:--------------|:------------------------|:------------------| +| my_list | Context | - | [1, 2, 'hello'] | - | +| my_value | Context | - | 2.0 | - | +| U | Field | 0 | XLinear(...) | NumPy | +| V | Field | 0 | XLinear(...) | NumPy | +| UV | VectorField | 0 | XLinear_Velocity(...) | - | +| U_wind | Field | 1 | XLinear(...) | NumPy | +| V_wind | Field | 1 | XLinear(...) | NumPy | +| UV_wind | VectorField | 1 | XLinear_Velocity(...) | - | +| constant_field | Field | 2 | XConstantField(...) | NumPy | mesh: flat time interval: (np.datetime64('2000-01-01T00:00:00.000000000'), np.datetime64('2001-01-01T00:00:00.000000000')) @@ -416,3 +418,75 @@ def test_fieldset_describe(fieldset_two_models: FieldSet): fieldset.describe(io) actual = io.getvalue() assert actual == expected + + +def test_fieldset_describe_backends(tmp_path): + ds_u = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/U") + ds_v = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/V") + ds_w = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/W") + ds_coords = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/mesh_mask")[["glamf", "gphif"]] + + ds_fset = convert.nemo_to_sgrid( + fields={"U": ds_u["uo"], "V": ds_v["vo"], "W": ds_w["wo"]}, + coords=ds_coords, + ) + fieldset = FieldSet.from_sgrid_conventions(ds_fset) + + io = StringIO() + expected = """\ +| Name | Type | Grid number | Interp method / value | Parcels backend | +|:-------|:------------|--------------:|:------------------------|:------------------| +| U | Field | 0 | XLinear(...) | Dask | +| V | Field | 0 | XLinear(...) | Dask | +| W | Field | 0 | XLinear(...) | Dask | +| UV | VectorField | 0 | CGrid_Velocity(...) | - | +| UVW | VectorField | 0 | CGrid_Velocity(...) | - | + +mesh: spherical +time interval: (np.datetime64('2000-01-02T12:00:00.000000000'), np.datetime64('2000-01-12T12:00:00.000000000')) +""" + fieldset.describe(io) + actual = io.getvalue() + assert actual == expected + + # Also run with WindowedArray backend + fieldset = fieldset.to_windowed_arrays() + + io = StringIO() + expected = """\ +| Name | Type | Grid number | Interp method / value | Parcels backend | +|:-------|:------------|--------------:|:------------------------|:------------------| +| U | Field | 0 | XLinear(...) | WindowedArray | +| V | Field | 0 | XLinear(...) | WindowedArray | +| W | Field | 0 | XLinear(...) | WindowedArray | +| UV | VectorField | 0 | CGrid_Velocity(...) | - | +| UVW | VectorField | 0 | CGrid_Velocity(...) | - | + +mesh: spherical +time interval: (np.datetime64('2000-01-02T12:00:00.000000000'), np.datetime64('2000-01-12T12:00:00.000000000')) +""" + fieldset.describe(io) + actual = io.getvalue() + assert actual == expected + + path = tmp_path / "ds.zarr" + ds_fset.to_zarr(path) + ds_zarr = open_raw_zarr(path) + fieldset = FieldSet.from_sgrid_conventions(ds_zarr) + + io = StringIO() + expected = """\ +| Name | Type | Grid number | Interp method / value | Parcels backend | +|:-------|:------------|--------------:|:------------------------|:------------------| +| 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(...) | - | + +mesh: spherical +time interval: (np.datetime64('2000-01-02T12:00:00.000000000'), np.datetime64('2000-01-12T12:00:00.000000000')) +""" + fieldset.describe(io) + actual = io.getvalue() + assert actual == expected