Adding support for 2D advection in from_croco

OceanParcels · Aug 7, 2024 · 9fefe5b · 9fefe5b
1 parent 0ced1ec
commit 9fefe5b
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 9 deletions.
diff --git a/parcels/fieldset.py b/parcels/fieldset.py
@@ -590,7 +590,9 @@ def from_croco(cls, filenames, variables, dimensions, indices=None, mesh='spheri
             kwargs['creation_log'] = 'from_croco'
         if kwargs.pop('gridindexingtype', 'croco') != 'croco':
             raise ValueError("gridindexingtype must be 'croco' in FieldSet.from_croco(). Use FieldSet.from_c_grid_dataset otherwise")
-        if 'h' not in variables:
+
+        dimsU = dimensions['U'] if 'U' in dimensions else dimensions
+        if ('depth' in dimsU) and ('h' not in variables):
             raise ValueError("FieldSet.from_croco() requires a field 'h' for the bathymetry")
 
         interp_method = {}

diff --git a/parcels/include/index_search.h b/parcels/include/index_search.h
@@ -105,9 +105,9 @@ static inline StatusCode search_indices_vertical_z(type_coord z, int zdim, float
       *zeta = z / zvals[0];
       return SUCCESS;
     }
-    if ((z > zvals[-1]) && (gridindexingtype == CROCO)){
-      *zi = -2;
-      *zeta = 0;
+    if ((z > zvals[zdim-1]) && (gridindexingtype == CROCO)){
+      *zi = zdim-2;
+      *zeta = 1;
       return SUCCESS;
     }
     if (z < zvals[0]) {return ERRORTHROUGHSURFACE;}

diff --git a/tests/test_advection.py b/tests/test_advection.py
@@ -191,16 +191,33 @@ def test_advection_RK45(lon, lat, mode, rk45_tol, npart=10):
 def test_advection_3DCROCO(mode, npart=10):
     data_path = os.path.join(os.path.dirname(__file__), 'test_data/')
     fieldset = FieldSet.from_modulefile(data_path + 'fieldset_CROCO3D.py')
-    print(fieldset.U.creation_log)
     assert fieldset.U.creation_log == 'from_croco'
 
-    # X, Z = np.meshgrid([40e3, 80e3, 120e3], [-10 -250, -400, -850, -1400, -1550])
-    X, Z = np.meshgrid([40e3], [-130])
+    runtime = 1e4
+    X, Z = np.meshgrid([40e3, 80e3, 120e3], [-10, -130])
     Y = np.ones(X.size) * 100e3
     pset = ParticleSet(fieldset=fieldset, pclass=ptype[mode], lon=X, lat=Y, depth=Z)
 
-    pset.execute([AdvectionRK4_3D], runtime=1e4, dt=100)
+    pset.execute([AdvectionRK4_3D], runtime=runtime, dt=100)
     assert np.allclose(pset.depth, Z.flatten(), atol=5)  # TODO lower this atol
+    assert np.allclose(pset.lon_nextloop, [x+runtime for x in X.flatten()], atol=1e-3)
+
+
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+def test_advection_2DCROCO(mode, npart=10):
+    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'
+
+    runtime = 1e4
+    X = np.array([40e3, 80e3, 120e3])
+    Y = np.ones(X.size) * 100e3
+    Z = np.zeros(X.size)
+    pset = ParticleSet(fieldset=fieldset, pclass=ptype[mode], lon=X, lat=Y, depth=Z)
+
+    pset.execute([AdvectionRK4], runtime=runtime, dt=100)
+    assert np.allclose(pset.depth, Z.flatten(), atol=1e-3)
+    assert np.allclose(pset.lon_nextloop, [x+runtime for x in X], atol=1e-3)
 
 
 def periodicfields(xdim, ydim, uvel, vvel):

diff --git a/tests/test_data/fieldset_CROCO2D.py b/tests/test_data/fieldset_CROCO2D.py
@@ -0,0 +1,22 @@
+import os
+
+import parcels
+
+
+def create_fieldset(indices=None):
+    file = os.path.join(os.path.dirname(__file__), "CROCO_idealized.nc")
+
+    variables = {"U": "u", "V": "v"}
+    dimensions = {
+        "U": {"lon": "x_rho", "lat": "y_rho", "time": "time"},
+        "V": {"lon": "x_rho", "lat": "y_rho", "time": "time"},
+    }
+    fieldset = parcels.FieldSet.from_croco(
+        file,
+        variables,
+        dimensions,
+        allow_time_extrapolation=True,
+        mesh="flat",
+    )
+
+    return fieldset
diff --git a/tests/test_data/fieldset_CROCO3D.py b/tests/test_data/fieldset_CROCO3D.py
@@ -18,7 +18,6 @@ def create_fieldset(indices=None):
         variables,
         dimensions,
         allow_time_extrapolation=True,
-        indices=indices,
         mesh="flat",
     )