ElegantTools/distribution.py

import pysdds
import copy
import subprocess
import numpy as np
import pandas as pd

class Distribution:
    def __init__(self,parent=None,filename=None):
        self.parent=parent
        self.filename=filename
        self.dist=None
        self.distCopy=None
        self.columns = None
        self.parameters = None
        self.loadDist()

    def getQ(self):
        return self.Q

    def plotLPS(self):
        self.parent.plot.LPS(self.t,self.p,self.Q)

    def loadDist(self):
        if self.filename is None:
            return
        self.dist = pysdds.read(self.filename)
        self.distCopy = copy.deepcopy(self.dist)
        self.Q = self.dist.par('Charge').data[0]
        self.x = self.dist.col('x').data[0]
        self.xp = self.dist.col('xp').data[0]
        self.y = self.dist.col('y').data[0]
        self.yp = self.dist.col('yp').data[0]
        self.t = self.dist.col('t').data[0]
        self.p = self.dist.col('p').data[0]

        if self.Q < 1e-15:
            self.Q = float(str(self.parent.UIDistCharge.text()))*1e-12
        self.plotLPS()

    def getTwiss(self,x,xp,p):
        x1 = np.mean(x)
        x2=  np.mean(x*x)
        xp1 = np.mean(xp)
        xp2 = np.mean(xp*xp)
        xxp1 = np.mean(x*xp)
        g1=np.mean(p)
        ex = np.sqrt((x2-x1*x1)*(xp2-xp1*xp1)-(xxp1-x1*xp1)**2)*g1
        bx = (x2-x1*x1)*g1/ex
        ax = - (xxp1-x1*xp1)*g1/ex

        return bx,ax,ex

    def analyseBeamShort(self,filename):
        dist = pysdds.read(filename)
        p = dist.col('p').data[0]
        x = dist.col('x').data[0]
        xp = dist.col('xp').data[0]
        y = dist.col('y').data[0]
        yp = dist.col('yp').data[0]
        bx, ax, ex = self.getTwiss(x,xp,p)
        by, ay, ey = self.getTwiss(y,yp,p)
        res={}
        res['betax']=bx
        res['alphax']=ax
        res['betay'] = by
        res['alphay'] = ay
        res['pAverage'] = np.mean(p)
        return res

    def analyseBeam(self,filename,Q):
        output = 'Runs/analysebeam.sdds'
        subprocess.run(['sddsanalyzebeam',filename,output])
        fields = ['enx','eny','betax','betay','alphax','alphay','St','Sdelta','pAverage']
        names =['Emittance in X','Emittance in Y','Beta in X','Beta in Y','Alpha in X','Alpha in Y',
                'RMD Duration','Energy Spread','Mean Energy']
        units = ['nm','nm','m','m','rad','rad','ps','%','MeV']
        scl =[1e9,1e9,1,1,1,1,1e12,1e2,0.511,1e12]
        out = sdds.load(output)
        self.parent.UIDistAnalysis.clear()
        self.parent.UIDistAnalysis.appendPlainText('Analysis of File:\n%s\n' % filename)
        ns = len(self.dist.columnData[self.columns.index('p')][0])
        self.parent.UIDistAnalysis.appendPlainText('Number of Particles: %d\n' % ns)
        for i ,field in enumerate(fields):
            idx = out.columnName.index(field)
            val=scl[i]*out.columnData[idx][0][0]
            self.parent.UIDistAnalysis.appendPlainText('%s: %7.3f (%s)' % (names[i],val,units[i]))
        self.parent.UIDistAnalysis.appendPlainText('Charge: %7.3f (pC)' % (Q*1e12))
        del out

    def matchDist(self):
        betax = float(str(self.parent.UIDistBetax.text()))
        betay = float(str(self.parent.UIDistBetay.text()))
        alphax = float(str(self.parent.UIDistAlphax.text()))
        alphay = float(str(self.parent.UIDistAlphay.text()))
        argx = '-xPlane=beta=%f,alpha=%f' % (betax,alphax)
        argy = '-yPlane=beta=%f,alpha=%f' % (betay, alphay)
        filename = 'Runs/inputdist.sdds'
        output = 'Runs/outputdist.sdds'
        self.saveDist(filename)
        subprocess.run(['sddsmatchtwiss', filename, output,argx,argy])
        temp = self.filename
        self.filename = output
        self.loadDist()
        self.filename = temp


    def addBlurr(self):
        ns=len(self.p)
        blurr = float(str(self.parent.UIDistSpread.text())) / 511.
        dg = np.random.normal(0, blurr, size=ns)
        self.p += dg
        self.plotLPS()

    def centerDist(self):
        self.Q = float(str(self.parent.UIDistCharge.text()))*1e-12
        pCenter = float(str(self.parent.UIDistEnergy.text()))/0.511
        self.p = self.p - np.mean(self.p)+pCenter
        self.plotLPS()


    def cutDist(self):
        sigma = float(str(self.parent.UIDistLength.text()))
        if sigma <=0:
            return
        rmsT=np.std(self.t)
        idx=np.argwhere(np.abs(self.t)< rmsT*sigma)
        N0 = float(len(self.t))
        self.x=self.x[idx]
        self.y=self.y[idx]
        self.xp=self.xp[idx]
        self.yp=self.yp[idx]
        self.p=self.p[idx]
        self.t=self.t[idx]
        N1 = float(len(self.t))
        self.Q = self.Q*N1/N0
        self.plotLPS()

    def revertDist(self):
        self.dist=copy.deepcopy(self.dist)

    def saveDist(self,filename):
        if len(self.t.shape) > 1:
            meas_df = {'t': self.t[:, 0], 'p': self.p[:, 0], 'x': self.x[:, 0], 'xp': self.xp[:, 0],
                       'y': self.y[:, 0], 'yp': self.yp[:, 0], }
        else:
            meas_df = {'t': self.t[:], 'p': self.p[:], 'x': self.x[:], 'xp': self.xp[:],
                       'y': self.y[:], 'yp': self.yp[:], }
        df_meas = pd.DataFrame.from_dict(meas_df)
        parameters = {'Charge': [float(self.Q)]}
        subprocess.run(['rm', filename])
        sdds = pysdds.SDDSFile.from_df([df_meas], parameter_dict=parameters, mode='binary')
        sdds.col('x').nm['units']= 'm'
        sdds.col('y').nm['units'] = 'm'
        sdds.col('t').nm['units'] = 's'
        sdds.col('p').nm['units'] = 'm$be$nc'
        sdds.par('Charge').nm['units']='C'
        sdds.validate_data()

        pysdds.write(sdds, filename)