From 99cf6d09e94675367482c070e88efcb96c6c97e4 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 23 Jul 2025 16:21:48 +0200 Subject: [PATCH 01/12] Create benchmark_moi_curvilinear.py --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 164 +++++++++++++++++++ 1 file changed, 164 insertions(+) create mode 100644 MOi-Curvilinear/benchmark_moi_curvilinear.py diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py new file mode 100644 index 0000000..8b60de6 --- /dev/null +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -0,0 +1,164 @@ +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(2, "D") +dt = np.timedelta64(15, "m") + +parcelsv4 = True +try: + from parcels.xgrid import _XGRID_AXES +except ImportError: + parcelsv4 = False + +DATA_ROOT = "/storage/shared/oceanparcels/input_data/MOi/GLO12" + +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"] + + 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] + + 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 + + + fileroot = f"{DATA_ROOT}/psy4v3r1-daily" + filenames = {"U": f"{fileroot}_U*.nc", "V": f"{fileroot}_V*.nc", "W": f"{fileroot}_W*1.nc"} + mesh_mask = f"{DATA_ROOT}/PSY4V3R1_mesh_hgr.nc" + + lon0_expected, lat0_expected = -9.820091, -30.106716 # values from v3 + if parcelsv4: + if interpolator == "BiRectiLinear": + interp_method = BiRectiLinear + elif interpolator == "PureXarrayInterp": + 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_u = xr.open_mfdataset(filenames["U"], concat_dim="time_counter", combine="nested", data_vars='minimal', coords='minimal', compat='override')[["vozocrtx"]].drop_vars( + ["nav_lon", "nav_lat"] + ) + ds_v = xr.open_mfdataset(filenames["V"], concat_dim="time_counter", combine="nested", data_vars='minimal', coords='minimal', compat='override')[["vomecrty"]].drop_vars( + ["nav_lon", "nav_lat"] + ) + ds_depth = xr.open_mfdataset(filenames["W"], concat_dim="time_counter", combine="nested", data_vars='minimal', coords='minimal', compat='override')[["depthw"]] + ds_mesh = xr.open_dataset(mesh_mask)[["glamf", "gphif"]].isel(t=0) + + ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) + + xgcm_grid = parcels.xgcm.Grid( + ds, + coords={ + "X": {"left": "x"}, + "Y": {"left": "y"}, + "Z": {"center": "deptht", "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: + filenames = { + "U": {"lon": mesh_mask, "lat": mesh_mask, "data": filenames["U"]}, + "V": {"lon": mesh_mask, "lat": mesh_mask, "data": filenames["V"]}, + } + interpolator = "v3_default" + fieldset = parcels.FieldSet.from_netcdf(filenames, variables={"U": "vozocrtx", "V": "vomecrty"}, dimensions={"time": "time_counter", "lat": "gphif", "lon": "glamf"}) + + pclass = parcels.Particle if parcelsv4 else parcels.ScipyParticle + + for npart in [1, 10, 100, 1000, 5000, 10000]: + lon = np.linspace(170, 190, 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(parcels.AdvectionEE, 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() + + 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 e9a9137977e4eaf350807dfa0726748e27a951f4 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Thu, 24 Jul 2025 08:50:01 +0200 Subject: [PATCH 02/12] Fixing moi-curvilinear benchmark to run --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 28 +++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 8b60de6..47ab2e7 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -10,6 +10,7 @@ runtime = np.timedelta64(2, "D") dt = np.timedelta64(15, "m") +depth_range = range(0, 2) # only taking upper-two depth levels parcelsv4 = True try: @@ -17,7 +18,7 @@ except ImportError: parcelsv4 = False -DATA_ROOT = "/storage/shared/oceanparcels/input_data/MOi/GLO12" +DATA_ROOT = "/storage/shared/oceanparcels/input_data/MOi" def run_benchmark(interpolator: str): if parcelsv4: @@ -78,11 +79,12 @@ def NoFieldAccess( return 0 - fileroot = f"{DATA_ROOT}/psy4v3r1-daily" - filenames = {"U": f"{fileroot}_U*.nc", "V": f"{fileroot}_V*.nc", "W": f"{fileroot}_W*1.nc"} - mesh_mask = f"{DATA_ROOT}/PSY4V3R1_mesh_hgr.nc" + lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) + + fileU = f"{DATA_ROOT}/GLO12/psy4v3r1-daily_U_2010-01-0[1-3].nc" + filenames = {"U": glob(fileU), "V": glob(fileU.replace("_U_", "_V_")), "W": glob(fileU.replace("_U_", "_W_"))} + mesh_mask = f"{DATA_ROOT}/domain_ORCA0083-N006/PSY4V3R1_mesh_hgr.nc" - lon0_expected, lat0_expected = -9.820091, -30.106716 # values from v3 if parcelsv4: if interpolator == "BiRectiLinear": interp_method = BiRectiLinear @@ -104,6 +106,7 @@ def NoFieldAccess( ds_mesh = xr.open_dataset(mesh_mask)[["glamf", "gphif"]].isel(t=0) ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) + ds = ds.isel(depth=depth_range) xgcm_grid = parcels.xgcm.Grid( ds, @@ -125,16 +128,21 @@ def NoFieldAccess( fieldset = parcels.FieldSet([U, V, UV]) else: filenames = { - "U": {"lon": mesh_mask, "lat": mesh_mask, "data": filenames["U"]}, - "V": {"lon": mesh_mask, "lat": mesh_mask, "data": filenames["V"]}, + "U": {"lon": mesh_mask, "lat": mesh_mask, "depth": filenames["W"][0], "data": filenames["U"]}, + "V": {"lon": mesh_mask, "lat": mesh_mask, "depth": filenames["W"][0], "data": filenames["V"]}, } interpolator = "v3_default" - fieldset = parcels.FieldSet.from_netcdf(filenames, variables={"U": "vozocrtx", "V": "vomecrty"}, dimensions={"time": "time_counter", "lat": "gphif", "lon": "glamf"}) + fieldset = parcels.FieldSet.from_netcdf( + filenames, + variables={"U": "vozocrtx", "V": "vomecrty"}, + dimensions={"time": "time_counter", "lat": "gphif", "lon": "glamf", "depth": "depthw"}, + indices={"depth": depth_range}, + ) - pclass = parcels.Particle if parcelsv4 else parcels.ScipyParticle + pclass = parcels.Particle if parcelsv4 else parcels.JITParticle for npart in [1, 10, 100, 1000, 5000, 10000]: - lon = np.linspace(170, 190, npart) + lon = np.linspace(-10, 10, npart) lat = np.linspace(-30, -20, npart) pset = parcels.ParticleSet(fieldset=fieldset, pclass=pclass, lon=lon, lat=lat) From b013cef260bd315d69a3ea4896fbbcb60c890429 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 13 Aug 2025 10:11:49 +0200 Subject: [PATCH 03/12] Updating moi curvilinear benchmark to also assess peak memory --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 108 ++++++------------- 1 file changed, 35 insertions(+), 73 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 47ab2e7..3f2a939 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.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 @@ -10,74 +12,17 @@ runtime = np.timedelta64(2, "D") dt = np.timedelta64(15, "m") -depth_range = range(0, 2) # only taking upper-two depth levels parcelsv4 = True try: from parcels.xgrid import _XGRID_AXES + from parcels.application_kernels.interpolation import XLinear except ImportError: parcelsv4 = False DATA_ROOT = "/storage/shared/oceanparcels/input_data/MOi" -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"] - - 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] - + 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 - +def run_benchmark(interpolator: str, trace_memory: bool = False): lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) @@ -86,13 +31,8 @@ def NoFieldAccess( mesh_mask = f"{DATA_ROOT}/domain_ORCA0083-N006/PSY4V3R1_mesh_hgr.nc" if parcelsv4: - if interpolator == "BiRectiLinear": - interp_method = BiRectiLinear - elif interpolator == "PureXarrayInterp": - interp_method = PureXarrayInterp - elif interpolator == "NoFieldAccess": - interp_method = NoFieldAccess - lon0_expected, lat0_expected = -10, -30 # Zero interpolation, so expect initial values + if interpolator == "XLinear": + interp_method = XLinear else: raise ValueError(f"Unknown interpolator: {interpolator}") @@ -106,7 +46,6 @@ def NoFieldAccess( ds_mesh = xr.open_dataset(mesh_mask)[["glamf", "gphif"]].isel(t=0) ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) - ds = ds.isel(depth=depth_range) xgcm_grid = parcels.xgcm.Grid( ds, @@ -116,6 +55,7 @@ def NoFieldAccess( "Z": {"center": "deptht", "left": "depth"}, "T": {"center": "time"}, }, + periodic=False, ) grid = parcels.xgrid.XGrid(xgcm_grid) @@ -136,19 +76,34 @@ def NoFieldAccess( filenames, variables={"U": "vozocrtx", "V": "vomecrty"}, dimensions={"time": "time_counter", "lat": "gphif", "lon": "glamf", "depth": "depthw"}, - indices={"depth": depth_range}, ) pclass = parcels.Particle if parcelsv4 else parcels.JITParticle - 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) 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(parcels.AdvectionEE, runtime=runtime, dt=dt) + + if trace_memory: + tracemalloc.start() + else: + start = time.time() + + pset.execute(parcels.AdvectionEE, 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}" @@ -160,12 +115,19 @@ def main(args=None): p.add_argument( "-i", "--Interpolator", - choices=("BiRectiLinear", "PureXarrayInterp", "NoFieldAccess"), - default="BiRectiLinear", + choices=("XLinear", "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__": From a7c2e21c05f9fcfb92e952f1f7fb7d7f4b120371 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 13 Aug 2025 13:50:49 +0200 Subject: [PATCH 04/12] improving output readibility --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 3f2a939..8ae7de9 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -86,7 +86,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False): 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") + print(f"Running {len(lon):_} particles with parcels v{4 if parcelsv4 else 3} and {interpolator} interpolator") if trace_memory: tracemalloc.start() @@ -101,7 +101,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False): 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(f"Execution time: {elapsed_time:.0f} seconds") print("") From 2f860291c5356331c9eda901846bc18ae673dcb8 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 25 Aug 2025 14:29:46 +0200 Subject: [PATCH 05/12] Adding support for choosing surface simulation only --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 42 ++++++++++++++------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 8ae7de9..12b23ef 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -22,7 +22,7 @@ DATA_ROOT = "/storage/shared/oceanparcels/input_data/MOi" -def run_benchmark(interpolator: str, trace_memory: bool = False): +def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulation=False): lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) @@ -47,16 +47,17 @@ def run_benchmark(interpolator: str, trace_memory: bool = False): ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) - xgcm_grid = parcels.xgcm.Grid( - ds, - coords={ - "X": {"left": "x"}, - "Y": {"left": "y"}, - "Z": {"center": "deptht", "left": "depth"}, - "T": {"center": "time"}, - }, - periodic=False, - ) + coords={ + "X": {"left": "x"}, + "Y": {"left": "y"}, + "T": {"center": "time"}, + } + if surface_simulation: + ds = ds.isel(depth=0, deptht=0) + else: + coords["Z"] = {"center": "deptht", "left": "depth"} + + xgcm_grid = parcels.xgcm.Grid(ds, coords=coords, periodic=False) grid = parcels.xgrid.XGrid(xgcm_grid) U = parcels.Field("U", ds["U"], grid, interp_method=interp_method) @@ -72,10 +73,16 @@ def run_benchmark(interpolator: str, trace_memory: bool = False): "V": {"lon": mesh_mask, "lat": mesh_mask, "depth": filenames["W"][0], "data": filenames["V"]}, } interpolator = "v3_default" + if surface_simulation: + indices={"depth": range(2)} + else: + indices=None + fieldset = parcels.FieldSet.from_netcdf( filenames, variables={"U": "vozocrtx", "V": "vomecrty"}, dimensions={"time": "time_counter", "lat": "gphif", "lon": "glamf", "depth": "depthw"}, + indices=indices, ) pclass = parcels.Particle if parcelsv4 else parcels.JITParticle @@ -93,6 +100,10 @@ def run_benchmark(interpolator: str, trace_memory: bool = False): else: start = time.time() + if surface_simulation and parcelsv4: + fieldset.U.data.load() + fieldset.V.data.load() + pset.execute(parcels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False) if trace_memory: @@ -126,8 +137,15 @@ def main(args=None): help="Enable memory tracing (default: False)", ) + p.add_argument( + "-s", + "--surface", + action="store_true", + help="Run surface simulation with only 1 or 2 depth levels (default: False)", + ) + args = p.parse_args(args) - run_benchmark(args.Interpolator, args.memory) + run_benchmark(args.Interpolator, args.memory, args.surface) if __name__ == "__main__": From aba0f88415b4b316af8ad84c5a8f8c074e03dc37 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 26 Aug 2025 14:29:10 +0200 Subject: [PATCH 06/12] Adding support for chunking cycling --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 206 +++++++++++-------- 1 file changed, 116 insertions(+), 90 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 12b23ef..22266e6 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -22,102 +22,121 @@ DATA_ROOT = "/storage/shared/oceanparcels/input_data/MOi" -def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulation=False): +def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulation: bool =False, cycle_chunks: bool = False): - lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) + if surface_simulation: + lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) + if cycle_chunks: + xy_chunks = [64, 128, 256, 512, 1024, 2084, 32, 18, 8, 4] + nparts = [10_000] + else: + xy_chunks = [64] + nparts = [1, 10, 100, 1_000, 5_000, 10_000, 50_000, 100_000, 500_000, 1_000_000] fileU = f"{DATA_ROOT}/GLO12/psy4v3r1-daily_U_2010-01-0[1-3].nc" filenames = {"U": glob(fileU), "V": glob(fileU.replace("_U_", "_V_")), "W": glob(fileU.replace("_U_", "_W_"))} mesh_mask = f"{DATA_ROOT}/domain_ORCA0083-N006/PSY4V3R1_mesh_hgr.nc" - if parcelsv4: - if interpolator == "XLinear": - interp_method = XLinear - else: - raise ValueError(f"Unknown interpolator: {interpolator}") - - ds_u = xr.open_mfdataset(filenames["U"], concat_dim="time_counter", combine="nested", data_vars='minimal', coords='minimal', compat='override')[["vozocrtx"]].drop_vars( - ["nav_lon", "nav_lat"] - ) - ds_v = xr.open_mfdataset(filenames["V"], concat_dim="time_counter", combine="nested", data_vars='minimal', coords='minimal', compat='override')[["vomecrty"]].drop_vars( - ["nav_lon", "nav_lat"] - ) - ds_depth = xr.open_mfdataset(filenames["W"], concat_dim="time_counter", combine="nested", data_vars='minimal', coords='minimal', compat='override')[["depthw"]] - ds_mesh = xr.open_dataset(mesh_mask)[["glamf", "gphif"]].isel(t=0) - - ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) - - coords={ - "X": {"left": "x"}, - "Y": {"left": "y"}, - "T": {"center": "time"}, - } - if surface_simulation: - ds = ds.isel(depth=0, deptht=0) - else: - coords["Z"] = {"center": "deptht", "left": "depth"} - - xgcm_grid = parcels.xgcm.Grid(ds, coords=coords, periodic=False) - 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: - filenames = { - "U": {"lon": mesh_mask, "lat": mesh_mask, "depth": filenames["W"][0], "data": filenames["U"]}, - "V": {"lon": mesh_mask, "lat": mesh_mask, "depth": filenames["W"][0], "data": filenames["V"]}, - } - interpolator = "v3_default" - if surface_simulation: - indices={"depth": range(2)} - else: - indices=None - - fieldset = parcels.FieldSet.from_netcdf( - filenames, - variables={"U": "vozocrtx", "V": "vomecrty"}, - dimensions={"time": "time_counter", "lat": "gphif", "lon": "glamf", "depth": "depthw"}, - indices=indices, - ) - - pclass = parcels.Particle if parcelsv4 else parcels.JITParticle - - 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) - - 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") - - if trace_memory: - tracemalloc.start() - else: - start = time.time() - - if surface_simulation and parcelsv4: - fieldset.U.data.load() - fieldset.V.data.load() - - pset.execute(parcels.AdvectionEE, 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") + for chunk in xy_chunks: + if parcelsv4: + if interpolator == "XLinear": + interp_method = XLinear + else: + raise ValueError(f"Unknown interpolator: {interpolator}") + + fileargs = {"concat_dim": "time_counter", + "combine": "nested", + "data_vars": 'minimal', + "coords": 'minimal', + "compat": 'override', + "chunks": {"time_counter": 1, "depth":2, "y": chunk, "x": chunk} + } + + ds_u = xr.open_mfdataset(filenames["U"], **fileargs)[["vozocrtx"]].drop_vars(["nav_lon", "nav_lat"]) + ds_v = xr.open_mfdataset(filenames["V"], **fileargs)[["vomecrty"]].drop_vars(["nav_lon", "nav_lat"]) + ds_depth = xr.open_mfdataset(filenames["W"], **fileargs)[["depthw"]] + ds_mesh = xr.open_dataset(mesh_mask)[["glamf", "gphif"]].isel(t=0) + + ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) + + coords={ + "X": {"left": "x"}, + "Y": {"left": "y"}, + "T": {"center": "time"}, + } + if surface_simulation: + ds = ds.isel(depth=0, deptht=0) + else: + coords["Z"] = {"center": "deptht", "left": "depth"} + print(ds) + + xgcm_grid = parcels.xgcm.Grid(ds, coords=coords, periodic=False) + 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: - elapsed_time = time.time() - start - print(f"Execution time: {elapsed_time:.0f} 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}" + filenames = { + "U": {"lon": mesh_mask, "lat": mesh_mask, "depth": filenames["W"][0], "data": filenames["U"]}, + "V": {"lon": mesh_mask, "lat": mesh_mask, "depth": filenames["W"][0], "data": filenames["V"]}, + } + interpolator = "v3_default" + if surface_simulation: + indices={"depth": range(2)} + else: + indices=None + + fieldset = parcels.FieldSet.from_netcdf( + filenames, + variables={"U": "vozocrtx", "V": "vomecrty"}, + dimensions={"time": "time_counter", "lat": "gphif", "lon": "glamf", "depth": "depthw"}, + indices=indices, + ) + + pclass = parcels.Particle if parcelsv4 else parcels.JITParticle + + for npart in nparts: + if cycle_chunks: + X, Y = np.meshgrid(np.linspace(-10, 10, int(np.sqrt(npart))), np.linspace(-30, -20, int(np.sqrt(npart)))) + lon = X.flatten() + lat = Y.flatten() + else: + 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 on {"surface" if surface_simulation else "3D"} with parcels v{4 if parcelsv4 else 3}, chunksize {chunk} and {interpolator} interpolator") + + if trace_memory: + tracemalloc.start() + else: + start = time.time() + + # if surface_simulation and parcelsv4: + # fieldset.U.data.load() + # fieldset.V.data.load() + + pset.execute(parcels.AdvectionEE, 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:.0f} seconds") + + print("") + + if not cycle_chunks: + 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): @@ -144,8 +163,15 @@ def main(args=None): help="Run surface simulation with only 1 or 2 depth levels (default: False)", ) + p.add_argument( + "-c", + "--chunks", + action="store_true", + help="Cycle through different chunk sizes (default: False)", + ) + args = p.parse_args(args) - run_benchmark(args.Interpolator, args.memory, args.surface) + run_benchmark(args.Interpolator, args.memory, args.surface, args.chunks) if __name__ == "__main__": From 4d778f1ac121ad39f7278e38bd638e55e0e171e9 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Fri, 29 Aug 2025 08:48:04 +0200 Subject: [PATCH 07/12] small fixes to benchmark --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 22266e6..481ea8f 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -24,11 +24,10 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulation: bool =False, cycle_chunks: bool = False): - if surface_simulation: - lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) + lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) if cycle_chunks: - xy_chunks = [64, 128, 256, 512, 1024, 2084, 32, 18, 8, 4] + xy_chunks = [64, 128, 256, 512, 1024, 2084, 32, 16, 8, 4] nparts = [10_000] else: xy_chunks = [64] From 4067f4c365177a68b3aaf22f155a0f9661389df4 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Sat, 27 Sep 2025 12:53:32 +0200 Subject: [PATCH 08/12] update benchmark to use xgcm package from conda --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 481ea8f..cc59e9a 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -7,6 +7,7 @@ from glob import glob import numpy as np +import xgcm import parcels @@ -30,7 +31,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat xy_chunks = [64, 128, 256, 512, 1024, 2084, 32, 16, 8, 4] nparts = [10_000] else: - xy_chunks = [64] + xy_chunks = ["auto"] nparts = [1, 10, 100, 1_000, 5_000, 10_000, 50_000, 100_000, 500_000, 1_000_000] fileU = f"{DATA_ROOT}/GLO12/psy4v3r1-daily_U_2010-01-0[1-3].nc" filenames = {"U": glob(fileU), "V": glob(fileU.replace("_U_", "_V_")), "W": glob(fileU.replace("_U_", "_W_"))} @@ -57,6 +58,11 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat ds_mesh = xr.open_dataset(mesh_mask)[["glamf", "gphif"]].isel(t=0) ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical").rename({"vozocrtx": "U", "vomecrty": "V"}).rename({"glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) + ds = xr.merge([ds_u, ds_v, ds_depth, ds_mesh], compat="identical") + ds = ds.rename({"vozocrtx": "U", "vomecrty": "V", "glamf": "lon", "gphif": "lat", "time_counter": "time", "depthw": "depth"}) + ds.deptht.attrs["c_grid_axis_shift"] = -0.5 + + coords={"X": {"left": "x"}, "Y": {"left": "y"}, "T": {"center": "time"}} coords={ "X": {"left": "x"}, @@ -69,8 +75,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat coords["Z"] = {"center": "deptht", "left": "depth"} print(ds) - xgcm_grid = parcels.xgcm.Grid(ds, coords=coords, periodic=False) - grid = parcels.xgrid.XGrid(xgcm_grid) + grid = parcels.xgrid.XGrid(xgcm.Grid(ds, coords=coords, autoparse_metadata=False, periodic=False)) U = parcels.Field("U", ds["U"], grid, interp_method=interp_method) V = parcels.Field("V", ds["V"], grid, interp_method=interp_method) From ae1c543f4f3e22c58cd8e530651678ae9c9f04ac Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 29 Sep 2025 07:47:41 +0200 Subject: [PATCH 09/12] Adding preloading to CLI --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 26 +++++++++++++------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index cc59e9a..43e9ad3 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -23,7 +23,7 @@ DATA_ROOT = "/storage/shared/oceanparcels/input_data/MOi" -def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulation: bool =False, cycle_chunks: bool = False): +def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulation: bool =False, preload: bool = False, cycle_chunks: bool = False): lon0_expected, lat0_expected = -10.128929, -29.721205 # values from v3 using from_netcf (so assuming A-grid!) @@ -31,7 +31,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat xy_chunks = [64, 128, 256, 512, 1024, 2084, 32, 16, 8, 4] nparts = [10_000] else: - xy_chunks = ["auto"] + xy_chunks = [256] nparts = [1, 10, 100, 1_000, 5_000, 10_000, 50_000, 100_000, 500_000, 1_000_000] fileU = f"{DATA_ROOT}/GLO12/psy4v3r1-daily_U_2010-01-0[1-3].nc" filenames = {"U": glob(fileU), "V": glob(fileU.replace("_U_", "_V_")), "W": glob(fileU.replace("_U_", "_W_"))} @@ -49,8 +49,9 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat "data_vars": 'minimal', "coords": 'minimal', "compat": 'override', - "chunks": {"time_counter": 1, "depth":2, "y": chunk, "x": chunk} } + if chunk: + fileargs["chunks"] = {"time_counter": 1, "depth":2, "y": chunk, "x": chunk} ds_u = xr.open_mfdataset(filenames["U"], **fileargs)[["vozocrtx"]].drop_vars(["nav_lon", "nav_lat"]) ds_v = xr.open_mfdataset(filenames["V"], **fileargs)[["vomecrty"]].drop_vars(["nav_lon", "nav_lat"]) @@ -104,6 +105,10 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat pclass = parcels.Particle if parcelsv4 else parcels.JITParticle + if parcelsv4 and preload: + fieldset.U.data.load() + fieldset.V.data.load() + for npart in nparts: if cycle_chunks: X, Y = np.meshgrid(np.linspace(-10, 10, int(np.sqrt(npart))), np.linspace(-30, -20, int(np.sqrt(npart)))) @@ -115,17 +120,13 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat pset = parcels.ParticleSet(fieldset=fieldset, pclass=pclass, lon=lon, lat=lat) - print(f"Running {len(lon):_} particles on {"surface" if surface_simulation else "3D"} with parcels v{4 if parcelsv4 else 3}, chunksize {chunk} and {interpolator} interpolator") + print(f"Running {len(lon):_} particles on {"surface" if surface_simulation else "3D"} with parcels v{4 if parcelsv4 else 3}, chunksize {chunk} ({'preloaded' if preload else 'not preloaded'}) and {interpolator} interpolator") if trace_memory: tracemalloc.start() else: start = time.time() - # if surface_simulation and parcelsv4: - # fieldset.U.data.load() - # fieldset.V.data.load() - pset.execute(parcels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False) if trace_memory: @@ -167,6 +168,13 @@ def main(args=None): help="Run surface simulation with only 1 or 2 depth levels (default: False)", ) + p.add_argument( + "-l", + "--preload", + action="store_true", + help="Preload data into memory (default: False)", + ) + p.add_argument( "-c", "--chunks", @@ -175,7 +183,7 @@ def main(args=None): ) args = p.parse_args(args) - run_benchmark(args.Interpolator, args.memory, args.surface, args.chunks) + run_benchmark(args.Interpolator, args.memory, args.surface, args.preload, args.chunks) if __name__ == "__main__": From 61cc582f48696dd860d534e0afa5f4b637219412 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 29 Sep 2025 08:08:21 +0200 Subject: [PATCH 10/12] Moving xgcm import to v4 testing logic --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 43e9ad3..a912793 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -7,7 +7,6 @@ from glob import glob import numpy as np -import xgcm import parcels @@ -16,6 +15,7 @@ parcelsv4 = True try: + import xgcm from parcels.xgrid import _XGRID_AXES from parcels.application_kernels.interpolation import XLinear except ImportError: From 2a3116b8ad16525fd492bad15609f03d651a8f06 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 29 Sep 2025 10:03:17 +0200 Subject: [PATCH 11/12] Making sure mesh=spherical This fixes https://github.com/OceanParcels/Parcels/issues/2254 --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 43e9ad3..82a200c 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -76,7 +76,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat coords["Z"] = {"center": "deptht", "left": "depth"} print(ds) - grid = parcels.xgrid.XGrid(xgcm.Grid(ds, coords=coords, autoparse_metadata=False, periodic=False)) + grid = parcels.xgrid.XGrid(xgcm.Grid(ds, coords=coords, autoparse_metadata=False, periodic=False), mesh="spherical") U = parcels.Field("U", ds["U"], grid, interp_method=interp_method) V = parcels.Field("V", ds["V"], grid, interp_method=interp_method) From 89dc3dc3f2e57648d60c0e9b04102555ec17b8fa Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Wed, 1 Oct 2025 13:50:02 +0200 Subject: [PATCH 12/12] Updating script to new v4 structure --- MOi-Curvilinear/benchmark_moi_curvilinear.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/MOi-Curvilinear/benchmark_moi_curvilinear.py b/MOi-Curvilinear/benchmark_moi_curvilinear.py index 7941e57..e693553 100644 --- a/MOi-Curvilinear/benchmark_moi_curvilinear.py +++ b/MOi-Curvilinear/benchmark_moi_curvilinear.py @@ -16,8 +16,7 @@ parcelsv4 = True try: import xgcm - from parcels.xgrid import _XGRID_AXES - from parcels.application_kernels.interpolation import XLinear + from parcels.interpolators import XLinear except ImportError: parcelsv4 = False @@ -76,7 +75,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat coords["Z"] = {"center": "deptht", "left": "depth"} print(ds) - grid = parcels.xgrid.XGrid(xgcm.Grid(ds, coords=coords, autoparse_metadata=False, periodic=False), mesh="spherical") + grid = parcels._core.xgrid.XGrid(xgcm.Grid(ds, coords=coords, autoparse_metadata=False, periodic=False), mesh="spherical") U = parcels.Field("U", ds["U"], grid, interp_method=interp_method) V = parcels.Field("V", ds["V"], grid, interp_method=interp_method) @@ -127,7 +126,7 @@ def run_benchmark(interpolator: str, trace_memory: bool = False, surface_simulat else: start = time.time() - pset.execute(parcels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False) + pset.execute(parcels.kernels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False) if trace_memory: current, peak = tracemalloc.get_traced_memory()