Merge branch 'master' into tutorials

Author: erikvansebille

parcels/examples/example_dask_chunk_OCMs.py

@@ -121,6 +121,23 @@ def fieldset_from_ofam(chunk_mode):
     return FieldSet.from_netcdf(filenames, variables, dimensions, allow_time_extrapolation=True, field_chunksize=chs, chunkdims_name_map=name_map)
 
 
+def fieldset_from_mitgcm(chunk_mode):
+    data_path = path.join(path.dirname(__file__), "MITgcm_example_data/")
+    filenames = {"U": data_path + "mitgcm_UV_surface_zonally_reentrant.nc",
+                 "V": data_path + "mitgcm_UV_surface_zonally_reentrant.nc"}
+    variables = {"U": "UVEL", "V": "VVEL"}
+    dimensions = {"U": {"lon": "XG", "lat": "YG", "time": "time"},
+                  "V": {"lon": "XG", "lat": "YG", "time": "time"}}
+
+    chs = False
+    name_map = {'lon': 'XG', 'lat': 'YG', 'time': 'time'}
+    if chunk_mode == 'auto':
+        chs = 'auto'
+    elif chunk_mode == 'specific':
+        chs = (1, 50, 100)
+    return FieldSet.from_mitgcm(filenames, variables, dimensions, mesh='flat', field_chunksize=chs, chunkdims_name_map=name_map)
+
+
 def compute_nemo_particle_advection(field_set, mode, lonp, latp):
 
     def periodicBC(particle, fieldSet, time):
@@ -300,6 +317,27 @@ def test_ofam_3D(mode, chunk_mode):
     assert(abs(pset[0].lat - 11) < 1)
 
 
+@pytest.mark.parametrize('mode', ['jit'])
+@pytest.mark.parametrize('chunk_mode', [False, 'auto', 'specific'])
+def test_mitgcm(mode, chunk_mode):
+    if chunk_mode == 'auto':
+        dask.config.set({'array.chunk-size': '1024KiB'})
+    else:
+        dask.config.set({'array.chunk-size': '128MiB'})
+    field_set = fieldset_from_mitgcm(chunk_mode)
+    lons, lats = 5e5, 5e5
+
+    pset = ParticleSet.from_list(fieldset=field_set, pclass=ptype[mode], lon=lons, lat=lats)
+    pset.execute(AdvectionRK4, runtime=delta(days=1), dt=delta(minutes=5))
+    # MITgcm sample file dimensions: time=10, XG=400, YG=200
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.V.grid.load_chunk))
+    if chunk_mode in [False, 'auto']:
+        assert (len(field_set.U.grid.load_chunk) == 1)
+    elif chunk_mode == 'specific':
+        assert (len(field_set.U.grid.load_chunk) == (1 * int(math.ceil(400.0/50.0)) * int(math.ceil(200.0/100.0))))
+    assert np.allclose(pset[0].lon, 5.27e5, atol=1e3)
+
+
 @pytest.mark.parametrize('mode', ['jit'])
 def test_diff_entry_dimensions_chunks(mode):
     data_path = path.join(path.dirname(__file__), 'NemoNorthSeaORCA025-N006_data/')

parcels/examples/example_nemo_curvilinear.py

@@ -29,7 +29,7 @@ def run_nemo_curvilinear(mode, outfile, advtype='RK4'):
                        'data': data_path + 'V_purely_zonal-ORCA025_grid_V.nc4'}}
     variables = {'U': 'U', 'V': 'V'}
     dimensions = {'lon': 'glamf', 'lat': 'gphif'}
-    field_chunksize = {'y': 2, 'x': 2}
+    field_chunksize = {'y': 256, 'x': 512}
     field_set = FieldSet.from_nemo(filenames, variables, dimensions, field_chunksize=field_chunksize)
     assert field_set.U.field_chunksize == field_chunksize
 

parcels/field.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/fieldfilebuffer.py

@@ -17,7 +17,7 @@ def add_grid(self, field):
         grid = field.grid
         existing_grid = False
         for g in self.grids:
-            if field.field_chunksize != grid.master_chunksize:
+            if field.field_chunksize != grid.master_chunksize and grid.master_chunksize not in [None, False]:
                 logger.warning_once("Field chunksize and Grid master chunksize are not equal - erroneous behaviour expected.")
                 break
             if g == grid:

parcels/fieldset.py

@@ -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");