changes to make the script marion's compatible + minor tweaks

alignment.py

@@ -56,7 +56,12 @@ g_fit_default_kw = dict(
     iblur1           = 0,
     limit_iblur1     = (0,20),
     error_iblur1     = 0.02,
-    fix_iblur1       = True
+    fix_iblur1       = True,
+    iblur2           = 0,
+    limit_iblur2     = (0,20),
+    error_iblur2     = 0.02,
+    fix_iblur2       = True
+
 )
 
 
@@ -65,15 +70,15 @@ def rebin1D(a,shape):
   sh = shape,n0
   return a[:n0*shape].reshape(sh).mean(1)
 
-
+cmap = plt.cm.viridis if hasattr(plt.cm,"viridis") else plt.cm.gray
 kw_2dplot = dict(
   interpolation = "none",
   aspect        = "auto",
-  cmap          = plt.cm.viridis
+  cmap          = cmap
   )
 
 fit_ret = collections.namedtuple("fit_ret",["fit_result","init_pars","final_pars",\
-  "final_transform1","final_transform2","im1","im2","p1","p2","fom","ratio","tneeded"] )
+  "final_transform1","final_transform2","im1","im2","E","p1","p2","fom","ratio","tneeded"] )
 
 def calcFOM(p1,p2,ratio):
   idx = ( p1>p1.max()/10 ) & (p2>p2.max()/10)
@@ -131,9 +136,9 @@ def plotShot(im1,im2,transf1=None,transf2=None,fig=None,ax=None,res=None,E=defau
     ax = fig.axes
   if E is None: E=np.arange(im1.shape[1])
   n = im1.shape[0]
-  ax[0].imshow(im1,**kw_2dplot,extent=(E[0],E[-1],0,n))
+  ax[0].imshow(im1,extent=(E[0],E[-1],0,n),**kw_2dplot)
-  ax[1].imshow(im2,**kw_2dplot,extent=(E[0],E[-1],0,n))
+  ax[1].imshow(im2,extent=(E[0],E[-1],0,n),**kw_2dplot)
-  ax[2].imshow(im1-im2,**kw_2dplot,extent=(E[0],E[-1],0,n))
+  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)
@@ -156,7 +161,7 @@ def plotRatios(r,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,**kw_2dplot,extent=(E[0],E[-1],0,n))
+  i = ax[0].imshow(r,extent=(E[0],E[-1],0,n),**kw_2dplot)
   i.set_clim(0,1.2)
   idx = np.random.random_integers(0,n-1)
   ax[1].plot(E,r[idx],label="Shot n %d"%idx)
@@ -260,17 +265,33 @@ def saveAlignment(fname,transform,roi1,roi2):
 def loadAlignment(fname):
   return np.load(fname).item()
 
-def unravel_results(res,getBest=False):
+def unravel_results(res):
-  out  = dict()
+  #out  = dict()
-  parnames = res[0].fit_result.parameters
+  #parnames = res[0].fit_result.parameters
-  out["parameters"] = dict()
+  #out["parameters"] = dict()
-  for n in parnames:
+  #for n in parnames:
-    out["parameters"][n] = np.asarray( [r.final_pars[n] for r in res])
+  #  out["parameters"][n] = np.asarray( [r.final_pars[n] for r in res])
-  out["ratio"] = np.asarray( [r.ratio for r in res])
+  #out["ratio"] = np.asarray( [r.ratio for r in res])
-  out["p1"] =    np.asarray( [r.p1 for r in res]  )
+  #out["p1"] =    np.asarray( [r.p1 for r in res]  )
-  out["p2"] =    np.asarray( [r.p2 for r in res]  )
+  #out["p2"] =    np.asarray( [r.p2 for r in res]  )
-  out["fom"] =   np.asarray( [r.fom for r in res] )
+  #out["fom"] =   np.asarray( [r.fom for r in res] )
-  return out
+  #out["E"]  =    res[0].E
+  return fit_ret(
+    fit_result   = [r.fit_result for r in res],
+    init_pars    = [r.init_pars  for r in res],
+    final_pars   = [r.final_pars for r in res],
+    final_transform1 = [r.final_transform1 for r in res],
+    final_transform2 = [r.final_transform2 for r in res],
+    im1          = np.asarray( [r.im1 for r in res]),
+    im2          = np.asarray( [r.im2 for r in res]),
+    E            = defaultE,
+    p1           = np.asarray( [r.p1 for r in res]),
+    p2           = np.asarray( [r.p2 for r in res]),
+    ratio        = np.asarray( [r.ratio for r in res]),
+    fom          = np.asarray( [r.fom for r in res] ),
+    tneeded      = np.asarray( [r.tneeded for r in res])
+  )
+
 
 def getTransform(translation=(0,0),scale=(1,1),rotation=0,shear=0):
   t= tf.AffineTransform(scale=scale,rotation=rotation,shear=shear,\
@@ -317,33 +338,33 @@ def transformIminuit(im1,im2,init_transform=dict(),show=False,verbose=True,zeroT
   def transforms(intensity,
     igauss1cen,igauss1sig,iblur1,
     scalex,scaley,rotation,transx,transy,shear,
-    igauss2cen,igauss2sig):
+    igauss2cen,igauss2sig,iblur2):
     t1 = SpecrometerTransformation(translation=(transx,transy),scale=(scalex,scaley),\
          rotation=rotation,shear=shear,intensity=intensity,igauss=(igauss1cen,igauss1sig),\
          iblur=iblur1)
-    t2 = SpecrometerTransformation(igauss=(igauss2cen,igauss2sig))
+    t2 = SpecrometerTransformation(igauss=(igauss2cen,igauss2sig),iblur=iblur2)
     return t1,t2
 
   def model(intensity,
     igauss1cen,igauss1sig,iblur1,
     scalex,scaley,rotation,transx,transy,shear,
-    igauss2cen,igauss2sig):
+    igauss2cen,igauss2sig,iblur2):
 
     t1,t2 = transforms(intensity,
     igauss1cen,igauss1sig,iblur1,
     scalex,scaley,rotation,transx,transy,shear,
-    igauss2cen,igauss2sig)
+    igauss2cen,igauss2sig,iblur2)
     return t1.transformImage(im1_toFit),t2.transformImage(im2_toFit)
   
   def chi2(intensity,
     igauss1cen,igauss1sig,iblur1,
     scalex,scaley,rotation,transx,transy,shear,
-    igauss2cen,igauss2sig):
+    igauss2cen,igauss2sig,iblur2):
 
     i1,i2 = model(intensity,     \
     igauss1cen,igauss1sig,iblur1, \
     scalex,scaley,rotation,transx,transy,shear, \
-    igauss2cen,igauss2sig)
+    igauss2cen,igauss2sig,iblur2)
     d = (i1-i2)/err
     return np.sum(d*d)
 
@@ -417,6 +438,7 @@ def transformIminuit(im1,im2,init_transform=dict(),show=False,verbose=True,zeroT
     final_transform2 = t2,
     im1          = i1,
     im2          = i2,
+    E            = defaultE,
     p1           = p1,
     p2           = p2,
     ratio        = r,

xanes_analyzeRun.py

@@ -11,10 +11,11 @@ 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          = plt.cm.viridis
+  cmap          = cmap
   )
 
 
@@ -27,10 +28,13 @@ 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":
@@ -96,11 +100,16 @@ def showShots(im1,im2):
       a[0].plot(p1)
       a[1].plot(p2)
 
+ 
 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: use None if you want default transformation parameters
+        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.d = readDataset(run)
     if isinstance(run,str):
@@ -162,7 +171,7 @@ class AnalyzeRun(object):
 
 
   def analyzeScan(self,initpars=None,nShotsPerCalib=20,nC=None,doFit=False,fitEveryCalib=False):
-    """ use xhift = None; in this way the fit routine does not try to automatically find the translationx parameter """
+    """ this is a comment """
     if initpars is None: initpars= self.initAlign
     d = self.d
     if nC is None: nC = d.opal1.nCalib
@@ -179,7 +188,7 @@ class AnalyzeRun(object):
       ret = alignment.doShots(s1,s2,initpars=initpars,doFit=doFit)
       res = alignment.unravel_results(ret)
       self.results[i] = res
-    return self.results.values()
+    return list(self.results.values())
 
   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
@@ -191,6 +200,7 @@ class AnalyzeRun(object):
     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()
@@ -200,16 +210,15 @@ class AnalyzeRun(object):
     if initpars is None: initpars= self.initAlign
     d = self.d
     s1,s2 = self.getShot(shots,calib=calib)
-    ret,t = alignment.doShots(s1,s2,initpars=initpars,doFit=doFit,\
+    ret,transformForBestFit = alignment.doShots(s1,s2,initpars=initpars,doFit=doFit,\
             returnBestTransform=True)
-    if doFit:
+    if doFit: self.initAlign = transformForBestFit
-      self.initAlign = t
     ret_unravel = alignment.unravel_results(ret)
     # keep it for later !
     self.results[calib] = ret_unravel
     if unravel: ret = ret_unravel
     if returnBestTransform:
-      return ret,t
+      return ret,transformForBestFit
     else:
       return ret
 
@@ -219,8 +228,7 @@ class AnalyzeRun(object):
     if os.path.exists(fname) and not overwrite:
       print("File %s exists, **NOT** saving, use overwrite=True is you want ..."%fname)
       return
-    if os.path.exists(fname) and overwrite:
+    if os.path.exists(fname) and overwrite: os.unlink(fname)
-      os.unlink(fname)
     print("Saving results to %s"%fname)
     h = h5py.File(fname)
     h["roi1"] = (self.roi1.start,self.roi1.stop)
@@ -250,6 +258,7 @@ class AnalyzeRun(object):
       fname = g_folder_init+"/run%04d_transform.npy" % self.run
     elif 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"]