243 lines
7.7 KiB
Python
243 lines
7.7 KiB
Python
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
|
|
import numbers
|
|
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 findSlice(array,lims):
|
|
start = np.ravel(np.argwhere(array>lims[0]))[0]
|
|
stop = np.ravel(np.argwhere(array<lims[1]))[-1]
|
|
return slice(int(start),int(stop))
|
|
|
|
def removeBackground(x,data,xlims=None,max_iter=100,background_regions=[],**kw):
|
|
from dualtree import dualtree
|
|
if data.ndim == 1: data = data[np.newaxis,:]
|
|
if xlims is not None:
|
|
idx = findSlice(x,xlims)
|
|
x = x[idx]
|
|
data = data[:,idx].copy()
|
|
else:
|
|
data = data.copy(); # create local copy
|
|
# has to be a list of lists ..
|
|
if background_regions != [] and isinstance(background_regions[0],numbers.Real):
|
|
background_regions = [background_regions,]
|
|
background_regions = [findSlice(x,brange) for brange in background_regions]
|
|
for i in range(len(data)):
|
|
data[i] = data[i] - dualtree.baseline(data[i],max_iter=max_iter,
|
|
background_regions=background_regions,**kw)
|
|
return x,np.squeeze(data)
|
|
|
|
|
|
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,lw=3,
|
|
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)
|
|
|
|
def radToQ(theta,**kw):
|
|
""" theta is the scattering angle (theta_out - theta_in);
|
|
kw should be have either E or wavelength
|
|
it returns the scattering vector in the units of wavelength """
|
|
# Energy or wavelength should be in kw
|
|
assert "E" in kw or "wavelength" in kw
|
|
# but not both
|
|
assert not ("E" in kw and "wavelength" in kw)
|
|
if "E" in kw: kw["wavelength"] = 12.398/kw["E"]
|
|
return 4*np.pi/kw["wavelength"]*np.sin(theta)
|
|
|
|
def degToQ(theta,**kw):
|
|
theta = theta/180.*np.pi
|
|
return radToQ(theta,**kw)
|
|
degToQ.__doc__ = radToQ.__doc__
|
|
|
|
def qToTheta(q,**kw):
|
|
""" Return scattering angle from q (given E or wavelength) """
|
|
# Energy or wavelength should be in kw
|
|
assert "E" in kw or "wavelength" in kw
|
|
# but not both
|
|
assert not ("E" in kw and "wavelength" in kw)
|
|
if "E" in kw: kw["wavelength"] = 12.398/kw["E"]
|
|
return np.arcsin(q*kw["wavelength"]/4/np.pi)
|