diff --git a/.gitignore b/.gitignore index c1a039b..333f7b8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +**.parquet +**.nc + # pixi environments .pixi/* !.pixi/config.toml diff --git a/benchmarks/copernicusstreaming.py b/benchmarks/copernicusstreaming.py new file mode 100644 index 0000000..28d95ee --- /dev/null +++ b/benchmarks/copernicusstreaming.py @@ -0,0 +1,79 @@ +import copernicusmarine +import parcels +import xarray as xr +import numpy as np +import time +import argparse + + +def run_copernicusmarine_benchmark(load_mode="as_file"): + copernicusmarine.login() + + ds = copernicusmarine.open_dataset( + dataset_id="cmems_mod_glo_phy-cur_anfc_0.083deg_P1D-m", + variables=["uo", "vo"], + minimum_longitude=-20, + maximum_longitude=20, + minimum_latitude=-20, + maximum_latitude=20, + start_datetime="2024-01-01", + end_datetime="2024-01-31", + minimum_depth=0.5, + maximum_depth=5, + service="arco-geo-series", + chunk_size_limit=1, + ) + ds = parcels.convert.copernicusmarine_to_sgrid( + fields={"U": ds["uo"], "V": ds["vo"]} + ) + + if load_mode == "as_file": + ds.to_netcdf("tmp.nc") + del ds + ds = xr.open_dataset("tmp.nc") + elif load_mode == "ds_load": + ds.load() + + fieldset = parcels.FieldSet.from_sgrid_conventions(ds) + + if load_mode == "streaming": + fieldset.to_windowed_arrays() + + fieldset.describe() + + pset = parcels.ParticleSet(fieldset, x=0, y=0, z=2) + oufile = parcels.ParticleFile( + "tmp.parquet", + outputdt=np.timedelta64(1, "D"), + mode="w", + ) + + pset.execute( + parcels.kernels.AdvectionRK2, + dt=np.timedelta64(1, "h"), + runtime=np.timedelta64(30, "D"), + output_file=oufile, + ) + + # check for correctness of the benchmark result + np.testing.assert_allclose(pset.x, -14.034891, atol=1e-2) + np.testing.assert_allclose(pset.y, -2.1642778, atol=1e-2) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--load_mode", + choices=["as_file", "ds_load", "streaming"], + default="streaming", + ) + args = parser.parse_args() + + time_start = time.time() + run_copernicusmarine_benchmark(load_mode=args.load_mode) + time_end = time.time() + print(f"Elapsed time: {(time_end - time_start):.2f}s") + + +if __name__ == "__main__": + main() diff --git a/pixi.toml b/pixi.toml index f067792..a7a9e61 100644 --- a/pixi.toml +++ b/pixi.toml @@ -66,3 +66,4 @@ netcdf4 = ">=1.7.4,<2" #curl = "*" # for some reason SSL key signing wasn't working on Lorenz with this curl - instead use curl which is installed on the system proj = "*" ty = "*" +copernicusmarine = ">=2.4.1,<3"