From 3d10ebb8525f598ff5f8c20247a60935872d94ef Mon Sep 17 00:00:00 2001
From: Marco Cammarata <marcocamma@gmail.com>
Date: Sun, 27 Nov 2016 23:52:54 +0100
Subject: [PATCH] first attempt to work with chunks to bypass memory limit

---
 xanes_analyzeRun.py | 45 ++++++++++++++++++++++++++++++---------------
 1 file changed, 30 insertions(+), 15 deletions(-)

diff --git a/xanes_analyzeRun.py b/xanes_analyzeRun.py
index 171ef37..d055edc 100644
--- a/xanes_analyzeRun.py
+++ b/xanes_analyzeRun.py
@@ -141,7 +141,8 @@ def medianRatio(p1,p2,threshold=0.03):
   ratio = p2/p1
   return np.ma.average(ratio,axis=0,weights=p1)
 
-
+def sliceToIndices(shot_slice,nShots):
+  return list(range(*shot_slice.indices(nShots)))
  
 class AnalyzeRun(object):
   def __init__(self,run,initAlign="auto",swapx=g_swapx,swapy=g_swapy):
@@ -221,27 +222,31 @@ class AnalyzeRun(object):
     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")  """
+  def analyzeScan(self,initpars=None,nShotsPerCalib="all",calibs="all",calibsToFit="all",nImagesToFit=0,nSaveImg=5):
+    """ nImagesToFit: number of images to Fit per calibcycle, (int or "all")
+        calibs to fit could be 'all','even','odd'
+    """
     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):
+      shots = list(range(self.nShotsPerCalib[calib]))
+      if nShotsPerCalib != "all": shots = shots[:nShotsPerCalib]
       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 calibsToFit == 'even' and (calib%2==1): nToFit=0
+      if calibsToFit == 'odd'  and (calib%2==0): nToFit=0
+      print("Calib %d, tofit %d"%(calib,nToFit)) 
       ret = None
-      if nImagesToFit > 0:
-        ret,bestTransf = alignment.doShots(s1[:nToFit],s2[:nToFit],doFit=True,\
+      if nToFit > 0:
+        ret,bestTransf = self.doShots(shots=shots[:nToFit],calib=calib,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 nToFit < len(shots):
+        ret2 = self.doShots(shots[nToFit:],initpars=initpars,doFit=False,nSaveImg=0)
         if ret is None:
           ret = ret2
         else:
@@ -268,17 +273,27 @@ class AnalyzeRun(object):
     if save: self.saveTransform()
     return r
 
-  def doShots(self,shots=slice(0,50),calib=None,initpars=None,doFit=False,returnBestTransform=False,nSaveImg='all'):
+  def doShots(self,shots=slice(0,50),calib=None,initpars=None,doFit=False,returnBestTransform=False,nSaveImg='all',nInChunks=250):
     """
         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 shots == "all": shots = list(range(slice(self.nShotsPerCalib[calib])))
+    if isinstance(shots,slice):
+      nmax = self.nShotsPerCalib[calib] if calib is not None else self.data.spec1.nShots
+      shots = sliceToIndices(shots,nmax)
+    chunks = x3py.toolsVarious.chunk(shots,nInChunks)
+    ret = []
+    for chunk in chunks:
+      s1,s2 = self.getShots(chunk,calib=calib)
+      temp = alignment.doShots(s1,s2,initpars=initpars,doFit=doFit,\
+            returnBestTransform=False,nSaveImg=nSaveImg)
+      ret.append(temp)
+    ret = alignment.unravel_results(ret)
+    idxBest = np.nanargmin(ret.fom)
+    transformForBestFit = dict( [ (p,ret.final_pars[p][idxBest]) for p in ret.final_pars.keys() ] )
     if doFit: self.initAlign = transformForBestFit
     # keep it for later !
     self.results[calib] = ret