attempt to write to zarr in separate thread

pfb/operators/gridder.py

@@ -3,6 +3,7 @@
 because we are not guaranteed that each imaging band has the same
 number of rows after BDA.
 """
+import concurrent.futures as cf
 from time import time
 import numpy as np
 import numba
@@ -757,17 +758,27 @@ def compute_residual(dsl,
     ti = time()
     ds['MODEL'] = (('x','y'), model)
     ds['RESIDUAL'] = (('x','y'), residual)
+    tassign = time() - ti
 
     # save
-    ds.to_zarr(output_name, mode='a')
+    ti = time()
+    with cf.ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(dataset_to_zarr, ds, output_name)
+    # future =None
+    # ti = time()
+    # ds.to_zarr(output_name, mode='a')
     twrite = time() - ti
 
     ttot = time() - tii
-    ttally = tread + tdegrid + tgrid + tdiff + twrite
+    ttally = tread + tdegrid + tgrid + tdiff + tassign + twrite
     print(f'tread = {tread/ttot}')
     print(f'tdegrid = {tdegrid/ttot}')
     print(f'tgrid = {tgrid/ttot}')
     print(f'tdiff = {tdiff/ttot}')
+    print(f'tassign = {tassign/ttot}')
     print(f'twrite = {twrite/ttot}')
     print(f'ttally = {ttally/ttot}')
-    return residual
+    return residual, future
+
+def dataset_to_zarr(ds, output_name):
+    ds.to_zarr(output_name, mode='a')

pfb/utils/dist.py

@@ -356,7 +356,7 @@ def set_residual(self, k, x=None):
             return np.zeros_like(x)
 
         # we don;t need counts here because we use the weights in the dds
-        residual = compute_residual(
+        residual, _ = compute_residual(
                                 self.cache_path,
                                 self.nx, self.ny,
                                 self.cell_rad, self.cell_rad,

pfb/workers/klean.py

@@ -273,7 +273,7 @@ def _klean(**kw):
         print(f'Computing residual', file=log)
         for ds_name, ds in zip(dds_list, dds):
             b = int(ds.bandid)
-            resid = compute_residual(ds_name,
+            resid, _ = compute_residual(ds_name,
                                      nx, ny,
                                      cell_rad, cell_rad,
                                      ds_name,
@@ -349,7 +349,7 @@ def _klean(**kw):
             print(f'Computing residual', file=log)
             for ds_name, ds in zip(dds_list, dds):
                 b = int(ds.bandid)
-                resid = compute_residual(ds_name,
+                resid, _ = compute_residual(ds_name,
                                         nx, ny,
                                         cell_rad, cell_rad,
                                         ds_name,

pfb/workers/sara.py

@@ -2,6 +2,7 @@
 import os
 import sys
 from contextlib import ExitStack
+import concurrent.futures as cf
 from pfb.workers.main import cli
 import click
 from omegaconf import OmegaConf
@@ -328,8 +329,7 @@ def _sara(**kw):
         residual *= beam  # avoid copy
         update = precond.idot(residual,
                               mode=opts.hess_approx,
-                              # only warm start close to convergence
-                              x0=update if eps < 0.1 else None)
+                              x0=update if update.any() else None)
         update_mfs = np.mean(update, axis=0)
         save_fits(update_mfs,
                   fits_oname + f'_{opts.suffix}_update_{k+1}.fits',
@@ -441,9 +441,10 @@ def _sara(**kw):
                   hdr_mfs)
 
         print(f'Computing residual', file=log)
+        write_futures = []
         for ds_name, ds in zip(dds_list, dds):
             b = int(ds.bandid)
-            residual[b] = compute_residual(
+            residual[b], fut = compute_residual(
                                     ds_name,
                                     nx, ny,
                                     cell_rad, cell_rad,
@@ -453,6 +454,8 @@ def _sara(**kw):
                                     epsilon=opts.epsilon,
                                     do_wgridding=opts.do_wgridding,
                                     double_accum=opts.double_accum)
+            write_futures.append(fut)
+
         residual /= wsum
         residual_mfs = np.sum(residual, axis=0)
         save_fits(residual_mfs,
@@ -532,4 +535,7 @@ def _sara(**kw):
                 print("Algorithm is diverging. Terminating.", file=log)
                 break
 
+    # make sure write futures have finished
+    cf.wait(write_futures)
+
     return