new helper function for masks; can interpret string like y>20

xray/azav.py

@@ -12,6 +12,7 @@ import glob
 import pathlib
 from . import storage
 from . import utils
+import re
 import fabio
 import pyFAI
 
@@ -130,6 +131,48 @@ def getAI(poni,folder=None):
     ai = pyFAI.load(fname)
   return ai
 
+g_mask_str = re.compile("(\w)\s*(<|>)\s*(\d+)")
+
+def interpretMask(mask,shape=None):
+  """
+    if mask is an existing filename, returns it
+    if mask is a string like [x|y] [<|>] int;
+      for example y>500 will dis-regard out for y>500
+  """
+  maskout = None
+  if isinstance(mask,str) and os.path.isfile(mask):
+    maskout = read(mask).astype(np.bool)
+  elif isinstance(mask,str) and not os.path.isfile(mask):
+    err_msg = ValueError("The string '%s' could not be interpreted as simple\
+              mask; it should be something like x>10"%mask)
+    assert shape is not None
+    # interpret string
+    maskout = np.zeros(shape,dtype=bool)
+    match = g_mask_str.match(mask)
+    if match is None: raise err_msg
+    (axis,sign,lim) = match.groups()
+    if axis not in ("x","y"): raise err_msg
+    if sign not in (">","<"): raise err_msg
+    lim = int(lim)
+    idx = slice(lim,None) if sign == ">" else slice(None,lim)
+    if axis == 'y':
+      maskout[idx,:] = True
+    else:
+      maskout[:,idx] = True
+  elif isinstance(mask,np.ndarray):
+    maskout = mask.astype(np.bool)
+  elif mask is None:
+    assert shape is not None
+    maskout = np.zeros(shape,dtype=bool)
+  else:
+    maskout = None
+    raise ValueError("Could not interpret %s as mask input"%mask)
+
+  if shape is not None and maskout.shape != shape:
+    raise ValueError("The mask shape %s does not match the shape given as\
+      argument %s"%(maskout.shape,shape))
+  return maskout
+
 
 def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
              saveChi=True,poni='pyfai.poni',storageFile='auto',diagnostic=None):
@@ -167,11 +210,8 @@ def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
     files = [f for f in files if utils.getBasename(f) not in saved["files"]]
 
   if len(files) > 0:
-    # work out mask to use
-    if isinstance(mask,np.ndarray):
-      mask = mask.astype(bool)
-    elif mask is not None:
-      mask = read(mask).astype(bool)
+    shape = read(files[0]).shape
+    mask = interpretMask(mask,shape)
       
     data   = np.empty( (len(files),nQ) )
     err    = np.empty( (len(files),nQ) )