Merge branch 'master' into new_xarray

delandmeterp · delandmeterp · commit e4f26059c8d4 · 2018-10-03T14:47:22.000+02:00
diff --git a/parcels/examples/example_nemo_curvilinear.py b/parcels/examples/example_nemo_curvilinear.py
@@ -12,9 +12,12 @@ def run_nemo_curvilinear(mode, outfile):
     """Function that shows how to read in curvilinear grids, in this case from NEMO"""
     data_path = path.join(path.dirname(__file__), 'NemoCurvilinear_data/')
 
-    filenames = {'U': data_path + 'U_purely_zonal-ORCA025_grid_U.nc4',
-                 'V': data_path + 'V_purely_zonal-ORCA025_grid_V.nc4',
-                 'mesh_mask': data_path + 'mesh_mask.nc4'}
+    filenames = {'U': {'lon': data_path + 'mesh_mask.nc4',
+                       'lat': data_path + 'mesh_mask.nc4',
+                       'data': data_path + 'U_purely_zonal-ORCA025_grid_U.nc4'},
+                 'V': {'lon': data_path + 'mesh_mask.nc4',
+                       'lat': data_path + 'mesh_mask.nc4',
+                       'data': data_path + 'V_purely_zonal-ORCA025_grid_V.nc4'}}
     variables = {'U': 'U', 'V': 'V'}
     dimensions = {'lon': 'glamf', 'lat': 'gphif'}
     field_set = FieldSet.from_nemo(filenames, variables, dimensions)
diff --git a/parcels/examples/tutorial_nemo_curvilinear.ipynb b/parcels/examples/tutorial_nemo_curvilinear.ipynb
@@ -33,7 +33,17 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We can create a `FieldSet` just like we do for normal grids. Note that we need to provide an extra filename for the `mesh_mask` file that holds information on the Curvilinear grid."
+    "We can create a `FieldSet` just like we do for normal grids.\n",
+    "Note that NEMO is discretised on a C-grid. U and V velocities are not located on the same nodes (see https://www.nemo-ocean.eu/doc/node19.html ).\n",
+    "\n",
+    "```\n",
+    " __V1__\n",
+    "|      |\n",
+    "U0     U1\n",
+    "|__V0__|\n",
+    "```\n",
+    "\n",
+    "To interpolate U, V velocities on the C-grid, Parcels needs to read the f-nodes, which are located on the corners of the cells (for indexing details: https://www.nemo-ocean.eu/doc/img360.png )."
    ]
   },
   {
@@ -52,9 +62,13 @@
     }
    ],
    "source": [
-    "filenames = {'U': 'NemoCurvilinear_data/U_purely_zonal-ORCA025_grid_U.nc4',\n",
-    "             'V': 'NemoCurvilinear_data/V_purely_zonal-ORCA025_grid_V.nc4',\n",
-    "             'mesh_mask': 'NemoCurvilinear_data/mesh_mask.nc4'}\n",
+    "data_path = 'NemoCurvilinear_data/'\n",
+    "filenames = {'U': {'lon': data_path + 'mesh_mask.nc4',\n",
+    "                   'lat': data_path + 'mesh_mask.nc4',\n",
+    "                   'data': data_path + 'U_purely_zonal-ORCA025_grid_U.nc4'},\n",
+    "             'V': {'lon': data_path + 'mesh_mask.nc4',\n",
+    "                   'lat': data_path + 'mesh_mask.nc4',\n",
+    "                   'data': data_path + 'V_purely_zonal-ORCA025_grid_V.nc4'}}\n",
     "variables = {'U': 'U',\n",
     "             'V': 'V'}\n",
     "dimensions = {'lon': 'glamf', 'lat': 'gphif', 'time': 'time_counter'}\n",
diff --git a/parcels/field.py b/parcels/field.py
@@ -124,11 +124,14 @@ def __init__(self, name, data, lon=None, lat=None, depth=None, time=None, grid=N
     @classmethod
     def from_netcdf(cls, filenames, variable, dimensions, indices=None, grid=None,
                     mesh='spherical', allow_time_extrapolation=None, time_periodic=False,
-                    full_load=False, dimension_filename=None, **kwargs):
+                    full_load=False, **kwargs):
         """Create field from netCDF file
 
         :param filenames: list of filenames to read for the field.
                Note that wildcards ('*') are also allowed
+               filenames can be a list [files]
+               or a dictionary {dim:[files]} (if lon, lat, depth and/or data not stored in same files as data)
+               time values are in filenames[data]
         :param variable: Name of the field to create. Note that this has to be a string
         :param dimensions: Dictionary mapping variable names for the relevant dimensions in the NetCDF file
         :param indices: dictionary mapping indices for each dimension to read from file.
@@ -152,25 +155,48 @@ def from_netcdf(cls, filenames, variable, dimensions, indices=None, grid=None,
 
         if not isinstance(filenames, Iterable) or isinstance(filenames, str):
             filenames = [filenames]
-        dimension_filename = dimension_filename if dimension_filename else filenames[0]
+
+        data_filenames = filenames['data'] if type(filenames) is dict else filenames
+        if type(filenames) == dict:
+            for k in filenames.keys():
+                assert k in ['lon', 'lat', 'depth', 'data'], \
+                    'filename dimension keys must be lon, lat, depth or data'
+            assert len(filenames['lon']) == 1
+            if filenames['lon'] != filenames['lat']:
+                raise NotImplementedError('longitude and latitude dimensions are currently processed together from one single file')
+            lonlat_filename = filenames['lon'][0]
+            if 'depth' in dimensions:
+                depth_filename = filenames['depth'][0]
+                if len(depth_filename) != 1:
+                    raise NotImplementedError('Vertically adaptive meshes not implemented for from_netcdf()')
+        else:
+            lonlat_filename = filenames[0]
+            depth_filename = filenames[0]
+
         indices = {} if indices is None else indices.copy()
-        with NetcdfFileBuffer(dimension_filename, dimensions, indices) as filebuffer:
+        with NetcdfFileBuffer(lonlat_filename, dimensions, indices) as filebuffer:
             lon, lat = filebuffer.read_lonlat
-            depth = filebuffer.read_depth
             indices = filebuffer.indices
             # Check if parcels_mesh has been explicitly set in file
             if 'parcels_mesh' in filebuffer.dataset.attrs:
                 mesh = filebuffer.dataset.attrs['parcels_mesh']
 
-        if len(filenames) > 1 and 'time' not in dimensions:
+        if 'depth' in dimensions:
+            with NetcdfFileBuffer(depth_filename, dimensions, indices) as filebuffer:
+                depth = filebuffer.read_depth
+        else:
+            indices['depth'] = [0]
+            depth = np.zeros(1)
+
+        if len(data_filenames) > 1 and 'time' not in dimensions:
             raise RuntimeError('Multiple files given but no time dimension specified')
 
         if grid is None:
             # Concatenate time variable to determine overall dimension
             # across multiple files
             timeslices = []
             dataFiles = []
-            for fname in filenames:
+            for fname in data_filenames:
                 with NetcdfFileBuffer(fname, dimensions, indices) as filebuffer:
                     ftime = filebuffer.time
                     timeslices.append(ftime)
@@ -204,7 +230,7 @@ def from_netcdf(cls, filenames, variable, dimensions, indices=None, grid=None,
             # Pre-allocate data before reading files into buffer
             data = np.empty((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32)
             ti = 0
-            for tslice, fname in zip(grid.timeslices, filenames):
+            for tslice, fname in zip(grid.timeslices, data_filenames):
                 with NetcdfFileBuffer(fname, dimensions, indices) as filebuffer:
                     # If Field.from_netcdf is called directly, it may not have a 'data' dimension
                     # In that case, assume that 'name' is the data dimension
@@ -905,7 +931,6 @@ def __init__(self, name, U, V, W=None):
             assert self.U.grid is self.V.grid
             if W:
                 assert self.W.interp_method == 'cgrid_linear'
-                assert self.U.grid is self.W.grid
 
     def dist(self, lon1, lon2, lat1, lat2, mesh):
         if mesh == 'spherical':
diff --git a/parcels/fieldset.py b/parcels/fieldset.py
@@ -140,6 +140,18 @@ def check_complete(self):
             else:
                 self.add_vector_field(VectorField('UVW', self.U, self.V, self.W))
 
+    @classmethod
+    def parse_wildcards(cls, paths, filenames, var):
+        if not isinstance(paths, list):
+            paths = sorted(glob(str(paths)))
+        if len(paths) == 0:
+            notfound_paths = filenames[var] if type(filenames) is dict and var in filenames else filenames
+            raise IOError("FieldSet files not found: %s" % str(notfound_paths))
+        for fp in paths:
+            if not path.exists(fp):
+                raise IOError("FieldSet file not found: %s" % str(fp))
+        return paths
+
     @classmethod
     def from_netcdf(cls, filenames, variables, dimensions, indices=None,
                     mesh='spherical', allow_time_extrapolation=None, time_periodic=False, full_load=False, **kwargs):
@@ -148,6 +160,10 @@ def from_netcdf(cls, filenames, variables, dimensions, indices=None,
         :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).
         :param dimensions: Dictionary mapping data dimensions (lon,
@@ -178,15 +194,12 @@ def from_netcdf(cls, filenames, variables, dimensions, indices=None,
         fields = {}
         for var, name in variables.items():
             # Resolve all matching paths for the current variable
-            paths = filenames[var] if type(filenames) is dict else filenames
-            if not isinstance(paths, list):
-                paths = sorted(glob(str(paths)))
-            if len(paths) == 0:
-                notfound_paths = filenames[var] if type(filenames) is dict else filenames
-                raise IOError("FieldSet files not found: %s" % str(notfound_paths))
-            for fp in paths:
-                if not path.exists(fp):
-                    raise IOError("FieldSet file not found: %s" % str(fp))
+            paths = filenames[var] if type(filenames) is dict and var in filenames else filenames
+            if type(paths) is not dict:
+                paths = cls.parse_wildcards(paths, filenames, var)
+            else:
+                for dim, p in paths.items():
+                    paths[dim] = cls.parse_wildcards(p, filenames, var)
 
             # Use dimensions[var] and indices[var] if either of them is a dict of dicts
             dims = dimensions[var] if var in dimensions else dimensions
@@ -198,11 +211,19 @@ def from_netcdf(cls, filenames, variables, dimensions, indices=None,
             for procvar, _ in fields.items():
                 procdims = dimensions[procvar] if procvar in dimensions else dimensions
                 procinds = indices[procvar] if (indices and procvar in indices) else indices
-                if (type(filenames) is not dict or filenames[procvar] == filenames[var]) \
-                        and procdims == dims and procinds == inds:
-                    grid = fields[procvar].grid
-                    kwargs['dataFiles'] = fields[procvar].dataFiles
-                    break
+                if procdims == dims and procinds == inds:
+                    sameGrid = False
+                    if (type(filenames) is not dict or filenames[procvar] == filenames[var]):
+                        sameGrid = True
+                    elif type(filenames[procvar]) == dict:
+                        sameGrid = True
+                        for dim in ['lon', 'lat', 'depth']:
+                            if dim in dimensions:
+                                sameGrid *= filenames[procvar][dim] == filenames[var][dim]
+                    if sameGrid:
+                        grid = fields[procvar].grid
+                        kwargs['dataFiles'] = fields[procvar].dataFiles
+                        break
             fields[var] = Field.from_netcdf(paths, var, dims, inds, grid=grid, mesh=mesh,
                                             allow_time_extrapolation=allow_time_extrapolation,
                                             time_periodic=time_periodic, full_load=full_load, **kwargs)
@@ -215,19 +236,29 @@ def from_nemo(cls, filenames, variables, dimensions, indices=None, mesh='spheric
                   allow_time_extrapolation=None, time_periodic=False,
                   tracer_interp_method='linear', **kwargs):
         """Initialises FieldSet object from NetCDF files of Curvilinear NEMO fields.
-        Note that this assumes there is a variable mesh_mask that is used for the dimensions
 
         :param filenames: Dictionary mapping variables to file(s). The
                filepath may contain wildcards to indicate multiple files,
-               or be a list of file. At least a 'mesh_mask' needs to be present
+               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). Must include a variable 'mesh_mask' that
-               holds the dimensions
+               names in the netCDF file(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
-               (e.g. dimensions['U'], dimensions['V'], etc).
+               dimension names are different for each variable.
+               Watch out: NEMO is discretised on a C-grid:
+               U and V velocities are not located on the same nodes (see https://www.nemo-ocean.eu/doc/node19.html ).
+                __V1__
+               |      |
+               U0     U1
+               |__V0__|
+               To interpolate U, V velocities on the C-grid, Parcels needs to read the f-nodes,
+               which are located on the corners of the cells.
+               (for indexing details: https://www.nemo-ocean.eu/doc/img360.png )
         :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.
@@ -247,7 +278,8 @@ def from_nemo(cls, filenames, variables, dimensions, indices=None, mesh='spheric
 
         """
 
-        dimension_filename = filenames.pop('mesh_mask') if type(filenames) is dict else filenames
+        if 'U' in dimensions and 'V' in dimensions and dimensions['U'] != dimensions['V']:
+            raise RuntimeError("On a c-grid discretisation like NEMO, U and V should have the same dimensions")
 
         interp_method = {}
         for v in variables:
@@ -257,8 +289,7 @@ def from_nemo(cls, filenames, variables, dimensions, indices=None, mesh='spheric
                 interp_method[v] = tracer_interp_method
 
         return cls.from_netcdf(filenames, variables, dimensions, mesh=mesh, indices=indices, time_periodic=time_periodic,
-                               allow_time_extrapolation=allow_time_extrapolation, interp_method=interp_method,
-                               dimension_filename=dimension_filename, **kwargs)
+                               allow_time_extrapolation=allow_time_extrapolation, interp_method=interp_method, **kwargs)
 
     @classmethod
     def from_parcels(cls, basename, uvar='vozocrtx', vvar='vomecrty', indices=None, extra_fields=None,
diff --git a/tests/test_fieldset_sampling.py b/tests/test_fieldset_sampling.py
@@ -118,7 +118,7 @@ def test_variable_init_from_field(mode, npart=9):
             'P': np.zeros(dims, dtype=np.float32)}
     data['P'][0, 0] = 1.
     fieldset = FieldSet.from_data(data, dimensions, mesh='flat', transpose=True)
-    xv, yv = np.meshgrid(np.linspace(0, 1, np.sqrt(npart)), np.linspace(0, 1, np.sqrt(npart)))
+    xv, yv = np.meshgrid(np.linspace(0, 1, int(np.sqrt(npart))), np.linspace(0, 1, int(np.sqrt(npart))))
 
     class VarParticle(pclass(mode)):
         a = Variable('a', dtype=np.float32, initial=fieldset.P)
@@ -161,7 +161,7 @@ def test_nearest_neighbour_interpolation2D(mode, k_sample_p, npart=81):
     data['P'][0, 1] = 1.
     fieldset = FieldSet.from_data(data, dimensions, mesh='flat', transpose=True)
     fieldset.P.interp_method = 'nearest'
-    xv, yv = np.meshgrid(np.linspace(0., 1.0, np.sqrt(npart)), np.linspace(0., 1.0, np.sqrt(npart)))
+    xv, yv = np.meshgrid(np.linspace(0., 1.0, int(np.sqrt(npart))), np.linspace(0., 1.0, int(np.sqrt(npart))))
     pset = ParticleSet(fieldset, pclass=pclass(mode),
                        lon=xv.flatten(), lat=yv.flatten())
     pset.execute(k_sample_p, endtime=1, dt=1)
@@ -181,7 +181,7 @@ def test_nearest_neighbour_interpolation3D(mode, k_sample_p, npart=81):
     data['P'][0, 1, 1] = 1.
     fieldset = FieldSet.from_data(data, dimensions, mesh='flat', transpose=True)
     fieldset.P.interp_method = 'nearest'
-    xv, yv = np.meshgrid(np.linspace(0, 1.0, np.sqrt(npart)), np.linspace(0, 1.0, np.sqrt(npart)))
+    xv, yv = np.meshgrid(np.linspace(0, 1.0, int(np.sqrt(npart))), np.linspace(0, 1.0, int(np.sqrt(npart))))
     # combine a pset at 0m with pset at 1m, as meshgrid does not do 3D
     pset = ParticleSet(fieldset, pclass=pclass(mode), lon=xv.flatten(), lat=yv.flatten(), depth=np.zeros(npart))
     pset2 = ParticleSet(fieldset, pclass=pclass(mode), lon=xv.flatten(), lat=yv.flatten(), depth=np.ones(npart))
diff --git a/tests/test_grids.py b/tests/test_grids.py
@@ -336,9 +336,12 @@ class MyParticle(ptype[mode]):
 def test_nemo_grid(mode):
     data_path = path.join(path.dirname(__file__), 'test_data/')
 
-    filenames = {'U': data_path + 'Uu_eastward_nemo_cross_180lon.nc',
-                 'V': data_path + 'Vv_eastward_nemo_cross_180lon.nc',
-                 'mesh_mask': data_path + 'mask_nemo_cross_180lon.nc'}
+    filenames = {'U': {'lon': data_path + 'mask_nemo_cross_180lon.nc',
+                       'lat': data_path + 'mask_nemo_cross_180lon.nc',
+                       'data': data_path + 'Uu_eastward_nemo_cross_180lon.nc'},
+                 'V': {'lon': data_path + 'mask_nemo_cross_180lon.nc',
+                       'lat': data_path + 'mask_nemo_cross_180lon.nc',
+                       'data': data_path + 'Vv_eastward_nemo_cross_180lon.nc'}}
     variables = {'U': 'U', 'V': 'V'}
     dimensions = {'lon': 'glamf', 'lat': 'gphif'}
     field_set = FieldSet.from_nemo(filenames, variables, dimensions)
@@ -367,9 +370,12 @@ class MyParticle(ptype[mode]):
 def test_advect_nemo(mode):
     data_path = path.join(path.dirname(__file__), 'test_data/')
 
-    filenames = {'U': data_path + 'Uu_eastward_nemo_cross_180lon.nc',
-                 'V': data_path + 'Vv_eastward_nemo_cross_180lon.nc',
-                 'mesh_mask': data_path + 'mask_nemo_cross_180lon.nc'}
+    filenames = {'U': {'lon': data_path + 'mask_nemo_cross_180lon.nc',
+                       'lat': data_path + 'mask_nemo_cross_180lon.nc',
+                       'data': data_path + 'Uu_eastward_nemo_cross_180lon.nc'},
+                 'V': {'lon': data_path + 'mask_nemo_cross_180lon.nc',
+                       'lat': data_path + 'mask_nemo_cross_180lon.nc',
+                       'data': data_path + 'Vv_eastward_nemo_cross_180lon.nc'}}
     variables = {'U': 'U', 'V': 'V'}
     dimensions = {'lon': 'glamf', 'lat': 'gphif'}
     field_set = FieldSet.from_nemo(filenames, variables, dimensions)
