Improve iterative storage (closes #239)

climix/datahandling.py

@@ -53,7 +53,7 @@ def prepare_input_data(datafiles):
 
 
 def save(result, output_filename, iterative_storage=False, client=None):
-    data = result.core_data()
+    data = result.core_data().rechunk()
     if iterative_storage:
         logger.info('Storing iteratively')
         logger.info('Creating empty data')
@@ -62,7 +62,6 @@ def save(result, output_filename, iterative_storage=False, client=None):
         iris.save(result, output_filename, fill_value=MISSVAL,
                   local_keys=['proposed_standard_name'])
         logger.info('Reopening output file and beginning storage')
-        result.data = data
         with netCDF4.Dataset(output_filename, 'a') as ds:
             var = ds[result.var_name]
             time_dim = result.coord_dims('time')[0]
@@ -70,19 +69,20 @@ def save(result, output_filename, iterative_storage=False, client=None):
 
             end = time.time()
             cumulative = 0.
-            for i, result_data in enumerate(result.core_data()):
-                result_id = f'{i+1}/{no_slices}'
-                r = result_data.persist()
-                progress(r)
-                print()
-                logging.info(f'Storing result {result_id}')
-                var[i, ...] = r
+            start_index = 0
+            for result_data in data.blocks:
+                result_id = f'{start_index+1}/{no_slices}'
                 logger.info(f'Storing partial result {result_id}')
+                end_index = start_index+result_data.shape[0]
+                logger.debug(f'{start_index}:{end_index}')
+                logger.debug(f'{result_data.shape}')
+                var[start_index:end_index, ...] = result_data
+                start_index = end_index
                 start = end
                 end = time.time()
                 last = end - start
                 cumulative += last
-                eta = cumulative/(i+1)*no_slices
+                eta = cumulative/(start_index+1)*no_slices
                 logger.info(f'Finished {result_id} in (last cum eta): '
                             f'{last:4.0f} {cumulative:4.0f} {eta:4.0f}')
     else: