Moving lon/lat loading to to_windowed_arrays

Author: erikvansebille

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,10 @@ 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", "depth"]:
+            # ensure lon and lat are loaded into memory for dask-backed datasets
+            if is_dask_collection(self.data[dim]):
+                self.data[dim].load()
         for name in self.scalar_field_names:
             current = windowed.get(name, self.data[name])
             windowed[name] = maybe_windowed(current, max_levels=max_levels)

src/parcels/_core/xgrid.py

@@ -7,7 +7,6 @@
 import numpy.typing as npt
 import xarray as xr
 import xgcm
-from dask import is_dask_collection
 
 import parcels._sgrid as sgrid
 import parcels._typing as ptyping
@@ -208,9 +207,6 @@ def lon(self):
             _ = self.xgcm_grid.axes["X"]
         except KeyError:
             return np.zeros(1)
-        # ensure lon is loaded into memory for dask-backed datasets, as it is used in the search method
-        if is_dask_collection(self._ds["lon"].data):
-            self._ds["lon"].load()
         return self._ds["lon"].values
 
     @property
@@ -225,9 +221,6 @@ def lat(self):
             _ = self.xgcm_grid.axes["Y"]
         except KeyError:
             return np.zeros(1)
-        # ensure lat is loaded into memory for dask-backed datasets, as it is used in the search method
-        if is_dask_collection(self._ds["lat"].data):
-            self._ds["lat"].load()
         return self._ds["lat"].values
 
     @property