Repository: OceanParcels/parcels
Pull request #498 (2018-11-29 13:59:00)

From data v2.0
solves issue #447


Patch:
diff --git a/parcels/fieldset.py b/parcels/fieldset.py
index 895262f2..b691aca6 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -1,6 +1,6 @@
 from parcels.field import Field, VectorField, SummedField, SummedVectorField, NestedField
 from parcels.gridset import GridSet
-from parcels.grid import RectilinearZGrid
+from parcels.grid import Grid
 from parcels.tools.loggers import logger
 from parcels.tools.converters import TimeConverter
 import numpy as np
@@ -78,7 +78,7 @@ class FieldSet(object):
             lat = dims['lat']
             depth = np.zeros(1, dtype=np.float32) if 'depth' not in dims else dims['depth']
             time = np.zeros(1, dtype=np.float64) if 'time' not in dims else dims['time']
-            grid = RectilinearZGrid(lon, lat, depth, time, time_origin=TimeConverter(), mesh=mesh)
+            grid = Grid.grid(lon, lat, depth, time, time_origin=TimeConverter(), mesh=mesh)
 
             fields[name] = Field(name, datafld, grid=grid, transpose=transpose,
                                  allow_time_extrapolation=allow_time_extrapolation, time_periodic=time_periodic, **kwargs)
diff --git a/parcels/grid.py b/parcels/grid.py
index 6ba7f1c3..17af8e52 100644
--- a/parcels/grid.py
+++ b/parcels/grid.py
@@ -51,6 +51,19 @@ class Grid(object):
         self.defer_load = False
         self.lonlat_minmax = np.array([np.nanmin(lon), np.nanmax(lon), np.nanmin(lat), np.nanmax(lat)], dtype=np.float32)
 
+    @classmethod
+    def grid(self, lon, lat, depth, time, time_origin, mesh, **kwargs):
+        if len(lon.shape) == 1:
+            if len(depth.shape) == 1:
+                return RectilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+            else:
+                return RectilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+        else:
+            if len(depth.shape) == 1:
+                return CurvilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+            else:
+                return CurvilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+
     @property
     def ctypes_struct(self):
         # This is unnecessary for the moment, but it could be useful when going will fully unstructured grids
@@ -166,7 +179,7 @@ class RectilinearGrid(Grid):
         if isinstance(time, np.ndarray):
             assert(len(time.shape) == 1), 'time is not a vector'
 
-        Grid.__init__(self, lon, lat, time, time_origin, mesh)
+        super(RectilinearGrid, self).__init__(lon, lat, time, time_origin, mesh)
         self.xdim = self.lon.size
         self.ydim = self.lat.size
         self.tdim = self.time.size
@@ -221,7 +234,7 @@ class RectilinearZGrid(RectilinearGrid):
     """
 
     def __init__(self, lon, lat, depth=None, time=None, time_origin=None, mesh='flat'):
-        RectilinearGrid.__init__(self, lon, lat, time, time_origin, mesh)
+        super(RectilinearZGrid, self).__init__(lon, lat, time, time_origin, mesh)
         if isinstance(depth, np.ndarray):
             assert(len(depth.shape) == 1), 'depth is not a vector'
 
@@ -258,7 +271,7 @@ class RectilinearSGrid(RectilinearGrid):
     """
 
     def __init__(self, lon, lat, depth, time=None, time_origin=None, mesh='flat'):
-        RectilinearGrid.__init__(self, lon, lat, time, time_origin, mesh)
+        super(RectilinearSGrid, self).__init__(lon, lat, time, time_origin, mesh)
         assert(isinstance(depth, np.ndarray) and len(depth.shape) in [3, 4]), 'depth is not a 3D or 4D numpy array'
 
         self.gtype = GridCode.RectilinearSGrid
@@ -290,7 +303,7 @@ class CurvilinearGrid(Grid):
 
         lon = lon.squeeze()
         lat = lat.squeeze()
-        Grid.__init__(self, lon, lat, time, time_origin, mesh)
+        super(CurvilinearGrid, self).__init__(lon, lat, time, time_origin, mesh)
         self.xdim = self.lon.shape[1]
         self.ydim = self.lon.shape[0]
         self.tdim = self.time.size
@@ -350,7 +363,7 @@ class CurvilinearZGrid(CurvilinearGrid):
     """
 
     def __init__(self, lon, lat, depth=None, time=None, time_origin=None, mesh='flat'):
-        CurvilinearGrid.__init__(self, lon, lat, time, time_origin, mesh)
+        super(CurvilinearZGrid, self).__init__(lon, lat, time, time_origin, mesh)
         if isinstance(depth, np.ndarray):
             assert(len(depth.shape) == 1), 'depth is not a vector'
 
@@ -386,7 +399,7 @@ class CurvilinearSGrid(CurvilinearGrid):
     """
 
     def __init__(self, lon, lat, depth, time=None, time_origin=None, mesh='flat'):
-        CurvilinearGrid.__init__(self, lon, lat, time, time_origin, mesh)
+        super(CurvilinearSGrid, self).__init__(lon, lat, time, time_origin, mesh)
         assert(isinstance(depth, np.ndarray) and len(depth.shape) in [3, 4]), 'depth is not a 4D numpy array'
 
         self.gtype = GridCode.CurvilinearSGrid
diff --git a/parcels/gridset.py b/parcels/gridset.py
index 4fd7c469..7d7b1821 100644
--- a/parcels/gridset.py
+++ b/parcels/gridset.py
@@ -45,8 +45,12 @@ class GridSet(object):
             if len(getattr(g, dim)) == 1:
                 continue  # not including grids where only one entry
             else:
-                maxleft = max(maxleft, getattr(g, dim)[0])
-                minright = min(minright, getattr(g, dim)[-1])
+                if dim == 'depth':
+                    maxleft = max(maxleft, np.min(getattr(g, dim)))
+                    minright = min(minright, np.max(getattr(g, dim)))
+                else:
+                    maxleft = max(maxleft, getattr(g, dim)[0])
+                    minright = min(minright, getattr(g, dim)[-1])
         maxleft = 0 if maxleft == -np.inf else maxleft  # if all len(dim) == 1
         minright = 0 if minright == np.inf else minright  # if all len(dim) == 1
         return maxleft, minright


Test patch:
diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py
index 82ad5f27..1579c512 100644
--- a/tests/test_fieldset.py
+++ b/tests/test_fieldset.py
@@ -413,3 +413,55 @@ def test_fieldset_from_xarray(maxlatind):
 
     pset.execute(AdvectionRK4, dt=1)
     assert pset[0].lon == 4.5 and pset[0].lat == 10
+
+
+def test_fieldset_from_data_gridtypes(xdim=20, ydim=10, zdim=4):
+    """ Simple test for fieldset initialisation from data. """
+    lon = np.linspace(0., 10., xdim, dtype=np.float32)
+    lat = np.linspace(0., 10., ydim, dtype=np.float32)
+    depth = np.linspace(0., 1., zdim, dtype=np.float32)
+    depth_s = np.ones((zdim, ydim, xdim))
+    U = np.ones((zdim, ydim, xdim))
+    V = np.ones((zdim, ydim, xdim))
+    dimensions = {'lat': lat, 'lon': lon, 'depth': depth}
+    data = {'U': np.array(U, dtype=np.float32), 'V': np.array(V, dtype=np.float32)}
+    lonm, latm = np.meshgrid(lon, lat)
+    for k in range(zdim):
+        data['U'][k, :, :] = lonm * (depth[k]+1) + .1
+        depth_s[k, :, :] = depth[k]
+
+    # Rectilinear Z grid
+    fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
+    pset = ParticleSet(fieldset, ScipyParticle, [0, 0], [0, 0], [0, .4])
+    pset.execute(AdvectionRK4, runtime=1, dt=.5)
+    plon = [p.lon for p in pset]
+    plat = [p.lat for p in pset]
+    # sol of  dx/dt = (init_depth+1)*x+0.1; x(0)=0
+    assert np.allclose(plon, [0.17173462592827032, 0.2177736932123214])
+    assert np.allclose(plat, [1, 1])
+
+    # Rectilinear S grid
+    dimensions['depth'] = depth_s
+    fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
+    pset = ParticleSet(fieldset, ScipyParticle, [0, 0], [0, 0], [0, .4])
+    pset.execute(AdvectionRK4, runtime=1, dt=.5)
+    assert np.allclose(plon, [p.lon for p in pset])
+    assert np.allclose(plat, [p.lat for p in pset])
+
+    # Curvilinear Z grid
+    dimensions['lon'] = lonm
+    dimensions['lat'] = latm
+    dimensions['depth'] = depth
+    fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
+    pset = ParticleSet(fieldset, ScipyParticle, [0, 0], [0, 0], [0, .4])
+    pset.execute(AdvectionRK4, runtime=1, dt=.5)
+    assert np.allclose(plon, [p.lon for p in pset])
+    assert np.allclose(plat, [p.lat for p in pset])
+
+    # Curvilinear S grid
+    dimensions['depth'] = depth_s
+    fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
+    pset = ParticleSet(fieldset, ScipyParticle, [0, 0], [0, 0], [0, .4])
+    pset.execute(AdvectionRK4, runtime=1, dt=.5)
+    assert np.allclose(plon, [p.lon for p in pset])
+    assert np.allclose(plat, [p.lat for p in pset])