KaptureControlGui/KCG/widgets/CorrelationCorrection.py

from PyQt4 import QtCore

from scipy.optimize import curve_fit, ridder
from scipy.stats import chisquare
from scipy.stats import chi2_contingency
from scipy.special import erfc

import h5py
import os
import ast
import sys
import traceback

from time import time, sleep

try:
    #for KCG
    from ..base import kcgwidget as kcgw
except:
    #for Analysis GUI
    import kcgwidget as kcgw


import matplotlib
matplotlib.use("Qt4Agg")
#matplotlib.use("Agg")
import numpy as np, matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def gauss(x, sigma, mu, A):
	return A*np.exp(-0.5*((x-mu)/sigma)**2)

def sinus(x,f):
    return np.sin(2*np.pi*f*x)

def linear(x, a,b):
    return a*x+b

def varianz(s,l):
    return np.sqrt(s**2 + 1/(l**2))

def skewness(s,l):
    return 2.0/(s**3*l**3) * (1+ 1.0/(s**2*l**2))**(-3.0/2)

def skewness2(s,v):
    return 2.0-2.0*(s/v)**3

def sigma(v,t): 
    return np.sqrt(v**2*(1-(t/2.0)**(2.0/3)))
  
def lamda(v,t):
    return (2.0/t)**(1.0/3)/v

def expgauss(x, s, m, a ,l):
    return a*l/2*np.exp(l/2*(2*m + l*s**2 - 2*x)) * erfc((m +l*s**2 -x)/(np.sqrt(2)*s)) 


def poly2(x, a,b,c):#,d):#,e):
    return a + b*x+ c*x**2# + d*x**3# +e*x**4

import itertools
def polyfit1d(x, y, order=3):
    ncols = (order + 1)**1
    G = np.zeros((x.size, ncols))
    i = itertools.product(range(order+1))
    for k, (i) in enumerate(i):
        G[:,k] = x**i 
    m, _, _, _ = np.linalg.lstsq(G, y)
    return m

def polyval1d(x, m):
    order = int(np.sqrt(len(m))) - 1
    i = itertools.product(range(order+1))
    z = np.zeros_like(x)
    for a, (i) in zip(m, i):
        z += a * x**i 
    return z

def polyfit2d(x, y, z, order=3):
    ncols = (order + 1)**2
    G = np.zeros((x.size, ncols))
    ij = itertools.product(range(order+1), range(order+1))
    for k, (i,j) in enumerate(ij):
        G[:,k] = x**i * y**j
    m, _, _, _ = np.linalg.lstsq(G, z, rcond=None)
    return m

def polyval2d(x, y, m):
    order = int(np.sqrt(len(m))) - 1
    ij = itertools.product(range(order+1), range(order+1),)
    z = np.zeros_like(x)
    for a, (i,j) in zip(m, ij):
        z += a * x**i * y**j
    return z

def polyfit3d(x, y, z,f, order=3):
    ncols = (order + 1)**3
    #print(ncols)
    G = np.zeros((x.size, ncols))
    ijk = itertools.product(range(order+1), range(order+1),range(order+1))
    for n, (i,j,k) in enumerate(ijk):
        G[:,n] = x**i * y**j * z**k
    m, _, _, _ = np.linalg.lstsq(G, f, rcond=None)
    return m

def polyval3d(x, y, z, m):
    order = int(len(m)**(1/3)) - 1
    #print(order)
    ijk = itertools.product(range(order+1), range(order+1), range(order+1))
    f = np.zeros_like(x)
    for a, (i,j,k) in zip(m, ijk):
        #print(i,j,k)
        f += a * x**i * y**j * z**k
    return f


class CorrelationCorrection(QtCore.QObject):
    """docstring for CorrelationCorrection

    """

    stopSignal = QtCore.pyqtSignal()

    def __init__(self, path, scanX, scanY, delays, workingChannels=[0,1,2,3,4,5,6,7]):
        super(CorrelationCorrection, self).__init__()
        self.host = "ibpt-kapture1"
        self.remotePath = "/home/kapture/GPU-Analyse/PeakReconstruction/"
        self.remoteUser = 'kapture'
        self.remotePW   = 'kapture123'

        #self.dirToReconstruction = '../PeakReconstruction/'

        self.workingChannels = workingChannels
        print(workingChannels)

        self.scanX  = scanX
        self.scanY  = scanY
        self.delays = delays

        self.path = os.path.join(path, 'simulationData' )
        if not os.path.exists(self.path):
            os.mkdir(self.path)


        self.nr=0
        filelist = np.sort(os.listdir(self.path))
        for file in filelist:
            if '.out' in file:
                if int(file[-6:-4]) >= self.nr:
                    self.nr = int(file[-6:-4])+1

        #self.nr-=1
        self.file = 'simu_{:02d}.out'.format(self.nr)
        self.fileMuster = 'simu_{:02d}.npy'.format(self.nr)


        self.process = QtCore.QProcess()

        self.process.readyRead.connect(self._processOut)
        self.process.readyReadStandardOutput.connect(self._processOut)
        self.process.finished.connect(self._processFinished)
        self.process.error.connect(self._processError)
    
        self.process.setWorkingDirectory(self.remotePath)
        self.process.setProcessChannelMode(QtCore.QProcess.MergedChannels)



    def run(self, plotting=False):

        startx = self.scanX[np.where(self.scanY==np.max(self.scanY))]

        p0  = [2.2099e-11, startx, 1.03567844e-8, 2.15525657e+10]
        self.popt, pcov = curve_fit(expgauss, self.scanX, self.scanY, p0=p0)

        self.meanSigma = self.popt[0] #22e-11
        self.meanMean  = self.popt[1]
        self.meanAmplitude = np.max(self.scanY)*2 #self.popt[2]
        self.meanVarianz  = 45
        self.deltaVarianz = 15
        self.deltaMean    = 10

        try:
            self._generateSimulationData(plotting)
            self._runReconstruction()
            self._calculateCorrections(plotting)
            if plotting:
                plt.show(block=False)
        except:
            traceback.print_exc()
            #popup = kcgw.PopupDialog("There was an error",
            #                    title='Reminder', okonly=True)
            #popup.exec_()
            #popup.deleteLater()
            #popup.get_return_value()

            
            pass
        self.stopSignal.emit()
        
        pass

    def save(self, file):
        with h5py.File(file, 'w') as f:
            f['Mean'] = self.poptMean
            f['Sigma'] = self.poptSigma
            f['Amplitude'] = self.poptAmpl



    def _generateSimulationData(self, plotting):
        print('_generateSimulationData')
        if plotting:
            self._plotScan()
            plt.show(block=False)

        fitx = np.linspace(np.min(self.delays)*0.8,np.max(self.delays)*1.2,2000)
        steps = 10
        steplength = 3000
        n = steps*steplength
        fh = np.memmap(os.path.join(self.path, self.file), mode='w+', dtype='uint32', shape=(n*184//2*8))
        
        ctr = np.left_shift(np.array(np.repeat(np.linspace(0,184-2,184//2),8), dtype=np.uint32), 24)
        
        master =[]

        for i in range(n):
            #print(i, end='\r')
            fh[i*184//2*8:(i+1)*184//2*8] = ctr + np.left_shift(2040,12) + 2040

            m = self.meanMean + self.deltaMean*sinus(i, 77.7777e-5)*1e-12
            v = (self.meanVarianz + self.deltaVarianz*sinus(i, 456.1e-5))*1e-12
            A = self.meanAmplitude +0.8*self.meanAmplitude*sinus(i, 3.33e-5) 

            s = self.meanSigma
            t = skewness2(s,v)
            #s = sigma(v,t)
            l = lamda(v,t)

            popt = [s, m, self.popt[2], l]

            fity = expgauss(fitx, *popt)  
            fityM = np.max(fity)
            m2 = fitx[np.where(fity==fityM)]
            m2 = m-(m2-m)

            popt = [s, m2, self.popt[2], l]

            #fity = expgauss(fitx, *popt)  
            #fityM = np.max(fity)
            #m = fitx[np.where(fity==fityM)]
            
            data = expgauss(self.delays, *popt)
            data = data/fityM*A
            #print(data)
            master.append([v/2.0, m-self.meanMean, A, l])

            fh[i*184//2*8:i*184//2*8+8] += np.array(data, dtype=np.uint32) 
            
        print('done')
        master = np.array(master)
        
        np.save(os.path.join(self.path, self.fileMuster), master)

        dataxb = self.delays*1e12
        #datax *= 1e12
        dataxb -= np.min(dataxb)-10
        f = open(os.path.join(self.path,'Measurement_board_6028.log'), 'a')
        f.write('\n---\n')
        f.write('2018-12-05 20:12:14\n')
        f.write('\tSimulation\n')
        f.write('\tFilename: {}\n'.format(self.file))
        f.write('\tNumber of Skipped Turns: 0\n')
        f.write('\tBeam Current (mA): 10\n')
        f.write('\tADC Delays: [0, 0, 0, 0, 0, 0, 0, 0]\n\tT/H Delay: 0\n\t25ps Delay: 0\n\tT/H Delay 2: 0\n\t25ps Delay 2: 0\n\tDelay Cascade: 0\n')
        f.write('\tdatax: ['+', '.join(str(x) for x in dataxb) + ']\n')
        f.write('\tWorking Channels: ['+', '.join(str(x) for x in self.workingChannels) + ']\n\n')
        f.close()



    def checkForCuda(self):
        #Check if CUDA available:
        cuda = True
        try:
            import pycuda.driver as cuda
            import pycuda.autoinit
            def printGPUstats(onlymem=True):
                (free,total)=cuda.mem_get_info()
                print("Global memory occupancy: %f%% free of %f"%(free*100/total, total/1024/1024))

                if onlymem:
                    return
                for devicenum in range(cuda.Device.count()):
                    device=cuda.Device(devicenum)
                    attrs=device.get_attributes()
                    
                    #Beyond this point is just pretty printing
                    print("\n===Attributes for device %d"%devicenum)
                    for (key,value) in attrs.items():
                        print("%s:%s"%(str(key),str(value)))
            printGPUstats()
        except Exception as e:
            print('no CUDA')
            cuda = False

        return cuda

    def _runReconstruction(self):
        print('_runReconstruction')

        cuda = self.checkForCuda()
        
        if cuda:
            param = ["-u", os.path.join("PeakReconstuction.py"), "-f", os.path.abspath(os.path.join(self.path, self.file)), "-c", 'd']

            self.process.start("python3", param)
                    # -u is for python to force it to unbuffered print
            self.process.waitForStarted()
            self.process.waitForFinished()

        else:
            
            remoteFile = self.remotePath + "simulationData/"+ self.file
            localPath  = os.path.abspath(self.path)
            localFile  = os.path.abspath(os.path.join(self.path, self.file))

            import paramiko
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(self.host, username=self.remoteUser, password=self.remotePW)
            
            print('put')

            # Define progress callback that prints the current percentage completed for the file
            def progress(sent, size):
                sys.stdout.write("progress: %.2f%%   \r" % (float(sent)/float(size)*100) )

            
            scp = ssh.open_sftp()
            scp.put(localFile, remoteFile)#, progress)
            scp.put(os.path.join(localPath, 'Measurement_board_6028.log'), self.remotePath+'simulationData/Measurement_board_6028.log', progress)

            print('calc')

            stdin, stdout, stderr = ssh.exec_command("cd {} ; python3 -u {} -f {} -c d -v -o '' ".format(self.remotePath, "PeakReconstruction.py", remoteFile))
            #stdout.channel.recv_exit_status()
            while not stdout.channel.exit_status_ready():
                o = stdout.readline()
                if o != '':
                    print(o[:-1])
            
            lines = stdout.readlines()
            for line in lines:
                print(line[:-1])

            lines = stderr.readlines()
            for line in lines:
                print('e', line[:-1])

            print('get')
            scp.get(remoteFile[:-4]+'_fit.hdf', localFile[:-4]+'_fit.hdf', progress)
            scp.close()
            ssh.close()

        
        pass


    def _calculateCorrections(self, plotting):
        print('_calculateCorrections')

        filename = os.path.join(self.path, self.file[:-4]+'_fit.hdf')
        fitData={}
        with h5py.File(filename,'r') as fh:
            fitData['Amplitude'] = np.reshape(fh['Amplitude'][()]/10.0, -1)
            fitData['Mean']  = np.reshape(fh['Mean'][()]/10.0, -1)
            fitData['Sigma'] = np.reshape(fh['Sigma'][()]/10.0, -1)
            try:
                fitData['Error'] = fh['Error'][()]/10.0
            except:
                pass

            fitData['log'] = ast.literal_eval(fh['fit_info']['log'][()])
            fitData['selector'] = fh['fit_info']['selector'][()]
            fitData['filename'] = filename


        m = fitData['Mean'][:]
        a = fitData['Amplitude'][:]
        s = fitData['Sigma'][:]

        selector = np.where(a>0)

        m = m[selector]
        a = a[selector]
        s = s[selector]

        #m = m-np.mean(m)

        original = np.load(os.path.join(self.path, self.fileMuster))
        #print(original.shape)
        mo = original[:,1]*1e12
        ao = original[:,2]
        so = original[:,0]*1e12

        mo = mo[selector]
        ao = ao[selector]
        so = so[selector]
        #print(mo.shape)


        w = np.array([ao,mo,so])
        pov = np.corrcoef(w)
        print('Korrelation Original:')
        print('A-M: {:.5f}'.format(pov[0,1]))
        print('A-S: {:.5f}'.format(pov[0,2]))
        print('M-S: {:.5f}'.format(pov[1,2]))   

        self.povOrg = pov

        w = np.array([a,m,s])
        pov = np.corrcoef(w)
        print('Korrelation unkorrigiert:')
        print('A-M: {:.5f}'.format(pov[0,1]))
        print('A-S: {:.5f}'.format(pov[0,2]))
        print('M-S: {:.5f}'.format(pov[1,2])) 
        self.povUnCorr = pov

        if plotting:
            self._plotComparison(mo,ao,so,m,a,s)


        self.poptMean, pcov = curve_fit(poly2, m,mo)
        mc = poly2(m, *self.poptMean)

        self.poptSigma = polyfit2d(mc, s, so, order=3)
        sc = polyval2d(mc, s, self.poptSigma)

        self.poptAmpl = polyfit3d(mc, sc, a, ao, order=2)
        ac = polyval3d(mc, sc, a, self.poptAmpl)

        print(self.poptAmpl)

        w = np.array([ac,mc,sc])
        pov = np.corrcoef(w)
        print('Korrelation korrigiert:')
        print('A-M: {:.5f}'.format(pov[0,1]))
        print('A-S: {:.5f}'.format(pov[0,2]))
        print('M-S: {:.5f}'.format(pov[1,2])) 

        self.povCorr = pov

        

        if plotting:
            self._plotComparison(mo,ao,so,mc,ac,sc)
            #plt.show()


        self.mo = mo
        self.so = so
        self.ao = ao 
        
        self.a = a
        self.s = s
        self.m = m 

        self.mc = mc
        self.sc = sc 
        self.ac = ac

        pass

    def test(self):
        param = ["-u", os.path.join("PeakReconstuction.py"), "-f", os.path.abspath(os.path.join(self.path, self.file)), "-c", 'd', '-c2', 'd']

        self.process.start("python3", param)
                # -u is for python to force it to unbuffered print
                
        self.process.waitForStarted()
        self.process.waitForFinished()


        filename = os.path.join(self.path, self.file[:-4]+'_fit.hdf')
        fitData={}
        with h5py.File(filename,'r') as fh:
            fitData['Amplitude'] = np.reshape(fh['Amplitude'][()]/10.0, -1)
            fitData['Mean']  = np.reshape(fh['Mean'][()]/10.0, -1)
            fitData['Sigma'] = np.reshape(fh['Sigma'][()]/10.0, -1)
            try:
                fitData['Error'] = fh['Error'][()]/10.0
            except:
                pass

            fitData['log'] = ast.literal_eval(fh['fit_info']['log'][()])
            fitData['selector'] = fh['fit_info']['selector'][()]
            fitData['filename'] = filename


        m = fitData['Mean'][:]
        a = fitData['Amplitude'][:]
        s = fitData['Sigma'][:]

        selector = np.where(a>0)

        m = m[selector]
        a = a[selector]
        s = s[selector]

        original = np.load(os.path.join(self.path, self.fileMuster))
        
        mo = original[:,1]*1e12
        ao = original[:,2]
        so = original[:,0]*1e12

        mo = mo[selector]
        ao = ao[selector]
        so = so[selector]

        self._plotComparison(mo,ao,so,m,a,s)

        plt.show()

        pass


    def _plotFitOverOrg(self, mo,ao,so,m,a,s, titel=""):

        fig4, axs3 = plt.subplots(1, 3, figsize=(14, 9))
        axs3[0].plot(mo,m,'.')
        axs3[1].plot(so,s,'.')
        axs3[2].plot(ao,a,'.')

        axs3[0].set_ylabel('Ankunftszeit\n(ps)')
        axs3[1].set_ylabel('Breite\n(ps)')
        axs3[2].set_ylabel('Amplitude\n(ADC counts)')
        plt.title('Fit over Org '+titel)
        plt.tight_layout()

    def _plotComparison(self, mo,ao,so,m,a,s, titel=""):
        
        fig4, axs3 = plt.subplots(3, 1, figsize=(14, 9),sharex=True)
        axs3[0].plot(mo)
        axs3[1].plot(so)
        axs3[2].plot(ao)

        axs3[0].plot(m)
        axs3[1].plot(s)
        axs3[2].plot(a)

        axs3[0].set_ylabel('Ankunftszeit\n(ps)')
        axs3[1].set_ylabel('Breite\n(ps)')
        axs3[2].set_ylabel('Amplitude\n(ADC counts)')
        plt.tight_layout()
        plt.title(titel)
        #plt.savefig('plots/simu2.png', dpi=200)  

    def _plotScan(self):
        plt.figure(figsize=(14,9))
        plt.plot(self.scanX*1e9, self.scanY, '.', markersize=16, label='Timescan')
        plt.plot(np.sort(self.scanX*1e9), expgauss(np.sort(self.scanX), *self.popt), linewidth=2, label='Fit')
        i=1
        ym = np.max(self.scanY)*1.06
        for xt in self.delays*1e9:
            if i in [5,6]:
                i += 1
                continue
            plt.plot([xt,xt],[0,ym], linewidth=1, color='black')#, label='ADC {}'.format(i))
            plt.text(xt, ym+1, '{}'.format(i), horizontalalignment='center',  verticalalignment='bottom')#, rotation='vertical')
            i += 1

        plt.ylabel('ADC count')
        plt.xlabel('t (ns)')
        plt.ylim(-2, 132)
        plt.tight_layout()


    def _processOut(self):
        #print('out')
        print(self.process.readAll().data().decode('utf8'))#,end='')

    def _processFinished(self):
        print('Process Finished!')

    def _processError(self):
        print('Process Error!')