diff --git a/docs/user_guide/examples/tutorial_Argofloats.ipynb b/docs/user_guide/examples/tutorial_Argofloats.ipynb index b5be395e0..dfd4f6a26 100644 --- a/docs/user_guide/examples/tutorial_Argofloats.ipynb +++ b/docs/user_guide/examples/tutorial_Argofloats.ipynb @@ -107,9 +107,6 @@ " \"CopernicusMarine_data_for_Argo_tutorial/data\"\n", ")\n", "\n", - "# TODO check how we can get good performance without loading full dataset in memory\n", - "ds_fields.load() # load the dataset into memory\n", - "\n", "# Select fields\n", "fields = {\n", " \"U\": ds_fields[\"uo\"],\n", @@ -120,6 +117,7 @@ "# Convert to SGRID-compliant dataset and create FieldSet\n", "ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n", "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n", + "fieldset.to_windowed_arrays()\n", "\n", "# Define a new Particle type including extra Variables\n", "ArgoParticle = parcels.Particle.add_variable(\n", diff --git a/src/parcels/_core/_windowed_array.py b/src/parcels/_core/_windowed_array.py new file mode 100644 index 000000000..77cb3cf1c --- /dev/null +++ b/src/parcels/_core/_windowed_array.py @@ -0,0 +1,111 @@ +"""Transparent rolling time-window cache for lazy (dask-backed) field data. + +Assumptions / current limits: + * ``time`` is the leading dimension of the field (true for both the SGRID and + UGRID ingestion paths; the structured path transposes to ``(time, ...)``). + * Valid while the requested time indices stay within the resident window + (i.e. all particles share the clock). A sample that requests time indices + spanning more than the retained levels would force reloads. + * The clock is assumed monotonic but may run in either direction: forward + (``dt > 0``) or backward (``dt < 0``). Eviction keeps only the levels each + ``isel`` actually requests, which is symmetric in time -- so direction never + enters the logic and no integration-direction flag is needed. +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr +from dask import is_dask_collection + +# xarray / uxarray ``isel`` keyword arguments that are NOT dimension indexers. +_NON_INDEXER_KWARGS = frozenset({"drop", "missing_dims", "ignore_grid"}) + + +class WindowedArray: + """Wrap a lazy DataArray so ``isel`` loads/caches/evicts time levels as NumPy.""" + + def __init__(self, data: xr.DataArray, time_dim: str = "time", max_levels: int | None = None): + if data.dims[0] != time_dim: + raise ValueError(f"WindowedArray expects {time_dim!r} as the leading dimension, got {data.dims}") + self._data = data + self._tdim = time_dim + self._cache: dict[int, np.ndarray] = {} # time index -> NumPy slab (remaining dims) + self._max = max_levels + # diagnostics + self.loads = 0 + self.bytes_read = 0 + self._slab_bytes = int(np.prod(data.isel({time_dim: 0}).shape)) * data.dtype.itemsize + + # -- transparency: forward everything we don't override ------------------- + def __getattr__(self, name): + # __getattr__ only fires for misses; reach _data without recursing. + return getattr(object.__getattribute__(self, "_data"), name) + + def __repr__(self): + return ( + f"WindowedArray(time_dim={self._tdim!r}, cached_levels={sorted(self._cache)}, " + f"loads={self.loads})\n{self._data!r}" + ) + + # -- window management ---------------------------------------------------- + def _read_level(self, lvl: int) -> np.ndarray: + """Bulk, sequential read of one time level into NumPy (the dask->NumPy step).""" + return np.asarray(self._data.isel({self._tdim: int(lvl)}).values) + + def _ensure(self, levels: np.ndarray) -> None: + for lvl in levels: + lvl = int(lvl) + if lvl not in self._cache: + self._cache[lvl] = self._read_level(lvl) + self.loads += 1 + self.bytes_read += self._slab_bytes + # retire cached levels outside the span this call requested. Direction never + # enters here: a forward (dt > 0) or backward (dt < 0) clock both shed their + # trailing edge. Consecutive brackets overlap on one endpoint (inside [lo, hi]), + # so it is retained and each level is still read at most once per pass. + lo, hi = int(np.min(levels)), int(np.max(levels)) + for old in [k for k in self._cache if k < lo or k > hi]: + del self._cache[old] + if self._max is not None and len(self._cache) > self._max: + for old in sorted(self._cache)[: len(self._cache) - self._max]: + del self._cache[old] + + # -- intercepted indexing ------------------------------------------------- + def isel(self, indexers: dict | None = None, **kwargs): + sel = dict(indexers) if indexers is not None else {} + sel.update({k: v for k, v in kwargs.items() if k not in _NON_INDEXER_KWARGS}) + + # no time selection -> nothing to window; preserve control kwargs + if self._tdim not in sel: + return self._data.isel(indexers, **kwargs) + + t_ind = sel[self._tdim] + t_vals = np.asarray(t_ind.values if isinstance(t_ind, xr.DataArray) else t_ind) + levels = np.unique(t_vals) + if levels.size == 0: + # empty selection (e.g. a kernel evaluating an empty particle subset): + # nothing to load or evict; gather from an empty NumPy block below + block = np.empty((0, *self._data.shape[1:]), dtype=self._data.dtype) + else: + self._ensure(levels) + # stack the resident levels into one small NumPy block; remap to local indices + block = np.stack([self._cache[int(lvl)] for lvl in levels]) # (nlevels, *rest) + nda = xr.DataArray(block, dims=self._data.dims) # NumPy-backed, original dim order + local = np.searchsorted(levels, t_vals) + sel[self._tdim] = xr.DataArray(local, dims=getattr(t_ind, "dims", ())) + return nda.isel(sel) # plain vectorised gather in NumPy (no ignore_grid needed) + + +def maybe_windowed(data: xr.DataArray, max_levels: int | None = None): + """Wrap dask-backed, field data in a ``WindowedArray``; else pass through. + + NumPy-backed fields (already resident) and fields without a leading ``time`` + dimension are returned unchanged, so existing eager workflows are unaffected. + Already-wrapped data is returned unchanged. + """ + if isinstance(data, WindowedArray): + return data + if data.dims and data.dims[0] == "time" and is_dask_collection(data.data): + return WindowedArray(data, max_levels=max_levels) + return data diff --git a/src/parcels/_core/field.py b/src/parcels/_core/field.py index 0853a05fa..701a90d91 100644 --- a/src/parcels/_core/field.py +++ b/src/parcels/_core/field.py @@ -101,7 +101,7 @@ def __init__( @property def data(self): - return self.model.data[self.name] + return self.model.field_data(self.name) @property def grid(self): # TODO PR: Remove in favour of referencing model grid directly diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index 9ed2f9460..940c98cf2 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -152,6 +152,39 @@ def add_field(self, field: Field, name: str | None = None): self.fields[name] = field + def to_windowed_arrays(self, *, max_levels: int | None = None): + """Wrap dask-backed field data in rolling time-window caches. + + Opt-in optimization for forward-marching simulations where all particles + share a single clock. Delegates to each underlying model; dask-backed, + time-leading fields are served through a resident NumPy window (each time + level loaded once and evicted as the clock advances) instead of re-reading + chunks on every kernel step. NumPy-backed (eager) and non-time-leading + fields are left unchanged, and re-invoking is idempotent, so this is safe + to call more than once. + + Parameters + ---------- + max_levels : int, optional + Hard cap on the number of time levels kept resident per field. + With the default ``None``, each interpolation call decides what + stays resident: the cache keeps exactly the span of time indices + that call requests and evicts every level outside it. During time + integration particles bracket the current time between two + adjacent levels, so the default keeps at most two levels resident. + Only when a single call requests a wider time span (e.g. particles + spread across many time levels) does the window grow beyond that, + and ``max_levels`` then bounds its size. + + Returns + ------- + FieldSet + ``self``, to allow chaining. + """ + for model in self.models: + model.to_windowed_arrays(max_levels=max_levels) + return self + def add_constant_field(self, name: str, value, mesh: ptyping.Mesh = "spherical"): """Wrapper function to add a Field that is constant in space, useful e.g. when using constant horizontal diffusivity diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index c44026370..93f403adb 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -10,6 +10,7 @@ import parcels._sgrid as sgrid import parcels._typing as ptyping +from parcels._core._windowed_array import maybe_windowed from parcels._core.basegrid import BaseGrid from parcels._core.field import Field, VectorField from parcels._core.utils.time import TimeInterval @@ -62,6 +63,50 @@ def assert_valid_model_data(self) -> None: raise e return + def field_data(self, name: str) -> Any: + """Return the array backing field ``name``. + + Normally this is the ``xr.DataArray`` held in the dataset. After + :meth:`to_windowed_arrays`, dask-backed fields are served through a + cached :class:`~parcels._core._windowed_array.WindowedArray` instead. + """ + windowed = self.__dict__.get("_windowed") + if windowed is not None and name in windowed: + return windowed[name] + return self.data[name] + + def to_windowed_arrays(self, *, max_levels: int | None = None) -> Self: + """Wrap dask-backed field data in rolling time-window caches. + + Opt-in optimization for forward-marching simulations where all particles + share a single clock. For each dask-backed, time-leading field, ``isel`` + then samples a resident NumPy window (each time level loaded once and + evicted as the clock advances) instead of re-reading chunks and paying the + dask scheduling overhead on every kernel step. NumPy-backed (eager) fields + and non-time-leading fields are left unchanged. + + Idempotent: re-invoking reuses the existing wrapper (and its warm cache) + rather than rebuilding it. + + Parameters + ---------- + max_levels : int, optional + Hard cap on the number of time levels kept resident per field. + With the default ``None``, each interpolation call decides what + stays resident: the cache keeps exactly the span of time indices + that call requests and evicts every level outside it. During time + integration particles bracket the current time between two + adjacent levels, so the default keeps at most two levels resident. + Only when a single call requests a wider time span (e.g. particles + spread across many time levels) does the window grow beyond that, + and ``max_levels`` then bounds its size. + """ + windowed = self.__dict__.setdefault("_windowed", {}) + for name in self.scalar_field_names: + current = windowed.get(name, self.data[name]) + windowed[name] = maybe_windowed(current, max_levels=max_levels) + return self + @property def time_interval(self) -> TimeInterval | None: try: diff --git a/tests/test_windowed_array.py b/tests/test_windowed_array.py new file mode 100644 index 000000000..47a85c8eb --- /dev/null +++ b/tests/test_windowed_array.py @@ -0,0 +1,164 @@ +"""Tests for the transparent rolling time-window cache (WindowedArray).""" + +import dask.array as da +import numpy as np +import pytest +import xarray as xr + +from parcels import FieldSet, ParticleSet +from parcels._core._windowed_array import WindowedArray, maybe_windowed +from parcels._datasets.structured.generated import simple_UV_dataset +from parcels.kernels import AdvectionRK2 + + +def test_windowed_isel_matches_dask_loads_once_and_evicts(): + """WindowedArray.isel must equal dask isel, load each level once, keep <=2 resident.""" + ntime, n, npart = 20, 64, 200 + rng = np.random.default_rng(0) + base = rng.standard_normal((ntime, 3, n, n)) + lazy = xr.DataArray(da.from_array(base, chunks=(1, 3, n, n)), dims=("time", "depth", "lat", "lon")) + win = WindowedArray(lazy) + + worst, max_cache = 0.0, 0 + for step in range(40): + ti = min(step // 2, ntime - 2) # advancing clock, 2 sub-steps per level + yi, xi = rng.integers(0, n, npart), rng.integers(0, n, npart) + zi = np.zeros(npart, dtype=int) + sel = dict( + time=xr.DataArray(np.r_[np.full(npart, ti), np.full(npart, ti + 1)], dims="p"), + depth=xr.DataArray(np.r_[zi, zi], dims="p"), + lat=xr.DataArray(np.r_[yi, yi], dims="p"), + lon=xr.DataArray(np.r_[xi, xi], dims="p"), + ) + got = win.isel(sel).data + ref = lazy.isel(sel).data.compute() + worst = max(worst, float(np.abs(got - ref).max())) + max_cache = max(max_cache, len(win._cache)) + + assert worst == 0.0 # byte-identical to dask + assert win.loads == ntime # each time level read exactly once + assert max_cache <= 2 # only the bracketing levels resident + + +def test_windowed_isel_backward_clock_loads_once_and_evicts(): + """A backward-running clock (dt < 0) must also load each level once and keep <=2 + resident: eviction is symmetric, so no integration-direction flag is needed. + """ + ntime, n, npart = 20, 64, 200 + rng = np.random.default_rng(0) + base = rng.standard_normal((ntime, 3, n, n)) + lazy = xr.DataArray(da.from_array(base, chunks=(1, 3, n, n)), dims=("time", "depth", "lat", "lon")) + win = WindowedArray(lazy) + + worst, max_cache = 0.0, 0 + for step in range(40): + ti = max(ntime - 2 - step // 2, 0) # receding clock, 2 sub-steps per level + yi, xi = rng.integers(0, n, npart), rng.integers(0, n, npart) + zi = np.zeros(npart, dtype=int) + sel = dict( + time=xr.DataArray(np.r_[np.full(npart, ti), np.full(npart, ti + 1)], dims="p"), + depth=xr.DataArray(np.r_[zi, zi], dims="p"), + lat=xr.DataArray(np.r_[yi, yi], dims="p"), + lon=xr.DataArray(np.r_[xi, xi], dims="p"), + ) + got = win.isel(sel).data + ref = lazy.isel(sel).data.compute() + worst = max(worst, float(np.abs(got - ref).max())) + max_cache = max(max_cache, len(win._cache)) + + assert worst == 0.0 # byte-identical to dask + assert win.loads == ntime # each time level read exactly once, going backward + assert max_cache <= 2 # only the bracketing levels resident + + +def test_to_windowed_arrays_wraps_dask_but_not_numpy(): + ds = simple_UV_dataset(mesh="flat") + fset_np = FieldSet.from_sgrid_conventions(ds, mesh="flat") + fset_dk = FieldSet.from_sgrid_conventions(ds.chunk({"time": 1}), mesh="flat") + + # construction is never windowing -- it is opt-in via the fieldset method + assert not isinstance(fset_np.U.data, WindowedArray) + assert not isinstance(fset_dk.U.data, WindowedArray) + assert isinstance(fset_dk.U.data.data, da.Array) # chunked input stays lazy (dask-backed) + + assert fset_np.to_windowed_arrays() is fset_np # chainable + fset_dk.to_windowed_arrays() + + # numpy-backed field is left eager; dask-backed field gets wrapped + assert not isinstance(fset_np.U.data, WindowedArray) + assert isinstance(fset_dk.U.data, WindowedArray) + # transparency: forwarded attributes still behave like the DataArray + assert fset_dk.U.data.dims == fset_np.U.data.dims + assert fset_dk.U.data.shape == fset_np.U.data.shape + + +def test_to_windowed_arrays_is_idempotent_and_forwards_max_levels(): + ds = simple_UV_dataset(mesh="flat") + fs = FieldSet.from_sgrid_conventions(ds.chunk({"time": 1}), mesh="flat") + + fs.to_windowed_arrays(max_levels=3) + first = fs.U.data + assert isinstance(first, WindowedArray) + assert first._max == 3 + + # re-wrapping returns the same object (idempotent, warm cache preserved) + fs.to_windowed_arrays(max_levels=3) + assert fs.U.data is first + + +def test_windowed_isel_empty_selection(): + """An empty pointwise selection (a kernel evaluating an empty particle subset, + as in the Argo floats tutorial's phase kernels) must return an empty result + without touching the cache, matching plain xarray/dask isel behaviour. + """ + ntime, n = 4, 8 + base = np.arange(ntime * 3 * n * n, dtype=float).reshape(ntime, 3, n, n) + lazy = xr.DataArray(da.from_array(base, chunks=(1, 3, n, n)), dims=("time", "depth", "lat", "lon")) + win = WindowedArray(lazy) + + empty = xr.DataArray(np.array([], dtype=int), dims="p") + sel = dict(time=empty, depth=empty, lat=empty, lon=empty) + got = win.isel(sel) + ref = lazy.isel(sel) + + assert got.shape == ref.shape == (0,) + assert got.dtype == base.dtype + assert win.loads == 0 # nothing read + assert win._cache == {} # nothing cached, nothing evicted + + # a warm cache must survive an interleaved empty call (no spurious eviction) + full = xr.DataArray(np.zeros(5, dtype=int), dims="p") + win.isel(dict(time=full, depth=full, lat=full, lon=full)) + assert sorted(win._cache) == [0] + win.isel(sel) + assert sorted(win._cache) == [0] + + +def test_maybe_windowed_passthrough_for_non_time_leading(): + da_no_time = xr.DataArray(da.zeros((3, 4), chunks=(3, 4)), dims=("lat", "lon")) + assert maybe_windowed(da_no_time) is da_no_time # not wrapped (no leading time dim) + + +@pytest.mark.parametrize("mesh", ["flat", "spherical"]) +@pytest.mark.parametrize("dt_minutes", [15, -15], ids=["forward", "backward"]) +def test_dask_advection_matches_numpy(mesh, dt_minutes): + """An identical advection must give identical trajectories whether the field + is numpy-backed or dask-backed (windowed) -- for both forward (dt > 0) and + backward (dt < 0) integration. + """ + ds = simple_UV_dataset(mesh=mesh) + ds["U"].data[:] = 1.0 # steady zonal flow -> in-bounds, deterministic + + def run(chunked): + d = ds.chunk({"time": 1}) if chunked else ds + fs = FieldSet.from_sgrid_conventions(d, mesh=mesh) + if chunked: + fs.to_windowed_arrays() + pset = ParticleSet(fs, x=np.zeros(10), y=np.linspace(-10, 10, 10)) + pset.execute(AdvectionRK2, runtime=7200, dt=np.timedelta64(dt_minutes, "m")) + return np.array(pset.x), np.array(pset.y) + + x_np, y_np = run(False) + x_dk, y_dk = run(True) + np.testing.assert_allclose(x_dk, x_np, atol=1e-9) + np.testing.assert_allclose(y_dk, y_np, atol=1e-9)