mcutils/xray/utils.py

from __future__ import print_function,division

import logging
log = logging.getLogger(__name__)  # __name__ is "foo.bar" here

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

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

_time_regex =      re.compile( "(-?\d+\.?\d*(?:ps|ns|us|ms)?)")
_timeInStr_regex = re.compile("_(-?\d+\.?\d*(?:ps|ns|us|ms)?)")

def getFiles(folder,basename="*.edf*",nFiles=None):
  files = glob.glob(folder + "/" + basename)
  files.sort()
  if nFiles is not None: files = files[:nFiles]
  return files

def getEdfFiles(folder,**kw):
  return getFiles(folder,basename="*.edf*",**kw)

def getDelayFromString(string) :
  match = _timeInStr_regex_regex.search(string)
  return match and match.group(1) or None

_time_regex = re.compile("(-?\d+\.?\d*)((?:s|fs|ms|ns|ps|us)?)")
def strToTime(delay) :
  _time2value = dict( fs = 1e-15, ps = 1e-12, ns = 1e-9, us = 1e-6, ms = 1e-3, s = 1)

  match = _time_regex.search(delay)
  if match:
    n,t = float(match.group(1)),match.group(2)
    value = _time2value.get(t,1)
    return n*value
  else:
    return None

def timeToStr(delay,fmt="%+.0f"):
  a_delay = abs(delay)
  if a_delay >= 1:
    ret = fmt % delay + "s"
  elif 1e-3 <= a_delay < 1: 
    ret = fmt % (delay*1e3) + "ms"
  elif 1e-6 <= a_delay < 1e-3: 
    ret = fmt % (delay*1e6) + "us"
  elif 1e-9 <= a_delay < 1e-6: 
    ret = fmt % (delay*1e9) + "ns"
  elif 1e-12 <= a_delay < 1e-9: 
    ret = fmt % (delay*1e12) + "ps"
  elif 1e-15 <= a_delay < 1e-12: 
    ret = fmt % (delay*1e12) + "fs"
  elif 1e-18 <= a_delay < 1e-15: 
    ret = fmt % (delay*1e12) + "as"
  else:
    ret = str(delay) +"s"
  return ret

def removeExt(fname):
  """ special remove extension meant to work with compressed files.edf and .edf.gz files """
  if fname[-3:] == ".gz": fname = fname[-3:]
  return os.path.splitext(fname)[0]

def getBasename(fname):
  return removeExt(os.path.basename(fname));

def plotdata(q,data,x=None,plot=True,showTrend=True,title=None,clim='auto'):
  if not (plot or showTrend): return
  if x is None: x = np.arange(data.shape[0])
  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(x,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 plotdiffs(q,diffs,t,select=None,err=None,absSignal=None,absSignalScale=10,
              showErr=False,cmap=plt.cm.jet):
  # this selection trick done in this way allows to keep the same colors when 
  # subselecting (because I do not change the size of diffs)
  if select is not None:
    indices = range(*select.indices(t.shape[0]))
  else:
    indices = range(len(t))
  lines = []
  if absSignal is not None:
    line = plt.plot(q,absSignal/absSignalScale,
                    color='k',label="absSignal/%s"%str(absSignalScale))[0]
    lines.append(line)
  for idiff in indices:
    color = cmap(idiff/(len(diffs)-1))
    label = timeToStr(t[idiff])
    kw = dict( color = color, label = label )
    if err is not None and showErr:
      line = plt.errorbar(q,diffs[idiff],err[idiff],**kw)[0]
    else:
      line = plt.plot(q,diffs[idiff],**kw)[0]
    lines.append(line)

  fig    = plt.gcf() 
  legend = plt.legend()
  plt.grid()
  plt.xlabel(r"q ($\AA^{-1}$)")
  # we will set up a dict mapping legend line to orig line, and enable
  # picking on the legend line
  lined = dict()
  for legline, origline in zip(legend.get_lines(), lines):
    legline.set_picker(5)  # 5 pts tolerance
    lined[legline] = origline

  def onpick(event):
    # on the pick event, find the orig line corresponding to the
    # legend proxy line, and toggle the visibility
    legline = event.artist
    origline = lined[legline]
    vis = not origline.get_visible()
    origline.set_visible(vis)
    # Change the alpha on the line in the legend so we can see what lines
    # have been toggled
    if vis:
        legline.set_alpha(1.0)
    else:
        legline.set_alpha(0.2)
    fig    = plt.gcf() 
    fig.canvas.draw()

  fig.canvas.mpl_connect('pick_event', onpick)


def saveTxt(fname,q,data,headerv=None,info=None,overwrite=True,columns=''):
  """ Write data to file 'fname' in text format.
      Inputs:
        q = x vector
        data = one or 2D array (first axis is q)
        info = dictionary (saved as '# key : value') or string
        headerv = vector to be used as header or string
  """
  if os.path.exists(fname) and not overwrite:
    log.warn("File %s exists, returning",fname)
    return
  if isinstance(info,dict):
    header = [ "# %s : %s" %(k,str(v)) for (k,v) in info.items() ]
    header = "\n".join(header); # skip first #, will be added by np
  elif isinstance(info,str):
    header = info
  else:
    header = ""
  if isinstance(headerv,str): header += "\n%s" % headerv
  if data.ndim == 1:
    x = np.vstack( (q,data) )
  elif data.ndim == 2:
    x = np.vstack( (q,data) )
    if headerv is not None:
      headerv = np.concatenate(( (data.shape[1],),headerv))
      x = np.hstack( (headerv[:,np.newaxis],x) )
  if columns != '':
    s = "#" + " ".join( [str(c).center(12) for c in columns] )
    header = header + "\n" + s if header != '' else s
  np.savetxt(fname,x.T,fmt="%+10.5e",header=header,comments='')

def reshapeToBroadcast(what,ref):
  """ expand the 1d array 'what' to allow broadbasting to match 
      multidimentional array 'ref'. The two arrays have to same the same 
      dimensions along the first axis
  """
  assert what.shape[0] == ref.shape[0]
  shape  = [ref.shape[0],] + [1,]*(ref.ndim-1)
  return what.reshape(shape)