Skip to content

Commit 091d2ae

Browse files
erikvansebilleerikvansebille
authored
Load no-time vars as numpy arrays in to_windowed_arrays (#2752)
* Compute variables that have a mockT dimension * Explicitly loading lon and lat * Adding Backend to fieldset repr To show which fields are WindowedArrays, which are Dask and which are NumPy * Adding unt test for backends * Moving lon/lat loading to to_windowed_arrays * Renaming Backend to Parcels backend Following suggestion by @VeckoTheGecko * Also adding support for zarr backend in describe() * Adding print statements for CI debugging As I can't reproduce the failing unit tests locally * fixing time_interval because of fewer NorthSea files Related to #2750 * Separate depth loading As not all fields have depth coordinate
1 parent 59c549e commit 091d2ae

4 files changed

Lines changed: 118 additions & 12 deletions

File tree

‎src/parcels/_core/_windowed_array.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,4 +108,6 @@ def maybe_windowed(data: xr.DataArray, max_levels: int | None = None):
108108
return data
109109
if data.dims and data.dims[0] == "time" and is_dask_collection(data.data):
110110
return WindowedArray(data, max_levels=max_levels)
111+
elif data.dims and data.dims[0] == "mockT" and is_dask_collection(data.data):
112+
return data.compute()
111113
return data

‎src/parcels/_core/model.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import cf_xarray # noqa: F401
88
import uxarray as ux
99
import xarray as xr
10+
from dask import is_dask_collection
1011

1112
import parcels._sgrid as sgrid
1213
import parcels._typing as ptyping
@@ -102,6 +103,12 @@ def to_windowed_arrays(self, *, max_levels: int | None = None) -> Self:
102103
and ``max_levels`` then bounds its size.
103104
"""
104105
windowed = self.__dict__.setdefault("_windowed", {})
106+
for dim in ["lon", "lat"]:
107+
# ensure lon and lat are loaded into memory for dask-backed datasets
108+
if is_dask_collection(self.data[dim]):
109+
self.data[dim].load()
110+
if "depth" in self.data.dims and is_dask_collection(self.data["depth"]):
111+
self.data["depth"].load()
105112
for name in self.scalar_field_names:
106113
current = windowed.get(name, self.data[name])
107114
windowed[name] = maybe_windowed(current, max_levels=max_levels)

‎src/parcels/_reprs.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88

99
import numpy as np
1010
import xarray as xr
11+
import zarr
12+
from dask.base import is_dask_collection
1113

14+
from parcels._core._windowed_array import WindowedArray
1215
from parcels._python import isinstance_noimport
1316

1417
if TYPE_CHECKING:
@@ -190,13 +193,15 @@ class _FieldSetDescriptionRow:
190193
model_id: int | None
191194
name: str
192195
interp_method_or_value: str
196+
backend: str | None = None
193197

194198
def to_dict(self) -> dict[str, str]:
195199
return {
196200
"Name": self.name,
197201
"Type": self.type_,
198202
"Grid number": str(self.model_id) if self.model_id is not None else "-",
199203
"Interp method / value": self.interp_method_or_value,
204+
"Parcels backend": self.backend if self.backend is not None else "-",
200205
}
201206

202207

@@ -213,6 +218,22 @@ def _print_time_interval(time_interval: TimeInterval | None) -> str:
213218
return repr((time_interval.left, time_interval.right))
214219

215220

221+
def _field_backend(field: Field | VectorField) -> str | None:
222+
if hasattr(field, "data"):
223+
if isinstance(field.data, WindowedArray):
224+
return "WindowedArray"
225+
elif is_dask_collection(field.data.data):
226+
return "Dask"
227+
elif isinstance(field.data.variable._data, zarr.Array):
228+
return "Zarr"
229+
elif isinstance(field.data.data, np.ndarray):
230+
return "NumPy"
231+
else:
232+
return type(field.data).__name__
233+
else:
234+
return None
235+
236+
216237
def fieldset_describe(fieldset: FieldSet) -> str:
217238
rows: list[_FieldSetDescriptionRow] = []
218239
models: dict[int, int] = {} # mapping of memory ID to a human readable ID
@@ -235,6 +256,7 @@ def fieldset_describe(fieldset: FieldSet) -> str:
235256
model_id=model_id,
236257
name=field.name,
237258
interp_method_or_value=repr(field.interp_method),
259+
backend=_field_backend(field),
238260
)
239261
)
240262
for k, v in fieldset.context.items():
@@ -244,6 +266,7 @@ def fieldset_describe(fieldset: FieldSet) -> str:
244266
model_id=None,
245267
name=k,
246268
interp_method_or_value=repr(v),
269+
backend=None,
247270
)
248271
)
249272
return (

‎tests/test_fieldset.py‎

Lines changed: 86 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
import pandas as pd
88
import pytest
99

10-
from parcels import Field, ParticleFile, ParticleSet, XGrid, convert
10+
import parcels.tutorial
11+
from parcels import Field, ParticleFile, ParticleSet, XGrid, convert, open_raw_zarr
1112
from parcels._core.fieldset import FieldSet, _datetime_to_msg
1213
from parcels._core.model import _default_vector_field_components
1314
from parcels._datasets.structured.generic import datasets as datasets_structured
@@ -398,21 +399,94 @@ def test_fieldset_describe(fieldset_two_models: FieldSet):
398399
fieldset = fieldset_two_models
399400
io = StringIO()
400401
expected = """\
401-
| Name | Type | Grid number | Interp method / value |
402-
|:---------------|:------------|:--------------|:------------------------|
403-
| my_list | Context | - | [1, 2, 'hello'] |
404-
| my_value | Context | - | 2.0 |
405-
| U | Field | 0 | XLinear(...) |
406-
| V | Field | 0 | XLinear(...) |
407-
| UV | VectorField | 0 | XLinear_Velocity(...) |
408-
| U_wind | Field | 1 | XLinear(...) |
409-
| V_wind | Field | 1 | XLinear(...) |
410-
| UV_wind | VectorField | 1 | XLinear_Velocity(...) |
411-
| constant_field | Field | 2 | XConstantField(...) |
402+
| Name | Type | Grid number | Interp method / value | Backend |
403+
| Name | Type | Grid number | Interp method / value | Parcels backend |
404+
|:---------------|:------------|:--------------|:------------------------|:------------------|
405+
| my_list | Context | - | [1, 2, 'hello'] | - |
406+
| my_value | Context | - | 2.0 | - |
407+
| U | Field | 0 | XLinear(...) | NumPy |
408+
| V | Field | 0 | XLinear(...) | NumPy |
409+
| UV | VectorField | 0 | XLinear_Velocity(...) | - |
410+
| U_wind | Field | 1 | XLinear(...) | NumPy |
411+
| V_wind | Field | 1 | XLinear(...) | NumPy |
412+
| UV_wind | VectorField | 1 | XLinear_Velocity(...) | - |
413+
| constant_field | Field | 2 | XConstantField(...) | NumPy |
412414
413415
mesh: flat
414416
time interval: (np.datetime64('2000-01-01T00:00:00.000000000'), np.datetime64('2001-01-01T00:00:00.000000000'))
415417
"""
416418
fieldset.describe(io)
417419
actual = io.getvalue()
418420
assert actual == expected
421+
422+
423+
def test_fieldset_describe_backends(tmp_path):
424+
ds_u = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/U")
425+
ds_v = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/V")
426+
ds_w = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/W")
427+
ds_coords = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/mesh_mask")[["glamf", "gphif"]]
428+
429+
ds_fset = convert.nemo_to_sgrid(
430+
fields={"U": ds_u["uo"], "V": ds_v["vo"], "W": ds_w["wo"]},
431+
coords=ds_coords,
432+
)
433+
fieldset = FieldSet.from_sgrid_conventions(ds_fset)
434+
435+
io = StringIO()
436+
expected = """\
437+
| Name | Type | Grid number | Interp method / value | Parcels backend |
438+
|:-------|:------------|--------------:|:------------------------|:------------------|
439+
| U | Field | 0 | XLinear(...) | Dask |
440+
| V | Field | 0 | XLinear(...) | Dask |
441+
| W | Field | 0 | XLinear(...) | Dask |
442+
| UV | VectorField | 0 | CGrid_Velocity(...) | - |
443+
| UVW | VectorField | 0 | CGrid_Velocity(...) | - |
444+
445+
mesh: spherical
446+
time interval: (np.datetime64('2000-01-02T12:00:00.000000000'), np.datetime64('2000-01-12T12:00:00.000000000'))
447+
"""
448+
fieldset.describe(io)
449+
actual = io.getvalue()
450+
assert actual == expected
451+
452+
# Also run with WindowedArray backend
453+
fieldset = fieldset.to_windowed_arrays()
454+
455+
io = StringIO()
456+
expected = """\
457+
| Name | Type | Grid number | Interp method / value | Parcels backend |
458+
|:-------|:------------|--------------:|:------------------------|:------------------|
459+
| U | Field | 0 | XLinear(...) | WindowedArray |
460+
| V | Field | 0 | XLinear(...) | WindowedArray |
461+
| W | Field | 0 | XLinear(...) | WindowedArray |
462+
| UV | VectorField | 0 | CGrid_Velocity(...) | - |
463+
| UVW | VectorField | 0 | CGrid_Velocity(...) | - |
464+
465+
mesh: spherical
466+
time interval: (np.datetime64('2000-01-02T12:00:00.000000000'), np.datetime64('2000-01-12T12:00:00.000000000'))
467+
"""
468+
fieldset.describe(io)
469+
actual = io.getvalue()
470+
assert actual == expected
471+
472+
path = tmp_path / "ds.zarr"
473+
ds_fset.to_zarr(path)
474+
ds_zarr = open_raw_zarr(path)
475+
fieldset = FieldSet.from_sgrid_conventions(ds_zarr)
476+
477+
io = StringIO()
478+
expected = """\
479+
| Name | Type | Grid number | Interp method / value | Parcels backend |
480+
|:-------|:------------|--------------:|:------------------------|:------------------|
481+
| U | Field | 0 | XLinear(...) | Zarr |
482+
| V | Field | 0 | XLinear(...) | Zarr |
483+
| W | Field | 0 | XLinear(...) | Zarr |
484+
| UV | VectorField | 0 | CGrid_Velocity(...) | - |
485+
| UVW | VectorField | 0 | CGrid_Velocity(...) | - |
486+
487+
mesh: spherical
488+
time interval: (np.datetime64('2000-01-02T12:00:00.000000000'), np.datetime64('2000-01-12T12:00:00.000000000'))
489+
"""
490+
fieldset.describe(io)
491+
actual = io.getvalue()
492+
assert actual == expected

0 commit comments

Comments
 (0)