Merge pull request #978 from OceanParcels/deferred_load_with_two_snapshots

Deferred load with two snapshots

Author: erikvansebille

parcels/examples/example_dask_chunk_OCMs.py

@@ -216,7 +216,7 @@ def compute_swash_particle_advection(field_set, mode, lonp, latp, depthp):
 
 def compute_ofam_particle_advection(field_set, mode, lonp, latp, depthp):
     pset = ParticleSet(field_set, pclass=ptype[mode], lon=lonp, lat=latp, depth=depthp)
-    pfile = ParticleFile("ofam_particles_chunk", pset, outputdt=delta(minutes=10))
+    pfile = ParticleFile("ofam_particles_chunk", pset, outputdt=delta(days=1))
     pset.execute(AdvectionRK4, runtime=delta(days=10), dt=delta(minutes=5), output_file=pfile)
     return pset
 
@@ -255,7 +255,7 @@ def test_nemo_3D(mode, chunk_mode):
 @pytest.mark.parametrize('chunk_mode', [False, 'auto', 'specific', 'failsafe'])
 def test_globcurrent_2D(mode, chunk_mode):
     if chunk_mode in ['auto', ]:
-        dask.config.set({'array.chunk-size': '32KiB'})
+        dask.config.set({'array.chunk-size': '16KiB'})
     else:
         dask.config.set({'array.chunk-size': '128MiB'})
     field_set = fieldset_from_globcurrent(chunk_mode)
@@ -325,7 +325,7 @@ def test_pop(mode, chunk_mode):
 @pytest.mark.parametrize('chunk_mode', [False, 'auto', 'specific', 'failsafe'])
 def test_swash(mode, chunk_mode):
     if chunk_mode in ['auto', ]:
-        dask.config.set({'array.chunk-size': '64KiB'})
+        dask.config.set({'array.chunk-size': '32KiB'})
     else:
         dask.config.set({'array.chunk-size': '128MiB'})
     field_set = fieldset_from_swash(chunk_mode)

parcels/field.py

@@ -200,7 +200,7 @@ def __init__(self, name, data, lon=None, lat=None, depth=None, time=None, grid=N
         self.c_data_chunks = []
         self.nchunks = []
         self.chunk_set = False
-        self.filebuffers = [None] * 3
+        self.filebuffers = [None] * 2
         if len(kwargs) > 0:
             raise SyntaxError('Field received an unexpected keyword argument "%s"' % list(kwargs.keys())[0])
 
@@ -392,7 +392,7 @@ def from_netcdf(cls, filenames, variable, dimensions, indices=None, grid=None,
         if 'full_load' in kwargs:  # for backward compatibility with Parcels < v2.0.0
             deferred_load = not kwargs['full_load']
 
-        if grid.time.size <= 3 or deferred_load is False:
+        if grid.time.size <= 2 or deferred_load is False:
             deferred_load = False
 
         if chunksize not in [False, None]:
@@ -1324,11 +1324,9 @@ def data_concatenate(self, data, data_to_concat, tindex):
         if tindex == 0:
             data = lib.concatenate([data_to_concat, data[tindex+1:, :]], axis=0)
         elif tindex == 1:
-            data = lib.concatenate([data[:tindex, :], data_to_concat, data[tindex+1:, :]], axis=0)
-        elif tindex == 2:
             data = lib.concatenate([data[:tindex, :], data_to_concat], axis=0)
         else:
-            raise ValueError("data_concatenate is used for computeTimeChunk, with tindex in [0, 1, 2]")
+            raise ValueError("data_concatenate is used for computeTimeChunk, with tindex in [0, 1]")
         return data
 
     def advancetime(self, field_new, advanceForward):

parcels/fieldset.py

@@ -993,7 +993,7 @@ def advancetime(self, fieldset_new):
         :param fieldset_new: FieldSet snapshot with which the oldest time has to be replaced"""
 
         logger.warning_once("Fieldset.advancetime() is deprecated.\n \
-                             Parcels deals automatically with loading only 3 time steps simustaneously\
+                             Parcels deals automatically with loading only 2 time steps simultaneously\
                              such that the total allocated memory remains limited.")
 
         advance = 0
@@ -1055,7 +1055,7 @@ def computeTimeChunk(self, time, dt):
                 else:
                     zd = g.zdim
                 data = lib.empty((g.tdim, zd, g.ydim-2*g.meridional_halo, g.xdim-2*g.zonal_halo), dtype=np.float32)
-                f.loaded_time_indices = range(3)
+                f.loaded_time_indices = range(2)
                 for tind in f.loaded_time_indices:
                     for fb in f.filebuffers:
                         if fb is not None:
@@ -1083,43 +1083,43 @@ def computeTimeChunk(self, time, dt):
                     zd = g.zdim
                 data = lib.empty((g.tdim, zd, g.ydim-2*g.meridional_halo, g.xdim-2*g.zonal_halo), dtype=np.float32)
                 if signdt >= 0:
-                    f.loaded_time_indices = [2]
+                    f.loaded_time_indices = [1]
                     if f.filebuffers[0] is not None:
                         f.filebuffers[0].close()
                         f.filebuffers[0] = None
-                    f.filebuffers[:2] = f.filebuffers[1:]
-                    data = f.computeTimeChunk(data, 2)
+                    f.filebuffers[0] = f.filebuffers[1]
+                    data = f.computeTimeChunk(data, 1)
                 else:
                     f.loaded_time_indices = [0]
-                    if f.filebuffers[2] is not None:
-                        f.filebuffers[2].close()
-                        f.filebuffers[2] = None
-                    f.filebuffers[1:] = f.filebuffers[:2]
+                    if f.filebuffers[1] is not None:
+                        f.filebuffers[1].close()
+                        f.filebuffers[1] = None
+                    f.filebuffers[1] = f.filebuffers[0]
                     data = f.computeTimeChunk(data, 0)
                 data = f.rescale_and_set_minmax(data)
                 if signdt >= 0:
-                    data = f.reshape(data)[2:, :]
+                    data = f.reshape(data)[1, :]
                     if lib is da:
-                        f.data = lib.concatenate([f.data[1:, :], data], axis=0)
+                        f.data = lib.stack([f.data[1, :], data], axis=0)
                     else:
                         if not isinstance(f.data, DeferredArray):
                             if isinstance(f.data, list):
                                 del f.data[0, :]
                             else:
                                 f.data[0, :] = None
-                        f.data[:2, :] = f.data[1:, :]
-                        f.data[2, :] = data
+                        f.data[0, :] = f.data[1, :]
+                        f.data[1, :] = data
                 else:
-                    data = f.reshape(data)[0:1, :]
+                    data = f.reshape(data)[0, :]
                     if lib is da:
-                        f.data = lib.concatenate([data, f.data[:2, :]], axis=0)
+                        f.data = lib.stack([data, f.data[0, :]], axis=0)
                     else:
                         if not isinstance(f.data, DeferredArray):
                             if isinstance(f.data, list):
-                                del f.data[2, :]
+                                del f.data[1, :]
                             else:
-                                f.data[2, :] = None
-                        f.data[1:, :] = f.data[:2, :]
+                                f.data[1, :] = None
+                        f.data[1, :] = f.data[0, :]
                         f.data[0, :] = data
                 g.load_chunk = np.where(g.load_chunk == g.chunk_loaded_touched,
                                         g.chunk_loading_requested, g.load_chunk)
@@ -1135,8 +1135,7 @@ def computeTimeChunk(self, time, dt):
                                     break
                                 block = f.get_block(block_id)
                                 f.data_chunks[block_id][0] = None
-                                f.data_chunks[block_id][:2] = f.data_chunks[block_id][1:]
-                                f.data_chunks[block_id][2] = np.array(f.data.blocks[(slice(3),)+block][2])
+                                f.data_chunks[block_id][1] = np.array(f.data.blocks[(slice(2),)+block][1])
                     else:
                         for block_id in range(len(g.load_chunk)):
                             if g.load_chunk[block_id] == g.chunk_loaded_touched:
@@ -1145,9 +1144,8 @@ def computeTimeChunk(self, time, dt):
                                     # happens when field not called by kernel, but shares a grid with another field called by kernel
                                     break
                                 block = f.get_block(block_id)
-                                f.data_chunks[block_id][2] = None
-                                f.data_chunks[block_id][1:] = f.data_chunks[block_id][:2]
-                                f.data_chunks[block_id][0] = np.array(f.data.blocks[(slice(3),)+block][0])
+                                f.data_chunks[block_id][1] = None
+                                f.data_chunks[block_id][0] = np.array(f.data.blocks[(slice(2),)+block][0])
         # do user-defined computations on fieldset data
         if self.compute_on_defer:
             self.compute_on_defer(self)

parcels/grid.py

@@ -191,45 +191,45 @@ def computeTimeChunk(self, f, time, signdt):
         prev_time_indices = self.time
         if self.update_status == 'not_updated':
             if self.ti >= 0:
-                if (time - periods*(self.time_full[-1]-self.time_full[0]) < self.time[0] or time - periods*(self.time_full[-1]-self.time_full[0]) > self.time[2]):
+                if time - periods*(self.time_full[-1]-self.time_full[0]) < self.time[0] or time - periods*(self.time_full[-1]-self.time_full[0]) > self.time[1]:
                     self.ti = -1  # reset
                 elif signdt >= 0 and (time - periods*(self.time_full[-1]-self.time_full[0]) < self.time_full[0] or time - periods*(self.time_full[-1]-self.time_full[0]) >= self.time_full[-1]):
                     self.ti = -1  # reset
                 elif signdt < 0 and (time - periods*(self.time_full[-1]-self.time_full[0]) <= self.time_full[0] or time - periods*(self.time_full[-1]-self.time_full[0]) > self.time_full[-1]):
                     self.ti = -1  # reset
-                elif signdt >= 0 and time - periods*(self.time_full[-1]-self.time_full[0]) >= self.time[1] and self.ti < len(self.time_full)-3:
+                elif signdt >= 0 and time - periods*(self.time_full[-1]-self.time_full[0]) >= self.time[1] and self.ti < len(self.time_full)-2:
                     self.ti += 1
-                    self.time = self.time_full[self.ti:self.ti+3]
+                    self.time = self.time_full[self.ti:self.ti+2]
                     self.update_status = 'updated'
-                elif signdt == -1 and time - periods*(self.time_full[-1]-self.time_full[0]) < self.time[1] and self.ti > 0:
+                elif signdt < 0 and time - periods*(self.time_full[-1]-self.time_full[0]) <= self.time[0] and self.ti > 0:
                     self.ti -= 1
-                    self.time = self.time_full[self.ti:self.ti+3]
+                    self.time = self.time_full[self.ti:self.ti+2]
                     self.update_status = 'updated'
             if self.ti == -1:
                 self.time = self.time_full
                 self.ti, _ = f.time_index(time)
                 periods = self.periods.value if isinstance(self.periods, c_int) else self.periods
                 if signdt == -1 and self.ti == 0 and (time - periods*(self.time_full[-1]-self.time_full[0])) == self.time[0] and f.time_periodic:
-                    self.ti = len(self.time)-2
+                    self.ti = len(self.time)-1
                     periods -= 1
                 if signdt == -1 and self.ti > 0:
                     self.ti -= 1
-                if self.ti >= len(self.time_full) - 2:
-                    self.ti = len(self.time_full) - 3
+                if self.ti >= len(self.time_full) - 1:
+                    self.ti = len(self.time_full) - 2
 
-                self.time = self.time_full[self.ti:self.ti+3]
-                self.tdim = 3
-                if prev_time_indices is None or len(prev_time_indices) != 3 or len(prev_time_indices) != len(self.time):
+                self.time = self.time_full[self.ti:self.ti+2]
+                self.tdim = 2
+                if prev_time_indices is None or len(prev_time_indices) != 2 or len(prev_time_indices) != len(self.time):
                     self.update_status = 'first_updated'
                 elif functools.reduce(lambda i, j: i and j, map(lambda m, k: m == k, self.time, prev_time_indices), True) and len(prev_time_indices) == len(self.time):
                     self.update_status = 'not_updated'
-                elif functools.reduce(lambda i, j: i and j, map(lambda m, k: m == k, self.time[:2], prev_time_indices[:2]), True) and len(prev_time_indices) == len(self.time):
+                elif functools.reduce(lambda i, j: i and j, map(lambda m, k: m == k, self.time[:1], prev_time_indices[:1]), True) and len(prev_time_indices) == len(self.time):
                     self.update_status = 'updated'
                 else:
                     self.update_status = 'first_updated'
-            if signdt >= 0 and (self.ti < len(self.time_full)-3 or not f.allow_time_extrapolation):
-                nextTime_loc = self.time[2] + periods*(self.time_full[-1]-self.time_full[0])
-            elif signdt == -1 and (self.ti > 0 or not f.allow_time_extrapolation):
+            if signdt >= 0 and (self.ti < len(self.time_full)-2 or not f.allow_time_extrapolation):
+                nextTime_loc = self.time[1] + periods*(self.time_full[-1]-self.time_full[0])
+            elif signdt < 0 and (self.ti > 0 or not f.allow_time_extrapolation):
                 nextTime_loc = self.time[0] + periods*(self.time_full[-1]-self.time_full[0])
         return nextTime_loc
 

tests/test_fieldset.py

@@ -521,10 +521,13 @@ def SampleUV2(particle, fieldset, time):
 
 @pytest.mark.parametrize('mode', ['scipy', 'jit'])
 @pytest.mark.parametrize('time_periodic', [4*86400.0, False])
+@pytest.mark.parametrize('dt', [-3600, 3600])
 @pytest.mark.parametrize('chunksize', [False, 'auto', {'time': ('time_counter', 1), 'lat': ('y', 32), 'lon': ('x', 32)}])
 @pytest.mark.parametrize('with_GC', [False, True])
 @pytest.mark.skipif(sys.platform.startswith("win"), reason="skipping windows test as windows memory leaks (#787)")
-def test_from_netcdf_memory_containment(mode, time_periodic, chunksize, with_GC):
+def test_from_netcdf_memory_containment(mode, time_periodic, dt, chunksize, with_GC):
+    if time_periodic and dt < 0:
+        return True  # time_periodic does not work in backward-time mode
     if chunksize == 'auto':
         dask.config.set({'array.chunk-size': '2MiB'})
     else:
@@ -575,7 +578,7 @@ def periodicBoundaryConditions(particle, fieldset, time):
     mem_0 = process.memory_info().rss
     mem_exhausted = False
     try:
-        pset.execute(pset.Kernel(AdvectionRK4)+periodicBoundaryConditions, dt=delta(hours=1), runtime=delta(days=7), postIterationCallbacks=postProcessFuncs, callbackdt=delta(hours=12))
+        pset.execute(pset.Kernel(AdvectionRK4)+periodicBoundaryConditions, dt=dt, runtime=delta(days=7), postIterationCallbacks=postProcessFuncs, callbackdt=delta(hours=12))
     except MemoryError:
         mem_exhausted = True
     mem_steps_np = np.array(perflog.memory_steps)
@@ -768,7 +771,7 @@ def DoNothing(particle, fieldset, time):
     assert np.allclose(fieldset.U.data, scale_fac*(zdim-1.)/zdim)
 
 
-@pytest.mark.parametrize('time2', [2, 7])
+@pytest.mark.parametrize('time2', [1, 7])
 def test_fieldset_initialisation_kernel_dask(time2, tmpdir, filename='test_parcels_defer_loading'):
     filepath = tmpdir.join(filename)
     data0, dims0 = generate_fieldset(3, 3, 4, 10)
@@ -789,7 +792,7 @@ class SampleParticle(JITParticle):
     pset = ParticleSet(fieldset, pclass=SampleParticle, time=[0, time2],
                        lon=[0.5, 0.5], lat=[0.5, 0.5], depth=[0.5, 0.5])
 
-    if time2 > 2:
+    if time2 > 1:
         failed = False
         try:
             pset.execute(SampleField, dt=0.)
@@ -905,3 +908,30 @@ def test_fieldset_from_data_gridtypes(xdim=20, ydim=10, zdim=4):
     pset.execute(AdvectionRK4, runtime=1, dt=.5)
     assert np.allclose(plon, pset.lon)
     assert np.allclose(plat, pset.lat)
+
+
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+@pytest.mark.parametrize('direction', [1, -1])
+@pytest.mark.parametrize('time_extrapolation', [True, False])
+def test_deferredload_simplefield(mode, direction, time_extrapolation, tmpdir, tdim=10):
+    filename = tmpdir.join("simplefield_deferredload.nc")
+    data = np.zeros((tdim, 2, 2))
+    for ti in range(tdim):
+        data[ti, :, :] = ti if direction == 1 else tdim-ti-1
+    ds = xr.Dataset({"U": (("t", "y", "x"), data), "V": (("t", "y", "x"), data)},
+                    coords={"x": [0, 1], "y": [0, 1], "t": np.arange(tdim)})
+    ds.to_netcdf(filename)
+
+    fieldset = FieldSet.from_netcdf(filename, {'U': 'U', 'V': 'V'}, {'lon': 'x', 'lat': 'y', 'time': 't'},
+                                    deferred_load=True, mesh='flat', allow_time_extrapolation=time_extrapolation)
+
+    class SamplingParticle(ptype[mode]):
+        p = Variable('p')
+    pset = ParticleSet(fieldset, SamplingParticle, lon=0.5, lat=0.5)
+
+    def SampleU(particle, fieldset, time):
+        particle.p = fieldset.U[particle]
+
+    runtime = tdim*2 if time_extrapolation else None
+    pset.execute(SampleU, dt=direction, runtime=runtime)
+    assert pset.p == tdim-1 if time_extrapolation else tdim-2