From 30622a0bd46dd5208e49e9254bf59da5f85b1c0f Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 21 Jul 2025 14:52:36 +0200 Subject: [PATCH 1/6] Adding benchmark test for CopernicusMarine with different Interpolators --- .../benchmark_CopernicusMarineService.py | 139 ++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 CopernicusMarineService/benchmark_CopernicusMarineService.py diff --git a/CopernicusMarineService/benchmark_CopernicusMarineService.py b/CopernicusMarineService/benchmark_CopernicusMarineService.py new file mode 100644 index 0000000..2e285aa --- /dev/null +++ b/CopernicusMarineService/benchmark_CopernicusMarineService.py @@ -0,0 +1,139 @@ +from argparse import ArgumentParser +from pathlib import Path +import xarray as xr + +from glob import glob + +import numpy as np + +import parcels + +runtime = np.timedelta64(10, "D") +dt = np.timedelta64(3, "h") + +parcelsv4 = True +try: + from parcels.xgrid import _XGRID_AXES +except ImportError: + parcelsv4 = False + +DATA_ROOT = "/storage/shared/oceanparcels/input_data/CopernicusMarineService/GLOBAL_ANALYSIS_FORECAST_PHY_001_024_SMOC/" + +def run_benchmark(interpolator: str): + if parcelsv4: + + def BiRectiLinear( + field: parcels.Field, + ti: int, + position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]], + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, + ): + """Bilinear interpolation on a rectilinear grid.""" + xi, xsi = position["X"] + yi, eta = position["Y"] + zi, zeta = position["Z"] + + data = field.data.isel({"time": slice(ti, ti + 2), "lat": slice(yi, yi + 2), "lon": slice(xi, xi + 2)}).data#.compute() + val_t0 =( + (1 - xsi) * (1 - eta) * data[0, 0, 0, 0] + + xsi * (1 - eta) * data[0, 0, 0, 1] + + xsi * eta * data[0, 0, 1, 1] + + (1 - xsi) * eta * data[0, 0, 1, 0] + ) + + val_t1 =( + (1 - xsi) * (1 - eta) * data[1, 0, 0, 0] + + xsi * (1 - eta) * data[1, 0, 0, 1] + + xsi * eta * data[1, 0, 1, 1] + + (1 - xsi) * eta * data[1, 0, 1, 0] + ) + return (val_t0 * (1 - tau) + val_t1 * tau) + + def PureXarrayInterp( + field: parcels.Field, + ti: int, + position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]], + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, + ): + return field.data.interp(time=t, lon=x, lat=y).values[0] + + + def NoFieldAccess( + field: parcels.Field, + ti: int, + position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]], + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, + ): + return 0 + + + files = f"{DATA_ROOT}/SMOC_202101*.nc" + + if parcelsv4: + if interpolator == "BiRectiLinear": + interp_method = BiRectiLinear + elif interpolator == "PureXarrayInterp": + interp_method = PureXarrayInterp + elif interpolator == "NoFieldAccess": + interp_method = NoFieldAccess + else: + raise ValueError(f"Unknown interpolator: {interpolator}") + ds = xr.open_mfdataset(files, concat_dim="time", combine="nested", data_vars='minimal', coords='minimal', compat='override') + ds = ds.rename({"uo": "U", "vo": "V", "longitude": "lon", "latitude": "lat"}) + + xgcm_grid = parcels.xgcm.Grid(ds, coords={"X": {"left": "lon"}, "Y": {"left": "lat"}, "Z": {"left": "depth"}, "T": {"center": "time"}}) + grid = parcels.xgrid.XGrid(xgcm_grid) + + U = parcels.Field("U", ds["U"], grid, interp_method=interp_method) + V = parcels.Field("V", ds["V"], grid, interp_method=interp_method) + U.units = parcels.GeographicPolar() + V.units = parcels.Geographic() + UV = parcels.VectorField("UV", U, V) + + fieldset = parcels.FieldSet([U, V, UV]) + else: + interpolator = "v3_default" + fieldset = parcels.FieldSet.from_netcdf(files, variables={"U": "uo", "V": "vo"}, dimensions={"time": "time", "lat": "latitude", "lon": "longitude", "depth": "depth"}) + + pclass = parcels.Particle if parcelsv4 else parcels.JITParticle + + kernel = parcels.AdvectionEE + + for npart in [1, 10, 30, 50]: + lon = np.linspace(-10, 10, npart) + lat = np.linspace(-30, -20, npart) + + pset = parcels.ParticleSet(fieldset=fieldset, pclass=pclass, lon=lon, lat=lat) + + print(f"Running {len(lon)} particles with parcels v{4 if parcelsv4 else 3} and {interpolator} interpolator") + pset.execute(kernel, runtime=runtime, dt=dt) + + +def main(args=None): + p = ArgumentParser() + + p.add_argument( + "-i", + "--Interpolator", + choices=("BiRectiLinear", "PureXarrayInterp", "NoFieldAccess"), + default="BiRectiLinear", + ) + + args = p.parse_args(args) + run_benchmark(args.Interpolator) + + +if __name__ == "__main__": + main() From 897e643096826e917ead446a7dcb7c73b958272d Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 22 Jul 2025 16:20:12 +0200 Subject: [PATCH 2/6] Updating Copernicus benchmark to use shorter dt (since field time interval is hourly) and more particles --- .../benchmark_CopernicusMarineService.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CopernicusMarineService/benchmark_CopernicusMarineService.py b/CopernicusMarineService/benchmark_CopernicusMarineService.py index 2e285aa..aeae6a9 100644 --- a/CopernicusMarineService/benchmark_CopernicusMarineService.py +++ b/CopernicusMarineService/benchmark_CopernicusMarineService.py @@ -8,8 +8,8 @@ import parcels -runtime = np.timedelta64(10, "D") -dt = np.timedelta64(3, "h") +runtime = np.timedelta64(2, "D") +dt = np.timedelta64(15, "m") parcelsv4 = True try: @@ -111,7 +111,7 @@ def NoFieldAccess( kernel = parcels.AdvectionEE - for npart in [1, 10, 30, 50]: + for npart in [1, 10, 100, 1000, 10000]: lon = np.linspace(-10, 10, npart) lat = np.linspace(-30, -20, npart) From 9599090cc30bebc07b30cb0020eb9956ff73f6a1 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 22 Jul 2025 17:02:12 +0200 Subject: [PATCH 3/6] updating Bilinear interpolation function --- CopernicusMarineService/benchmark_CopernicusMarineService.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/CopernicusMarineService/benchmark_CopernicusMarineService.py b/CopernicusMarineService/benchmark_CopernicusMarineService.py index aeae6a9..db8da89 100644 --- a/CopernicusMarineService/benchmark_CopernicusMarineService.py +++ b/CopernicusMarineService/benchmark_CopernicusMarineService.py @@ -35,9 +35,8 @@ def BiRectiLinear( """Bilinear interpolation on a rectilinear grid.""" xi, xsi = position["X"] yi, eta = position["Y"] - zi, zeta = position["Z"] - data = field.data.isel({"time": slice(ti, ti + 2), "lat": slice(yi, yi + 2), "lon": slice(xi, xi + 2)}).data#.compute() + data = field.data.data[:, :, yi:yi + 2, xi:xi + 2] val_t0 =( (1 - xsi) * (1 - eta) * data[0, 0, 0, 0] + xsi * (1 - eta) * data[0, 0, 0, 1] From 31be70c04840c616cb360776cb982c2e355acbef Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 23 Jul 2025 14:48:39 +0200 Subject: [PATCH 4/6] Adding correctness test to CopernicusMarine benchmark --- .../benchmark_CopernicusMarineService.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/CopernicusMarineService/benchmark_CopernicusMarineService.py b/CopernicusMarineService/benchmark_CopernicusMarineService.py index db8da89..faf1ea9 100644 --- a/CopernicusMarineService/benchmark_CopernicusMarineService.py +++ b/CopernicusMarineService/benchmark_CopernicusMarineService.py @@ -80,6 +80,7 @@ def NoFieldAccess( files = f"{DATA_ROOT}/SMOC_202101*.nc" + lon0_expected, lat0_expected = -9.820091, -30.106716 # values from v3 if parcelsv4: if interpolator == "BiRectiLinear": interp_method = BiRectiLinear @@ -87,9 +88,11 @@ def NoFieldAccess( interp_method = PureXarrayInterp elif interpolator == "NoFieldAccess": interp_method = NoFieldAccess + lon0_expected, lat0_expected = -10, -30 # Zero interpolation, so expect initial values else: raise ValueError(f"Unknown interpolator: {interpolator}") ds = xr.open_mfdataset(files, concat_dim="time", combine="nested", data_vars='minimal', coords='minimal', compat='override') + ds = ds.drop_vars(["vsdx", "vsdy", "utide", "vtide", "utotal", "vtotal"]) ds = ds.rename({"uo": "U", "vo": "V", "longitude": "lon", "latitude": "lat"}) xgcm_grid = parcels.xgcm.Grid(ds, coords={"X": {"left": "lon"}, "Y": {"left": "lat"}, "Z": {"left": "depth"}, "T": {"center": "time"}}) @@ -110,7 +113,7 @@ def NoFieldAccess( kernel = parcels.AdvectionEE - for npart in [1, 10, 100, 1000, 10000]: + for npart in [1, 10, 100, 1000, 5000, 10000]: lon = np.linspace(-10, 10, npart) lat = np.linspace(-30, -20, npart) @@ -119,6 +122,9 @@ def NoFieldAccess( print(f"Running {len(lon)} particles with parcels v{4 if parcelsv4 else 3} and {interpolator} interpolator") pset.execute(kernel, runtime=runtime, dt=dt) + assert np.allclose(pset[0].lon, lon0_expected, atol=1e-5), f"Expected lon {lon0_expected}, got {pset[0].lon}" + assert np.allclose(pset[0].lat, lat0_expected, atol=1e-5), f"Expected lat {lat0_expected}, got {pset[0].lat}" + def main(args=None): p = ArgumentParser() From 05ebc3e213459a08e9e83df59a97425271982fbe Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 5 Aug 2025 09:52:55 +0200 Subject: [PATCH 5/6] Updating benchmark to use XLinear interpolation as default --- .../benchmark_CopernicusMarineService.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/CopernicusMarineService/benchmark_CopernicusMarineService.py b/CopernicusMarineService/benchmark_CopernicusMarineService.py index faf1ea9..81dfabc 100644 --- a/CopernicusMarineService/benchmark_CopernicusMarineService.py +++ b/CopernicusMarineService/benchmark_CopernicusMarineService.py @@ -14,6 +14,7 @@ parcelsv4 = True try: from parcels.xgrid import _XGRID_AXES + from parcels.application_kernels.interpolation import XLinear except ImportError: parcelsv4 = False @@ -82,7 +83,9 @@ def NoFieldAccess( lon0_expected, lat0_expected = -9.820091, -30.106716 # values from v3 if parcelsv4: - if interpolator == "BiRectiLinear": + if interpolator == "XLinear": + interp_method = XLinear + elif interpolator == "BiRectiLinear": interp_method = BiRectiLinear elif interpolator == "PureXarrayInterp": interp_method = PureXarrayInterp @@ -95,7 +98,7 @@ def NoFieldAccess( ds = ds.drop_vars(["vsdx", "vsdy", "utide", "vtide", "utotal", "vtotal"]) ds = ds.rename({"uo": "U", "vo": "V", "longitude": "lon", "latitude": "lat"}) - xgcm_grid = parcels.xgcm.Grid(ds, coords={"X": {"left": "lon"}, "Y": {"left": "lat"}, "Z": {"left": "depth"}, "T": {"center": "time"}}) + xgcm_grid = parcels.xgcm.Grid(ds, coords={"X": {"left": "lon"}, "Y": {"left": "lat"}, "Z": {"left": "depth"}, "T": {"center": "time"}}, periodic=False) grid = parcels.xgrid.XGrid(xgcm_grid) U = parcels.Field("U", ds["U"], grid, interp_method=interp_method) @@ -113,7 +116,7 @@ def NoFieldAccess( kernel = parcels.AdvectionEE - for npart in [1, 10, 100, 1000, 5000, 10000]: + for npart in [1, 10, 100, 1_000, 5_000, 10_000, 50_000, 100_000, 500_000, 1_000_000]: lon = np.linspace(-10, 10, npart) lat = np.linspace(-30, -20, npart) @@ -133,7 +136,7 @@ def main(args=None): "-i", "--Interpolator", choices=("BiRectiLinear", "PureXarrayInterp", "NoFieldAccess"), - default="BiRectiLinear", + default="XLinear", ) args = p.parse_args(args) From 75df1253876333c9a1c5b26543827a010790e766 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 7 Aug 2025 13:30:19 +0200 Subject: [PATCH 6/6] Adding support for tracing memory load --- .../benchmark_CopernicusMarineService.py | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/CopernicusMarineService/benchmark_CopernicusMarineService.py b/CopernicusMarineService/benchmark_CopernicusMarineService.py index 81dfabc..7798213 100644 --- a/CopernicusMarineService/benchmark_CopernicusMarineService.py +++ b/CopernicusMarineService/benchmark_CopernicusMarineService.py @@ -1,6 +1,8 @@ from argparse import ArgumentParser from pathlib import Path import xarray as xr +import tracemalloc +import time from glob import glob @@ -20,7 +22,7 @@ DATA_ROOT = "/storage/shared/oceanparcels/input_data/CopernicusMarineService/GLOBAL_ANALYSIS_FORECAST_PHY_001_024_SMOC/" -def run_benchmark(interpolator: str): +def run_benchmark(interpolator: str, trace_memory: bool = False): if parcelsv4: def BiRectiLinear( @@ -123,8 +125,23 @@ def NoFieldAccess( pset = parcels.ParticleSet(fieldset=fieldset, pclass=pclass, lon=lon, lat=lat) print(f"Running {len(lon)} particles with parcels v{4 if parcelsv4 else 3} and {interpolator} interpolator") - pset.execute(kernel, runtime=runtime, dt=dt) + if trace_memory: + tracemalloc.start() + else: + start = time.time() + + pset.execute(kernel, runtime=runtime, dt=dt, verbose_progress=False) + + if trace_memory: + current, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + print(f"Memory usage: current={current / 1e6:.0f} MB, peak={peak / 1e6:.0f} MB") + else: + elapsed_time = time.time() - start + print(f"Execution time: {elapsed_time:.2f} seconds") + + print("") assert np.allclose(pset[0].lon, lon0_expected, atol=1e-5), f"Expected lon {lon0_expected}, got {pset[0].lon}" assert np.allclose(pset[0].lat, lat0_expected, atol=1e-5), f"Expected lat {lat0_expected}, got {pset[0].lat}" @@ -138,9 +155,15 @@ def main(args=None): choices=("BiRectiLinear", "PureXarrayInterp", "NoFieldAccess"), default="XLinear", ) + p.add_argument( + "-m", + "--memory", + action="store_true", + help="Enable memory tracing (default: False)", + ) args = p.parse_args(args) - run_benchmark(args.Interpolator) + run_benchmark(args.Interpolator, args.memory) if __name__ == "__main__":