Merge pull request #660 from OceanParcels/4d_depth_sgrids

4D (time evolving) depth S-grids

parcels/examples/example_dask_chunk_OCMs.py

@@ -78,6 +78,30 @@ def fieldset_from_pop_1arcs(chunk_mode):
     return fieldset
 
 
+def fieldset_from_swash(chunk_mode):
+    filenames = path.join(path.join(path.dirname(__file__), 'SWASH_data'), 'field_*.nc')
+    variables = {'U': 'cross-shore velocity',
+                 'V': 'along-shore velocity',
+                 'W': 'vertical velocity',
+                 'depth': 'time varying depth',
+                 'depth_u': 'time varying depth_u'}
+    dimensions = {'U': {'lon': 'x', 'lat': 'y', 'depth': 'not_yet_set', 'time': 't'},
+                  'V': {'lon': 'x', 'lat': 'y', 'depth': 'not_yet_set', 'time': 't'},
+                  'W': {'lon': 'x', 'lat': 'y', 'depth': 'not_yet_set', 'time': 't'},
+                  'depth': {'lon': 'x', 'lat': 'y', 'depth': 'not_yet_set', 'time': 't'},
+                  'depth_u': {'lon': 'x', 'lat': 'y', 'depth': 'not_yet_set', 'time': 't'}}
+    chs = False
+    if chunk_mode == 'auto':
+        chs = 'auto'
+    elif chunk_mode == 'specific':
+        chs = (1, 7, 4, 4)
+    fieldset = FieldSet.from_netcdf(filenames, variables, dimensions, mesh='flat', allow_time_extrapolation=True, field_chunksize=chs)
+    fieldset.U.set_depth_from_field(fieldset.depth_u)
+    fieldset.V.set_depth_from_field(fieldset.depth_u)
+    fieldset.W.set_depth_from_field(fieldset.depth)
+    return fieldset
+
+
 def compute_nemo_particle_advection(field_set, mode, lonp, latp):
 
     def periodicBC(particle, fieldSet, time):
@@ -111,6 +135,13 @@ def compute_pop_particle_advection(field_set, mode, lonp, latp):
     return pset
 
 
+def compute_swash_particle_advection(field_set, mode, lonp, latp, depthp):
+    pset = ParticleSet.from_list(field_set, ptype[mode], lon=lonp, lat=latp, depth=depthp)
+    pfile = ParticleFile("swash_particles_chunk", pset, outputdt=delta(seconds=0.05))
+    pset.execute(AdvectionRK4, runtime=delta(seconds=0.2), dt=delta(seconds=0.005), output_file=pfile)
+    return pset
+
+
 @pytest.mark.parametrize('mode', ['jit'])
 @pytest.mark.parametrize('chunk_mode', [False, 'auto', 'specific'])
 def test_nemo_3D(mode, chunk_mode):
@@ -157,6 +188,32 @@ def test_pop(mode, chunk_mode):
         assert (len(field_set.U.grid.load_chunk) == (int(math.ceil(21.0/3.0)) * int(math.ceil(60.0/8.0)) * int(math.ceil(60.0/8.0))))
 
 
+@pytest.mark.parametrize('mode', ['jit'])
+@pytest.mark.parametrize('chunk_mode', [False, 'auto', 'specific'])
+def test_swash(mode, chunk_mode):
+    if chunk_mode == 'auto':
+        dask.config.set({'array.chunk-size': '32KiB'})
+    else:
+        dask.config.set({'array.chunk-size': '128MiB'})
+    field_set = fieldset_from_swash(chunk_mode)
+    npart = 20
+    lonp = [i for i in 9.5 + (-0.2 + np.random.rand(npart) * 2.0 * 0.2)]
+    latp = [i for i in np.arange(start=12.3, stop=13.1, step=0.04)[0:20]]
+    depthp = [-0.1, ] * npart
+    compute_swash_particle_advection(field_set, mode, lonp, latp, depthp)
+    # SWASH sample file dimensions: t=1, z=7, z_u=6, y=21, x=51
+    assert (len(field_set.U.grid.load_chunk) == len(field_set.V.grid.load_chunk))
+    if chunk_mode != 'auto':
+        assert (len(field_set.U.grid.load_chunk) == len(field_set.W.grid.load_chunk))
+    if chunk_mode is False:
+        assert (len(field_set.U.grid.load_chunk) == 1)
+    elif chunk_mode == '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(6.0 / 7.0)) * int(math.ceil(21.0 / 4.0)) * int(math.ceil(51.0 / 4.0))))
+        assert (len(field_set.U.grid.load_chunk) == (1 * int(math.ceil(7.0 / 7.0)) * int(math.ceil(21.0 / 4.0)) * int(math.ceil(51.0 / 4.0))))
+
+
 @pytest.mark.parametrize('mode', ['jit'])
 @pytest.mark.parametrize('chunk_mode', [False, 'auto', 'specific'])
 def test_globcurrent_2D(mode, chunk_mode):
@@ -213,7 +270,6 @@ def test_3d_2dfield_sampling(mode):
 
     filenames = {'U': {'lon': mesh_mask, 'lat': mesh_mask, 'data': ufiles},
                  'V': {'lon': mesh_mask, 'lat': mesh_mask, 'data': vfiles},
-                 # 'nav_lon': {'lon': mesh_mask, 'lat': mesh_mask, 'data': ufiles[0]}}
                  'nav_lon': {'lon': mesh_mask, 'lat': mesh_mask, 'data': [ufiles[0], ]}}
     variables = {'U': 'uo',
                  'V': 'vo',
@@ -333,7 +389,8 @@ def test_diff_entry_chunksize_error_nemo_complex_conform_depth(mode):
 
 @pytest.mark.parametrize('mode', ['jit'])
 def test_diff_entry_chunksize_error_nemo_complex_nonconform_depth(mode):
-    # ==== this test is expected to fall-back to a pre-defined minimal chunk as the requested chunks don't match, or throw a value error ==== #
+    # ==== this test is expected to fall-back to a pre-defined minimal chunk as the ==== #
+    # ==== requested chunks don't match, or throw a value error                     ==== #
     data_path = path.join(path.dirname(__file__), 'NemoNorthSeaORCA025-N006_data/')
     ufiles = sorted(glob(data_path + 'ORCA*U.nc'))
     vfiles = sorted(glob(data_path + 'ORCA*V.nc'))
@@ -354,7 +411,9 @@ def test_diff_entry_chunksize_error_nemo_complex_nonconform_depth(mode):
     latp = [i for i in 52.0+(-1e-3+np.random.rand(npart)*2.0*1e-3)]
     try:
         compute_nemo_particle_advection(fieldset, mode, lonp, latp)
-    except IndexError:
+    except IndexError:  # incorrect data access, in case grids were created
+        return True
+    except AssertionError:  # U-V grids are not equal to one another, throwing assertion errors
         return True
     return False
 

parcels/field.py

@@ -168,6 +168,10 @@ def __init__(self, name, data, lon=None, lat=None, depth=None, time=None, grid=N
         self.loaded_time_indices = []
         self.creation_log = kwargs.pop('creation_log', '')
         self.field_chunksize = kwargs.pop('field_chunksize', None)
+        self.grid.depth_field = kwargs.pop('depth_field', None)
+
+        if self.grid.depth_field == 'not_yet_set':
+            assert self.grid.z4d, 'Providing the depth dimensions from another field data is only available for 4d S grids'
 
         # data_full_zdim is the vertical dimension of the complete field data, ignoring the indices.
         # (data_full_zdim = grid.zdim if no indices are used, for A- and C-grids and for some B-grids). It is used for the B-grid,
@@ -286,7 +290,11 @@ def from_netcdf(cls, filenames, variable, dimensions, indices=None, grid=None,
         if 'depth' in dimensions:
             with NetcdfFileBuffer(depth_filename, dimensions, indices, netcdf_engine, interp_method=interp_method, field_chunksize=False, lock_file=False) as filebuffer:
                 filebuffer.name = filebuffer.parse_name(variable[1])
-                depth = filebuffer.read_depth
+                if dimensions['depth'] == 'not_yet_set':
+                    depth = filebuffer.read_depth_dimensions
+                    kwargs['depth_field'] = 'not_yet_set'
+                else:
+                    depth = filebuffer.read_depth
                 data_full_zdim = filebuffer.data_full_zdim
         else:
             indices['depth'] = [0]
@@ -494,6 +502,9 @@ def set_scaling_factor(self, factor):
         if not self.grid.defer_load:
             self.data *= factor
 
+    def set_depth_from_field(self, field):
+        self.grid.depth_field = field
+
     def __getitem__(self, key):
         return self.eval(*key)
 
@@ -528,15 +539,26 @@ def cell_areas(self):
     def search_indices_vertical_z(self, z):
         grid = self.grid
         z = np.float32(z)
-        if z < grid.depth[0]:
-            raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
-        elif z > grid.depth[-1]:
-            raise FieldOutOfBoundError(0, 0, z, field=self)
-        depth_index = grid.depth <= z
-        if z >= grid.depth[-1]:
-            zi = len(grid.depth) - 2
+        if grid.depth[-1] > grid.depth[0]:
+            if z < grid.depth[0]:
+                raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
+            elif z > grid.depth[-1]:
+                raise FieldOutOfBoundError(0, 0, z, field=self)
+            depth_indices = grid.depth <= z
+            if z >= grid.depth[-1]:
+                zi = len(grid.depth) - 2
+            else:
+                zi = depth_indices.argmin() - 1 if z >= grid.depth[0] else 0
         else:
-            zi = depth_index.argmin() - 1 if z >= grid.depth[0] else 0
+            if z > grid.depth[0]:
+                raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
+            elif z < grid.depth[-1]:
+                raise FieldOutOfBoundError(0, 0, z, field=self)
+            depth_indices = grid.depth >= z
+            if z <= grid.depth[-1]:
+                zi = len(grid.depth) - 2
+            else:
+                zi = depth_indices.argmin() - 1 if z <= grid.depth[0] else 0
         zeta = (z-grid.depth[zi]) / (grid.depth[zi+1]-grid.depth[zi])
         return (zi, zeta)
 
@@ -567,15 +589,27 @@ def search_indices_vertical_s(self, x, y, z, xi, yi, xsi, eta, ti, time):
                 xsi*eta * grid.depth[:, yi+1, xi+1] + \
                 (1-xsi)*eta * grid.depth[:, yi+1, xi]
         z = np.float32(z)
-        depth_index = depth_vector <= z
-        if z >= depth_vector[-1]:
-            zi = len(depth_vector) - 2
+
+        if depth_vector[-1] > depth_vector[0]:
+            depth_indices = depth_vector <= z
+            if z >= depth_vector[-1]:
+                zi = len(depth_vector) - 2
+            else:
+                zi = depth_indices.argmin() - 1 if z >= depth_vector[0] else 0
+            if z < depth_vector[zi]:
+                raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
+            elif z > depth_vector[zi+1]:
+                raise FieldOutOfBoundError(x, y, z, field=self)
         else:
-            zi = depth_index.argmin() - 1 if z >= depth_vector[0] else 0
-        if z < depth_vector[zi]:
-            raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
-        elif z > depth_vector[zi+1]:
-            raise FieldOutOfBoundError(x, y, z, field=self)
+            depth_indices = depth_vector >= z
+            if z <= depth_vector[-1]:
+                zi = len(depth_vector) - 2
+            else:
+                zi = depth_indices.argmin() - 1 if z <= depth_vector[0] else 0
+            if z > depth_vector[zi]:
+                raise FieldOutOfBoundSurfaceError(0, 0, z, field=self)
+            elif z < depth_vector[zi+1]:
+                raise FieldOutOfBoundError(x, y, z, field=self)
         zeta = (z - depth_vector[zi]) / (depth_vector[zi+1]-depth_vector[zi])
         return (zi, zeta)
 
@@ -1618,7 +1652,8 @@ def __getitem__(self, key):
 class NetcdfFileBuffer(object):
     _name_maps = {'lon': ['lon', 'nav_lon', 'x', 'longitude', 'lo', 'ln', 'i'],
                   'lat': ['lat', 'nav_lat', 'y', 'latitude', 'la', 'lt', 'j'],
-                  'depth': ['depth', 'depthu', 'depthv', 'depthw', 'depths', 'deptht', 'depthx', 'depthy', 'depthz', 'z', 'd', 'k', 'w_dep', 'w_deps'],
+                  'depth': ['depth', 'depthu', 'depthv', 'depthw', 'depths', 'deptht', 'depthx', 'depthy', 'depthz',
+                            'z', 'z_u', 'z_v', 'z_w', 'd', 'k', 'w_dep', 'w_deps'],
                   'time': ['time', 'time_count', 'time_counter', 'timer_count', 't']}
     _min_dim_chunksize = 16
 
@@ -1961,12 +1996,20 @@ def read_depth(self):
             elif len(depth.shape) == 3:
                 return np.array(depth[self.indices['depth'], self.indices['lat'], self.indices['lon']])
             elif len(depth.shape) == 4:
-                raise NotImplementedError('Time varying depth data cannot be read in netcdf files yet')
                 return np.array(depth[:, self.indices['depth'], self.indices['lat'], self.indices['lon']])
         else:
             self.indices['depth'] = [0]
             return np.zeros(1)
 
+    @property
+    def read_depth_dimensions(self):
+        if 'depth' in self.dimensions:
+            data = self.dataset[self.name]
+            depthsize = data.shape[-3]
+            self.data_full_zdim = depthsize
+            self.indices['depth'] = self.indices['depth'] if 'depth' in self.indices else range(depthsize)
+            return np.empty((0, len(self.indices['depth']), len(self.indices['lat']), len(self.indices['lon'])))
+
     @property
     def data(self):
         return self.data_access()

parcels/fieldset.py

@@ -251,6 +251,16 @@ def check_complete(self):
                 counter += 1
             ccode_fieldnames.append(fld.ccode_name)
 
+        for f in self.get_fields():
+            if type(f) in [VectorField, NestedField, SummedField] or f.dataFiles is None:
+                continue
+            if f.grid.depth_field is not None:
+                if f.grid.depth_field == 'not_yet_set':
+                    raise ValueError("If depth dimension is set at 'not_yet_set', it must be added later using Field.set_depth_from_field(field)")
+                if not f.grid.defer_load:
+                    depth_data = f.grid.depth_field.data
+                    f.grid.depth = depth_data if isinstance(depth_data, np.ndarray) else np.array(depth_data)
+
     @classmethod
     def parse_wildcards(cls, paths, filenames, var):
         if not isinstance(paths, list):
@@ -346,6 +356,8 @@ def from_netcdf(cls, filenames, variables, dimensions, indices=None, fieldtype=N
                 procpaths = filenames[procvar] if isinstance(filenames, dict) and procvar in filenames else filenames
                 nowpaths = filenames[var] if isinstance(filenames, dict) and var in filenames else filenames
                 if procdims == dims and procinds == inds:
+                    if 'depth' in dims and dims['depth'] == 'not_yet_set':
+                        break
                     processedGrid = False
                     if ((not isinstance(filenames, dict)) or filenames[procvar] == filenames[var]):
                         processedGrid = True
@@ -963,6 +975,14 @@ def computeTimeChunk(self, time, dt):
         if self.compute_on_defer:
             self.compute_on_defer(self)
 
+        # update time varying grid depth
+        for f in self.get_fields():
+            if type(f) in [VectorField, NestedField, SummedField] or not f.grid.defer_load or f.dataFiles is None:
+                continue
+            if f.grid.depth_field is not None:
+                depth_data = f.grid.depth_field.data
+                f.grid.depth = depth_data if isinstance(depth_data, np.ndarray) else np.array(depth_data)
+
         if abs(nextTime) == np.infty or np.isnan(nextTime):  # Second happens when dt=0
             return nextTime
         else:

parcels/grid.py

@@ -65,6 +65,7 @@ def __init__(self, lon, lat, time, time_origin, mesh):
         self.chunk_info = None
         self.master_chunksize = None
         self._add_last_periodic_data_timestep = False
+        self.depth_field = None
 
     @staticmethod
     def create_grid(lon, lat, depth, time, time_origin, mesh, **kwargs):
@@ -342,9 +343,10 @@ def __init__(self, lon, lat, depth, time=None, time_origin=None, mesh='flat'):
         self.zdim = self.depth.shape[-3]
         self.z4d = len(self.depth.shape) == 4
         if self.z4d:
-            assert self.tdim == self.depth.shape[0], 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
-            assert self.xdim == self.depth.shape[-1], 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
-            assert self.ydim == self.depth.shape[-2], 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
+            # self.depth.shape[0] is 0 for S grids loaded from netcdf file
+            assert self.tdim == self.depth.shape[0] or self.depth.shape[0] == 0, 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
+            assert self.xdim == self.depth.shape[-1] or self.depth.shape[-1] == 0, 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
+            assert self.ydim == self.depth.shape[-2] or self.depth.shape[-2] == 0, 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
         else:
             assert self.xdim == self.depth.shape[-1], 'depth dimension has the wrong format. It should be [zdim, ydim, xdim]'
             assert self.ydim == self.depth.shape[-2], 'depth dimension has the wrong format. It should be [zdim, ydim, xdim]'
@@ -472,9 +474,10 @@ def __init__(self, lon, lat, depth, time=None, time_origin=None, mesh='flat'):
         self.zdim = self.depth.shape[-3]
         self.z4d = len(self.depth.shape) == 4
         if self.z4d:
-            assert self.tdim == self.depth.shape[0], 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
-            assert self.xdim == self.depth.shape[-1], 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
-            assert self.ydim == self.depth.shape[-2], 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
+            # self.depth.shape[0] is 0 for S grids loaded from netcdf file
+            assert self.tdim == self.depth.shape[0] or self.depth.shape[0] == 0, 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
+            assert self.xdim == self.depth.shape[-1] or self.depth.shape[-1] == 0, 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
+            assert self.ydim == self.depth.shape[-2] or self.depth.shape[-2] == 0, 'depth dimension has the wrong format. It should be [tdim, zdim, ydim, xdim]'
         else:
             assert self.xdim == self.depth.shape[-1], 'depth dimension has the wrong format. It should be [zdim, ydim, xdim]'
             assert self.ydim == self.depth.shape[-2], 'depth dimension has the wrong format. It should be [zdim, ydim, xdim]'

parcels/gridset.py

@@ -39,9 +39,6 @@ def add_grid(self, field):
                 continue
             sameGrid &= (grid.master_chunksize == g.master_chunksize) or (grid.master_chunksize in [False, None] and g.master_chunksize in [False, None])
             if not sameGrid and sameDims and grid.master_chunksize is not None:
-                print(field.field_chunksize)
-                print(grid.master_chunksize)
-                print(g.master_chunksize)
                 res = False
                 if (isinstance(grid.master_chunksize, tuple) and isinstance(g.master_chunksize, tuple)) or \
                         (isinstance(grid.master_chunksize, dict) and isinstance(g.master_chunksize, dict)):

parcels/include/index_search.h

@@ -54,10 +54,18 @@ typedef enum
 
 static inline ErrorCode search_indices_vertical_z(type_coord z, int zdim, float *zvals, int *zi, double *zeta)
 {
-  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);
-  while (*zi > 0 && z < zvals[*zi]) --(*zi);
+  if (zvals[zdim-1] > zvals[0]){
+    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);
+    while (*zi > 0 && z < zvals[*zi]) --(*zi);
+  }
+  else{
+    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);
+    while (*zi > 0 && z > zvals[*zi]) --(*zi);
+  }
   if (*zi == zdim-1) {--*zi;}
 
   *zeta = (z - zvals[*zi]) / (zvals[*zi+1] - zvals[*zi]);
@@ -103,10 +111,18 @@ static inline ErrorCode search_indices_vertical_s(type_coord z, int xdim, int yd
     }
   }
 
-  if (z < zcol[0]) {return ERROR_THROUGH_SURFACE;}
-  if (z > zcol[zdim-1]) {return ERROR_OUT_OF_BOUNDS;}
-  while (*zi < zdim-1 && z > zcol[*zi+1]) ++(*zi);
-  while (*zi > 0 && z < zcol[*zi]) --(*zi);
+  if (zcol[zdim-1] > zcol[0]){
+    if (z < zcol[0]) {return ERROR_THROUGH_SURFACE;}
+    if (z > zcol[zdim-1]) {return ERROR_OUT_OF_BOUNDS;}
+    while (*zi < zdim-1 && z > zcol[*zi+1]) ++(*zi);
+    while (*zi > 0 && z < zcol[*zi]) --(*zi);
+  }
+  else{
+    if (z > zcol[0]) {return ERROR_THROUGH_SURFACE;}
+    if (z < zcol[zdim-1]) {return ERROR_OUT_OF_BOUNDS;}
+    while (*zi < zdim-1 && z < zcol[*zi+1]) ++(*zi);
+    while (*zi > 0 && z > zcol[*zi]) --(*zi);
+  }
   if (*zi == zdim-1) {--*zi;}
 
   *zeta = (z - zcol[*zi]) / (zcol[*zi+1] - zcol[*zi]);