Fixing review comments

Author: erikvansebille

parcels/_typing.py

@@ -37,7 +37,7 @@ class ParcelsAST(ast.AST):
 Mesh = Literal["spherical", "flat"]  # corresponds with `mesh`
 VectorType = Literal["3D", "3DSigma", "2D"] | None  # corresponds with `vector_type`
 ChunkMode = Literal["auto", "specific", "failsafe"]  # corresponds with `chunk_mode`
-GridIndexingType = Literal["pop", "mom5", "mitgcm", "nemo"]  # corresponds with `gridindexingtype`
+GridIndexingType = Literal["pop", "mom5", "mitgcm", "nemo", "croco"]  # corresponds with `gridindexingtype`
 UpdateStatus = Literal["not_updated", "first_updated", "updated"]  # corresponds with `_update_status`
 TimePeriodic = float | datetime.timedelta | Literal[False]  # corresponds with `time_periodic`
 NetcdfEngine = Literal["netcdf4", "xarray", "scipy"]

parcels/compilation/codegenerator.py

@@ -426,7 +426,7 @@ def __init__(self, fieldset=None, ptype=JITParticle):
         self.fieldset = fieldset
         self.ptype = ptype
         self.field_args = collections.OrderedDict()
-        if isinstance(fieldset.U, Field) and fieldset.U.gridindexingtype == "croco" and hasattr(fieldset, "H"):
+        if fieldset.U.gridindexingtype == "croco" and hasattr(fieldset, "H"):
             self.field_args["H"] = fieldset.H  # CROCO requires H field
         self.vector_field_args = collections.OrderedDict()
         self.const_args = collections.OrderedDict()
@@ -827,7 +827,7 @@ def visit_FieldEvalNode(self, node):
         self.visit(node.args)
         args = self._check_FieldSamplingArguments(node.args.ccode)
         statements_croco = []
-        if "croco" in node.field.obj.gridindexingtype and node.field.obj.name != "H":  # TODO needs to be sigma
+        if "croco" in node.field.obj.gridindexingtype and node.field.obj.name != "H":
             statements_croco.append(
                 c.Assign(
                     "parcels_interp_state",

tests/test_advection.py

@@ -1,5 +1,4 @@
 import math
-import os
 from datetime import timedelta
 
 import numpy as np
@@ -22,6 +21,7 @@
     StatusCode,
     Variable,
 )
+from tests.utils import TEST_DATA
 
 ptype = {"scipy": ScipyParticle, "jit": JITParticle}
 kernel = {
@@ -197,9 +197,7 @@ def test_advection_RK45(lon, lat, mode, rk45_tol):
 
 @pytest.mark.parametrize("mode", ["scipy", "jit"])
 def test_advection_3DCROCO(mode):
-    data_path = os.path.join(os.path.dirname(__file__), "test_data/")
-    fieldset = FieldSet.from_modulefile(data_path + "fieldset_CROCO3D.py")
-    assert fieldset.U.creation_log == "from_croco"
+    fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py")
 
     runtime = 1e4
     X, Z = np.meshgrid([40e3, 80e3, 120e3], [-10, -130])
@@ -218,9 +216,7 @@ def SampleW(particle, fieldset, time):
 
 @pytest.mark.parametrize("mode", ["scipy", "jit"])
 def test_advection_2DCROCO(mode):
-    data_path = os.path.join(os.path.dirname(__file__), "test_data/")
-    fieldset = FieldSet.from_modulefile(data_path + "fieldset_CROCO2D.py")
-    assert fieldset.U.creation_log == "from_croco"
+    fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO2D.py")
 
     runtime = 1e4
     X = np.array([40e3, 80e3, 120e3])

tests/test_data/CROCO_idealized.nc

@@ -4,7 +4,8 @@
 
 
 def create_fieldset(indices=None):
-    file = os.path.join(os.path.dirname(__file__), "CROCO_idealized.nc")
+    example_dataset_folder = parcels.download_example_dataset("CROCOidealized_data")
+    file = os.path.join(example_dataset_folder, "CROCO_idealized.nc")
 
     variables = {"U": "u", "V": "v"}
     dimensions = {
@@ -17,6 +18,7 @@ def create_fieldset(indices=None):
         dimensions,
         allow_time_extrapolation=True,
         mesh="flat",
+        indices=indices,
     )
 
     return fieldset

tests/test_data/fieldset_CROCO2D.py

@@ -4,7 +4,8 @@
 
 
 def create_fieldset(indices=None):
-    file = os.path.join(os.path.dirname(__file__), "CROCO_idealized.nc")
+    example_dataset_folder = parcels.download_example_dataset("CROCOidealized_data")
+    file = os.path.join(example_dataset_folder, "CROCO_idealized.nc")
 
     variables = {"U": "u", "V": "v", "W": "w", "H": "h"}
     dimensions = {
@@ -19,6 +20,7 @@ def create_fieldset(indices=None):
         dimensions,
         allow_time_extrapolation=True,
         mesh="flat",
+        indices=indices,
     )
 
     return fieldset

tests/test_data/fieldset_CROCO3D.py

@@ -618,7 +618,7 @@ def test_sampling_out_of_bounds_time(mode, allow_time_extrapolation):
 
 
 @pytest.mark.parametrize("mode", ["scipy", "jit"])
-def test_sampling_3DCROCO(mode, npart=10):
+def test_sampling_3DCROCO(mode):
     data_path = os.path.join(os.path.dirname(__file__), "test_data/")
     fieldset = FieldSet.from_modulefile(data_path + "fieldset_CROCO3D.py")