now mask can be a list of masks + add removebackground + bugfix

xray/azav.py

@@ -12,6 +12,7 @@ import glob
 import pathlib
 from . import storage
 from . import utils
+from . import filters
 import re
 import fabio
 import pyFAI
@@ -108,7 +109,7 @@ def getAI(poni,folder=None):
   elif isinstance(poni,str):
     # look is file exists in cwd
     if os.path.isfile(poni):
-      ai = pyFAI.load(poni)
+      fname = poni
     # if file does not exist look for one with that name around
     else:
       # build search paths
@@ -133,7 +134,7 @@ def getAI(poni,folder=None):
 
 g_mask_str = re.compile("(\w)\s*(<|>)\s*(\d+)")
 
-def interpretMask(mask,shape=None):
+def _interpretMask(mask,shape=None):
   """
     if mask is an existing filename, returns it
     if mask is a string like [x|y] [<|>] int;
@@ -174,14 +175,49 @@ def interpretMask(mask,shape=None):
   return maskout
 
 
-def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
-             saveChi=True,poni='pyfai.poni',storageFile='auto',diagnostic=None):
+def interpretMask(masks,shape=None):
+  """
+    if masks is a list of masks, eachone can be:
+    *  an existing filename
+    *  a string like [x|y] [<|>] int;
+  """
+  if not isinstance( masks, (list,tuple,np.ndarray) ):
+    masks = (masks,)
+  masks = [_interpretMask(mask,shape) for mask in masks]
+  # put them all together
+  mask = masks[0]
+  for m in masks[1:]:
+    mask = np.logical_or(mask,m)
+  return mask
+
+def removeBackground(data,qlims=(0,10),max_iter=30,background_regions=[],force=False,
+      storageFile=None,save=True,**removeBkg):
+  """ similar function to the zray.utils one, this works on dataset created by
+      doFolder """
+  idx = utils.findSlice(data.q_orig,qlims)
+  # see if there are some to do ...
+  if force:
+    idx_start = 0
+  else:
+    idx_start = len(data.data)
+  if idx_start < len(data.data_orig):
+    _q,_data = utils.removeBackground(data.q_orig[idx],data.data_orig[:,idx],
+               max_iter=max_iter,background_regions=background_regions,**removeBkg)
+    data.q = _q
+    data.data  = np.concatenate( (data.data,_data  ) )
+    data.err   = np.concatenate( (data.err ,data.err[idx_start,idx]   ) )
+  if save: data.save(storageFile); # if None uses .filename
+  return data
+
+
+def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,dark=10,norm='auto',
+             saveChi=True,poni='pyfai.poni',storageFile='auto',save=True,diagnostic=None):
   """ calc 1D curves from files in folder, returning a dictionary of stuff
       nQ    : number of Q-points (equispaced)
       force : if True, redo from beginning even if previous data are found
               if False, do only new files
-      mask  : can be a filename or an array of booleans; pixels that are True
-              are dis-regarded
+      mask  : can be a list of [filenames|array of booleans|mask string]
+              pixels that are True are dis-regarded
       saveChi: self-explanatory
       poni  : could be:
               → an AzimuthalIntegrator instance
@@ -199,17 +235,25 @@ def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
 
   if os.path.isfile(storageFile) and not force:
     saved = storage.DataStorage(storageFile)
+    log.info("Found %d images in storage file"%saved.data.shape[0])
   else:
     saved = None
 
-  # which poni file to use:
-  ai = getAI(poni,folder)
-
   files = utils.getFiles(folder,files)
   if saved is not None:
     files = [f for f in files if utils.getBasename(f) not in saved["files"]]
 
+  if diagnostic is not None:
+    key = list( diagnostic.keys() )[0]
+    files_diagnostic = list(diagnostic[key].keys())
+    files = [ f for f in files if utils.getBasename(f) in files_diagnostic ]
+  log.info("Will do azimuthal integration for %d files"%(len(files)))
+
   if len(files) > 0:
+    # which poni file to use:
+    ai = getAI(poni,folder)
+
+
     shape = read(files[0]).shape
     mask = interpretMask(mask,shape)
       
@@ -217,7 +261,7 @@ def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
     err    = np.empty( (len(files),nQ) )
     for ifname,fname in enumerate(files):
       img = read(fname)
-      q,i,e = do1d(ai,img,mask=mask,npt_radial=nQ)
+      q,i,e = do1d(ai,img,mask=mask,npt_radial=nQ,dark=dark)
       data[ifname] = i
       err[ifname]  = e
       if saveChi:
@@ -231,6 +275,7 @@ def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
       data  = np.concatenate( (saved["data"]  ,data  ) )
       err   = np.concatenate( (saved["err"]   ,err   ) )
     ret = dict(q=q,folder=folder,files=files,data=data,err=err,
+          data_orig=data.copy(),err_orig=err.copy(),q_orig=q.copy(),
           pyfai=ai_as_dict(ai),pyfai_info=ai_as_str(ai),mask=mask)
 
     # add info from diagnostic if provided
@@ -238,12 +283,24 @@ def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
       for k in diagnostic:
         ret[k] = np.asarray( [diagnostic[k][f] for f in ret['files']] )
     ret = storage.DataStorage(ret)
-    if storageFile is not None: ret.save(storageFile)
+    # sometime saving is not necessary (if one has to do it after subtracting background
+    if storageFile is not None and save: ret.save(storageFile)
   else:
     ret = saved
   return ret
 
 
+def removeBackground(data,qlims=(0,10),max_iter=30,background_regions=[],
+      storageFile=None,save=True,**removeBkg):
+  """ similar function to the zray.utils one, this works on dataset created by
+      doFolder """
+  idx = utils.findSlice(data.q,qlims)
+  data.q,data.data = utils.removeBackground(data.q[idx],data.data[:,idx],
+              max_iter=max_iter,background_regions=background_regions,**removeBkg)
+  data.err  = data.err[idx]
+  if save: data.save(storageFile); # if None uses .filename
+  return data
+
 def _calc_R(x,y, xc, yc):
   """ calculate the distance of each 2D points from the center (xc, yc) """
   return np.sqrt((x-xc)**2 + (y-yc)**2)
@@ -295,6 +352,7 @@ def find_center(img,psize=100e-6,dist=0.1,wavelength=0.8e-10,center=None,referen
     if reference is not None: ax_pyfai.axvline(reference)
     plt.pause(0.01)
     plt.draw()
+    plt.draw_all()
     ans=input("Enter to continue with clinking or enter xc,yc values ")
     if ans == '':
       center = None
@@ -361,7 +419,7 @@ def chiAverage(folder,basename="",scale=1,norm=None,returnAll=False,plot=False,s
     idx  = ( q>norm[0] ) & (q<norm[1])
     norm = np.nanmean(i[:,idx],axis=1)
     i = i/norm[:,np.newaxis]
-  if isinstance(norm,np.ndarray):
+  elif isinstance(norm,np.ndarray):
     i = i/norm[:,np.newaxis]
   title = "%s %s" % (folder,basename)
   utils.plotdata(q,i,plot=plot,showTrend=showTrend,title=title,clim=clim)