145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
from __future__ import print_function,division
|
|
import mcutils as mc
|
|
import joblib
|
|
import numpy as np
|
|
|
|
# /--------\
|
|
# | |
|
|
# | UTILS |
|
|
# | |
|
|
# \--------/
|
|
|
|
|
|
def rebin1D(a,shape):
|
|
n0 = a.shape[0]//shape
|
|
sh = shape,n0
|
|
return a[:n0*shape].reshape(sh).mean(1)
|
|
|
|
def calcFOM(p1,p2,ratio,threshold=0.1):
|
|
idx = ( p1>p1.max()*threshold )# & (p2>p2.max()/10)
|
|
ratio = ratio[idx]
|
|
return ratio.std()/np.abs(ratio.mean())
|
|
|
|
def getCenterOfMass(img,x=None,axis=0,threshold=0.05):
|
|
img = img.copy()
|
|
if img.ndim == 1: img=img[np.newaxis,:]; axis=1
|
|
img[img<img.max()*threshold] = 0
|
|
if axis == 1: img=img.T
|
|
p = img.mean(1)
|
|
if x is None: x = np.arange(img.shape[0])
|
|
return np.sum(x*p)/np.sum(p)
|
|
|
|
|
|
def maskLowIntensity(p1,p2,threshold=0.03,squeeze=True):
|
|
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
|
|
if squeeze:
|
|
p1 = np.squeeze(p1);
|
|
p2 = np.squeeze(p2)
|
|
return p1,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 """
|
|
p1,p2 = maskLowIntensity(p1,p2,threshold=threshold,squeeze=False)
|
|
# using masked array because some pixel will have zero shots contributing
|
|
av1 = np.ma.average(p1,axis=0,weights=p1)
|
|
av2 = np.ma.average(p2,axis=0,weights=p2)
|
|
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 """
|
|
p1,p2 = maskLowIntensity(p1,p2,threshold=threshold,squeeze=False)
|
|
ratio = p2/p1
|
|
return np.ma.average(ratio,axis=0,weights=p1)
|
|
|
|
|
|
# /--------------------\
|
|
# | |
|
|
# | PLOTS & CO. |
|
|
# | |
|
|
# \--------------------/
|
|
|
|
def plotShot(im1,im2,transf1=None,transf2=None,fig=None,ax=None,res=None,E=defaultE,save=None):
|
|
if transf1 is not None: im1 = transf1.transformImage(im1)
|
|
if transf2 is not None: im2 = transf2.transformImage(im2)
|
|
if fig is None and ax is None:
|
|
fig = plt.subplots(2,3,figsize=[7,5],sharex=True)[0]
|
|
ax = fig.axes
|
|
elif fig is not None:
|
|
ax = fig.axes
|
|
if E is None: E=np.arange(im1.shape[1])
|
|
n = im1.shape[0]
|
|
ax[0].imshow(im1,extent=(E[0],E[-1],0,n),**kw_2dplot)
|
|
ax[1].imshow(im2,extent=(E[0],E[-1],0,n),**kw_2dplot)
|
|
ax[2].imshow(im1-im2,extent=(E[0],E[-1],0,n),**kw_2dplot)
|
|
if res is None:
|
|
p1 = np.nansum(im1,axis=0)
|
|
p2 = np.nansum(im2,axis=0)
|
|
pr = p2/p1
|
|
else:
|
|
p1 = res.p1; p2 = res.p2; pr = res.ratio
|
|
ax[3].plot(E,p1,lw=3)
|
|
ax[4].plot(E,p1,lw=1)
|
|
ax[4].plot(E,p2,lw=3)
|
|
idx = (p1>p1.max()/10.)
|
|
ax[5].plot(E[idx],pr[idx])
|
|
if res is not None:
|
|
ax[5].set_title("FOM: %.2f"%res.fom)
|
|
else:
|
|
ax[5].set_title("FOM: %.2f"% calcFOM(p1,p2,pr))
|
|
if (save is not None) and (save is not False): plt.savefig(save,transparent=True,dpi=500)
|
|
return fig
|
|
|
|
def plotRatios(r,shot='random',fig=None,E=defaultE,save=None):
|
|
if fig is None: fig = plt.subplots(2,1,sharex=True)[0]
|
|
ax = fig.axes
|
|
n = r.shape[0]
|
|
i = ax[0].imshow(r,extent=(E[0],E[-1],0,n),**kw_2dplot)
|
|
i.set_clim(0,1.2)
|
|
if shot == 'random' : shot = np.random.random_integers(0,n-1)
|
|
ax[1].plot(E,r[shot],label="Shot n %d"%shot)
|
|
ax[1].plot(E,np.nanmedian(r[:10],axis=0),label="median 10 shots")
|
|
ax[1].plot(E,np.nanmedian(r,axis=0),label="all shots")
|
|
ax[1].legend()
|
|
ax[1].set_ylim(0,1.5)
|
|
ax[1].set_xlabel("Energy")
|
|
ax[1].set_ylabel("Transmission")
|
|
ax[0].set_ylabel("Shot num")
|
|
if (save is not None) and (save is not False): plt.savefig(save,transparent=True,dpi=500)
|
|
|
|
def plotSingleShots(r,nShots=10,fig=None,E=defaultE,save=None,ErangeForStd=(7090,7150)):
|
|
if fig is None: fig = plt.subplots(2,1,sharex=True)[0]
|
|
ax = fig.axes
|
|
for i in range(nShots):
|
|
ax[0].plot(E,r[i]+i)
|
|
ax[0].set_ylim(0,nShots+0.5)
|
|
av = (1,3,10,30,100)
|
|
good = np.nanmedian(r,0)
|
|
for i,a in enumerate(av):
|
|
m = np.nanmedian(r[:a],0)
|
|
idx = (E>ErangeForStd[0]) & (E<ErangeForStd[1])
|
|
fom = np.nanstd( m[idx]/good[idx] )
|
|
print("n shots %d, std %.2f"%(a,fom) )
|
|
ax[1].plot(E,m+i,label="%d shots, std :%.2f"%(a,fom))
|
|
ax[1].legend()
|
|
ax[1].set_ylim(0,len(av)+0.5)
|
|
ax[1].set_xlabel("Energy")
|
|
ax[1].set_ylabel("Transmission")
|
|
if (save is not None) and (save is not False): plt.savefig(save,transparent=True,dpi=500)
|
|
|