Compatability updates for parcelsv4

benchmarks/fesom2.py

@@ -4,6 +4,7 @@
     FieldSet,
     Particle,
     ParticleSet,
+    StatusCode,
     convert,
 )
 from parcels.kernels import AdvectionRK2_3D
@@ -14,6 +15,12 @@
 dt = np.timedelta64(2400, "s")
 
 
+def _delete_error_particles(particles, fieldset):
+    """Delete any particle that errors out (e.g. runs aground or off the grid)"""
+    any_error = particles.state >= 50
+    particles[any_error].state = StatusCode.Delete
+
+
 def _load_ds():
     """Helper function to load uxarray dataset from datapath"""
 
@@ -39,12 +46,15 @@ def pset_execute(self, npart, integrator):
         ds = _load_ds()
         ds = convert.fesom_to_ugrid(ds)
         fieldset = FieldSet.from_ugrid_conventions(ds)
+        fieldset = fieldset.to_windowed_arrays()
 
         lon = np.linspace(2.0, 15.0, npart)
         lat = np.linspace(32.0, 19.0, npart)
 
-        pset = ParticleSet(fieldset=fieldset, pclass=Particle, lon=lon, lat=lat)
-        pset.execute(kernels=integrator, runtime=runtime, dt=dt)
+        pset = ParticleSet(fieldset=fieldset, pclass=Particle, x=lon, y=lat)
+        pset.execute(
+            kernels=[integrator, _delete_error_particles], runtime=runtime, dt=dt
+        )
 
     def time_pset_execute(self, npart, integrator):
         self.pset_execute(npart, integrator)

benchmarks/moi_curvilinear.py

@@ -1,13 +1,18 @@
 import numpy as np
 import parcels
-from parcels.interpolators import XLinear
 
 from .catalogs import Catalogs
 
 runtime = np.timedelta64(2, "D")
 dt = np.timedelta64(15, "m")
 
 
+def _delete_error_particles(particles, fieldset):
+    """Delete any particle that errors out (e.g. runs aground or off the grid)"""
+    any_error = particles.state >= 50
+    particles[any_error].state = parcels.StatusCode.Delete
+
+
 def _load_ds(chunk):
     """Helper function to load xarray dataset from catalog with or without chunking"""
     cat = Catalogs.CAT_BENCHMARKS
@@ -19,6 +24,11 @@ def _load_ds(chunk):
     ds_v = (
         cat.moi_v(chunks=chunks).to_dask()[["vomecrty"]].rename_vars({"vomecrty": "V"})
     )
+
+    # undo the .to_dask float32 -> float64 conversion
+    ds_u["U"] = ds_u["U"].astype("float32")
+    ds_v["V"] = ds_v["V"].astype("float32")
+
     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
@@ -31,49 +41,45 @@ class MOICurvilinear:
     """Mercator Ocean International model data based benchmark on a curvilinear grid."""
 
     params = (
-        ["XLinear"],
         [256],
         [10000],
     )
     param_names = [
-        "interpolator",
         "chunk",
         "npart",
     ]
 
-    def time_load_data_3d(self, interpolator, chunk, npart):
+    def time_load_data_3d(self, 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(chunk)
-        for j in range(min(ds.coords["deptht"].size, 2)):
+        for j in range(min(ds.coords["depth_center"].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()
+                _u = ds["U"].isel(depth_center=j, time=i).compute()
+                _v = ds["V"].isel(depth_center=j, time=i).compute()
 
-    def pset_execute_3d(self, interpolator, chunk, npart):
+    def pset_execute_3d(self, chunk, npart):
         ds = _load_ds(chunk)
         fieldset = parcels.FieldSet.from_sgrid_conventions(ds)
-        if interpolator == "XLinear":
-            fieldset.U.interp_method = XLinear
-            fieldset.V.interp_method = XLinear
-        else:
-            raise ValueError(f"Unknown interpolator: {interpolator}")
-
+        fieldset = fieldset.to_windowed_arrays()
         pclass = parcels.Particle
 
         lon = np.linspace(-10, 10, npart)
         lat = np.linspace(-30, -20, npart)
 
-        pset = parcels.ParticleSet(fieldset=fieldset, pclass=pclass, lon=lon, lat=lat)
+        pset = parcels.ParticleSet(fieldset=fieldset, pclass=pclass, x=lon, y=lat)
 
         pset.execute(
-            parcels.kernels.AdvectionEE, runtime=runtime, dt=dt, verbose_progress=False
+            [parcels.kernels.AdvectionEE, _delete_error_particles],
+            runtime=runtime,
+            dt=dt,
+            verbose_progress=False,
         )
 
-    def time_pset_execute_3d(self, interpolator, chunk, npart):
-        self.pset_execute_3d(interpolator, chunk, npart)
+    def time_pset_execute_3d(self, chunk, npart):
+        self.pset_execute_3d(chunk, npart)
 
-    def peakmem_pset_execute_3d(self, interpolator, chunk, npart):
-        self.pset_execute_3d(interpolator, chunk, npart)
+    def peakmem_pset_execute_3d(self, chunk, npart):
+        self.pset_execute_3d(chunk, npart)