From 2465fc86c64010ed5ab758bc8937c6b260982eb5 Mon Sep 17 00:00:00 2001
From: Thierry Zamofing <thierry.zamofing@psi.ch>
Date: Wed, 18 Apr 2018 15:50:43 +0200
Subject: [PATCH] progress in findxtal.py

---
 python/findxtal.py | 184 +++++++++++++++++++++++++++++++++------------
 1 file changed, 135 insertions(+), 49 deletions(-)

diff --git a/python/findxtal.py b/python/findxtal.py
index 9682854..800a037 100755
--- a/python/findxtal.py
+++ b/python/findxtal.py
@@ -12,50 +12,77 @@ implements an image alalyser for ESB MX
 from scipy import fftpack, ndimage
 import scipy.ndimage as ndi
 import matplotlib.pyplot as plt
+import matplotlib as mpl
 import numpy as np
 #plt.ion()
 
 def ffttest():
+  for phi in np.arange(0.,180.,10.):
+    ffttest2(phi=phi)
+    plt.show()
+    pass
+
+def ffttest2(phi=45.,frq=4.2,amp=1.,n=256.):
   #find the main frequency and phase in 1-D
+  #plt.ion()
+  plt.figure(1)
   d2r=np.pi/180.
-  n=16.
   x=np.arange(n)
-  amp=1.
-  frq=4.5
-  phi=40.*d2r
-  y=amp*np.cos(phi+frq*x/n*2.*np.pi)
+  phi=phi*d2r
+  y=amp*np.cos(frq*x/n*2.*np.pi-phi)
+  plt.plot(x,y,'y')
+  #y[np.where(y<-.5)]=0
   #y*=np.hamming(n)
-  y*=np.hanning(n)
+  w=np.hanning(n)
+  plt.plot(x,w,'y')
+  y*=w
   #y*=1.-np.cos(x/(n-1.)*2.*np.pi)
 
   #y=[1,-1,1,-1,1,-1,1,-1,1,-1,1,-1,1,-1,1,-1]
-  plt.ion()
-  plt.figure()
   plt.stem(x,y)
 
   fy=np.fft.fft(y)
   fya=np.abs(fy)
-  plt.figure()
+  plt.figure(2)
   plt.subplot(211)
   plt.stem(x,fya)
   plt.subplot(212)
   plt.stem(x,np.angle(fy)/d2r)
   print(np.angle(fy[frq])/d2r)
 
+  fya[0]=0
   i=(fya.reshape(2,-1)[0,:]).argmax()
   (vn,v0,vp)=fya[i-1:i+2]
-  frq2=i+(vn-vp)/(2.*(vp+vn-2*v0))
-  print('freq: %g phase %g'%(frq2,np.angle(fy[i])/d2r))
-  #TODO: THE PHASE CALCULATION IS NOT YET WORKING!!!
+  frq_=i+(vn-vp)/(2.*(vp+vn-2*v0))
+  print('freq: %g phase %g %g %g'%((frq_,)+tuple(np.angle(fy[i-1:i+2])/d2r)))
 
-  #y_inv=np.fft.fft(fy)
-  #plt.figure()
-  #plt.plot(x,np.abs(y_inv))
+  #PHASE CALCULATION
+  plt.figure(1)
+  y=np.zeros(x.shape)
+  z=np.zeros(x.shape)
+  for ii in (i,i-1,i+1):
+    y+=np.abs(fy[ii])/n*np.cos(ii*x/n*2.*np.pi+np.angle(fy[ii]))
+    z+=np.abs(fy[ii])/n*np.sin(ii*x/n*2.*np.pi+np.angle(fy[ii]))
+  y*=2 #double because of conjugate part
+  z*=2 #double because of conjugate part
+  plt.plot(x,y,'r')
+  amp_=fya[i-1:i+2].sum()/n*2.
+  t=int(n/2)-1
+  #->phase: find maximum or where the sin is 0
+  w=np.arccos(y[t]/amp_)
+  if(z[t]<0): w=-w
+  print('amplitude %g, value at middle (%d) cos->%g sin->%g -> acos %g deg'%(amp_,t,y[t],z[t],w/d2r))
+  #rot angle at middle x=127 =frq_*t/n*2.*np.pi-phi_=w '%(amp_,t,y[t],phi_/d2r))
+  phi_=frq_*t/n*2.*np.pi-w
+  y=amp_*np.cos(frq_*x/n*2.*np.pi-phi_)
+  print('y[%d] %g'%(t,y[t]))
 
+  plt.plot(x,y,'g')
 
   pass
 
 def findGrid(image,numPeak=2):
+  d2r=np.pi/180.
   #image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/grid_20180409_115332.png', flatten=True)     # flatten=True gives a greyscale image
   s=image.shape
   w1=np.hamming(s[0]).reshape((-1,1))
@@ -67,36 +94,83 @@ def findGrid(image,numPeak=2):
   #plt.figure(num='hamming window*img')
   #plt.imshow(image, interpolation="nearest")
 
-  fft2 = fftpack.fft2(image)
+  fft2 = np.fft.fft2(image)
   fft2=np.fft.fftshift(fft2)
   fa =abs(fft2)
-  fal =np.log(fa)
-  #img=fft3;
-  ofs=int(fal.min())
-  mx2=(image.shape[0]/2,image.shape[1]/2)
-  fal[mx2[0] - 1:mx2[0] + 2, mx2[1] - 1:mx2[1] + 2]=ofs
-  for i in range(numPeak*2):
-    mx=fal .argmax()
-    mx2=divmod(mx,fal .shape[1])
-    peakPos=(mx2[0] - image.shape[0] / 2, mx2[1] - image.shape[1] / 2)
-    peak=fal[mx2[0] - 1:mx2[0] + 2, mx2[1] - 1:mx2[1] + 2]
+  plt.figure(num='log of fft: hamming wnd*image')
+  hi=plt.imshow(fa, interpolation="nearest",norm=mpl.colors.LogNorm(vmin=.1, vmax=fa.max()))
+  #plt.xlim(s[1]/2-50, s[1]/2+50);plt.ylim(s[0]/2-50, s[0]/2+50)
+
+  ctr=np.array(image.shape,dtype=np.int16)/2
+  fa[ctr[0] - 1:ctr[0] + 2, ctr[1] - 1:ctr[1] + 2]=0 # set dc to 0
+  hi.set_data(fa)
+  gen = np.zeros(fft2.shape)
+  x=np.arange(s[1])/float(s[1])*2.*np.pi
+  y=np.arange(s[0])/float(s[0])*2.*np.pi
+  #x=np.linspace(0,2*np.pi,s[1],endpoint=False)
+  #y=np.linspace(0,2*np.pi,s[0],endpoint=False)
+  my=int(s[0]/2)
+  mx=int(s[1]/2)
+  xx, yy = np.meshgrid(x, y)
+
+  res=[] #list of tuples (freq_x,freq_y, phase)
+  for i in range(numPeak):
+    maxAmpIdx=fa.argmax()
+    maxAmpPos=np.array(divmod(maxAmpIdx,fa.shape[1]),dtype=np.int16)
+    peakPos=maxAmpPos-ctr
+    peak=fft2[maxAmpPos[0] - 1:maxAmpPos[0] + 2, maxAmpPos[1] - 1:maxAmpPos[1] + 2]
     print(peakPos)
-    print(peak)
-    (vn, v0, vp)=fal[mx2[0], mx2[1] - 1:mx2[1] + 2]
-    frq1x=peakPos[1]+(vn-vp)/(2.*(vp+vn-2*v0))
-    (vn, v0, vp)=fal[mx2[0]-1:mx2[0]+2,mx2[1]]
-    frq1y=peakPos[0]+(vn-vp)/(2.*(vp+vn-2*v0))
-    print((frq1x,frq1y))
-    fal[mx2[0]-1:mx2[0]+2,mx2[1]-1:mx2[1]+2]=i+ofs
+    print(abs(peak))
+    (vn, v0, vp)=np.log(fa[maxAmpPos[0], maxAmpPos[1] - 1:maxAmpPos[1] + 2]) #using log for interpolation is more precise
+    freq_x=peakPos[1]+(vn-vp)/(2.*(vp+vn-2*v0))
+    (vn, v0, vp)=np.log(fa[maxAmpPos[0]-1:maxAmpPos[0]+2,maxAmpPos[1]])
+    freq_y=peakPos[0]+(vn-vp)/(2.*(vp+vn-2*v0))
+    #calculate phase
+    sumCos=0.
+    sumSin=0.
+    sumAmp=0.
+    n=np.prod(s)
+    for iy in (-1,0,1):#(0,):
+      for ix in (-1,0,1):
+        v=peak[iy+1,ix+1]
+        fx=peakPos[1]+ix
+        fy=peakPos[0]+iy
+        amp=np.abs(v)/n; ang=np.angle(v)
+        sumAmp+=amp
+        sumCos+=amp*np.cos(fx*x[mx] + fy*y[my] + ang)
+        sumSin+=amp*np.sin(fx*x[mx] + fy*y[my] + ang)
+        gen+=amp*np.cos(fx*xx + fy*yy + ang)
+    sumAmp*=2. #double because of conjugate part
+    sumCos*=2.
+    sumSin*=2.
+    #sumAmp=np.abs(peak).sum()/n*2 #double because of conjugate part
+    w=np.arccos(sumCos/sumAmp)
+    if sumSin<0: w=-w
+    phi_= freq_x*x[mx]+freq_y*y[my]-w
+    phi_%=(np.pi*2)
+    res.append((freq_x,freq_y,phi_))
+    fa[maxAmpPos[0]-1:maxAmpPos[0]+2,maxAmpPos[1]-1:maxAmpPos[1]+2]=0 # clear peak
+    maxAmpPos_=2*ctr-maxAmpPos
+    fa[maxAmpPos_[0]-1:maxAmpPos_[0]+2,maxAmpPos_[1]-1:maxAmpPos_[1]+2]=0 # clear conjugated peak
+    hi.set_data(fa)
+
+  gen*=2. # double because of conjugate part
+  for fx,fy,phase in res:
+    print('fx: %g fy: %g phase: %g deg'%(fx,fy,phase/d2r))
 
 
-
-
-  plt.figure(num='fft of hamming window*img')
-  plt.imshow(fal, interpolation="nearest")
   plt.xlim(s[1]/2-50, s[1]/2+50)
   plt.ylim(s[0]/2-50, s[0]/2+50)
-  plt.show()
+  plt.figure('image*wnd')
+  plt.imshow(image,interpolation="nearest")
+  plt.figure('reconstruct')
+  plt.imshow(gen,interpolation="nearest")
+
+  plt.figure()
+  x=range(s[1])
+  y=int(s[0]/2)-1
+  plt.plot(x,image[y,:],'r')
+  plt.plot(x,gen[y,:],'g')
   pass
 
 def findObj(image,objSize=150,tol=0,viz=0):
@@ -196,31 +270,43 @@ def findObj(image,objSize=150,tol=0,viz=0):
 def genImg(shape,*args):
   '''args is a list of tuples (freq_x,freq_y, phase) multiple args can be added'''
   image=np.ndarray(shape)#,dtype=np.uint8)
-  x=np.linspace(0,2*np.pi,shape[1])
-  y=np.linspace(0,2*np.pi,shape[0])
+  x=np.linspace(0,2*np.pi,shape[1],endpoint=False)
+  y=np.linspace(0,2*np.pi,shape[0],endpoint=False)
   #xx, yy = np.meshgrid(x, y, sparse=True)
   xx, yy = np.meshgrid(x, y)
-
+  dc=0
   for i,f in enumerate(args):
     (freq_x, freq_y, phase)=f
     if i==0:
-      image = np.cos(freq_x*xx + freq_y*yy + phase)
+      image = dc+np.cos(freq_x*xx + freq_y*yy - phase)
     else:
-      image += np.cos(freq_x * xx + freq_y * yy + phase)
+      image += np.cos(freq_x * xx + freq_y * yy - phase)
 
   plt.imshow(image, interpolation="nearest")
-  plt.show()
   return image
 
 if __name__ == '__main__':
+  #plt.ion()
   #ffttest()
-  #image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/grid_20180409_115332.png')
+  image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/grid_20180409_115332.png')
   #image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/grid_20180409_115332_45deg.png')
-  #image=-image
-  image=genImg((600,800),(9.5,.2,0))
-  #image=genImg((600,800),(9.5,.2,0),(.4,5.2,0))
+  image=-image
+  #image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/honeycomb.png')
+
+  d2r=np.pi/180.
+  #for phi in np.arange(0.,180.,10.):
+  #  image = genImg((600, 800), (4.5, .2, phi*d2r))
+  #  plt.show()
+  #image=genImg((600,800),(4.,1.0,10.*d2r))
+  #image=genImg((600,800),(4.5,.2,20.*d2r))
+  #findGrid(image,numPeak=1)
+
+  #image=genImg((600,800),(9.5,.2,20.*d2r),(.4,5.2,60.*d2r))
+  #image=genImg((600,800),(3.5,0.,0.*d2r),(0,5.,0.*d2r))
+  findGrid(image,numPeak=2)
+
   #image=genImg((600,800),(9.5,.2,0),(.4,5.2,0),(4,8,0))
-  findGrid(image,numPeak=1)
   #findObj(image,viz=1)
   #findObj(image,viz=255)
   #print(findObj(image))
+  plt.show()