Repository: OceanParcels/parcels
Pull request #383 (2018-06-12 11:42:10)

Field.gradient for deferred loading
This PR fixes the `Field.gradient()` function for `Fields` that are "deferred_loaded"

Note that this requires the output of `Field.gradient()` to be a `Field`, rather than a `numpy.array`, as it was before

This PR should fix Issue #371 reported by @JoostSH


Patch:
diff --git a/parcels/examples/example_globcurrent.py b/parcels/examples/example_globcurrent.py
index 6d68e050..432435e1 100644
--- a/parcels/examples/example_globcurrent.py
+++ b/parcels/examples/example_globcurrent.py
@@ -97,3 +97,18 @@ def test_globcurrent_variable_fromfield(mode, dt):
     pset = ParticleSet(fieldset, pclass=MyParticle, lon=[25], lat=[-35], time=time)
 
     pset.execute(AdvectionRK4, runtime=delta(days=1), dt=dt)
+
+
+@pytest.mark.parametrize('full_load', [True, False])
+def test_globcurrent_deferred_fieldset_gradient(full_load):
+    fieldset = set_globcurrent_fieldset(full_load=full_load)
+    (dU_dx, dU_dy) = fieldset.U.gradient()
+    fieldset.add_field(dU_dy)
+
+    pset = ParticleSet(fieldset, pclass=JITParticle, lon=25, lat=-35)
+    pset.execute(AdvectionRK4, runtime=delta(days=1), dt=delta(days=1))
+
+    tdim = 365 if full_load else 3
+    assert(dU_dx.data.shape == (tdim, 41, 81))
+    assert(fieldset.dU_dy.data.shape == (tdim, 41, 81))
+    assert(dU_dx is fieldset.U.gradientx)
diff --git a/parcels/field.py b/parcels/field.py
index dccf854a..fb39d43a 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -230,6 +230,8 @@ class Field(object):
                 self.data[self.data > self.vmax] = 0.
 
         self._scaling_factor = None
+        (self.gradientx, self.gradienty) = (None, None)  # to store if Field is a gradient() of another field
+        self.is_gradient = False
 
         # Variable names in JIT code
         self.ccode_data = self.name
@@ -446,15 +448,29 @@ class Field(object):
             self.calc_cell_edge_sizes()
         return self.grid.cell_edge_sizes['x'] * self.grid.cell_edge_sizes['y']
 
-    def gradient(self):
+    def gradient(self, update=False):
         """Method to calculate horizontal gradients of Field.
-                Returns two numpy arrays: the zonal and meridional gradients,
-                on the same Grid as the original Field, using numpy.gradient() method"""
+                Returns two Fields: the zonal and meridional gradients,
+                on the same Grid as the original Field, using numpy.gradient() method
+                Names of these grids are dNAME_dx and dNAME_dy, where NAME is the name
+                of the original Field"""
         if not self.grid.cell_edge_sizes:
             self.calc_cell_edge_sizes()
-        dFdy = np.gradient(self.data, axis=-2) / self.grid.cell_edge_sizes['y']
-        dFdx = np.gradient(self.data, axis=-1) / self.grid.cell_edge_sizes['x']
-        return dFdx, dFdy
+        if self.grid.defer_load and self.data is None:
+            (dFdx, dFdy) = (None, None)
+        else:
+            dFdy = np.gradient(self.data, axis=-2) / self.grid.cell_edge_sizes['y']
+            dFdx = np.gradient(self.data, axis=-1) / self.grid.cell_edge_sizes['x']
+        if update:
+            self.gradientx.data = dFdx
+            self.gradienty.data = dFdy
+        else:
+            dFdx_fld = Field('d%s_dx' % self.name, dFdx, grid=self.grid)
+            dFdy_fld = Field('d%s_dy' % self.name, dFdy, grid=self.grid)
+            dFdx_fld.is_gradient = True
+            dFdy_fld.is_gradient = True
+            (self.gradientx, self.gradienty) = (dFdx_fld, dFdy_fld)
+            return (dFdx_fld, dFdy_fld)
 
     def interpolator2D_scipy(self, ti, z_idx=None):
         """Provide a SciPy interpolator for spatial interpolation
diff --git a/parcels/fieldset.py b/parcels/fieldset.py
index 5abd72f3..d3dc0bd2 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -384,7 +384,7 @@ class FieldSet(object):
             nextTime = min(nextTime, nextTime_loc) if signdt >= 0 else max(nextTime, nextTime_loc)
 
         for f in self.fields:
-            if f.name == 'UV' or not f.grid.defer_load:
+            if f.name == 'UV' or not f.grid.defer_load or f.is_gradient:
                 continue
             g = f.grid
             if g.update_status == 'first_updated':  # First load of data
@@ -406,6 +406,8 @@ class FieldSet(object):
                 if f._scaling_factor:
                     data *= f._scaling_factor
                 f.data[tindex, :] = f.reshape(data)[tindex, :]
+            if f.gradientx is not None and g.update_status in ['first_updated', 'updated']:
+                f.gradient(update=True)
 
         if abs(nextTime) == np.infty or np.isnan(nextTime):  # Second happens when dt=0
             return nextTime


Test patch:
diff --git a/tests/test_diffusion.py b/tests/test_diffusion.py
index a0fb7c53..126e4575 100644
--- a/tests/test_diffusion.py
+++ b/tests/test_diffusion.py
@@ -73,8 +73,8 @@ def test_fieldKh_SpatiallyVaryingBrownianMotion(mesh, mode, xdim=200, ydim=100):
 
     dKh_zonal_dx, _ = fieldset.Kh_zonal.gradient()
     _, dKh_meridional_dy = fieldset.Kh_meridional.gradient()
-    fieldset.add_field(Field('dKh_zonal_dx', dKh_zonal_dx, grid=grid))
-    fieldset.add_field(Field('dKh_meridional_dy', dKh_meridional_dy, grid=grid))
+    fieldset.add_field(dKh_zonal_dx)
+    fieldset.add_field(dKh_meridional_dy)
     if mesh is 'spherical':
         fieldset.dKh_zonal_dx.units = GeographicPolar()
         fieldset.dKh_meridional_dy.units = Geographic()
diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py
index 0f141324..c49800e6 100644
--- a/tests/test_fieldset.py
+++ b/tests/test_fieldset.py
@@ -177,8 +177,8 @@ def test_fieldset_gradient(mesh):
         for y in range(np_dFdx.shape[0]):
             np_dFdx[:, y] /= math.cos(fieldset.V.grid.lat[y] * math.pi / 180.)
 
-    assert np.allclose(dFdx, np_dFdx, rtol=5e-2)  # Field gradient dx.
-    assert np.allclose(dFdy, np_dFdy, rtol=5e-2)  # Field gradient dy.
+    assert np.allclose(dFdx.data, np_dFdx, rtol=5e-2)  # Field gradient dx.
+    assert np.allclose(dFdy.data, np_dFdy, rtol=5e-2)  # Field gradient dy.
 
 
 def addConst(particle, fieldset, time, dt):