bench_windowed_parcels.py

"""End-to-end Parcels benchmark: WindowedArray vs raw dask, on the 20 GB Atlantic.

Unlike bench_windowed.py (which benchmarks the loader strategy in isolation), this
runs the *actual* Parcels advection through `pset.execute`, comparing the new
transparent rolling-window cache against the old per-step `isel().compute()` path.
The only difference between the two runs is whether `FieldSet.to_windowed_arrays()`
(the opt-in that wraps dask-backed fields in `WindowedArray`) has been applied.

REQUIRES a Parcels build that includes the WindowedArray change (Parcels issue
#2656 / the issue-2656-windowed-array branch). Run it against that environment,
e.g.:

    PARCELS_PYTHON=/path/to/parcels/.pixi/envs/default/bin/python
    "$PARCELS_PYTHON" bench_windowed_parcels.py --dir ../04_atlantic/data/atlantic

(The stage-05 pixi env here intentionally has no Parcels; see ../03_parcels.)
"""

from __future__ import annotations

import argparse
import time
import warnings
from glob import glob
from pathlib import Path

import numpy as np
import xarray as xr

warnings.filterwarnings("ignore")

import parcels  # noqa: E402
from parcels._core._windowed_array import WindowedArray  # noqa: E402


def build_fieldset(files):
    ds = xr.open_mfdataset(files, chunks={"time": 1}, combine="by_coords")
    dsf = parcels.convert.copernicusmarine_to_sgrid(fields={"U": ds["uo"], "V": ds["vo"]})
    return parcels.FieldSet.from_sgrid_conventions(dsf), ds


def run(files, windowed, npart, days, dt_min):
    fieldset, ds = build_fieldset(files)
    if windowed:  # opt-in: wrap dask-backed fields in the rolling NumPy window
        fieldset.to_windowed_arrays()
    backing = "WindowedArray (NumPy window)" if isinstance(fieldset.U.data, WindowedArray) else "raw dask (per-step)"

    t0 = ds["time"].values[0]
    z0 = float(ds["depth"].values[0])
    rng = np.random.default_rng(0)
    lon = rng.uniform(-60.0, -20.0, npart)
    lat = rng.uniform(10.0, 40.0, npart)
    pset = parcels.ParticleSet(fieldset, lon=lon, lat=lat,
                               z=np.full(npart, z0), time=np.repeat(t0, npart))

    start = time.perf_counter()
    pset.execute(parcels.kernels.AdvectionRK2,
                 runtime=np.timedelta64(int(days * 86400), "s"),
                 dt=np.timedelta64(int(dt_min * 60), "s"))
    secs = time.perf_counter() - start
    nsteps = int(days * 24 * 60 / dt_min)
    print(f"  {backing:30s}: {secs:8.2f} s  "
          f"({npart} particles x {nsteps} steps, {npart*nsteps/secs:,.0f} steps/s)")
    return secs, np.array(pset.lon), np.array(pset.lat)


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--dir", type=Path, default=Path("../04_atlantic/data/atlantic"))
    p.add_argument("--days", type=float, default=2.0)
    p.add_argument("--dt-min", type=float, default=15.0)
    p.add_argument("--npart", type=int, default=500)
    args = p.parse_args()

    files = sorted(glob(str(args.dir / "*.nc")))
    if not files:
        raise SystemExit(f"no NetCDF files in {args.dir}")
    gb = sum(Path(f).stat().st_size for f in files) / 1024**3
    print(f"=== Parcels advection on {len(files)} files ({gb:.0f} GiB), "
          f"{args.days:g} d @ dt={args.dt_min:g} min ===")

    s_old, lon_old, lat_old = run(files, False, args.npart, args.days, args.dt_min)
    s_new, lon_new, lat_new = run(files, True, args.npart, args.days, args.dt_min)

    dmax = max(float(np.abs(lon_new - lon_old).max()), float(np.abs(lat_new - lat_old).max()))
    print(f"\n  trajectories identical: max|d| = {dmax:.2e}  ({'OK' if dmax < 1e-9 else 'MISMATCH'})")
    print(f"  WindowedArray is {s_old / s_new:.1f}x faster than raw per-step dask")


if __name__ == "__main__":
    main()