Update fieldset ingestion to use convert modules (#40)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: VeckoTheGecko

Parcels

@@ -1,6 +1,5 @@
 import numpy as np
 import uxarray as ux
-import xarray as xr
 from parcels import (
     FieldSet,
     Particle,
@@ -9,7 +8,7 @@
 )
 from parcels.kernels import AdvectionRK2_3D
 
-from . import PARCELS_BENCHMARKS_DATA_FOLDER
+from .catalogs import Catalogs
 
 runtime = np.timedelta64(1, "D")
 dt = np.timedelta64(2400, "s")
@@ -18,12 +17,9 @@
 def _load_ds():
     """Helper function to load uxarray dataset from datapath"""
 
-    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"
-    )
+    cat = Catalogs.CAT_BENCHMARKS
+    grid_file = cat.fesom_baroclinic_gyre_mesh().to_dask()
+    data_files = cat.fesom_baroclinic_gyre_data().to_dask()
 
     grid = ux.open_grid(grid_file)
     return ux.UxDataset(data_files, uxgrid=grid)
@@ -48,7 +44,7 @@ def pset_execute(self, npart, integrator):
         lat = np.linspace(32.0, 19.0, npart)
 
         pset = ParticleSet(fieldset=fieldset, pclass=Particle, lon=lon, lat=lat)
-        pset.execute(runtime=runtime, dt=dt, pyfunc=integrator)
+        pset.execute(kernels=integrator, runtime=runtime, dt=dt)
 
     def time_pset_execute(self, npart, integrator):
         self.pset_execute(npart, integrator)

benchmarks/fesom2.py

@@ -1,62 +1,28 @@
-from glob import glob
-
 import numpy as np
 import parcels
-import xarray as xr
-import xgcm
 from parcels.interpolators import XLinear
 
+from .catalogs import Catalogs
+
 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"""
-
-    fileU = f"{datapath}/psy4v3r1-daily_U_2025-01-0[1-3].nc"
-    filenames = {
-        "U": glob(fileU),
-        "V": glob(fileU.replace("_U_", "_V_")),
-        "W": glob(fileU.replace("_U_", "_W_")),
-    }
-    mesh_mask = f"{datapath}/PSY4V3R1_mesh_hgr.nc"
-    fileargs = {
-        "concat_dim": "time_counter",
-        "combine": "nested",
-        "data_vars": "minimal",
-        "coords": "minimal",
-        "compat": "override",
-    }
-    if chunk:
-        fileargs["chunks"] = {"time_counter": 1, "depth": 2, "y": chunk, "x": chunk}
+def _load_ds(chunk):
+    """Helper function to load xarray dataset from catalog with or without chunking"""
+    cat = Catalogs.CAT_BENCHMARKS
+    chunks = {"time_counter": 1, "depth": 2, "y": chunk, "x": chunk} if chunk else None
 
-    ds_u = xr.open_mfdataset(filenames["U"], **fileargs)[["vozocrtx"]].drop_vars(
-        ["nav_lon", "nav_lat"]
+    ds_u = (
+        cat.moi_u(chunks=chunks).to_dask()[["vozocrtx"]].rename_vars({"vozocrtx": "U"})
     )
-    ds_v = xr.open_mfdataset(filenames["V"], **fileargs)[["vomecrty"]].drop_vars(
-        ["nav_lon", "nav_lat"]
+    ds_v = (
+        cat.moi_v(chunks=chunks).to_dask()[["vomecrty"]].rename_vars({"vomecrty": "V"})
     )
-    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")
-    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
+    da_depth = cat.moi_w(chunks=chunks).to_dask()["depthw"]
+    ds_mesh = cat.moi_mesh(chunks=None).read()[["glamf", "gphif"]].isel(t=0)
+    ds_mesh["depthw"] = da_depth
+    ds = parcels.convert.nemo_to_sgrid(fields=dict(U=ds_u, V=ds_v), coords=ds_mesh)
 
     return ds
 
@@ -75,47 +41,26 @@ class MOICurvilinear:
         "npart",
     ]
 
-    def setup(self, interpolator, chunk, npart):
-        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"""
 
         # To have a reasonable runtime, we only consider the time it takes to load two time levels
         # and two depth levels (at most)
-        ds = _load_ds(self.datapath, chunk)
+        ds = _load_ds(chunk)
         for j in range(min(ds.coords["deptht"].size, 2)):
             for i in range(min(ds.coords["time"].size, 2)):
                 _u = ds["U"].isel(deptht=j, time=i).compute()
                 _v = ds["V"].isel(deptht=j, time=i).compute()
 
     def pset_execute_3d(self, interpolator, chunk, npart):
-        ds = _load_ds(self.datapath, chunk)
-        coords = {
-            "X": {"left": "x"},
-            "Y": {"left": "y"},
-            "Z": {"center": "deptht", "left": "depth"},
-            "T": {"center": "time"},
-        }
-
-        grid = parcels._core.xgrid.XGrid(
-            xgcm.Grid(ds, coords=coords, autoparse_metadata=False, periodic=False),
-            mesh="spherical",
-        )
-
+        ds = _load_ds(chunk)
+        fieldset = parcels.FieldSet.from_sgrid_conventions(ds)
         if interpolator == "XLinear":
-            interp_method = XLinear
+            fieldset.U.interp_method = XLinear
+            fieldset.V.interp_method = XLinear
         else:
             raise ValueError(f"Unknown interpolator: {interpolator}")
 
-        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])
-
         pclass = parcels.Particle
 
         lon = np.linspace(-10, 10, npart)