mcutils/xray/azav.py

from __future__ import print_function,division

import logging as log
log.basicConfig(level=log.INFO)

import numpy as np
np.seterr(all='ignore')
import os
import collections
import glob
import pathlib
from . import storage
from . import utils
import pyFAI

try:
  import matplotlib.pyplot as plt
except ImportError:
  log.warn("Can't import matplotlib !")

def pyFAIread(fname):
  """ read data from file using fabio """
  import fabio
  f = fabio.open(fname)
  data = f.data
  del f
  return data

def pyFAI_dict(ai):
  """ ai is a pyFAI azimuthal intagrator"""
  methods = dir(ai)
  methods = [m for m in methods if m.find("get_") == 0]
  names   = [m[4:] for m in methods]
  values  = [getattr(ai,m)() for m in methods]
  ret = dict( zip(names,values) )
  ret["detector"] = ai.detector.get_name()
  return ret

def pyFAI_info(ai):
  """ ai is a pyFAI azimuthal intagrator"""
  return "#" + str(ai).replace("\n","\n#")

def pyFAI1d(ai, imgs, mask = None, npt_radial = 600, method = 'csr',safe=True,dark=10., polCorr = 1):
    """ ai is a pyFAI azimuthal intagrator 
              it can be defined with pyFAI.load(ponifile)
        mask: True are points to be masked out """
    # force float to be sure of type casting for img
    if isinstance(dark,int): dark = float(dark);
    if imgs.ndim == 2: imgs = (imgs,)
    out_i = np.empty( ( len(imgs), npt_radial) )
    out_s = np.empty( ( len(imgs), npt_radial) )
    for _i,img in enumerate(imgs):
      q,i, sig = ai.integrate1d(img-dark, npt_radial, mask= mask, safe = safe,\
                 unit="q_A^-1", method = method, error_model = "poisson",
                 polarization_factor = polCorr)
      out_i[_i] = i
      out_s[_i] = sig
    return q,np.squeeze(out_i),np.squeeze(out_s)

def pyFAI2d(ai, imgs, mask = None, npt_radial = 600, npt_azim=360,method = 'csr',safe=True,dark=10., polCorr = 1):
    """ ai is a pyFAI azimuthal intagrator 
              it can be defined with pyFAI.load(ponifile)
        mask: True are points to be masked out """
    # force float to be sure of type casting for img
    if isinstance(dark,int): dark = float(dark);
    if imgs.ndim == 2: imgs = (imgs,)
    out = np.empty( ( len(imgs), npt_azim,npt_radial) )
    for _i,img in enumerate(imgs):
      i2d,q,azTheta = ai.integrate2d(img-dark, npt_radial, npt_azim=npt_azim,
                      mask= mask, safe = safe,unit="q_A^-1", method = method,
                      polarization_factor = polCorr )
      out[_i] = i2d
    return q,azTheta,np.squeeze(out)

def _getAI(poni,folder):
  if isinstance(poni,pyFAI.azimuthalIntegrator.AzimuthalIntegrator):
    ai = poni
  elif isinstance(poni,dict):
    ai = pyFAI.azimuthalIntegrator.AzimuthalIntegrator(**poni)
  else:
    if poni == 'auto':
      temp = os.path.abspath(folder)
      path = pathlib.Path(temp)
      folders = [ str(path), ]
      for p in path.parents: folders.append(str(p))
      folders.append( "./" )
      folders.append( os.path.expanduser("~/") )
      for path in folders:
        poni = path + "/" + "pyfai.poni"
        if os.path.exists(poni):
          log.info("Found pyfai.poni in %s",path)
          break
        else:
          log.debug("Could not find pyfai.poni in %s",path)
    ai = pyFAI.load(poni)
  return ai


def doFolder(folder,files='*.edf*',nQ = 1500,force=False,mask=None,
             saveChi=True,poni='auto',storageFile='auto',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
      saveChi: self-explanatory
      poni  : could be:
              → an AzimuthalIntegrator instance
              → a filename
              → a dictionary (use to bootstrap an AzimuthalIntegrator using 
                AzimuthalIntegrator(**poni)
              → the string 'auto' that will look for the file 'pyfai.poni' in:
                 1 'folder' first
                 2 in ../folder
                 3 in ../../folder
                 ....
                 n-1 in pwd
                 n   in homefolder
  """
  if storageFile == 'auto': storageFile = folder + "/" + "pyfai_1d.h5"

  if os.path.exists(storageFile) and not force:
    saved = utils.data_storage(storageFile)
  else:
    saved = None

  # which poni file to use:
  ai = _getAI(poni,folder)

  files = glob.glob("%s/%s"%(folder,files))
  files.sort()
  if saved is not 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 = pyFAIread(mask).astype(bool)
      
    data   = np.empty( (len(files),nQ),dtype=np.float32 )
    err    = np.empty( (len(files),nQ),dtype=np.float32 )
    for ifname,fname in enumerate(files):
      img = pyFAIread(fname)
      q,i,e = pyFAI1d(ai,img,mask=mask,npt_radial=nQ)
      data[ifname] = i
      err[ifname]  = e
      if saveChi:
        chi_fname = utils.removeExt(fname) + ".chi"
        utils.saveTxt(chi_fname,q,i,e,info=pyFAI_dict(ai),overwrite=True)

    files = [ utils.getBasename(f) for f in files ]
    files = np.asarray(files)
    if saved is not None:
      files = np.concatenate( (saved["files"] ,files ) )
      data  = np.concatenate( (saved["data"]  ,data  ) )
      err   = np.concatenate( (saved["err"]   ,err   ) )
    ret = dict(q=q,folder=folder,files=files,data=data,err=err,
          pyfai=pyFAI_dict(ai),pyfai_info=pyFAI_info(ai),mask=mask)

    # add info from diagnostic if provided
    if diagnostic is not None:
      for k in diagnostic:
        ret[k] = np.asarray( [diagnostic[k][f] for f in ret['files']] )
    ret = utils.data_storage(ret)
    if storageFile is not None: ret.save(storageFile)
  else:
    ret = saved
  return ret


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)

def _chi2(c, x, y):
  """ calculate the algebraic distance between the data points and the mean
      circle centered at c=(xc, yc) """
  Ri = _calc_R(x, y, *c)
  return Ri - Ri.mean()

def leastsq_circle(x,y):
  from scipy import optimize
  # coordinates of the barycenter
  center_estimate = np.nanmean(x), np.nanmean(y)
  center, ier = optimize.leastsq(_chi2, center_estimate, args=(x,y))
  xc, yc = center
  Ri       = _calc_R(x, y, *center)
  R        = Ri.mean()
  residu   = np.sum((Ri - R)**2)
  return xc, yc, R

def pyFAI_find_center(img,psize=100e-6,dist=0.1,wavelength=0.8e-10,**kwargs):
  plt.ion()
  kw = dict( pixel1 = psize, pixel2 = psize, dist = dist,wavelength=wavelength )
  kw.update(kwargs)
  ai =  pyFAI.azimuthalIntegrator.AzimuthalIntegrator(**kw)
  fig_img,ax_img = plt.subplots(1,1)
  fig_pyfai,ax_pyfai = plt.subplots(1,1)
  fig_pyfai = plt.figure(2)
  ax_img.imshow(img)
  plt.sca(ax_img); # set figure to use for mouse interaction
  ans = ""
  print("Enter 'end' when done")
  while ans != "end":
    if ans == "":
      print("Click on beam center:")
      plt.sca(ax_img); # set figure to use for mouse interaction
      xc,yc = plt.ginput()[0]
    else:
      xc,yc = map(float,ans.split(","))
    print("Selected center:",xc,yc)
    ai.set_poni1(xc*psize)
    ai.set_poni2(yc*psize)
    q,az,i = pyFAI2d(ai,img)
    ax_pyfai.pcolormesh(q,az,i)
    ax_pyfai.set_title(str( (xc,yc) ))
    plt.pause(0.01)
    plt.draw()
    ans=input("Enter to continue with clinking or enter xc,yc values")
  print("Final values: (in pixels) %.3f %.3f"%(xc,yc))
  return ai

def plotdata(q,data,plot=True,showTrend=True,title=None,clim='auto'):
  if not (plot or showTrend): return
  if clim == 'auto': clim = np.nanpercentile(data,(1.5,98.5))
  one_plot = showTrend or plot
  two_plot = showTrend and plot
  if one_plot and not two_plot:
    fig,ax = plt.subplots(1,1)
  if two_plot:
    fig,ax = plt.subplots(1,2,sharey=True)
  ax = np.atleast_1d(ax)
  if showTrend:
    plt.sca(ax[0])
    plt.pcolormesh(np.arange(data.shape[0]),q,data.T)
    plt.xlabel("image number, 0 being older")
    plt.ylabel(r"q ($\AA^{-1}$)")
    plt.clim( *clim )
  if plot:
    if showTrend:
      ax[1].plot(data.mean(axis=0),q)
    else:
      ax[0].plot(q,data.mean(axis=0))
  if (plot or showTrend) and title is not None:
    plt.title(title)

def average(fileOrFolder,delays=slice(None),scale=1,norm=None,returnAll=False,plot=False,
  showTrend=False):
  data = utils.data_storage(fileOrFolder)
  if isinstance(delays,slice):
    idx = np.arange(data.delays.shape[0])[delays]
  elif isinstance(delays,(int,float)):
    idx = data.delays == float(delays)
  else:
    idx = data.delays < 0
  if idx.sum() == 0:
    print("No data with the current filter")
    return None
  i   = data.data[idx]
  q   = data.q
  if isinstance(norm,(tuple,list)):
    idx  = ( q>norm[0] ) & (q<norm[1])
    norm = np.nanmean(i[:,idx],axis=1)
    i = i/norm[:,np.newaxis]
  if isinstance(norm,np.ndarray):
    i = i/norm[:,np.newaxis]
  title = "%s %s" % (fileOrFolder,str(delays))
  plotdata(q,i*scale,showTrend=showTrend,plot=plot,title=title)
  if returnAll:
    return q,i.mean(axis=0)*scale,i
  else:
    return q,i.mean(axis=0)*scale

#### Utilities for chi files ####
def chiRead(fname,scale=1):
  q,i = np.loadtxt(fname,unpack=True,usecols=(0,1))
  return q,i*scale

def chiPlot(fname,useTheta=False,E=12.4):
  q,i = chiRead(fname)
  lam = 12.4/E
  theta = 2*180/3.14*np.arcsin(q*lam/4/3.14)
  x = theta if useTheta else q
  plt.plot(x,i,label=fname)



def chiAverage(folder,basename="",scale=1,norm=None,returnAll=False,plot=False,showTrend=False,clim='auto'):
  files = glob.glob("%s/%s*chi"%(folder,basename))
  files.sort()
  print(files)
  if len(files) == 0:
    print("No file found (basename %s)" % basename)
    return None
  q,_ = chiRead(files[0])
  i   = np.asarray( [ chiRead(f)[1] for f in files ] )
  if isinstance(norm,(tuple,list)):
    idx  = ( q>norm[0] ) & (q<norm[1])
    norm = np.nanmean(i[:,idx],axis=1)
    i = i/norm[:,np.newaxis]
  if isinstance(norm,np.ndarray):
    i = i/norm[:,np.newaxis]
  title = "%s %s" % (folder,basename)
  plotdata(q,i,plot=plot,showTrend=showTrend,title=title,clim=clim)
  if (showTrend and plot): plt.subplot(1,2,1)
  if showTrend:
    plt.pcolormesh(np.arange(i.shape[0]),q,i.T)
    plt.xlabel("image number, 0 being older")
    plt.ylabel(r"q ($\AA^{-1}$)")
  if (showTrend and plot): plt.subplot(1,2,2)
  if plot:
    plt.plot(q,i.mean(axis=0)*scale)
  if (plot or showTrend):
    plt.title(folder+"/"+basename)
  if returnAll:
    return q,i.mean(axis=0)*scale,i
  else:
    return q,i.mean(axis=0)*scale