From 8e48d72fa6a356dcc29b3ee72dd62da57ea168b3 Mon Sep 17 00:00:00 2001 From: marco cammarata Date: Thu, 12 May 2016 17:37:35 +0200 Subject: [PATCH] first git.ipr commit --- mcutils.py | 1293 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 1293 insertions(+) create mode 100644 mcutils.py diff --git a/mcutils.py b/mcutils.py new file mode 100644 index 0000000..8ceaf4f --- /dev/null +++ b/mcutils.py @@ -0,0 +1,1293 @@ +from __future__ import print_function,division +#import sys +#if (sys.version_info[0] < 3): +import numpy as np +from numpy import exp +import re +import codecs +import string +import scipy.signal +import functools +import types +import os +import pylab as plt +from itertools import chain +sqrt2=np.sqrt(2) + + +### COLORS, ETC ### +colors = None + +nice_colors = ( +# from http://www.mulinblog.com/a-color-palette-optimized-for-data-visualization/ + "#4D4D4D", # (gray) + "#5DA5DA", # (blue) + "#FAA43A", # (orange) + "#60BD68", # (green) + "#F17CB0", # (pink) + "#B2912F", # (brown) + "#B276B2", # (purple) + "#DECF3F", # (yellow) + "#F15854", # (red) +) + +def colormap( list_of_colors ): + from matplotlib.colors import colorConverter,LinearSegmentedColormap + c = [ colorConverter.to_rgba(l) for l in list_of_colors ] + # make the colormaps + cmap = LinearSegmentedColormap.from_list('cmap',c,256) + return cmap + +def simpleaxis(ax=plt.gca()): + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + ax.get_xaxis().tick_bottom() + ax.get_yaxis().tick_left() + +def noaxis(ax=plt.gca()): + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + ax.spines['left'].set_visible(False) + ax.spines['bottom'].set_visible(False) + +def xrayAttLenght(*args,**kw): + from periodictable import xsf + n = xsf.index_of_refraction(*args,**kw) + if not "wavelength" in kw: + wavelength = 12.398/np.asarray(kw["energy"]) + else: + wavelength = np.asarray(kw["wavelength"]) + attenuation_length = np.abs( (wavelength*1e-10)/ (4*np.pi*np.imag(n)) ) + return attenuation_length + + +def xrayFluo(atom,density,energy=7.,length=30.,I0=1e10,\ + det_radius=1.,det_dist=10.,det_material="Si",det_thick=300,verbose=True): + """ compound: anything periodictable would understand + density: in mM + length: sample length in um + energy: in keV, could be array + """ + import periodictable + from periodictable import xsf + wavelength = 12.398/energy + atom = periodictable.__dict__[atom] + # 1e3 -> from mM to M + # 1e3 -> from L to cm3 + # so 1 mM = 1e-3M/L = 1e-6 M/cm3 + density_g_cm3 = density*1e-6*atom.mass + n = xsf.index_of_refraction(atom,density=density_g_cm3,wavelength=wavelength) + attenuation_length = xrayAttLenght( atom,density=density_g_cm3,energy=energy ) + # um-> m: 1e-6 + fraction_absorbed = 1.-np.exp(-length*1e-6/attenuation_length) + if verbose: + print("Fraction of x-ray photons absorbed by the sample:",fraction_absorbed) + ## detector ## + det_area = np.pi*det_radius**2 + det_solid_angle = det_area/(4*np.pi*det_dist**2) + if verbose: + print("Detector fraction of solid angle:",det_solid_angle) + det_att_len = xrayAttLenght(det_material,wavelength=atom.K_alpha) + det_abs = 1-np.exp(-det_thick*1e-6/det_att_len) + if verbose: + print("Detector efficiency (assuming E fluo = E_K_alpha):",det_abs) + eff = fraction_absorbed*det_solid_angle*det_abs + if verbose: + print("Overall intensity (as ratio to incoming one):",eff) + return eff + + +def pulseDuration(t0,L,GVD): + return t0*np.sqrt( 1.+(L/(t0**2/2/np.abs(GVD)))) + +class MCError(Exception): + def __init__(self, value): + self.value = value + def __str__(self): + return repr(self.value) + + +### VECTOR ... ETC. ### + +def vectorLenght(v): + """ assuming axis -1 as coord """ + return np.sqrt(np.sum(v*v,axis=-1)) + +def versor(v): + return v/vectorLenght(v)[:,np.newaxis] + + +### LIST,INDEXING ... ETC. ### + +def smartIdx(idx,forceContigous=False): + """ Try to interpret an array of bool as slice; + this allows selecting a subarray alot more efficient + since array[slice] it returns a view and not a copy """ + if (isinstance(idx,int)): + ret = slice(idx,idx+1) + elif (isinstance(idx,slice)): + ret = idx + else: + idx = np.asarray(idx) + if idx.dtype == np.bool: + i = np.where(idx)[0] + else: + i = idx + # in case there is only one + if (len(i) == 1): + ret = slice(i[0],i[0]+1) + return ret + if forceContigous: + ret = slice(i[0],i[-1]) + else: + d = i[1:]-i[0:-1] + dmean = int(d.mean()) + if np.all(d==dmean): + ret = slice(i[0],i[-1]+1,dmean) + else: + ret = idx + return ret + + +### CONVOLUTION,INTERPOLATION,SMOOTHING ... ETC. ### + +def poly_approximant(x,y,order=10,allowExtrapolation=False,fill_value=0): + """ return a polinomial view """ + poly = np.polyfit(x,y,order) + def f(xx): + res = np.polyval(poly,xx) + if allowExtrapolation: + return res + else: + if np.isscalar(xx) and ( ( xx>x.max()) or (xxx.max() ) + res[idx] = fill_value + return res + return f + + +def smoothing(x,y,err=None,k=5,s=None,newx=None,derivative_order=0): + idx = np.isnan(x)|np.isnan(y) + idx = ~ idx + if newx is None: newx=x + if idx.sum() > 0: + x=x[idx] + y=y[idx] + if idx.sum() < 3: + return np.ones(len(newx)) + if err is None: + w=None + elif err == "auto": + n=len(x) + imin = max(0,n/2-20) + imax = min(n,n/2+20) + idx = range(imin,imax) + p = np.polyfit(x[idx],y[idx],4) + e = np.std( y[idx] - np.polyval(p,x[idx] ) ) + w = np.ones_like(x)/e + else: + w=np.ones_like(x)/err + from scipy.interpolate import UnivariateSpline + if (s is not None): + s = len(x)*s + s = UnivariateSpline(x, y,w=w, k=k,s=s) + if (derivative_order==0): + return s(newx) + else: + try: + len(derivative_order) + return [s.derivative(d)(newx) for d in derivative_order] + except: + return s.derivative(derivative_order)(newx) + + +def interpolator(x,y,kind='linear',axis=-1, copy=False, bounds_error=False, fill_value=np.nan): + from scipy import interpolate + if (kind != "linear"): + print("Warning interp1d can be VERY SLOW when using something that is not liear") + f = interpolate.interp1d(x,y,kind=kind,axis=axis,copy=copy,bounds_error=bounds_error,fill_value=fill_value) + return f + +def interpolator_spl(x,y,kind="cubic"): + from scipy import interpolate as itp + if kind == "linear": kind=1 + if kind == "cubic" : kind=3 + splinepars = itp.splrep(x,y,k=kind) + def f(x,der=0,): + """f(x) returns values for x[i]. f(x,order) return order-th derivative""" + return itp.splev(x,splinepars,der=der) + return f + +def interpolate_fast(x,y,newx,kind='cubic'): + f = interpolator_spl(x,y,kind=kind) + return f(newx) + +def interpolate(x,y,newx,kind='linear',axis=-1, copy=False, bounds_error=False, fill_value=np.nan): + f = interpolator(x,y,kind=kind,axis=axis,copy=copy,bounds_error=bounds_error,fill_value=fill_value) + return f(newx) + +def getElement(a,i,axis=-1): + nDim = a.ndim + if (axis<0): axis = nDim+axis + colon = (slice(None),) + return a[colon*axis+(i,)+colon*(nDim-axis-1)] + +def setElement(a,i,res,axis=-1): + temp = getElement(a,i,axis=axis) + np.copyto(temp,res) + +def invertedView(x): + return x[ slice(None,None,-1) ] + +def convolveQuad(x,y,xres,yres,useAutoCrop=True, + fill_value=0.,axis=-1): + print("NOT FINISHED") + return + from scipy import integrate + if (useAutoCrop): + idxRes = smartIdx( yres>(1e-4*yres.max()) ) + #print("Autocropping",idxRes,xres.shape) + else: + idxRes = slice(None) + xresInv = -invertedView( xres[idxRes] ) + yresInv = invertedView( yres[idxRes] ) + area = integrate.simps(yres,x=xres) + + f = interpolator(x,y) + r = interpolator(xres,yres) + + if y.ndim < 2: y = y[np.newaxis,:] + nDim = y.ndim + if (axis<0): axis = nDim+axis + + ## expand yres to allow broadcasting + sh = [1,]*y.ndim + sh[axis] = len(yres) + yres = yres.reshape(sh) + + ## prepare output + out = np.empty_like(y) + nP = len(x) + for i in range(nP): + # interpolate x,y on xres+x[i] + x_integral = xresInv+x[i] + ytemp = interpolate(x,y,x_integral,fill_value=fill_value,axis=axis)/area + # do integration + res = integrate.simps(ytemp*yresInv,x=xresInv+x[i],axis=axis) + colon = (slice(None),) + setElement(out,i,res,axis=axis) + return out + + +def convolve(x,y,xres,yres,useAutoCrop=True,approximantOrder=None,fill_value=0.,axis=-1): + """ if approximantOrder is not None, use interpolating polynomial of order + approximantOrder to perform integration """ + from scipy import integrate + import copy + if (useAutoCrop): + idxRes = smartIdx( yres>(1e-4*yres.max()) ) + #print("Autocropping",idxRes,xres.shape) + else: + idxRes = slice(None) + xresInv = -invertedView( xres[idxRes] ) + yresInv = invertedView( yres[idxRes] ) + area = integrate.simps(yres,x=xres) + + if approximantOrder is not None: + #print("Using approximant!!",xresInv.shape) + approx = poly_approximant(xresInv,yresInv/area,approximantOrder, + allowExtrapolation=False,fill_value=0) + #return approx,xresInv,yresInv + return convolveFunc(x,y,approx,fill_value=fill_value,axis=axis,) + + if y.ndim < 2: y = y[np.newaxis,:] + nDim = y.ndim + if (axis<0): axis = nDim+axis + + ## expand yres to allow broadcasting + sh = [1,]*y.ndim + sh[axis] = len(yres) + yres = yres.reshape(sh) + + ## prepare output + out = np.empty_like(y) + nP = len(x) + f = interpolator(x,y,fill_value=fill_value,axis=axis) + for i in range(nP): + # interpolate x,y on xres+x[i] + x_integral = xresInv+x[i] + ytemp = f(x_integral)/area + #ytemp = interpolate(x,y,x_integral,fill_value=fill_value,axis=axis)/area + #ytemp = f(x_integral)/area + # do integration + res = integrate.simps( np.transpose(ytemp.T*yresInv),x=xresInv+x[i],axis=axis) + colon = (slice(None),) + setElement(out,i,res,axis=axis) + return out + + +def fftconvolve(x,y,yres,xres=None,normalize=False): + if (xres is not None): + yres = interpolate(xres,yres,x,fill_value=0) + fft = scipy.signal.fftconvolve(y,yres,"full") + _idx = np.argmin( np.abs(x) ); # find t=0 + fft = fft[_idx:_idx+len(x)] + if normalize: + norm = fftconvolve_find_norm(x,yres,xres=None) + else: + norm = 1 + return fft/norm + +def fftconvolve_find_norm(x,res,xres=None): + step = np.ones_like(x) + n = int( len(step)/2 ) + step[:n] = 0 + norm = fftconvolve(x,step,res,xres=xres,normalize=False).max() + return norm + +def convolveGaussian(x,y,sig=1.,nPointGaussian=51,fill_value=0.,axis=-1): + xG = np.linspace(-5*sig,5*sig,nPointGaussian) + G = gaussian(xG,x0=0,sig=sig) + return convolve(x,y,xG,G,fill_value=fill_value,axis=axis) + + +def convolveFunc(x,y,func_res,fill_value=0.,axis=-1): + from scipy import integrate + if y.ndim < 2: y = y[np.newaxis,:] + nDim = y.ndim + if (axis<0): axis = nDim+axis + + ## prepare output + out = np.empty_like(y) + + nP = len(x) + for i in range(nP): + # calculate the values of the resolution function on x-x[i] + ytemp = func_res(x-x[i]) + # do integration + res = integrate.simps(y*ytemp,x=x-x[i],axis=axis) + setElement(out,i,res,axis=axis) + return out + +def convolveFuncParams(x,y,func_res,func_res_pars,fill_value=0.,axis=-1): + def tempFunc(xx): + return func_res(xx,*func_res_pars) + return convolveFunc(x,y,tempFunc,fill_value=fill_value,axis=axis) + +def convolve_test(nG=51): + import time + x=np.arange(100) + y=(x>50).astype(np.float) + sig = 3 + xG = np.linspace(-5*sig,5*sig,nG) + G = gaussian(xG,x0=0,sig=sig) + Gpoly = np.polyfit(xG,G,20) + t0 = time.time() + conv_num = convolve(x,y,xG,G,fill_value=0.) + print("Num:",time.time()-t0) + t0 = time.time() + conv_poly = convolve(x,y,xG,G,approximantOrder=10,fill_value=0.) + print("Num:",time.time()-t0) + t0 = time.time() + conv_fun = convolveFuncParams(x,y,gaussian,(0,sig),fill_value=0.) + print("Fun:",time.time()-t0) + import pylab as plt + plt.plot(conv_num.T,label="Numerical %d points" % nG) + plt.plot(conv_fun.T,"o",label="Gaussian F") + plt.plot(conv_poly.T,label="Gaussian (poly)") + plt.plot(conv_num.T-conv_fun.T) + return conv_num,conv_fun + + +def conv_gauss_and_const(x,sig): + from scipy.special import erf + return 0.5*(1-erf(-x/sqrt2/sig)) + +def conv_gauss_and_exp(x,sig,tau): + from scipy.special import erf + #from mpmath import erf + #http://www.numberempire.com/integralcalculator.php?function=exp%28-x%2Fl%29*exp%28-%28t-x%29**2%2F2%2Fs**2%29%2FSqrt%282%29%2FSqrt%28pi%29%2Fs&var=x&answers= + # actually calculated with sympy + #return -(erf(sqrt2*(sig**2 - x*tau)/(2*sig*tau)) - 1)*exp(sig**2/(2*tau**2) - x/tau)/2 + return 0.5*np.exp(-(2*tau*x-sig**2)/2/tau**2)*(1-erf( (-tau*x+sig**2)/sqrt2/tau/sig)) + +def gaussian(x,x0=0,sig=1,normalize=True): + g = np.exp(-(x-x0)**2/2/sig**2) + if normalize: + return 1/np.sqrt(2*np.pi)/sig*g + else: + return g + + +## FFT filter ## +class FFTfilter(object): + def __init__(self,s,dx=1,wins=((0.024,0.01),),wintype="gauss"): + f,ffts = fft(s,dx=dx) + filter = np.ones_like(f) + for win in wins: + if wintype == "gauss": + filter *= (1-gaussian(f,win[0],win[1],normalize=False)) + filter *= (1-gaussian(f,-win[0],win[1],normalize=False)) + else: + print("Not implemented") + + self.filter=filter + + def apply(self,s): + s = np.fft.fft(s) + return np.fft.ifft(s*self.filter) + + +### PROMPT,PROCESS,TIME,DATE ... ETC. ### + +class procCom(object): + def __init__(self,cmd): + import pexpect + self.proc = pexpect.spawn(cmd) + + def get(self,timeout=None,waitFor=None): + import pexpect + if (waitFor is not None): + s = "" + try: + while s.find(waitFor)<0: + s+=self.proc.read_nonblocking(timeout=timeout) + except (pexpect.TIMEOUT,pexpect.EOF): + pass + else: + s="" + try: + while 1: + s+=self.proc.read_nonblocking(timeout=timeout) + except (pexpect.TIMEOUT,pexpect.EOF): + pass + #print "got",s + return s + + def send(self,what): + self.proc.write(what) + self.proc.flush() + #print "send",what + + def query(self,what,timeout=None,waitFor=None): + self.send(what) + return self,get(timeout=timeout,waitFor=waitFor) + +def getCMD(cmd,strip=True): + import os + shell = os.popen(cmd) + ret = shell.readlines() + shell.close() + if (strip): + ret = [x.strip() for x in ret] + return ret + +def lsdir(path,withQuotes=False,recursive=False): + if recursive: + dirs = [] + for (dir, _, file) in os.walk(path): dirs.append(dir) + else: + content = getCMD("ls -1 %s" % path) + content = ["%s/%s" % (path,x) for x in content] + dirs = [x for x in content if os.path.isdir(x)] + if (withQuotes): + dirs = [ "'%s'" % x for x in dirs ] + return dirs + + +def lsfiles(path,withQuotes=False,recursive=False): + if recursive: + print("Not sure is working") + files = [] + for (dir, _, file) in os.walk(path): files.append(file) + else: + content = getCMD("ls -1 %s" % path) + content = ["%s/%s" % (path,x) for x in content] + files = [x for x in content if os.path.isfile(x)] + if (withQuotes): + files = [ "'%s'" % x for x in files ] + return files + +def downloadPDB(pdbID,outFileName=None): + import urllib2 + import os + address = "http://www.rcsb.org/pdb/download/downloadFile.do?fileFormat=pdb&compression=NO&structureId=%s" % pdbID + p = urllib2.urlopen(address) + lines = p.readlines() + if (outFileName is None): + outFileName = pdbID+".pdb" + folder = os.path.dirname(outFileName) + if (folder != '' and not os.path.exists(folder)): + os.makedirs(folder) + f=open(outFileName,"w") + f.write( "".join(lines) ) + f.close() + + +def dateStringToObj(s,format="%Y.%m.%d %H:%M:%S"): + import datetime + return datetime.datetime.strptime(s,format) + +def now(): + import time + return time.strftime("%Y.%m.%d %H:%M:%S",time.localtime()) + +def mytimer(func,args): + import time + t0=time.time() + ret=func(*args) + return time.time()-t0,ret + +### TXT I/O ### + +def lineToVals(line): + return map(eval,string.split(line)) + +class DataFile(object): + filename="" + lines=[] + comments=[] + head={} + Data=[] + Ndata=0 + + def __init__(self,f): + self.filename=f + self.Read() + + def Read(self): + f=open(self.filename,"r") + temp = f.readlines(); + f.close(); + r=re.compile('\s+'); + c=re.compile('^\s*#'); + for l in temp: + if (re.match(c,l)): + self.comments.append(l) + else: + self.lines.append(l) + # first of non commented line might be header + # try to understand it + keys = [] + try: + v=lineToVals(self.lines[0]) + for i in range(len(v)): + keys.append(i) + except: + keys=self.lines[0].split() + self.lines.remove( self.lines[0] ) + datatemp = [] + for l in self.lines: + datatemp.append( lineToVals(l) ) + self.Data = np.asarray(datatemp) + for i in range(len(keys)): + self.head[keys[i]] = self.Data[:,i] + self.__dict__[keys[i]] = self.Data[:,i] + (self.Ndata,self.Ncol) = self.Data.shape + + def clean(self): + del self.Data + + +def writev(fname,x,Ys,form="%+.6g",sep=" ",header=None,headerv=None): + """ Write data to file 'fname' in text format. + Inputs: + x = x vector + Ys = vector(or array or list of vectors) for the Ys + form = format to use + sep = separator + header = text header (must be a string) + headerv = vector to be used as header, it is convienient when + the output must be of the form + Ncol 252 253 254 + x1 y11 y12 y13 + ....... + In this case headerv should be [252,253,254] + """ + if (type(x) != np.ndarray): x=np.array(x) + if (type(Ys) != np.ndarray): Ys=np.array(Ys) + if (len(Ys.shape)==1): + Ys=Ys.reshape(Ys.shape[0],1) + nx = len(x) + if (Ys.shape[0] == nx): + ny=Ys.shape[1] + elif (Ys.shape[1] == nx): + ny=Ys.shape[0] + Ys=np.transpose(Ys) + else: + raise MCError("dimension of x (%d) does not match any of the dimensions of Ys (%d,%d)" % (nx,Ys.shape[0],Ys.shape[1])) + f=codecs.open(fname,encoding='utf-8',mode="w") + if (header is not None): + f.write(header.strip()+"\n") + if (headerv is not None): + f.write("%d" % (ny+1)) + for i in range(ny): + f.write(sep) + f.write(form % headerv[i]) + f.write("\n") + for i in range(nx): + f.write(form % x[i]) + f.write(sep) + for j in range(ny-1): + f.write(form % Ys[i,j]) + f.write(sep) + f.write(form % Ys[i,-1]) + f.write("\n") + f.close() + + +def writev(fname,x,Ys,form="%+.6g",sep=" ",header=None,headerv=None): + """ Write data to file 'fname' in text format. + Inputs: + x = x vector + Ys = vector(or array or list of vectors) for the Ys + form = format to use + sep = separator + header = text header (must be a string) + headerv = vector to be used as header, it is convienient when + the output must be of the form + Ncol 252 253 254 + x1 y11 y12 y13 + ....... + In this case headerv should be [252,253,254] + """ + if (type(x) != np.ndarray): x=np.array(x) + if (type(Ys) != np.ndarray): Ys=np.array(Ys) + if (len(Ys.shape)==1): + Ys=Ys.reshape(Ys.shape[0],1) + nx = len(x) + if (Ys.shape[1] == nx): + ny=Ys.shape[0] + elif (Ys.shape[0] == nx): + ny=Ys.shape[1] + Ys=np.transpose(Ys) + else: + raise MCError("dimension of x (%d) does not match any of the dimensions of Ys (%d,%d)" % (nx,Ys.shape[0],Ys.shape[1])) + f=codecs.open(fname,encoding='utf-8',mode="w") + if (header is not None): + f.write(header.strip()+"\n") + if (headerv is not None): + f.write("%d" % (ny+1)) + for i in range(ny): + f.write(sep) + f.write(form % headerv[i]) + f.write("\n") + out = np.vstack( (x,Ys) ) + np.savetxt(f,np.transpose(out),fmt=form,delimiter=sep) + +def loadtxt(fname,hasVectorHeader=True,asObj=False,isRecArray=False): + if (isRecArray): + return loadRecArray(fname,asObj=asObj) + a=np.loadtxt(fname) + if (not hasVectorHeader): + x = a[:,0] + y = a[:,1:].T + t = None + else: + x = a[1:,0] + y = a[1:,1:].T + t = a[0,1:] + if (asObj): + class txtFile(object): + def __init__(self,x,y,t): + self.x = x + self.y = y + self.t = t + return txtFile(x,y,t) + else: + return x,y,t + +def loadRecArray(fname,hasVectorHeader=True,asObj=False): + a=np.loadtxt(fname,skiprows=1) + # find header + f=open(fname,"r") + found = None + while found is None: + s=f.readline().strip() + if s[0] != "#": + found = s + names = found.split() + if (asObj): + class txtFile(object): + def __init__(self): + pass + ret = txtFile() + for i in range(len(names)): + ret.__dict__[names[i]] = a[:,i] + else: + ret = np.core.records.fromarrays(a.transpose(),names=",".join(names)) + return ret + + +def prepareRecArrayFromDict( mydict,n=1,leaveEmpty=True ): + names = mydict.keys() + formats = [ type(mydict[p]) for p in names ] + if leaveEmpty: + array_kind = np.empty + else: + array_kind = np.zeros + return array_kind(n, dtype={'names':names, 'formats':formats}) + +def prepareRecArrayFromNamesAndArray( names,ar,n=1,leaveEmpty=True ): + formats = [ ar.dtype.type for p in names ] + if leaveEmpty: + array_kind = np.empty + else: + array_kind = np.zeros + return array_kind(n, dtype={'names':names, 'formats':formats}) + + + +def writeMatrix(fout,M,x,y,form="%+.6g",sep=" ",header=None): + (nx,ny) = M.shape + if ( (nx == len(x)) and (ny==len(y)) ): + pass + elif ( (nx == len(y)) and (ny==len(x)) ): + M=M.transpose() + (nx,ny) = M.shape + else: + e = "Dimensions of matrix and the x or y vectors don't match" + e += "shapes of M, x, y: " + str(M.shape) + " " +str(len(x))+ " " +str(len(y)) + raise MCError(e) + temp = np.zeros( (nx+1,ny) ) + temp[1:,:] = M + temp[0,:] = y + writev(fout,np.hstack( (0,x) ),temp,form=form,sep=sep,header=header) + + +### MATPLOTLIB ... ETC. ### + +def lt(i,style="-"): + colors="rkbgy" + i = i%len(colors) + color = colors[i] + if (style is not None): color += style + return color + +def displayFig(i,x=None,y=None,roi=None): + import pylab as plt + from matplotlib.widgets import Slider, Button, RadioButtons + fig=plt.figure() + n1,n2=i.shape + if x is None: + x = np.arange(n1) + if y is None: + y = np.arange(n2) + + ax = fig.add_subplot(111) + if roi is not None: + (x1,x2,y1,y2) = roi + else: + x1,x2,y1,y2 = (0,n1,0,n2) + xm=x[x1]; xM=x[x2-1] + ym=y[y1]; yM=y[y2-1] + def _format_coord(x, y): + col = int(x+0.5) + row = int(y+0.5) + if col>=0 and col=0 and rownhalf): + vec = np.concatenate( (vec[at_idx-nhalf:],vec[:at_idx-nhalf] )) + else: + vec = np.concatenate( (vec[at_idx+nhalf:],vec[:at_idx+nhalf])) + return vec + +### REBIN, RUNNING AVG, ... ETC ### + +def rebinOLD(bins_edges,x,*Y): + """ rebins a list of Y based on x using bin_edges + returns + center of bins (vector) + rebinned Y (list of vectors) + Simga's (list of vectors), note: just std not std of average + N (list of vectors), number of occurrances + """ + n=len(bins_edges)-1 + outX = [] + outY = [] + outS = [] + outN = [] + for j in range(len(Y)): + outY.append(np.empty(n)) + outS.append(np.empty(n)) + outN.append(np.empty(n)) + outX = np.empty(n) + for i in range(n): + idx = (x>=bins_edges[i]) & (x 0): + for j in range(len(Y)): + outN[j][i] = idx.sum() + outY[j][i] = Y[j][idx].mean() + outS[j][i] = Y[j][idx].std() + else: + for j in range(len(Y)): + outN[j][i] = 0 + outY[j][i] = np.nan + outS[j][i] = np.nan + return outX,outY,outS,outN + +def rebin1D(a,shape): + sh = shape,a.shape[0]//shape + return a.reshape(sh).mean(1) + +def rebin1Dnew(a,shape): + n0 = a.shape[0]//shape + sh = shape,n0 + return a[:n0*shape].reshape(sh).mean(1) + +def rebin2D(a, shape): + sh = shape[0],a.shape[0]//shape[0],shape[1],a.shape[1]//shape[1] + return a.reshape(sh).mean(-1).mean(1) + +def rebin2Dnew(a, shape): + # // means floor + n0 = a.shape[0]//shape[0] + n1 = a.shape[1]//shape[1] + crop = shape[0]*n0,shape[1]*n1 + sh = shape[0],n0,shape[1],n1 + #print a[:n0*shape[0],:n1*shape[1]].reshape(sh) + return a[:crop[0],:crop[1]].reshape(sh).mean(-1).mean(1) + +def rebin3Dnew(a, shape): + # // means floor + n0 = a.shape[0]//shape[0] + n1 = a.shape[1]//shape[1] + n2 = a.shape[2]//shape[2] + sh = shape[0],n0,shape[1],n1,shape[2],n2 + crop = n0*shape[0],n1*shape[1],n2*shape[2] + #print a[:n0*shape[0],:n1*shape[1]].reshape(sh) + return a[:crop[0],:crop[1],:crop[2]].reshape(sh).mean(-1).mean(-2).mean(-3) + +def rebin(a, shape): + a = np.asarray(a) + ndim = a.ndim + if (ndim == 1): + return rebin1Dnew(a,shape) + elif (ndim == 2): + return rebin2Dnew(a,shape) + elif (ndim == 3): + return rebin3Dnew(a,shape) + else: + print("Can't do rebin of",a) + + +def rebinTODO(a, shape): + ndim = a.ndim + if (len(shape) != ndim): + print("Error, asked to rebin a %d dimensional array but provided shape with %d lengths" % (ndim,len(shape))) + return None + nout = [] + for n in range(ndim): + nout.append( a.shape[n]%shape[n] ) + print("Not implemented ...") + +### DISTRIBUTION, MEDIAN, SIGMA, ETC. ### + +def idx_within_std_from_center(vector,range): + (m,s) = MedianAndSigma(vector) + return np.abs(vector-m)<(range*s) + +def MedianAndSigma(a): + median = np.median(a) + MAD = np.median(np.abs(a-median)) + sigma = 1.4826*MAD; # this assumes gauss distribution + return (median,sigma) + +def weigthed_average(y,e=None,axis=0): + if e is None: + e=np.ones_like(y) + if (axis != 0): + y = y.transpose() + e = e.transpose() + (n0,n1) = y.shape + yout = np.empty(n1) + eout = np.empty(n1) + for i in range(n1): + toav = y[:,i] + valid = np.isfinite(toav) + toav = toav[valid] + w = 1/e[valid,i]**2 + yout[i] = np.sum(toav*w)/np.sum(w) + eout[i] = np.sqrt(1/np.sum(w)) + return yout,eout + +### CONVERSION ### +def convert(value=1,unit="eV",out="nm"): + if (unit == "eV") & (out =="nm"): + return 1239.8/value + + +### IMINUIT RELATED ### + +def iminuitClass(modelFunc): + """ build a class to help fitting, first argumes of modelFunc should be x""" + import inspect + import iminuit + import iminuit.util + args = inspect.getargspec(modelFunc).args + defaults = inspect.getargspec(modelFunc).defaults + if defaults is not None: + nDef = len( defaults ) + args = args[:-nDef] + args_str = ",".join(args[1:]) + + + class iminuitFit(object): + def __init__(self,x,data,init_pars,err=1.): + self.x = x + self.data = data + self.err = err + self.init_pars=init_pars + self.func_code = iminuit.util.make_func_code(args[1:])#docking off independent variable + self.func_defaults = None #this keeps np.vectorize happy + + def __call__(self,*arg): + return self.chi2(*arg) + + def model(self,*arg): + return modelFunc(self.x,*arg) + + def chi2(self,*arg): + c2 = (self.model(*arg)-self.data)/self.err + return np.sum(c2*c2) + + def fit(self,showInit=True,showPlot=True,doFit=True,doMinos=False): + import pylab as plt + p = self.init_pars + if "errordef" not in p: + p["errordef"] = 1. + #m = iminuit.Minuit(self,print_level=0,pedantic=False,**p) + m = iminuit.Minuit(self,**p) + if showInit: + model = self.model(*m.args) + plt.figure("initial pars") + plt.grid() + plt.plot(self.x,self.data,"o") + plt.plot(self.x,model,'r-',linewidth=2) + raw_input() + if doFit: + m.migrad() + if doMinos: + m.minos() + for p in m.parameters: + err = m.get_merrors()[p] + err = "+ %.4f - %.4f" % (np.abs(err["upper"]),np.abs(err["lower"])) + print("%10s %.4f %s"%(p,m.values[p],err)) + else: + for p in m.parameters: + err = m.errors[p] + err = "+/- %.4f" % (err) + print("%10s %.4f %s"%(p,m.values[p],err)) + + model = self.model(*m.args) + if (showPlot): + plt.figure("final fit") + plt.grid() + plt.plot(self.x,self.data,"o") + plt.plot(self.x,model,'r-',linewidth=2) + self.m = m + return m,self.x,self.data,model + + return iminuitFit + + +def iminuitParsToStr(iminuit,withErrs=True,withFixed=True): + values = iminuit.values + errs = iminuit.errors + pnames = values.keys() + lenParNames = max( [len(p) for p in pnames] ) + fmt = "%%%ds" % lenParNames + pnames.sort() + res = [] + for p in pnames: + v = values[p] + e = errs[p] + isf = iminuit.is_fixed(p) + if not withFixed and isf: + continue + v,e = approx_err(v,e,asstring=True) + if isf: e="fixed" + s = fmt % p + if withErrs: + s += " = %s +/- %s" % (v,e) + else: + s += " = %s" % (v) + res.append(s) + return res + + +### VARIOUS ### + +def myProgressBar(N,title="Percentage"): + import progressbar as pb + widgets = [title, pb.Percentage(), ' ', pb.Bar(), ' ', pb.ETA()] + pbar = pb.ProgressBar(widgets=widgets, maxval=N) + return pbar + +def chunk(iterableOrNum, size): + temp = [] + try: + n = len(iterableOrNum) + except TypeError: + n = iterableOrNum + nGroups = int(np.ceil(float(n)/size)) + for i in range(nGroups): + m = i*size + M = (i+1)*size; M=min(M,n) + if (m>=n): + break + temp.append( slice(m,M) ) + try: + ret = [iterableOrNum[x] for x in temp] + except TypeError: + ret = temp + return ret + +def timeres(*K): + """ return sqrt(a1**2+a2**2+...) """ + s = 0 + for k in K: + s += k**2 + return np.sqrt(s) + +def approx_err(value,err,asstring=False): + if (not (np.isfinite(err))): + err = np.abs(value/1e3) + if (err != 0): + ndigerr = -int(np.floor(np.log10(err))) + if (ndigerr<1): ndigerr=2 + e =round(err,ndigerr) + ndigerr = -int(np.floor(np.log10(err))) + v =round(value,ndigerr) + else: + v=value + e=err + if (asstring): + return "%s" % v,"%s" % e + else: + return v,e + +def linFitOld(A,B): + """ solve Ax = B, returning x """ + temp = np.dot(A.T,B) + square = np.dot(A.T,A) + if (np.asarray(square).ndim==0): + inv = 1./square + else: + inv = np.linalg.inv(square) + x = np.dot(inv,temp) + return x + + + +def linFit(A,B,cond=None): + """ solve Ax = B, returning x """ + from scipy import linalg + temp = np.dot(A.T,B) + square = np.dot(A.T,A) + if (np.asarray(square).ndim==0): + inv = 1./square + else: + inv = linalg.pinvh(square,cond=cond) + x = np.dot(inv,temp) + return x + #return np.linalg.lstsq(A,B)[0] + +def insertInSortedArray(a,v): + if v>a.max(): return a + idx = np.argmin(a