Start implementing fieldset.from_croco()

Also added idealised unit test

parcels/application_kernels/advection.py

@@ -4,7 +4,7 @@
 from parcels.tools.statuscodes import StatusCode
 
 __all__ = ['AdvectionRK4', 'AdvectionEE', 'AdvectionRK45', 'AdvectionRK4_3D',
-           'AdvectionAnalytical']
+           'AdvectionAnalytical', 'AdvectionRK4_3D_CROCO']
 
 
 def AdvectionRK4(particle, fieldset, time):
@@ -46,7 +46,7 @@ def AdvectionRK4_3D(particle, fieldset, time):
     particle_ddepth += (w1 + 2*w2 + 2*w3 + w4) / 6. * particle.dt  # noqa
 
 
-def AdvectionRK4_3DCroco(particle, fieldset, time):
+def AdvectionRK4_3D_CROCO(particle, fieldset, time):
     """Advection of particles using fourth-order Runge-Kutta integration including vertical velocity.
     This kernel assumes the vertical velocity is the 'w' field from CROCO output and works on sigma-layers.
     """

parcels/field.py

@@ -1520,7 +1520,7 @@ def __init__(self, name, U, V, W=None):
             assert self._check_grid_dimensions(U.grid, V.grid), (
                 'Dimensions of U and V are not the same.')
             if self.vector_type == '3D':
-                assert self.W.interp_method == 'cgrid_velocity', (
+                assert (self.W.interp_method == 'cgrid_velocity') or (self.gridindexingtype == 'croco'), (
                     'Interpolation methods of U and W are not the same.')
                 assert self._check_grid_dimensions(U.grid, W.grid), (
                     'Dimensions of U and W are not the same.')
@@ -1572,7 +1572,7 @@ def spatial_c_grid_interpolation2D(self, ti, z, y, x, time, particle=None, apply
         c3 = self.dist(px[2], px[3], py[2], py[3], grid.mesh, np.dot(i_u.phi2D_lin(xsi, 1.), py))
         c4 = self.dist(px[3], px[0], py[3], py[0], grid.mesh, np.dot(i_u.phi2D_lin(0., eta), py))
         if grid.zdim == 1:
-            if self.gridindexingtype == 'nemo':
+            if self.gridindexingtype in ['nemo', 'croco']:
                 U0 = self.U.data[ti, yi+1, xi] * c4
                 U1 = self.U.data[ti, yi+1, xi+1] * c2
                 V0 = self.V.data[ti, yi, xi+1] * c1
@@ -1583,7 +1583,7 @@ def spatial_c_grid_interpolation2D(self, ti, z, y, x, time, particle=None, apply
                 V0 = self.V.data[ti, yi, xi] * c1
                 V1 = self.V.data[ti, yi + 1, xi] * c3
         else:
-            if self.gridindexingtype == 'nemo':
+            if self.gridindexingtype in ['nemo', 'croco']:
                 U0 = self.U.data[ti, zi, yi+1, xi] * c4
                 U1 = self.U.data[ti, zi, yi+1, xi+1] * c2
                 V0 = self.V.data[ti, zi, yi, xi+1] * c1

parcels/fieldset.py

@@ -247,8 +247,8 @@ def check_velocityfields(U, V, W):
                 if U.grid.xdim == 1 or U.grid.ydim == 1 or V.grid.xdim == 1 or V.grid.ydim == 1:
                     raise NotImplementedError('C-grid velocities require longitude and latitude dimensions at least length 2')
 
-            if U.gridindexingtype not in ['nemo', 'mitgcm', 'mom5', 'pop']:
-                raise ValueError("Field.gridindexing has to be one of 'nemo', 'mitgcm', 'mom5' or 'pop'")
+            if U.gridindexingtype not in ['nemo', 'mitgcm', 'mom5', 'pop', 'croco']:
+                raise ValueError("Field.gridindexing has to be one of 'nemo', 'mitgcm', 'mom5', 'pop' or 'croco'")
 
             if V.gridindexingtype != U.gridindexingtype or (W and W.gridindexingtype != U.gridindexingtype):
                 raise ValueError('Not all velocity Fields have the same gridindexingtype')
@@ -579,6 +579,32 @@ def from_mitgcm(cls, filenames, variables, dimensions, indices=None, mesh='spher
                                            chunksize=chunksize, gridindexingtype='mitgcm', **kwargs)
         return fieldset
 
+    @classmethod
+    def from_croco(cls, filenames, variables, dimensions, indices=None, mesh='spherical',
+                   allow_time_extrapolation=None, time_periodic=False,
+                   tracer_interp_method='cgrid_tracer', chunksize=None, **kwargs):
+        """Initialises FieldSet object from NetCDF files of CROCO fields.
+        All parameters and keywords are exactly the same as for FieldSet.from_nemo().
+        """  # TODO expand this docstring
+        if 'creation_log' not in kwargs.keys():
+            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:
+            raise ValueError("FieldSet.from_croco() requires a field 'h' for the bathymetry")
+
+        interp_method = {}
+        for v in variables:
+            if v in ['U', 'V']:
+                interp_method[v] = 'cgrid_velocity'
+            elif v in ['W', 'h']:
+                interp_method[v] = 'linear'
+
+        fieldset = cls.from_netcdf(filenames, variables, dimensions, mesh=mesh, indices=indices, time_periodic=time_periodic,
+                                   allow_time_extrapolation=allow_time_extrapolation, interp_method=interp_method,
+                                   chunksize=chunksize, gridindexingtype='croco', **kwargs)
+        return fieldset
+
     @classmethod
     def from_c_grid_dataset(cls, filenames, variables, dimensions, indices=None, mesh='spherical',
                             allow_time_extrapolation=None, time_periodic=False,

parcels/include/index_search.h

@@ -26,7 +26,7 @@ typedef enum
 
 typedef enum
   {
-    NEMO = 0, MITGCM = 1, MOM5 = 2, POP = 3
+    NEMO = 0, MITGCM = 1, MOM5 = 2, POP = 3, CROCO = 4
   } GridIndexingType;
 
 typedef struct

parcels/include/parcels.h

@@ -505,7 +505,7 @@ static inline StatusCode temporal_interpolation_structured_grid(type_coord x, ty
       (interp_method == BGRID_VELOCITY) || (interp_method == BGRID_W_VELOCITY)) {
     // adjust the normalised coordinate for flux-based interpolation methods
     if ((interp_method == CGRID_VELOCITY) || (interp_method == BGRID_W_VELOCITY)) {
-      if ((gridindexingtype == NEMO)   || (gridindexingtype == MOM5) || (gridindexingtype == POP)) {
+      if ((gridindexingtype == NEMO) || (gridindexingtype == MOM5) || (gridindexingtype == POP) || (gridindexingtype == CROCO)) {
         // velocity is on the northeast of a tracer cell
         xsi = 1;
         eta = 1;
@@ -666,7 +666,7 @@ static inline StatusCode temporal_interpolationUV_c_grid(type_coord x, type_coor
     status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zeta, gtype, ti[igrid], time, t0, t1, CGRID_VELOCITY, gridindexingtype); CHECKSTATUS(status);
     if (grid->zdim==1){
       float data2D_U[2][2][2], data2D_V[2][2][2];
-      if (gridindexingtype == NEMO) {
+      if ((gridindexingtype == NEMO)  || (gridindexingtype == CROCO)){
         status = getCell2D(U, xi[igrid], yi[igrid], ti[igrid], data2D_U, 0); CHECKSTATUS(status);
         status = getCell2D(V, xi[igrid], yi[igrid], ti[igrid], data2D_V, 0); CHECKSTATUS(status);
       }
@@ -679,7 +679,7 @@ static inline StatusCode temporal_interpolationUV_c_grid(type_coord x, type_coor
 
     } else {
       float data3D_U[2][2][2][2], data3D_V[2][2][2][2];
-      if (gridindexingtype == NEMO) {
+      if ((gridindexingtype == NEMO)  || (gridindexingtype == CROCO)){
         status = getCell3D(U, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D_U, 0); CHECKSTATUS(status);
         status = getCell3D(V, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D_V, 0); CHECKSTATUS(status);
       }
@@ -698,7 +698,7 @@ static inline StatusCode temporal_interpolationUV_c_grid(type_coord x, type_coor
     status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zeta, gtype, ti[igrid], t0, t0, t0+1, CGRID_VELOCITY, gridindexingtype); CHECKSTATUS(status);
     if (grid->zdim==1){
       float data2D_U[2][2][2], data2D_V[2][2][2];
-      if (gridindexingtype == NEMO) {
+      if ((gridindexingtype == NEMO)  || (gridindexingtype == CROCO)){
         status = getCell2D(U, xi[igrid], yi[igrid], ti[igrid], data2D_U, 1); CHECKSTATUS(status);
         status = getCell2D(V, xi[igrid], yi[igrid], ti[igrid], data2D_V, 1); CHECKSTATUS(status);
       }
@@ -710,7 +710,7 @@ static inline StatusCode temporal_interpolationUV_c_grid(type_coord x, type_coor
     }
     else{
       float data3D_U[2][2][2][2], data3D_V[2][2][2][2];
-      if (gridindexingtype == NEMO) {
+      if ((gridindexingtype == NEMO)  || (gridindexingtype == CROCO)){
         status = getCell3D(U, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D_U, 1); CHECKSTATUS(status);
         status = getCell3D(V, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D_V, 1); CHECKSTATUS(status);
       }

parcels/kernel.py

@@ -27,6 +27,7 @@
 from parcels.application_kernels.advection import (
     AdvectionAnalytical,
     AdvectionRK4_3D,
+    AdvectionRK4_3D_CROCO,
     AdvectionRK45,
 )
 from parcels.compilation.codegenerator import KernelGenerator, LoopGenerator
@@ -161,6 +162,10 @@ def __init__(self, fieldset, ptype, pyfunc=None, funcname=None, funccode=None, p
         # Derive meta information from pyfunc, if not given
         self.check_fieldsets_in_kernels(pyfunc)
 
+        if (pyfunc is AdvectionRK4_3D) and fieldset.U.gridindexingtype == 'croco':
+            pyfunc = AdvectionRK4_3D_CROCO
+            self.funcname = "AdvectionRK4_3D_CROCO"
+
         if funcvars is not None:
             self.funcvars = funcvars
         elif hasattr(pyfunc, '__code__'):

tests/test_advection.py

@@ -1,4 +1,5 @@
 import math
+import os
 from datetime import timedelta
 
 import numpy as np
@@ -186,6 +187,22 @@ def test_advection_RK45(lon, lat, mode, rk45_tol, npart=10):
     print(fieldset.RK45_tol)
 
 
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+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])
+    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)
+    assert np.allclose(pset.depth, Z.flatten(), atol=5)  # TODO lower this atol
+
+
 def periodicfields(xdim, ydim, uvel, vvel):
     dimensions = {'lon': np.linspace(0., 1., xdim+1, dtype=np.float32)[1:],  # don't include both 0 and 1, for periodic b.c.
                   'lat': np.linspace(0., 1., ydim+1, dtype=np.float32)[1:]}

tests/test_data/fieldset_CROCO3D.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', 'W': 'w', 'h': 'h'}
+    dimensions = {'U': {'lon': 'x_rho', 'lat': 'y_rho', 'depth': 's_w', 'time': 'time'},
+                  'V': {'lon': 'x_rho', 'lat': 'y_rho', 'depth': 's_w', 'time': 'time'},
+                  'W': {'lon': 'x_rho', 'lat': 'y_rho', 'depth': 's_w', 'time': 'time'},
+                  'h': {'lon': 'x_rho', 'lat': 'y_rho'}}
+    indices = {'lon': range(61),
+               'lat': range(51),
+               'depth': range(10)}  # TODO make this work under-the-hood in Parcels
+    fieldset = parcels.FieldSet.from_croco(file, variables, dimensions,
+                                           allow_time_extrapolation=True,
+                                           indices=indices,
+                                           mesh='flat')
+
+    return fieldset