Streamline parcels-benchmarks

.env.example

@@ -0,0 +1,2 @@
+# Target location for the data to be downloaded to
+PARCELS_BENCHMARKS_DATA_FOLDER='./data'

.gitignore

@@ -10,6 +10,9 @@ credentials.json
 *.egg-info
 __pycache__
 build/
-parcels/
 .asv/
 html/
+.DS_Store
+
+data
+.env

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "Parcels"]
+	path = Parcels
+	url = git@github.com:Parcels-code/Parcels

README.md

@@ -6,15 +6,27 @@ This repository houses performance benchmarks for [Parcels](https://github.com/O
 
 ## Development instructions
 
+This project uses a combination of [Pixi](https://pixi.sh/dev/installation/), [ASV](https://asv.readthedocs.io/), and [intake-xarray](https://github.com/intake/intake-xarray) to coordinate the setting up and running of benchmarks.
+
+- Scripts are used to download the datasets required into the correct location
+- intake-xarray is used to define data catalogues which can be easily accessed from within benchmark scripts
+- ASV is used to run the benchmarks (see the [Writing the benchmarks](#writing-the-benchmarks) section).
+- Pixi is used to orchestrate all the above into a convenient, user friendly workflow
+
+You can run `pixi task list` to see the list of available tasks in the workspace.
+
+In brief, you can set up the data and run the benchmarks by doing:
+
 - [install Pixi](https://pixi.sh/dev/installation/) `curl -fsSL https://pixi.sh/install.sh | bash`
 - `pixi install`
-- `pixi run asv run`
-
-You can run the linting with `pixi run lint`
+- `PARCELS_BENCHMARKS_DATA_FOLDER=./data pixi run benchmarks`
 
 > [!IMPORTANT]
-> The default path for the benchmark data is set by [pooch.os_cache](https://www.fatiando.org/pooch/latest/api/generated/pooch.os_cache.html), which typically is a subdirectory of your home directory. Currently, you will need at least 50GB of disk space available to store the benchmark data.
-> To change the location of the benchmark data cache, you can set the environment variable `PARCELS_DATADIR` to a preferred location to store the benchmark data.
+> Currently, you will need at least 50GB of disk space available to store the benchmark data.
+> You need to be explicit to determine where the benchmark data will be saved by
+> setting the `PARCELS_BENCHMARKS_DATA_FOLDER` environment variable. This
+> environment variable is used in the downloading of the data and definition of
+> the benchmarks.
 
 To view the benchmark data
 
@@ -34,7 +46,7 @@ Members of the Parcels community can contribute benchmark data using the followi
 2. Clone your fork onto your system
 
 ```
-git clone git@github.com:<your-github-handle>/parcels-benchmarks.git ~/parcels-benchmarks
+git clone --recurse-submodules git@github.com:<your-github-handle>/parcels-benchmarks.git
 ```
 
 3. Run the benchmarks
@@ -61,13 +73,9 @@ Adding benchmarks for parcels typically involves adding a dataset and defining t
 ### Adding new data
 
 Data is hosted remotely on a SurfDrive managed by the Parcels developers. You will need to open an issue on this repository to start the process of getting your data hosted in the shared SurfDrive.
-Once your data is hosted in the shared SurfDrive, you can easily add your dataset to the benchmark dataset manifest using
-
-```
-pixi run benchmark-setup pixi add-dataset --name "Name for your dataset" --file "Path to ZIP archive in the SurfDrive"
-```
+Once your data is hosted in the shared SurfDrive, you can easily add your dataset to the benchmark dataset catalogue by modifying `catalogs/parcels-benchmarks/catalog.yml`.
 
-During this process, the dataset will be downloaded and a complete entry will be added to the [parcels_benchmarks/benchmarks.json](./parcels_benchmarks/benchmarks.json) manifest file. Once updated, this file can be committed to this repository and contributed via a pull request.
+In the benchmark you can now use this catalogue entry.
 
 ### Writing the benchmarks
 

asv.conf.json

@@ -2,7 +2,8 @@
   "version": 1,
   "project": "parcels",
   "project_url": "https://github.com/Parcels-Code/parcels",
-  "repo": "https://github.com/Parcels-Code/parcels.git",
+  "repo": "./Parcels",
+  "dvcs": "git",
   "branches": ["main"],
   "environment_type": "rattler",
   "conda_channels": [
@@ -14,12 +15,17 @@
   "env_dir": ".asv/env",
   "results_dir": "results",
   "html_dir": "html",
-  "build_command": [],
-  "install_command": [
-    "in-dir={conf_dir} python -m pip install .",
-    "in-dir={build_dir} python -m pip install ."
-  ],
-  "uninstall_command": [
-    "return-code=any python -m pip uninstall -y parcels parcels_benchmarks"
-  ]
+  "build_command": ["python -m build --wheel -o {build_cache_dir} {build_dir}"],
+  // "install_command": [
+  //   "in-dir={conf_dir} python -m pip install .",
+  //   "in-dir={build_dir} python -m pip install ."
+  // ],
+  // "uninstall_command": [
+  //   "return-code=any python -m pip uninstall -y parcels parcels_benchmarks"
+  // ]
+  "matrix": {
+    "req": {
+      "intake-xarray": []
+    }
+  }
 }

benchmarks/__init__.py

@@ -0,0 +1,24 @@
+import os
+from pathlib import Path
+
+import intake
+
+PARCELS_BENCHMARKS_DATA_FOLDER: str
+PROJECT_ROOT = Path(__file__).parent.parent.resolve()
+
+try:
+    PARCELS_BENCHMARKS_DATA_FOLDER = Path(os.environ["PARCELS_BENCHMARKS_DATA_FOLDER"])
+except KeyError as e:
+    raise RuntimeError(
+        "Set the PARCELS_BENCHMARKS_DATA_FOLDER environment variable to specify where the data is/should be downloaded."
+    ) from e
+
+if not PARCELS_BENCHMARKS_DATA_FOLDER.is_absolute():
+    PARCELS_BENCHMARKS_DATA_FOLDER = PROJECT_ROOT / str(PARCELS_BENCHMARKS_DATA_FOLDER)
+
+CAT_EXAMPLES = intake.open_catalog(
+    f"{PARCELS_BENCHMARKS_DATA_FOLDER}/surf-data/parcels-examples/catalog.yml"
+)
+CAT_BENCHMARKS = intake.open_catalog(
+    f"{PARCELS_BENCHMARKS_DATA_FOLDER}/surf-data/parcels-benchmarks/catalog.yml"
+)

benchmarks/fesom2.py

@@ -1,5 +1,6 @@
 import numpy as np
 import uxarray as ux
+import xarray as xr
 from parcels import (
     FieldSet,
     Particle,
@@ -8,39 +9,38 @@
 )
 from parcels.kernels import AdvectionRK2_3D
 
-from parcels_benchmarks.benchmark_setup import PARCELS_DATADIR, download_example_dataset
+from . import PARCELS_BENCHMARKS_DATA_FOLDER
 
 runtime = np.timedelta64(1, "D")
 dt = np.timedelta64(2400, "s")
 
 
-def _load_ds(datapath):
+def _load_ds():
     """Helper function to load uxarray dataset from datapath"""
 
-    grid_file = f"{datapath}/mesh/fesom.mesh.diag.nc"
-    data_files = f"{datapath}/*.nc"
-    return ux.open_mfdataset(grid_file, data_files, combine="by_coords")
+    grid_file = xr.open_mfdataset(
+        f"{PARCELS_BENCHMARKS_DATA_FOLDER}/surf-data/parcels-benchmarks/data/Parcelsv4_Benchmarking_data/Parcels_Benchmarks_FESOM-baroclinic-gyre/data/mesh/fesom.mesh.diag.nc"
+    )
+    data_files = xr.open_mfdataset(
+        f"{PARCELS_BENCHMARKS_DATA_FOLDER}/surf-data/parcels-benchmarks/data/Parcelsv4_Benchmarking_data/Parcels_Benchmarks_FESOM-baroclinic-gyre/data/*.nc"
+    )
+
+    grid = ux.open_grid(grid_file)
+    return ux.UxDataset(data_files, uxgrid=grid)
 
 
 class FESOM2:
     params = ([10000], [AdvectionRK2_3D])
     param_names = ["npart", "integrator"]
 
-    def setup(self, npart, integrator):
-        # Ensure the dataset is downloaded in the desired data_home
-        # and obtain the path to the dataset
-        self.datapath = download_example_dataset(
-            "FESOM-baroclinic-gyre", data_home=PARCELS_DATADIR
-        )
-
     def time_load_data(self, npart, integrator):
-        ds = _load_ds(self.datapath)
+        ds = _load_ds()
         for i in range(min(ds.coords["time"].size, 2)):
             _u = ds["u"].isel(time=i).compute()
             _v = ds["v"].isel(time=i).compute()
 
     def pset_execute(self, npart, integrator):
-        ds = _load_ds(self.datapath)
+        ds = _load_ds()
         ds = convert.fesom_to_ugrid(ds)
         fieldset = FieldSet.from_ugrid_conventions(ds)
 

benchmarks/moi_curvilinear.py

@@ -6,12 +6,16 @@
 import xgcm
 from parcels.interpolators import XLinear
 
-from parcels_benchmarks.benchmark_setup import PARCELS_DATADIR, download_example_dataset
-
 runtime = np.timedelta64(2, "D")
 dt = np.timedelta64(15, "m")
 
 
+PARCELS_DATADIR = ...  # TODO: Replace with intake
+
+
+def download_dataset(*args, **kwargs): ...  # TODO: Replace with intake
+
+
 def _load_ds(datapath, chunk):
     """Helper function to load xarray dataset from datapath with or without chunking"""
 
@@ -72,9 +76,7 @@ class MOICurvilinear:
     ]
 
     def setup(self, interpolator, chunk, npart):
-        self.datapath = download_example_dataset(
-            "MOi-curvilinear", data_home=PARCELS_DATADIR
-        )
+        self.datapath = download_dataset("MOi-curvilinear", data_home=PARCELS_DATADIR)
 
     def time_load_data_3d(self, interpolator, chunk, npart):
         """Benchmark that times loading the 'U' and 'V' data arrays only for 3-D"""

catalogs/parcels-benchmarks/catalog.yml

@@ -0,0 +1,64 @@
+# zip_url: https://surfdrive.surf.nl/index.php/s/7xlfdOFaUGDEmpD/download?path=%2F&files=
+# ^ Do not remove this line! Used by the download script to find the data source
+plugins:
+  source:
+    - module: intake_xarray
+sources: #!TODO Update
+  croco:
+    description: CROCO_idealized
+    driver: netcdf
+    #cache:
+    #  - argkey: urlpath
+    #    regex: ''
+    #    type: file
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/CROCOidealized_data/CROCO_idealized.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  GlobCurrent_example_data:
+    description: GlobCurrent_example_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/GlobCurrent_example_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  MITgcm_example_data:
+    description: MITgcm_example_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/MITgcm_example_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  MovingEddies_data:
+    description: MovingEddies_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/MovingEddies_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+
+  # NemoCurvilinear_data:
+  # NemoNorthSeaORCA025-N006_data:
+  # OFAM_example_data
+  # Peninsula_data
+  # SWASH_data
+  SWASH_data:
+    description: SWASH_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/SWASH_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  WOA_data:
+    description: WOA_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/WOA_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"

catalogs/parcels-examples/catalog.yml

@@ -0,0 +1,64 @@
+# zip_url: https://surfdrive.surf.nl/index.php/s/cmdSy8wBtCLDaGJ/download?path=%2F&files=
+# ^ Do not remove this line! Used by the download script to find the data source
+plugins:
+  source:
+    - module: intake_xarray
+sources:
+  croco:
+    description: CROCO_idealized
+    driver: netcdf
+    #cache:
+    #  - argkey: urlpath
+    #    regex: ''
+    #    type: file
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/CROCOidealized_data/CROCO_idealized.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  GlobCurrent_example_data:
+    description: GlobCurrent_example_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/GlobCurrent_example_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  MITgcm_example_data:
+    description: MITgcm_example_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/MITgcm_example_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  MovingEddies_data:
+    description: MovingEddies_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/MovingEddies_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+
+  # NemoCurvilinear_data:
+  # NemoNorthSeaORCA025-N006_data:
+  # OFAM_example_data
+  # Peninsula_data
+  # SWASH_data
+  SWASH_data:
+    description: SWASH_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/SWASH_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"
+  WOA_data:
+    description: WOA_data
+    driver: netcdf
+    args:
+      urlpath: "{{ CATALOG_DIR }}/data/WOA_data/*.nc"
+      chunks: {}
+      xarray_kwargs:
+        engine: "netcdf4"