Merge branch 'main' into support_some_fields_without_time_in_fieldset

Author: erikvansebille

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",

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

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

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

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: