diff --git a/benchmark_chunk_cache.png b/benchmark_chunk_cache.png new file mode 100644 index 000000000..0f9ee9edc Binary files /dev/null and b/benchmark_chunk_cache.png differ diff --git a/benchmark_chunk_cache.py b/benchmark_chunk_cache.py new file mode 100644 index 000000000..d21047e22 --- /dev/null +++ b/benchmark_chunk_cache.py @@ -0,0 +1,139 @@ +"""Benchmark: plain dask vs windowed arrays vs cached chunk arrays. + +Scales particle count from 10 to 1,000,000 on ds_2d_left_agrid.zarr. +""" + +import time as time_mod + +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +import parcels +import parcels._sgrid as sgrid + + +def make_fieldset(ds: xr.Dataset) -> parcels.FieldSet: + """Build a FieldSet from the 2D left A-grid zarr dataset.""" + ds = ds.copy() + ds["lon"].attrs["units"] = "m" + ds["lat"].attrs["units"] = "m" + ds = ds.pipe( + sgrid._attach_sgrid_metadata, + sgrid.SGrid2DMetadata( + cf_role="grid_topology", + topology_dimension=2, + node_dimensions=("XG", "YG"), + node_coordinates=("lon", "lat"), + face_dimensions=( + sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW), + sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.LOW), + ), + vertical_dimensions=(sgrid.FaceNodePadding("ZC", "ZG", sgrid.Padding.LOW),), + ), + ) + return parcels.FieldSet.from_sgrid_conventions( + ds, + vector_fields={"UV": ("U_A_grid", "V_A_grid")}, + skip_field_data_validation=True, + ) + + +def delete_on_boundary(particles, fieldset): + """Delete particles that hit the boundary instead of erroring.""" + particles.state = np.where( + particles.state == parcels.StatusCode.ErrorOutOfBounds, + parcels.StatusCode.Delete, + particles.state, + ) + + +def run_simulation(fieldset, ds, npart): + """Run a simulation and return elapsed time in seconds.""" + np.random.seed(42) + pset = parcels.ParticleSet( + fieldset=fieldset, + pclass=parcels.Particle, + t=np.full(npart, ds.time.values[0]), + z=np.full(npart, 1), + y=np.random.uniform(1.0, 5.0, npart), + x=np.random.uniform(1.0, 5.0, npart), + ) + + t0 = time_mod.perf_counter() + pset.execute( + [parcels.kernels.AdvectionRK2, delete_on_boundary], + runtime=np.timedelta64(100, "ms"), + dt=np.timedelta64(10, "ms"), + ) + return time_mod.perf_counter() - t0 + + +def main(): + zarr_path = "./datasets/ds_2d_left_agrid.zarr" + particle_counts = [10, 100, 1_000, 10_000, 100_000, 1_000_000] + + print(f"Loading dataset from {zarr_path}") + ds = xr.open_zarr(zarr_path, consolidated=False) + print(f" shape: {dict(ds.dims)}") + print(f" chunks: U_A_grid {ds['U_A_grid'].encoding.get('chunks', 'N/A')}") + + # methods = { + # "plain dask": lambda ds: make_fieldset(ds), + # } + methods = { + "windowed": lambda ds: make_fieldset(ds).to_windowed_arrays(), + "cached chunks": lambda ds: make_fieldset(ds).to_chunk_cached_arrays(), + } + + results = {name: [] for name in methods} + + for npart in particle_counts: + print(f"\n--- {npart:,} particles ---") + for name, build_fn in methods.items(): + fieldset = build_fn(ds) + elapsed = run_simulation(fieldset, ds, npart) + results[name].append(elapsed) + print(f" {name}: {elapsed:.3f}s") + + # --- Print results table --- + print("\n" + "=" * 60) + print("Results summary") + print("=" * 60) + header = f"{'N particles':>12s}" + for name in methods: + header += f" {name:>15s}" + print(header) + print("-" * len(header)) + for i, npart in enumerate(particle_counts): + row = f"{npart:>12,d}" + for name in methods: + row += f" {results[name][i]:>14.3f}s" + print(row) + + # --- Plot --- + fig, ax = plt.subplots(figsize=(9, 6)) + markers = ["o", "s", "^", "D"] + for j, (name, times) in enumerate(results.items()): + ax.loglog( + particle_counts, + times, + marker=markers[j % len(markers)], + linewidth=2, + markersize=7, + label=name, + ) + + ax.set_xlabel("Number of particles") + ax.set_ylabel("Wall-clock time (s)") + ax.set_title("Parcels simulation scaling: windowed vs cached chunk arrays") + ax.legend() + ax.grid(True, which="both", alpha=0.3) + fig.tight_layout() + fig.savefig("benchmark_chunk_cache.png", dpi=150) + print("\nPlot saved to benchmark_chunk_cache.png") + plt.show() + + +if __name__ == "__main__": + main() diff --git a/data-generation.py b/data-generation.py new file mode 100644 index 000000000..28c09a48e --- /dev/null +++ b/data-generation.py @@ -0,0 +1,294 @@ +import operator +from functools import reduce +from pathlib import Path + +import dask.array as da +import numpy as np +import xarray as xr +from zarr.storage import LocalStore + +X = 1000 +Y = 1000 +Z = 50 +T = 30 + +X_SMALL = 200 +Y_SMALL = 200 + +TIME = xr.date_range("2000", "2001", T) + + +def prod(seq): + return reduce(operator.mul, seq, 1) + + +def _rotated_curvilinear_grid(): + XG = np.arange(X) + YG = np.arange(Y) + LON, LAT = np.meshgrid(XG, YG) + + angle = -np.pi / 24 + rotation = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) + + # rotate the LON and LAT grids + LON, LAT = np.einsum("ji, mni -> jmn", rotation, np.dstack([LON, LAT])) + + return xr.Dataset( + { + "data_g": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "data_c": ( + ["time", "ZC", "YC", "XC"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "U_A_grid": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "V_A_grid": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "U_C_grid": ( + ["time", "ZG", "YC", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "V_C_grid": ( + ["time", "ZG", "YG", "XC"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + }, + coords={ + "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}), + "YG": (["YG"], YG, {"axis": "Y", "c_grid_axis_shift": -0.5}), + "XC": (["XC"], XG + 0.5, {"axis": "X"}), + "YC": (["YC"], YG + 0.5, {"axis": "Y"}), + "ZG": ( + ["ZG"], + np.arange(Z), + {"axis": "Z", "c_grid_axis_shift": -0.5}, + ), + "ZC": ( + ["ZC"], + np.arange(Z) + 0.5, + {"axis": "Z"}, + ), + "depth": (["ZG"], np.arange(Z), {"axis": "Z"}), + "time": (["time"], TIME, {"axis": "T"}), + "lon": ( + ["YG", "XG"], + LON, + {"axis": "X", "c_grid_axis_shift": -0.5}, # ? Needed? + ), + "lat": ( + ["YG", "XG"], + LAT, + {"axis": "Y", "c_grid_axis_shift": -0.5}, # ? Needed? + ), + }, + ) + + +def random_dask_array(shape, scaling=1): + return da.random.uniform(scaling, size=prod(shape)).reshape(shape) + + +def _cartesion_to_polar(x, y): + r = np.sqrt(x**2 + y**2) + theta = np.arctan2(y, x) + return r, theta + + +def _polar_to_cartesian(r, theta): + x = r * np.cos(theta) + y = r * np.sin(theta) + return x, y + + +def _unrolled_cone_curvilinear_grid(): + # Not a great unrolled cone, but this is good enough for testing + # you can use matplotlib pcolormesh to plot + XG = np.arange(X) + YG = np.arange(Y) * 0.25 + + pivot = -10, 0 + LON, LAT = np.meshgrid(XG, YG) + + new_lon_lat = [] + + min_lon = np.min(XG) + for lon, lat in zip(LON.flatten(), LAT.flatten(), strict=True): + r, _ = _cartesion_to_polar(lon - pivot[0], lat - pivot[1]) + _, theta = _cartesion_to_polar(min_lon - pivot[0], lat - pivot[1]) + theta *= 1.2 + r *= 1.2 + lon, lat = _polar_to_cartesian(r, theta) + new_lon_lat.append((lon + pivot[0], lat + pivot[1])) + + new_lon, new_lat = zip(*new_lon_lat, strict=True) + LON, LAT = ( + np.array(new_lon).reshape(LON.shape), + np.array(new_lat).reshape(LAT.shape), + ) + + return xr.Dataset( + { + "data_g": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "data_c": ( + ["time", "ZC", "YC", "XC"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "U_A_grid": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "V_A_grid": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "U_C_grid": ( + ["time", "ZG", "YC", "XG"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + "V_C_grid": ( + ["time", "ZG", "YG", "XC"], + random_dask_array(scaling=5, shape=(T, Z, Y, X)), + ), + }, + coords={ + "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}), + "YG": (["YG"], YG, {"axis": "Y", "c_grid_axis_shift": -0.5}), + "XC": (["XC"], XG + 0.5, {"axis": "X"}), + "YC": (["YC"], YG + 0.5, {"axis": "Y"}), + "ZG": ( + ["ZG"], + np.arange(Z), + {"axis": "Z", "c_grid_axis_shift": -0.5}, + ), + "ZC": ( + ["ZC"], + np.arange(Z) + 0.5, + {"axis": "Z"}, + ), + "depth": (["ZG"], np.arange(Z), {"axis": "Z"}), + "time": (["time"], TIME, {"axis": "T"}), + "lon": ( + ["YG", "XG"], + LON, + {"axis": "X", "c_grid_axis_shift": -0.5}, # ? Needed? + ), + "lat": ( + ["YG", "XG"], + LAT, + {"axis": "Y", "c_grid_axis_shift": -0.5}, # ? Needed? + ), + }, + ) + + +def _ds_2d_left(x, y, z, t, time): + """MITgcm indexing style dataset.""" + return xr.Dataset( + { + "data_g": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(t, z, y, x)), + ), + "data_c": ( + ["time", "ZC", "YC", "XC"], + random_dask_array(scaling=5, shape=(t, z, y, x)), + ), + "U_A_grid": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(t, z, y, x)), + ), + "V_A_grid": ( + ["time", "ZG", "YG", "XG"], + random_dask_array(scaling=5, shape=(t, z, y, x)), + ), + "U_C_grid": ( + ["time", "ZG", "YC", "XG"], + random_dask_array(scaling=5, shape=(t, z, y, x)), + ), + "V_C_grid": ( + ["time", "ZG", "YG", "XC"], + random_dask_array(scaling=5, shape=(t, z, y, x)), + ), + }, + coords={ + "XG": ( + ["XG"], + 2 * np.pi / x * np.arange(0, x), + {"axis": "X", "c_grid_axis_shift": -0.5}, + ), + "XC": (["XC"], 2 * np.pi / x * (np.arange(0, x) + 0.5), {"axis": "X"}), + "YG": ( + ["YG"], + 2 * np.pi / y * np.arange(0, y), + {"axis": "Y", "c_grid_axis_shift": -0.5}, + ), + "YC": ( + ["YC"], + 2 * np.pi / y * (np.arange(0, y) + 0.5), + {"axis": "Y"}, + ), + "ZG": ( + ["ZG"], + np.arange(z), + {"axis": "Z", "c_grid_axis_shift": -0.5}, + ), + "ZC": ( + ["ZC"], + np.arange(z) + 0.5, + {"axis": "Z"}, + ), + "lon": (["XG"], 2 * np.pi / x * np.arange(0, x)), + "lat": (["YG"], 2 * np.pi / y * np.arange(0, y)), + "depth": (["ZG"], np.arange(z)), + "time": (["time"], time, {"axis": "T"}), + }, + ) + + +datasets = { + "2d_left_rotated": _rotated_curvilinear_grid(), + "ds_2d_left": _ds_2d_left(X, Y, Z, T, TIME), + "2d_left_unrolled_cone": _unrolled_cone_curvilinear_grid(), +} + + +def save(ds: xr.Dataset, path: str, chunks: dict) -> None: + """Save dataset to zarr with specified chunking.""" + store = LocalStore(path) + ds.chunk(chunks).to_zarr(store, mode="w", encoding=None, consolidated=False) + size_mb = sum(ds[v].nbytes for v in ds.data_vars) / 1e6 + print(f" {path} dims={dict(ds.sizes)} ~{size_mb:.0f} MB uncompressed") + + +if __name__ == "__main__": + dataset_path = "datasets/ds_2d_left_agrid.zarr" + print("Generating ds_2d_left...") + if Path(dataset_path).exists(): + print(f"Dataset {dataset_path} already exists") + else: + save( + datasets["ds_2d_left"][["U_A_grid", "V_A_grid"]], + dataset_path, + {"time": 15, "XG": 40, "YG": 40, "ZG": 8}, + ) + + dataset_path_small = "datasets/ds_2d_left_agrid_small.zarr" + print("Generating ds_2d_left_small...") + if Path(dataset_path_small).exists(): + print(f"Dataset {dataset_path_small} already exists") + else: + save( + _ds_2d_left(X_SMALL, Y_SMALL, Z, T, TIME)[["U_A_grid", "V_A_grid"]], + dataset_path_small, + {"time": 15, "XG": 40, "YG": 40, "ZG": 8}, + ) diff --git a/src/parcels/_chunk_cached_array/__init__.py b/src/parcels/_chunk_cached_array/__init__.py new file mode 100644 index 000000000..89920eefe --- /dev/null +++ b/src/parcels/_chunk_cached_array/__init__.py @@ -0,0 +1,4 @@ +from .core import ChunkCachedArray, wrap_dataset +from .lru_cache import ByteBoundedLRUCache + +__all__ = ["ByteBoundedLRUCache", "ChunkCachedArray", "wrap_dataset"] diff --git a/src/parcels/_chunk_cached_array/core.py b/src/parcels/_chunk_cached_array/core.py new file mode 100644 index 000000000..9d545c1ad --- /dev/null +++ b/src/parcels/_chunk_cached_array/core.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from xarray.core.indexing import BasicIndexer, ExplicitlyIndexedNDArrayMixin, OuterIndexer, VectorizedIndexer +from xarray.namedarray.pycompat import is_duck_array + +from .lru_cache import ByteBoundedLRUCache + +if TYPE_CHECKING: + import dask.array + import xarray as xr + + +def wrap_dataset(ds: xr.Dataset, max_cache_bytes: int) -> xr.Dataset: + """Replace all dask-backed data variables with ChunkCachedArray wrappers. + + Returns a shallow copy of the dataset. Each dask-backed data variable's + internal ``variable._data`` is swapped for a ``ChunkCachedArray`` that + caches chunks on vectorized indexing. Coordinate variables are loaded + eagerly into memory to avoid dask task-graph overhead on every + ``.isel()`` call. + + Parameters + ---------- + ds : xr.Dataset + Source dataset (not modified). + max_cache_bytes : int + Maximum cache size in bytes, per variable. + + Returns + ------- + xr.Dataset + Copy with dask arrays wrapped in ChunkCachedArray. + """ + from dask.base import is_dask_collection + + ds = ds.copy() + # Load coordinates eagerly — they are small 1D arrays and keeping them + # as dask arrays causes expensive task-graph construction on every .isel(). + for name in list(ds.coords): + ds[name].load() + for name in ds.data_vars: + var = ds[name].variable + if is_duck_array(var._data) and is_dask_collection(var._data): + var._data = ChunkCachedArray(var._data, max_cache_bytes) # type: ignore[assignment, arg-type] + return ds + + +class ChunkCachedArray(ExplicitlyIndexedNDArrayMixin): + """Chunk-level LRU cache on top of a dask array for vectorized indexing. + + Implements xarray's ExplicitlyIndexed protocol so it can be used as + a drop-in replacement for the dask array in ``da.data``. Xarray's + ``.isel()`` with vectorized indexers will route through ``_vindex_get``, + which uses the chunk cache. Other indexing modes delegate to the + underlying dask array. + + On each vectorized index: + 1. Maps global indices -> (chunk_coord, local_index) per dimension. + 2. Fetches missing chunks via dask_array.blocks[...].compute(). + 3. Assembles the result from cached numpy arrays. + """ + + def __init__(self, dask_array: dask.array.Array, max_cache_bytes: int) -> None: + self.array = dask_array + self.cache = ByteBoundedLRUCache(max_cache_bytes) + + # Precompute chunk boundaries per dimension. + # _boundaries[d] is a 1D array of cumulative chunk sizes, e.g., [0, 15, 30]. + self._boundaries: list[np.ndarray] = [] + for dim_chunks in dask_array.chunks: + self._boundaries.append(np.concatenate(([0], np.cumsum(dim_chunks)))) + + def get_duck_array(self): + return self.array.compute() + + def _raw_vindex(self, *indices: np.ndarray) -> np.ndarray: + """Vectorized indexing with chunk caching. + + Parameters + ---------- + *indices : np.ndarray + One 1D integer index array per dimension. All must have the same length N. + + Returns + ------- + np.ndarray + 1D array of length N with the selected values. + """ + ndim = len(self.array.chunks) + assert len(indices) == ndim + n_points = len(indices[0]) + + # Step 1: Map global indices to chunk coords and local indices. + # Normalize negative indices (e.g. -1 → last element) to positive, + # matching standard numpy fancy-indexing semantics. + indices = tuple(np.where(idx < 0, idx + self.array.shape[d], idx) for d, idx in enumerate(indices)) + chunk_ids = np.empty((ndim, n_points), dtype=np.intp) + local_indices = np.empty((ndim, n_points), dtype=np.intp) + for d in range(ndim): + cid = np.searchsorted(self._boundaries[d], indices[d], side="right") - 1 + chunk_ids[d] = cid + local_indices[d] = indices[d] - self._boundaries[d][cid] + + # Step 2: Group points by chunk using a structured array for vectorized grouping. + # Encode each point's chunk coords as a single int for fast grouping. + # Use np.ravel_multi_index on chunk_ids to get a flat chunk key per point. + numblocks = np.array(self.array.numblocks, dtype=np.intp) + flat_keys = np.ravel_multi_index(chunk_ids, numblocks) + + # Sort points by flat chunk key to group them. + sort_order = np.argsort(flat_keys, kind="mergesort") + sorted_flat_keys = flat_keys[sort_order] # type: ignore[index] + + # Find group boundaries. + boundaries = np.concatenate(([0], np.flatnonzero(np.diff(sorted_flat_keys)) + 1, [n_points])) + + out = np.empty(n_points, dtype=self.array.dtype) + for g in range(len(boundaries) - 1): + grp_slice = slice(boundaries[g], boundaries[g + 1]) + grp_indices = sort_order[grp_slice] + + # Recover the chunk key tuple from any point in this group. + key = tuple(int(chunk_ids[d, grp_indices[0]]) for d in range(ndim)) + + chunk_data = self.cache.get(key) + if chunk_data is None: + chunk_data = self.array.blocks[key].compute() + self.cache.put(key, chunk_data) + + # Vectorized fancy-index: extract all points from this chunk at once. + local_idx = tuple(local_indices[d, grp_indices] for d in range(ndim)) + out[grp_indices] = chunk_data[local_idx] + + return out + + # --- ExplicitlyIndexed protocol --- + + def _vindex_get(self, indexer: VectorizedIndexer): + key = indexer.tuple + return self._raw_vindex(*key) + + def _oindex_get(self, indexer: OuterIndexer): + # Delegate to dask for orthogonal indexing + return self.array[indexer.tuple] + + def __getitem__(self, indexer): + if isinstance(indexer, VectorizedIndexer): + return self._vindex_get(indexer) + if isinstance(indexer, OuterIndexer): + return self._oindex_get(indexer) + if isinstance(indexer, BasicIndexer): + return self.array[indexer.tuple] + return self.array[indexer] diff --git a/src/parcels/_chunk_cached_array/lru_cache.py b/src/parcels/_chunk_cached_array/lru_cache.py new file mode 100644 index 000000000..117593439 --- /dev/null +++ b/src/parcels/_chunk_cached_array/lru_cache.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Hashable + +import numpy as np + + +class ByteBoundedLRUCache: + """LRU cache bounded by total stored bytes. + + Keys are hashable (typically chunk coordinate tuples). + Values are numpy arrays whose .nbytes drives eviction. + """ + + def __init__(self, max_bytes: int) -> None: + assert max_bytes > 0 + self._max_bytes = max_bytes + self._cache: OrderedDict[Hashable, np.ndarray] = OrderedDict() + self._current_bytes = 0 + + @property + def current_bytes(self) -> int: + return self._current_bytes + + def get(self, key: Hashable) -> np.ndarray | None: + if key not in self._cache: + return None + self._cache.move_to_end(key) + return self._cache[key] + + def put(self, key: Hashable, value: np.ndarray) -> None: + nbytes = value.nbytes + if nbytes > self._max_bytes: + return + # Remove existing entry if present (will be re-inserted at end) + if key in self._cache: + self._current_bytes -= self._cache.pop(key).nbytes + # Evict LRU entries until there's room + while self._current_bytes + nbytes > self._max_bytes: + _, evicted = self._cache.popitem(last=False) + self._current_bytes -= evicted.nbytes + self._cache[key] = value + self._current_bytes += nbytes + + def clear(self) -> None: + self._cache.clear() + self._current_bytes = 0 diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index 696aa6544..c4865dad0 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -186,6 +186,30 @@ def to_windowed_arrays(self, *, max_levels: int | None = None): model.to_windowed_arrays(max_levels=max_levels) return self + def to_chunk_cached_arrays(self, *, max_cache_bytes: int = 600_000_000): + """Wrap dask-backed field data in chunk-level LRU caches. + + Opt-in optimization that replaces each dask-backed data variable's + internal storage with a :class:`~parcels._chunk_cached_array.ChunkCachedArray`. + Delegates to each underlying model; repeated vectorized ``.isel()`` + calls then hit an in-memory LRU cache instead of recomputing dask task + graphs. NumPy-backed (eager) fields are left unchanged, and re-invoking + is idempotent. + + Parameters + ---------- + max_cache_bytes : int, optional + Maximum cache size in bytes, per variable. Defaults to 600 MB. + + Returns + ------- + FieldSet + ``self``, to allow chaining. + """ + for model in self.models: + model.to_chunk_cached_arrays(max_cache_bytes=max_cache_bytes) + return self + def add_constant_field(self, name: str, value, mesh: ptyping.TMesh = "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 cdb2c0800..36e904843 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -12,6 +12,7 @@ import parcels._sgrid as sgrid import parcels._typing as ptyping +from parcels._chunk_cached_array import wrap_dataset from parcels._core._windowed_array import maybe_windowed from parcels._core.basegrid import BaseGrid from parcels._core.field import Field, VectorField @@ -114,6 +115,29 @@ def to_windowed_arrays(self, *, max_levels: int | None = None) -> Self: windowed[name] = maybe_windowed(current, max_levels=max_levels) return self + def to_chunk_cached_arrays(self, *, max_cache_bytes: int = 600_000_000) -> Self: + """Wrap dask-backed field data in chunk-level LRU caches. + + Opt-in optimization that replaces each dask-backed data variable's + internal storage with a :class:`~parcels._chunk_cached_array.ChunkCachedArray`. + Repeated vectorized ``.isel()`` calls then hit an in-memory LRU cache + keyed by chunk coordinates instead of recomputing dask task graphs. + + Coordinate variables are loaded eagerly into memory (they are small 1D + arrays) to avoid dask task-graph construction overhead on every + ``.isel()`` call. + + Idempotent: re-invoking is safe — ``wrap_dataset`` skips variables + whose storage is already a ``ChunkCachedArray``. + + Parameters + ---------- + max_cache_bytes : int, optional + Maximum cache size in bytes, per variable. Defaults to 600 MB. + """ + self.data = wrap_dataset(self.data, max_cache_bytes=max_cache_bytes) + return self + @property def time_interval(self) -> TimeInterval | None: try: diff --git a/tests/test_backends.py b/tests/test_backends.py new file mode 100644 index 000000000..31286cc7f --- /dev/null +++ b/tests/test_backends.py @@ -0,0 +1,115 @@ +import io +from pathlib import Path +from typing import Literal + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +import parcels +import parcels.tutorial +from parcels.kernels import AdvectionRK4 + +BackendT = Literal["WindowedArray", "Dask", "Zarr", "NumPy", "CachedChunkArray"] +BACKENDS = {"WindowedArray", "Dask", "Zarr", "NumPy", "CachedChunkArray"} + + +@pytest.fixture(scope="module") +def nemo_dataset() -> xr.Dataset: + ds_u = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/U") + ds_v = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/V") + ds_w = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/W") + ds_coords = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/mesh_mask")[["glamf", "gphif"]] + + ds_fset = parcels.convert.nemo_to_sgrid( + fields={"U": ds_u["uo"], "V": ds_v["vo"], "W": ds_w["wo"]}, + coords=ds_coords, + ) + return ds_fset + + +@pytest.fixture(scope="module") +def nemo_results(tmp_path_factory, nemo_dataset) -> tuple[xr.Dataset, Path]: + ref_parquet = tmp_path_factory.mktemp("nemo_ref") / "ref.parquet" + run_simulation(nemo_dataset, ref_parquet, "NumPy") + return nemo_dataset, ref_parquet + + +def assert_fieldset_backend(fset: parcels.FieldSet, backend: BackendT): + # a bit of a hacky way to check for the backend.... probably better for us to change how backends are stored + buf = io.StringIO() + fset.describe(buf) + return backend in buf.getvalue() + + +def run_simulation(ds: xr.Dataset, output_path: Path, backend: BackendT) -> Path: + if backend == "Zarr": + raise NotImplementedError("Doesn't work at this level of execution. Also will likely remove Zarr backend.") + + if backend == "NumPy": + ds.load() + + fset = parcels.FieldSet.from_sgrid_conventions(ds) + + if backend == "WindowedArray": + fset.to_windowed_arrays() + if backend == "CachedChunkArray": + fset.to_chunk_cached_arrays() + + assert_fieldset_backend(fset, backend) + + npart = 1000 + lons = np.linspace(1.9, 3.4, npart) + lats = np.linspace(51.6, 52.5, npart) + z = np.ones(npart) + pset = parcels.ParticleSet(fset, x=lons, y=lats, z=z) + + def delete_particle(particles, fieldset): + error_states = ( + parcels.StatusCode.ErrorOutOfBounds, + parcels.StatusCode.ErrorGridSearching, + ) + for error in error_states: + particles.state = np.where( + particles.state == error, + parcels.StatusCode.Delete, + particles.state, + ) + + pfile = parcels.ParticleFile(output_path, outputdt=np.timedelta64(6, "h")) + pset.execute( + [AdvectionRK4, delete_particle], + runtime=np.timedelta64(3, "D"), + dt=np.timedelta64(5, "m"), + output_file=pfile, + ) + + return output_path + + +@pytest.mark.parametrize( + "backend", + BACKENDS + - { + "NumPy", # reference point + "Zarr", # not supported + }, +) +def test_nemo_identical_across_backends(nemo_results, tmp_parquet, backend): + ds = nemo_results[0] + ref_parquet = nemo_results[1] + + assert str(ref_parquet) != str(tmp_parquet) # just covering my bases with Pytest fixture usage + + run_simulation(ds, tmp_parquet, backend) + + ref_df = pd.read_parquet(ref_parquet) + test_df = pd.read_parquet(tmp_parquet) + + ref_df = ref_df.sort_values(["particle_id", "t"]).reset_index(drop=True) + test_df = test_df.sort_values(["particle_id", "t"]).reset_index(drop=True) + + np.testing.assert_allclose(test_df["x"].values, ref_df["x"].values, atol=1e-5) + np.testing.assert_allclose(test_df["y"].values, ref_df["y"].values, atol=1e-5) + np.testing.assert_allclose(test_df["z"].values, ref_df["z"].values, atol=1e-5) diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index b1b6b16f6..eea5ce07d 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -540,3 +540,5 @@ def test_fieldset_describe_backends(tmp_path): fieldset.describe(io) actual = io.getvalue() assert actual == expected + + # TODO: Add test for the ChunkedArray backend (can also refactor this test at the same time)