Merge pull request #923 from OceanParcels/mom5_gridindexing

Implementing gridindexing for MOM5 B-grids

Author: erikvansebille

parcels/field.py

@@ -570,7 +570,11 @@ def search_indices_vertical_z(self, z):
         z = np.float32(z)
         if grid.depth[-1] > grid.depth[0]:
             if z < grid.depth[0]:
-                raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
+                # Since MOM5 is indexed at cell bottom, allow z at depth[0] - dz where dz = (depth[1] - depth[0])
+                if self.gridindexingtype == "mom5" and z > 2*grid.depth[0] - grid.depth[1]:
+                    return (-1, z / grid.depth[0])
+                else:
+                    raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
             elif z > grid.depth[-1]:
                 raise FieldOutOfBoundError(0, 0, z, field=self)
             depth_indices = grid.depth <= z
@@ -932,7 +936,10 @@ def interpolator3D(self, ti, z, y, x, time):
                 return (1 - zeta) * f0 + zeta * f1
         elif self.interp_method in ['linear', 'bgrid_velocity', 'bgrid_w_velocity']:
             if self.interp_method == 'bgrid_velocity':
-                zeta = 0.
+                if self.gridindexingtype == 'mom5':
+                    zeta = 1.
+                else:
+                    zeta = 0.
             elif self.interp_method == 'bgrid_w_velocity':
                 eta = 1.
                 xsi = 1.
@@ -946,7 +953,11 @@ def interpolator3D(self, ti, z, y, x, time):
                 xsi*(1-eta) * data[yi, xi+1] + \
                 xsi*eta * data[yi+1, xi+1] + \
                 (1-xsi)*eta * data[yi+1, xi]
-            return (1-zeta) * f0 + zeta * f1
+            if self.interp_method == 'bgrid_w_velocity' and self.gridindexingtype == 'mom5' and zi == -1:
+                # Since MOM5 is indexed at cell bottom, allow linear interpolation of W to zero in upper cell)
+                return zeta * f1
+            else:
+                return (1-zeta) * f0 + zeta * f1
         elif self.interp_method in ['cgrid_tracer', 'bgrid_tracer']:
             return self.data[ti, zi, yi+1, xi+1]
         else:

parcels/fieldset.py

@@ -234,8 +234,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']:
-                raise ValueError("Field.gridindexing has to be one of 'nemo' or 'mitgcm'")
+            if U.gridindexingtype not in ['nemo', 'mitgcm', 'mom5']:
+                raise ValueError("Field.gridindexing has to be one of 'nemo', 'mitgcm' or 'mom5")
 
             if U.gridindexingtype == 'mitgcm' and U.grid.gtype in [GridCode.CurvilinearZGrid, GridCode.CurvilinearZGrid]:
                 raise NotImplementedError('Curvilinear Grids are not implemented for mitgcm-style grid indexing.'
@@ -680,6 +680,72 @@ def from_pop(cls, filenames, variables, dimensions, indices=None, mesh='spherica
                 raise SyntaxError("'depth_units' has to be 'm' or 'cm'")
         return fieldset
 
+    @classmethod
+    def from_mom5(cls, filenames, variables, dimensions, indices=None, mesh='spherical',
+                  allow_time_extrapolation=None, time_periodic=False,
+                  tracer_interp_method='bgrid_tracer', field_chunksize='auto', **kwargs):
+        """Initialises FieldSet object from NetCDF files of MOM5 fields.
+
+        :param filenames: Dictionary mapping variables to file(s). The
+               filepath may contain wildcards to indicate multiple files,
+               or be a list of file.
+               filenames can be a list [files], a dictionary {var:[files]},
+               a dictionary {dim:[files]} (if lon, lat, depth and/or data not stored in same files as data),
+               or a dictionary of dictionaries {var:{dim:[files]}}
+               time values are in filenames[data]
+        :param variables: Dictionary mapping variables to variable names in the netCDF file(s).
+               Note that the built-in Advection kernels assume that U and V are in m/s
+        :param dimensions: Dictionary mapping data dimensions (lon,
+               lat, depth, time, data) to dimensions in the netCF file(s).
+               Note that dimensions can also be a dictionary of dictionaries if
+               dimension names are different for each variable.
+               U[k,j+1,i],V[k,j+1,i] ____________________U[k,j+1,i+1],V[k,j+1,i+1]
+               |                                         |
+               |      W[k-1:k+1,j+1,i+1],T[k,j+1,i+1]      |
+               |                                         |
+               U[k,j,i],V[k,j,i] ________________________U[k,j,i+1],V[k,j,i+1]
+               In 2D: U and V nodes are on the cell vertices and interpolated bilinearly as a A-grid.
+                      T node is at the cell centre and interpolated constant per cell as a C-grid.
+               In 3D: U and V nodes are at the midlle of the cell vertical edges,
+                      They are interpolated bilinearly (independently of z) in the cell.
+                      W nodes are at the centre of the horizontal interfaces, but below the U and V.
+                      They are interpolated linearly (as a function of z) in the cell.
+                      Note that W is normally directed upward in MOM5, but Parcels requires W
+                      in the positive z-direction (downward) so W is multiplied by -1.
+                      T node is at the cell centre, and constant per cell.
+        :param indices: Optional dictionary of indices for each dimension
+               to read from file(s), to allow for reading of subset of data.
+               Default is to read the full extent of each dimension.
+               Note that negative indices are not allowed.
+        :param fieldtype: Optional dictionary mapping fields to fieldtypes to be used for UnitConverter.
+               (either 'U', 'V', 'Kh_zonal', 'Kh_meridional' or None)
+        :param mesh: String indicating the type of mesh coordinates and
+               units used during velocity interpolation, see also https://nbviewer.jupyter.org/github/OceanParcels/parcels/blob/master/parcels/examples/tutorial_unitconverters.ipynb:
+
+               1. spherical (default): Lat and lon in degree, with a
+                  correction for zonal velocity U near the poles.
+               2. flat: No conversion, lat/lon are assumed to be in m.
+        :param allow_time_extrapolation: boolean whether to allow for extrapolation
+               (i.e. beyond the last available time snapshot)
+               Default is False if dimensions includes time, else True
+        :param time_periodic: To loop periodically over the time component of the Field. It is set to either False or the length of the period (either float in seconds or datetime.timedelta object). (Default: False)
+               This flag overrides the allow_time_interpolation and sets it to False
+        :param tracer_interp_method: Method for interpolation of tracer fields. It is recommended to use 'bgrid_tracer' (default)
+               Note that in the case of from_mom5() and from_bgrid(), the velocity fields are default to 'bgrid_velocity'
+        :param field_chunksize: size of the chunks in dask loading
+
+
+        """
+
+        if 'creation_log' not in kwargs.keys():
+            kwargs['creation_log'] = 'from_mom5'
+        fieldset = cls.from_b_grid_dataset(filenames, variables, dimensions, mesh=mesh, indices=indices, time_periodic=time_periodic,
+                                           allow_time_extrapolation=allow_time_extrapolation, tracer_interp_method=tracer_interp_method,
+                                           field_chunksize=field_chunksize, gridindexingtype='mom5', **kwargs)
+        if hasattr(fieldset, 'W'):
+            fieldset.W.set_scaling_factor(-1)
+        return fieldset
+
     @classmethod
     def from_b_grid_dataset(cls, filenames, variables, dimensions, indices=None, mesh='spherical',
                             allow_time_extrapolation=None, time_periodic=False,

parcels/include/index_search.h

@@ -23,7 +23,7 @@ typedef enum
 
 typedef enum
   {
-    NEMO = 0, MITGCM = 1
+    NEMO = 0, MITGCM = 1, MOM5 = 2
   } GridIndexingType;
 
 typedef struct
@@ -56,9 +56,14 @@ typedef enum
     RECTILINEAR_Z_GRID=0, RECTILINEAR_S_GRID=1, CURVILINEAR_Z_GRID=2, CURVILINEAR_S_GRID=3
   } GridCode;
 
-static inline StatusCode search_indices_vertical_z(type_coord z, int zdim, float *zvals, int *zi, double *zeta)
+static inline StatusCode search_indices_vertical_z(type_coord z, int zdim, float *zvals, int *zi, double *zeta, int gridindexingtype)
 {
   if (zvals[zdim-1] > zvals[0]){
+    if ((z < zvals[0]) && (gridindexingtype == MOM5) && (z > 2 * zvals[0] - zvals[1])){
+      *zi = -1;
+      *zeta = z / zvals[0];
+      return SUCCESS;
+    }
     if (z < zvals[0]) {return ERROR_THROUGH_SURFACE;}
     if (z > zvals[zdim-1]) {return ERROR_OUT_OF_BOUNDS;}
     while (*zi < zdim-1 && z > zvals[*zi+1]) ++(*zi);
@@ -162,7 +167,8 @@ static inline void reconnect_bnd_indices(int *xi, int *yi, int xdim, int ydim, i
 
 static inline StatusCode search_indices_rectilinear(type_coord x, type_coord y, type_coord z, CStructuredGrid *grid, GridCode gcode,
                                                    int *xi, int *yi, int *zi, double *xsi, double *eta, double *zeta,
-                                                   int ti, double time, double t0, double t1, int interp_method)
+                                                   int ti, double time, double t0, double t1, int interp_method,
+                                                   int gridindexingtype)
 {
   int xdim = grid->xdim;
   int ydim = grid->ydim;
@@ -239,7 +245,7 @@ static inline StatusCode search_indices_rectilinear(type_coord x, type_coord y,
   if (zdim > 1){
     switch(gcode){
       case RECTILINEAR_Z_GRID:
-        status = search_indices_vertical_z(z, zdim, zvals, zi, zeta);
+        status = search_indices_vertical_z(z, zdim, zvals, zi, zeta, gridindexingtype);
         break;
       case RECTILINEAR_S_GRID:
         status = search_indices_vertical_s(z, xdim, ydim, zdim, zvals,
@@ -264,7 +270,8 @@ static inline StatusCode search_indices_rectilinear(type_coord x, type_coord y,
 
 static inline StatusCode search_indices_curvilinear(type_coord x, type_coord y, type_coord z, CStructuredGrid *grid, GridCode gcode,
                                                    int *xi, int *yi, int *zi, double *xsi, double *eta, double *zeta,
-                                                   int ti, double time, double t0, double t1, int interp_method)
+                                                   int ti, double time, double t0, double t1, int interp_method,
+                                                   int gridindexingtype)
 {
   int xi_old = *xi;
   int yi_old = *yi;
@@ -375,7 +382,7 @@ static inline StatusCode search_indices_curvilinear(type_coord x, type_coord y,
   if (zdim > 1){
     switch(gcode){
       case CURVILINEAR_Z_GRID:
-        status = search_indices_vertical_z(z, zdim, zvals, zi, zeta);
+        status = search_indices_vertical_z(z, zdim, zvals, zi, zeta, gridindexingtype);
         break;
       case CURVILINEAR_S_GRID:
         status = search_indices_vertical_s(z, xdim, ydim, zdim, zvals,
@@ -402,18 +409,19 @@ static inline StatusCode search_indices_curvilinear(type_coord x, type_coord y,
  * */
 static inline StatusCode search_indices(type_coord x, type_coord y, type_coord z, CStructuredGrid *grid,
                                        int *xi, int *yi, int *zi, double *xsi, double *eta, double *zeta,
-                                       GridCode gcode, int ti, double time, double t0, double t1, int interp_method)
+                                       GridCode gcode, int ti, double time, double t0, double t1, int interp_method,
+                                       int gridindexingtype)
 {
   switch(gcode){
     case RECTILINEAR_Z_GRID:
     case RECTILINEAR_S_GRID:
       return search_indices_rectilinear(x, y, z, grid, gcode, xi, yi, zi, xsi, eta, zeta,
-                                   ti, time, t0, t1, interp_method);
+                                   ti, time, t0, t1, interp_method, gridindexingtype);
       break;
     case CURVILINEAR_Z_GRID:
     case CURVILINEAR_S_GRID:
       return search_indices_curvilinear(x, y, z, grid, gcode, xi, yi, zi, xsi, eta, zeta,
-                                   ti, time, t0, t1, interp_method);
+                                   ti, time, t0, t1, interp_method, gridindexingtype);
       break;
     default:
       printf("Only RECTILINEAR_Z_GRID, RECTILINEAR_S_GRID, CURVILINEAR_Z_GRID and CURVILINEAR_S_GRID grids are currently implemented\n");

parcels/include/parcels.h

@@ -147,6 +147,19 @@ static inline StatusCode spatial_interpolation_trilinear(double xsi, double eta,
   return SUCCESS;
 }
 
+/* Trilinear interpolation routine for MOM surface 3D grid */
+static inline StatusCode spatial_interpolation_trilinear_surface(double xsi, double eta, double zeta,
+                                                                 float data[2][2][2], float *value)
+{
+  float f1;
+  f1 = (1-xsi)*(1-eta) * data[0][0][0]
+     +    xsi *(1-eta) * data[0][0][1]
+     +    xsi *   eta  * data[0][1][1]
+     + (1-xsi)*   eta  * data[0][1][0];
+  *value = zeta * f1;
+  return SUCCESS;
+}
+
 /* Trilinear interpolation routine for 3D grid for tracers with squared inverse distance weighting near land*/
 static inline StatusCode spatial_interpolation_trilinear_invdist_land(double xsi, double eta, double zeta, float data[2][2][2], float *value)
 {
@@ -475,15 +488,19 @@ static inline StatusCode temporal_interpolation_structured_grid(type_coord x, ty
 
   status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid],
 			  &xsi, &eta, &zeta, gcode, ti[igrid],
-			  tsrch, t0, t1, interp_method);
+			  tsrch, t0, t1, interp_method, gridindexingtype);
   CHECKSTATUS(status);
 
   if (grid->zdim == 1) {
     // last param is a flag, which denotes that we only want the first timestep
     // (rather than both)
     status = getCell2D(f, xi[igrid], yi[igrid], ti[igrid], data2D, tii == 1); CHECKSTATUS(status);
   } else {
-    status = getCell3D(f, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D, tii == 1); CHECKSTATUS(status);
+    if ((gridindexingtype == MOM5) && (zi[igrid] == -1)){
+      status = getCell3D(f, xi[igrid], yi[igrid], 0, ti[igrid], data3D, tii == 1); CHECKSTATUS(status);
+    } else {
+      status = getCell3D(f, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D, tii == 1); CHECKSTATUS(status);
+    }
   }
 
   // define a helper macro that will select the appropriate interpolation method
@@ -507,7 +524,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) {
+      if ((gridindexingtype == NEMO)   || (gridindexingtype == MOM5)) {
         // velocity is on the northeast of a tracer cell
         xsi = 1;
         eta = 1;
@@ -517,10 +534,17 @@ static inline StatusCode temporal_interpolation_structured_grid(type_coord x, ty
         eta = 0;
       }
     } else if (interp_method == BGRID_VELOCITY) {
-      zeta = 0;
+      if (gridindexingtype == MOM5) {
+        zeta = 1;
+      } else {
+        zeta = 0;
+      }
+    }
+    if ((gridindexingtype == MOM5) && (zi[igrid] == -1)){
+      INTERP(spatial_interpolation_bilinear, spatial_interpolation_trilinear_surface);
+    } else {
+      INTERP(spatial_interpolation_bilinear, spatial_interpolation_trilinear);
     }
-
-    INTERP(spatial_interpolation_bilinear, spatial_interpolation_trilinear);
   } else if (interp_method == NEAREST) {
     INTERP(spatial_interpolation_nearest2D, spatial_interpolation_nearest3D);
   } else if ((interp_method == CGRID_TRACER) || (interp_method == BGRID_TRACER)) {
@@ -652,7 +676,7 @@ static inline StatusCode temporal_interpolationUV_c_grid(type_coord x, type_coor
     float u0, u1, v0, v1;
     double t0 = grid->time[ti[igrid]]; double t1 = grid->time[ti[igrid]+1];
     /* Identify grid cell to sample through local linear search */
-    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zeta, gcode, ti[igrid], time, t0, t1, CGRID_VELOCITY); CHECKSTATUS(status);
+    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zeta, gcode, 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) {
@@ -684,7 +708,7 @@ static inline StatusCode temporal_interpolationUV_c_grid(type_coord x, type_coor
     return SUCCESS;
   } else {
     double t0 = grid->time[ti[igrid]];
-    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zeta, gcode, ti[igrid], t0, t0, t0+1, CGRID_VELOCITY); CHECKSTATUS(status);
+    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zeta, gcode, 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) {
@@ -873,7 +897,7 @@ static inline StatusCode temporal_interpolationUVW_c_grid(type_coord x, type_coo
     float u0, u1, v0, v1, w0, w1;
     double t0 = grid->time[ti[igrid]]; double t1 = grid->time[ti[igrid]+1];
     /* Identify grid cell to sample through local linear search */
-    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zet, gcode, ti[igrid], time, t0, t1, CGRID_VELOCITY); CHECKSTATUS(status);
+    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zet, gcode, ti[igrid], time, t0, t1, CGRID_VELOCITY, gridindexingtype); CHECKSTATUS(status);
     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);
     status = getCell3D(W, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D_W, 0); CHECKSTATUS(status);
@@ -889,7 +913,7 @@ static inline StatusCode temporal_interpolationUVW_c_grid(type_coord x, type_coo
     return SUCCESS;
   } else {
     double t0 = grid->time[ti[igrid]];
-    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zet, gcode, ti[igrid], t0, t0, t0+1, CGRID_VELOCITY); CHECKSTATUS(status);
+    status = search_indices(x, y, z, grid, &xi[igrid], &yi[igrid], &zi[igrid], &xsi, &eta, &zet, gcode, ti[igrid], t0, t0, t0+1, CGRID_VELOCITY, gridindexingtype); CHECKSTATUS(status);
     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);
     status = getCell3D(W, xi[igrid], yi[igrid], zi[igrid], ti[igrid], data3D_W, 1); CHECKSTATUS(status);