ElegantTools/distribution.py

import sdds
import copy
import subprocess
import numpy as np

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.dist.parameterData[self.parameters.index('Charge')][0]

    def plotLPS(self):
        x = self.dist.columnData[self.columns.index('t')][0]
        y = self.dist.columnData[self.columns.index('p')][0]
        Q = self.dist.parameterData[self.parameters.index('Charge')][0]
        self.parent.plot.LPS(x,y,Q)

    def loadDist(self):
        if self.filename is None:
            return
        self.dist = sdds.load(self.filename)
        self.distcopy = copy.deepcopy(self.dist)
        self.columns = self.dist.columnName
        self.parameters =self.dist.parameterName
        Q = self.dist.parameterData[self.parameters.index('Charge')][0]
        if Q < 1e-15:
            Q = float(str(self.parent.UIDistCharge.text()))*1e-12
        self.dist.setParameterValue('Charge',Q)
        self.analyseBeam(self.filename,Q)

        self.plotLPS()


    def analyseBeamShort(self,filename):
        output = 'Runs/analysebeam.sdds'
        subprocess.run(['sddsanalyzebeam', filename, output])
        fields = ['betax', 'betay', 'alphax', 'alphay', 'pAverage']
        out = sdds.load(output)
        res={}
        for i ,field in enumerate(fields):
            idx = out.columnName.index(field)
            res[field]=out.columnData[idx][0][0]
        res['Q']=self.dist.parameterData[self.parameters.index('Charge')][0]
        del out
        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):
        p=self.dist.columnData[self.columns.index('p')][0]
        ns=len(p)
        blurr = float(str(self.parent.UIDistSpread.text())) / 511.
        dg = np.random.normal(0, blurr, size=ns)
        self.dist.setColumnValueList('p', [p[i]+dg[i] for i in range(ns)])
        self.plotLPS()

    def centerDist(self):
        Q = float(str(self.parent.UIDistCharge.text()))*1e-12
        self.dist.setParameterValue('Charge', Q)
        pCenter = float(str(self.parent.UIDistEnergy.text()))/0.511
        p = self.dist.columnData[self.columns.index('p')][0]
        p = p - np.mean(p)+pCenter
        self.dist.columnData[self.columns.index('p')][0] = [p0 for p0 in p]
        #self.dist.setColumnValueList('p', p)
        self.plotLPS()

    def cutDist(self):
        cut = float(str(self.parent.UIDistLength.text()))
        t = self.dist.columnData[self.columns.index('t')][0]
        nti=len(t)
        t0 = np.mean(t)
        trms = np.std(t)
        tmin = t0-cut*trms
        tmax = t0+cut*trms
        idx = np.argwhere(np.logical_and(t > tmin,t < tmax))
        ntf = len(idx)
        for i, fld in enumerate(self.columns):
            dat = np.array(self.dist.columnData[i][0])
            self.dist.setColumnValueList(fld, [d for d in dat[idx].ravel()])
        Q = self.dist.parameterData[self.parameters.index('Charge')][0]*float(ntf)/float(nti)
        self.dist.setParameterValue('Charge', Q)
        self.plotLPS()

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

    def saveDist(self,filename):
        x=sdds.SDDS()
        x.mode = sdds.SDDS_BINARY
        x.description[0]="test"
        x.description[1]="test"
        x.parameterName=self.dist.parameterName
        x.parameterData=self.dist.parameterData
        x.parameterDefinition = self.dist.parameterDefinition
        x.columnName=self.dist.columnName
        x.columnData=[self.dist.columnData[i] for i in range(len(self.dist.columnName))]
        x.columnDefinition=self.dist.columnDefinition
        x.save(filename)
        del x