372 lines
13 KiB
Python
372 lines
13 KiB
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
np.warnings.simplefilter('ignore')
|
|
import time
|
|
import matplotlib.pyplot as plt
|
|
import h5py
|
|
import collections
|
|
import re
|
|
|
|
from x3py import x3py
|
|
import alignment
|
|
import mcutils as mc
|
|
|
|
cmap = plt.cm.viridis if hasattr(plt.cm,"viridis") else plt.cm.gray
|
|
kw_2dplot = dict(
|
|
interpolation = "none",
|
|
aspect = "auto",
|
|
cmap = cmap
|
|
)
|
|
|
|
|
|
g_exp = "mecl3616"
|
|
g_exp = "xppl3716"
|
|
g_bml = g_exp[:3]
|
|
|
|
x3py.config.updateBeamline(g_bml)
|
|
|
|
g_folder_init = g_exp+"_init_pars/"
|
|
g_folder_out = g_exp+"_output/"
|
|
g_folder_data = "/reg/d/psdm/"+g_bml+"/"+ g_exp +"/hdf5/"
|
|
|
|
import socket
|
|
hostname = socket.gethostname()
|
|
if hostname == "x1":
|
|
g_folder_data = "/home/marco/temp"
|
|
if hostname == "apcluster0":
|
|
g_folder_data = "/data/marcoc/singleShotXanes/"+ g_exp +"/hdf5/"
|
|
|
|
# set defaults based on experiment
|
|
if g_bml == "xpp":
|
|
g_roi_height = 200
|
|
g_swapx = False
|
|
g_swapy = False
|
|
else:
|
|
g_roi_height = 100
|
|
g_swapx = True
|
|
g_swapy = False
|
|
|
|
print("Working on experiment",g_exp,"(beamline %s)"%g_bml)
|
|
print(" folder data →",g_folder_data)
|
|
print(" folder init_pars →",g_folder_init)
|
|
print(" folder outout →",g_folder_out)
|
|
|
|
#g_folder = "/reg/d/psdm/xpp/xppl3716/ftc/hdf5/"
|
|
|
|
def readDataset(fnameOrRun=7,
|
|
force=False,
|
|
doBkgSub=False):
|
|
if isinstance(fnameOrRun,str) and (fnameOrRun[-3:]=="npz"):
|
|
d = x3py.toolsVarious.DropObject()
|
|
temp = np.load(fnameOrRun)
|
|
spec1 = temp["spec1"]
|
|
spec2 = temp["spec2"]
|
|
nS = spec1.shape[0]
|
|
d.spec1 = x3py.toolsDetectors.wrapArray("spec1",spec1,time=np.arange(nS))
|
|
d.spec2 = x3py.toolsDetectors.wrapArray("spec2",spec2,time=np.arange(nS))
|
|
else:
|
|
if isinstance(fnameOrRun,int):
|
|
fnameOrRun=g_folder_data+"/"+g_exp+"-r%04d.h5" % fnameOrRun
|
|
d = x3py.Dataset(fnameOrRun,detectors=["opal0","opal1","fee_spec","opal2","ebeam"])
|
|
if g_bml == "xpp":
|
|
d.spec1 = d.opal0
|
|
d.spec2 = d.opal1
|
|
else:
|
|
d.spec1 = d.fee_spec
|
|
d.spec2 = d.opal2
|
|
if not hasattr(d,"scan"):
|
|
d.scan = x3py.toolsVarious.DropObject()
|
|
d.scan.scanmotor0_values = [0,]
|
|
return d
|
|
|
|
|
|
def getCenter(img,axis=0,threshold=0.05):
|
|
img = img.copy()
|
|
img[img<img.max()*threshold] = 0
|
|
if axis == 1: img=img.T
|
|
p = img.mean(1)
|
|
x = np.arange(img.shape[0])
|
|
return int(np.sum(x*p)/np.sum(p))
|
|
|
|
def showShots(im1,im2):
|
|
nS = im1.shape[0]
|
|
fig,ax = plt.subplots(2,nS,sharex=True,sharey=True)
|
|
if im1.ndim == 3:
|
|
for a,i1,i2 in zip(ax.T,im1,im2):
|
|
a[0].imshow(i1.T,**kw_2dplot)
|
|
a[1].imshow(i2.T,**kw_2dplot)
|
|
else:
|
|
for a,p1,p2 in zip(ax.T,im1,im2):
|
|
a[0].plot(p1)
|
|
a[1].plot(p2)
|
|
|
|
def ratioOfAverage(p1,p2,threshold=0.03):
|
|
"""
|
|
p1 and p2 are the energy spectrum. if 2D the first index has to be the shot number
|
|
calculate median ratio taking into account only regions where p1 and p2 are > 5% of the max """
|
|
# check if they are 2D
|
|
if p1.ndim == 1:
|
|
p1 = p1[np.newaxis,:]
|
|
p2 = p2[np.newaxis,:]
|
|
# w1 and w2 are the weights
|
|
w1 = p1.copy(); w2 = p2.copy()
|
|
if threshold is not None:
|
|
# weights will be set to zero if intensity is smaller than 5% of max
|
|
# for each shots, get maximum
|
|
m1 = np.nanmax(p1,axis=1); m2 = np.nanmax(p2,axis=1)
|
|
# find where each spectrum is smaller than threshold*max_for_that_shot; they will be masked out
|
|
idx1 = p1 < (m1[:,np.newaxis]*threshold)
|
|
idx2 = p2 < (m2[:,np.newaxis]*threshold)
|
|
w1[idx1]=0
|
|
w2[idx2]=0
|
|
# using masked array because some pixel will have zero shots contributing
|
|
av1 = np.ma.average(p1,axis=0,weights=w1)
|
|
av1[av1.mask] = np.nan
|
|
av2 = np.ma.average(p2,axis=0,weights=w2)
|
|
av2[av2.mask] = np.nan
|
|
return av2/av1
|
|
|
|
def medianRatio(p1,p2,threshold=0.03):
|
|
"""
|
|
p1 and p2 are the energy spectrum. if 2D the first index has to be the shot number
|
|
calculate median ratio taking into account only regions where p1 and p2 are > 5% of the max """
|
|
# check if they are 2D
|
|
if p1.ndim == 1:
|
|
p1 = p1[np.newaxis,:]
|
|
p2 = p2[np.newaxis,:]
|
|
p1 = np.ma.asarray( p1.copy() )
|
|
p2 = np.ma.asarray( p2.copy() )
|
|
if threshold is not None:
|
|
m1 = np.nanmax(p1,axis=1); m2 = np.nanmax(p2,axis=1)
|
|
# find where each spectrum is smaller than threshold*max_for_that_shot; they will be masked out
|
|
idx1 = p1 < (m1[:,np.newaxis]*threshold)
|
|
idx2 = p2 < (m2[:,np.newaxis]*threshold)
|
|
idx = idx1 & idx2
|
|
p1.mask = idx
|
|
p2.mask = idx
|
|
ratio = p2/p1
|
|
return np.ma.average(ratio,axis=0,weights=p1)
|
|
|
|
|
|
|
|
class AnalyzeRun(object):
|
|
def __init__(self,run,initAlign="auto",swapx=g_swapx,swapy=g_swapy):
|
|
""" swapx → swap x axis of first spectrometer
|
|
swapy → swap y axis of first spectrometer
|
|
initAlign: could be:
|
|
1. None if you want default transformation parameters
|
|
2. a dict if you want to overwrite certain parameters of the default ones
|
|
3. an integer (to look for xppl3716_init_pars/run????_transform.npy)
|
|
4. a file name (that has been previosly saved with r.saveTransform(fname)
|
|
"""
|
|
self.data = readDataset(run)
|
|
self.scanpos = self.data.scan.scanmotor0_values
|
|
self.nCalib = self.data.spec1.nCalib
|
|
self.nShotsPerCalib = self.data.spec1.lens
|
|
if isinstance(run,str):
|
|
run = int( re.search("\d{3,4}",run).group() )
|
|
self.run = run
|
|
self.results = collections.OrderedDict()
|
|
self.swap = (swapx,swapy)
|
|
#self.clearCache()
|
|
|
|
d = self.data
|
|
self.spec1 = d.spec1 ; # spec1 is the one that is moved
|
|
self.spec2 = d.spec2 ;
|
|
self.E = alignment.defaultE
|
|
|
|
try:
|
|
self.loadTransform(initAlign)
|
|
except (AttributeError,FileNotFoundError):
|
|
if initAlign is None:
|
|
print("Set to default transform")
|
|
self.initAlign = self.setDefaultTransform()
|
|
else:
|
|
self.initAlign = initAlign
|
|
|
|
def getShots(self,shots=0,calib=None,bkgSub="line",roi=g_roi_height):
|
|
if shots == "all":
|
|
if calib != None:
|
|
shots = slice(0,self.nShotsPerCalib[calib])
|
|
else:
|
|
shots = slice(0,self.data.spec1.nShots)
|
|
# read data
|
|
im1 = self.spec1.getShots(shots,calib=calib)
|
|
im2 = self.spec2.getShots(shots,calib=calib)
|
|
# subtractBkg bkg
|
|
im1 = alignment.subtractBkg(im1,bkg_type=bkgSub)
|
|
im2 = alignment.subtractBkg(im2,bkg_type=bkgSub)
|
|
# rebin and swap im1 if necessary
|
|
if im1.shape[-1] != 1024:
|
|
im1 = mc.rebin(im1, (im1.shape[0],im1.shape[1],1024) )
|
|
if self.swap[0]:
|
|
im1 = im1[:,:,::-1]
|
|
if self.swap[1]:
|
|
im1 = im1[:,::-1,:]
|
|
if roi is None:
|
|
pass
|
|
elif isinstance(roi,slice):
|
|
im1 = im1[:,roi,:]
|
|
im2 = im2[:,roi,:]
|
|
elif isinstance(roi,int):
|
|
if not hasattr(self,"roi1"): self.roi1 = alignment.findRoi(im1[0],roi)
|
|
if not hasattr(self,"roi2"): self.roi2 = alignment.findRoi(im2[0],roi)
|
|
im1 = im1[:,self.roi1,:]; im2 = im2[:,self.roi2,:]
|
|
return im1,im2
|
|
|
|
def guiAlign(self,shot=0,save="auto"):
|
|
im1,im2 = self.getShot(shot)
|
|
gui = alignment.GuiAlignment(im1[0],im2[0])
|
|
input("Enter to start")
|
|
gui.start()
|
|
if save == "auto":
|
|
fname = g_folder_init+"/run%04d_gui_align.npy" % self.run
|
|
else:
|
|
fname = save
|
|
self.initAlign = gui.transform
|
|
gui.save(fname)
|
|
|
|
|
|
def analyzeScan(self,initpars=None,shots=slice(0,30),calibs="all",nImagesToFit=0,nSaveImg=5):
|
|
""" nImagesToFit: number of images to Fit per calibcycle, (int or "all") """
|
|
if initpars is None: initpars= self.initAlign
|
|
if calibs == "all": calibs=list(range(self.nCalib))
|
|
if isinstance(calibs,slice): calibs=list(range(self.nCalib))[calibs]
|
|
nC = len(calibs)
|
|
for ic,calib in enumerate(calibs):
|
|
if nImagesToFit == "all":
|
|
nToFit = self.nShotsPerCalib[calib]
|
|
else:
|
|
nToFit = nImagesToFit
|
|
#print("Memory available 1",x3py.toolsOS.memAvailable())
|
|
s1,s2 = self.getShots(shots,calib=calib)
|
|
#print("Memory available 2",x3py.toolsOS.memAvailable())
|
|
if nImagesToFit > 0:
|
|
ret,bestTransf = alignment.doShots(s1[:nToFit],s2[:nToFit],doFit=True,\
|
|
initpars=initpars,nSaveImg=nSaveImg,returnBestTransform=True);
|
|
initpars = bestTransf; self.initAlign=bestTransf
|
|
if nToFit < s1.shape[0]:
|
|
ret2 = alignment.doShots(s1[nToFit:],s2[nToFit:],initpars=initpars,doFit=False,nSaveImg=0)
|
|
if ret is None:
|
|
ret = ret2
|
|
else:
|
|
ret = alignment.unravel_results( (ret,ret2) )
|
|
#print("Memory available 3",x3py.toolsOS.memAvailable())
|
|
self.results[calib] = ret
|
|
#print("Memory available 4",x3py.toolsOS.memAvailable())
|
|
print("Calib cycle %d/%d -> %.3f (best FOM: %.2f)" % (ic,nC,self.scanpos[ic],np.nanmin(ret.fom)))
|
|
return [self.results[c] for c in calibs]
|
|
|
|
def doShot(self,shot=0,calib=None,initpars=None,im1=None,im2=None,doFit=True,show=False,showInit=False,save=False,savePlot="auto"):
|
|
if initpars is None: initpars= self.initAlign
|
|
if (im1 is None) or (im2 is None):
|
|
im1,im2 = self.getShots(shot,calib=calib); im1=im1[0]; im2=im2[0]
|
|
r = alignment.doShot(im1,im2,initpars,doFit=doFit,show=showInit)
|
|
im1 = r.im1
|
|
im2 = r.im2
|
|
self.initAlign = r.final_pars
|
|
if show:
|
|
if savePlot == "auto":
|
|
if not os.path.isdir(g_folder_out): os.makedirs(g_folder_out)
|
|
savePlot = g_folder_out+"/run%04d_calib%s_shot%04d_fit.png" % (self.run,calib,shot)
|
|
alignment.plotShot(im1,im2,res=r,save=savePlot)
|
|
if save: self.saveTransform()
|
|
return r
|
|
|
|
def doShots(self,shots=slice(0,50),calib=None,initpars=None,doFit=False,returnBestTransform=False,nSaveImg='all'):
|
|
"""
|
|
shots : slice to define shots to read, use 'all' for all shots in calibcycle
|
|
nSaveImg : save saveImg images in memory (self.results), use 'all' for all
|
|
useful for decreasing memory footprint
|
|
"""
|
|
if initpars is None: initpars= self.initAlign
|
|
if shots == "all": shots = slice(self.nShotsPerCalib[calib])
|
|
s1,s2 = self.getShots(shots,calib=calib)
|
|
ret,transformForBestFit = alignment.doShots(s1,s2,initpars=initpars,doFit=doFit,\
|
|
returnBestTransform=True,nSaveImg=nSaveImg)
|
|
if doFit: self.initAlign = transformForBestFit
|
|
# keep it for later !
|
|
self.results[calib] = ret
|
|
if returnBestTransform:
|
|
return ret,transformForBestFit
|
|
else:
|
|
return ret
|
|
|
|
def save(self,fname="auto",overwrite=False):
|
|
if len(self.results) == 0: print("self.results are empty, returning without saving")
|
|
if not os.path.isdir(g_folder_out): os.makedirs(g_folder_out)
|
|
if fname == "auto":
|
|
fname = g_folder_out+"/run%04d_analysis" % self.run
|
|
if os.path.exists(fname) and not overwrite:
|
|
print("File %s exists, **NOT** saving, use overwrite=True is you want ..."%fname)
|
|
return
|
|
|
|
def load(self,fname="auto"):
|
|
if fname == "auto": fname = g_folder_out+"/run%04d_analysis.npz" % self.run
|
|
temp = np.load(fname)
|
|
self.results = temp["results"].item()
|
|
temp.close()
|
|
|
|
def _auto_transform_name(self,run=None,calib=None):
|
|
if run is None: run = self.run
|
|
fname = g_folder_init+"/run%04d_transform" % run
|
|
if calib is not None:
|
|
fname = fname + "_c%03d" % calib
|
|
return fname + ".npy"
|
|
|
|
def saveTransform(self,fname="auto",calib=None,transform=None):
|
|
if transform is None: transform = self.initAlign
|
|
if fname == "auto": fname = self._auto_transform_name(calib=calib)
|
|
print("Saving roi and transformation parameter to %s"%fname)
|
|
alignment.saveAlignment(fname,self.initAlign,self.roi1,self.roi2)
|
|
|
|
def loadTransform(self,fname="auto", calib=None):
|
|
if isinstance(fname,dict): raise FileNotFoundError
|
|
if fname == "auto": fname = self._auto_transform_name(calib=calib)
|
|
if isinstance(fname,int):
|
|
fname = g_folder_init+"/run%04d_transform.npy" % fname
|
|
if not os.path.exists(fname): print("Asked to read %s, but it does not exist"%fname)
|
|
temp = np.load(fname).item()
|
|
self.initAlign = temp["transform"]
|
|
self.roi1 = temp["roi1"]
|
|
self.roi2 = temp["roi2"]
|
|
print("init transform and ROIs from %s"%fname)
|
|
|
|
|
|
def clearCache(self):
|
|
del self.roi1
|
|
del self.roi2
|
|
alignment.clearCache(); # nedded for multiprocessing can leave bad parameters in the cache
|
|
|
|
def setDefaultTransform( self ):
|
|
#dict( scalex=0.65,rotation=0.0,transx=90, iblur1=4.3,fix_iblur1=False )
|
|
t = alignment.g_fit_default_kw
|
|
self.initAlign = t
|
|
return t
|
|
|
|
def quick_mec(run,ref=236,divideByRef=False,returnRes=False):
|
|
""" useful to analyze the runs around 140 (done with the focusing """
|
|
ref_run = 236
|
|
h=h5py.File("mecl3616_output/run%04d_analysis.h5" %ref,"r")
|
|
ref = np.nanmean(h["calibNone"]["ratio"][...],axis=0)
|
|
r = AnalyzeRun(run,initAlign=ref,swapx=True,swapy=False)
|
|
res=r.doShots(slice(5),doFit=False)
|
|
ret = res["ratio"]/ref if divideByRef else res["ratio"]
|
|
if returnRes:
|
|
return ret,res
|
|
else:
|
|
return ret
|
|
|
|
def quickAndDirty(run,nShots=300,returnAll=True,doFit=False):
|
|
""" useful to analyze the runs around 140 (done with the focusing """
|
|
r = AnalyzeRun(run,swap=True,initAlign=g_folder_init+"/run0144_transform.npy")
|
|
res=r.doShots(slice(nShots),doFit=doFit)
|
|
o = alignment.unravel_results(res)
|
|
ref = np.nanmedian(o["ratio"][:40],0)
|
|
sam = np.nanmedian(o["ratio"][50:],0)
|
|
if returnAll:
|
|
return sam/ref,o["ratio"]/ref
|
|
else:
|
|
return sam/ref
|