Windowed Array (#2718)

* Carry over windowed_array implementation to new main branch with "model" concept

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Change caching pattern to be based on spanning time index range

This supports both forward and backward time stepping.

Additional tests have been added to test backwards in time integration.
This includes tests verifying expected data is held in the windowed cache
when dt is negative and a test that compares backwards integration with
purely numpy backed fieldset.

* Breaking tutorial

* Return empty numpy array for empty selection dictionaries

* Change fs_* to fset_* per review; consistent with other naming conventions in tests

* Flesh out docstring on how the default windowedarray caching works

* Assert that dask array is indeed initially a dask array in test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Erik van Sebille <erikvansebille@gmail.com>

Author: 3 people

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: