Clean up schema and dependencies

README.rst

@@ -75,7 +75,7 @@ By default only a single worker is used so that datasets will be processed one a
 This results in the smallest memory footprint but won't always result in optimal resource utilisation.
 It also allows for easy debugging as there is no Dask cluster involved in this case.
 However, for better resource utilisation and or distributing computations over a cluster, you may wish to set `--nworkers` larger than one.
-This uses multiple Dask workers (processes) to process chunks in parallel and is especially useful for the `init`, `grid` and `fluxmop` applications.
+This uses multiple Dask workers (processes) to process chunks in parallel and is especially useful for the `init`, `grid` and `fluxtractor` applications.
 It is usually advisable to set `--nworkers` to the number of desired imaging bands which is set by the `--channels-per-image` parameter when initialising corrected Stokes visibilities with the `init` application.
 The product of `--nworkers` and `--nthreads-per-worker` should not exceed the available resources.
 

pfb/__init__.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import logging
 # import psutil
 # import resource
 
@@ -14,7 +15,9 @@ def set_envs(nthreads, ncpu):
     os.environ["OPENBLAS_NUM_THREADS"] = str(nthreads)
     os.environ["MKL_NUM_THREADS"] = str(nthreads)
     os.environ["VECLIB_MAXIMUM_THREADS"] = str(nthreads)
-    # os.environ["NUMBA_NUM_THREADS"] = str(nthreads)
+    os.environ["NPY_NUM_THREADS"] = str(nthreads)
+    os.environ["NUMBA_NUM_THREADS"] = str(nthreads)
+    os.environ["JAX_PLATFORMS"] = 'cpu'
     os.environ["JAX_ENABLE_X64"] = 'True'
     # this may be required for numba parallelism
     # find python and set LD_LIBRARY_PATH
@@ -30,23 +33,45 @@ def set_envs(nthreads, ncpu):
     os.environ["NUMEXPR_NUM_THREADS"] = str(ne_threads)
 
 
-def set_client(nworkers, log, stack=None,
-               host_address=None, direct_to_workers=False):
+def set_client(nworkers, log, stack=None, host_address=None,
+               direct_to_workers=False, client_log_level=None):
+
+    import warnings
+    warnings.filterwarnings("ignore", message="Port 8787 is already in use")
+    if client_log_level == 'error':
+        logging.getLogger("distributed").setLevel(logging.ERROR)
+        logging.getLogger("bokeh").setLevel(logging.ERROR)
+        logging.getLogger("tornado").setLevel(logging.CRITICAL)
+    elif client_log_level == 'warning':
+        logging.getLogger("distributed").setLevel(logging.WARNING)
+        logging.getLogger("bokeh").setLevel(logging.WARNING)
+        logging.getLogger("tornado").setLevel(logging.WARNING)
+    elif client_log_level == 'info':
+        logging.getLogger("distributed").setLevel(logging.INFO)
+        logging.getLogger("bokeh").setLevel(logging.INFO)
+        logging.getLogger("tornado").setLevel(logging.INFO)
+    elif client_log_level == 'debug':
+        logging.getLogger("distributed").setLevel(logging.DEBUG)
+        logging.getLogger("bokeh").setLevel(logging.DEBUG)
+        logging.getLogger("tornado").setLevel(logging.DEBUG)
+
     import dask
-    dask.config.set({'distributed.comm.compression': 'lz4'})
     # set up client
     host_address = host_address or os.environ.get("DASK_SCHEDULER_ADDRESS")
     if host_address is not None:
         from distributed import Client
         print("Initialising distributed client.", file=log)
-        client = stack.enter_context(Client(host_address))
+        if stack is not None:
+            client = stack.enter_context(Client(host_address))
+        else:
+            client = Client(host_address)
     else:
         from dask.distributed import Client, LocalCluster
         print("Initialising client with LocalCluster.", file=log)
         dask.config.set({
                 'distributed.comm.compression': {
                     'on': True,
-                    'type': 'blosc'
+                    'type': 'lz4'
                 }
         })
         cluster = LocalCluster(processes=True,

pfb/deconv/clark.py

@@ -116,9 +116,9 @@ def clark(ID,
     model = np.zeros((nband, nx, ny), dtype=ID.dtype)
     IR = ID.copy()
     # pre-allocate arrays for doing FFT's
-    xout = empty_noncritical(ID.shape, dtype=ID.dtype)
-    xpad = empty_noncritical(PSF.shape, dtype=ID.dtype)
-    xhat = empty_noncritical(PSFHAT.shape, dtype=PSFHAT.dtype)
+    xout = empty_noncritical(ID.shape, dtype='f8')
+    xpad = empty_noncritical(PSF.shape, dtype='f8')
+    xhat = empty_noncritical(PSFHAT.shape, dtype='c16')
     # square avoids abs of full array
     IRsearch = np.sum(IR, axis=0)**2
     pq = IRsearch.argmax()