diff --git a/scripts/convertAstra2Elegant.py b/scripts/convertAstra2Elegant.py index e1c20b3..dae1830 100644 --- a/scripts/convertAstra2Elegant.py +++ b/scripts/convertAstra2Elegant.py @@ -171,13 +171,22 @@ class Astra2ElegantConverter: return bx,ax,ex def exportElegantDistribution(self,output = 'SF_input_dist.sdds'): - 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],} + 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',output]) sdds = pysdds.SDDSFile.from_df([df_meas], parameter_dict=parameters, mode='binary') sdds.validate_data() + 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' pysdds.write(sdds, output) def plotLPS(self): @@ -217,6 +226,9 @@ def modifyElegantDist(input,output,cut,slice,scale): if slice > 1: a2e.analyseBeamSlice(slice) a2e.exportElegantDistribution(output) + a2e.importElegantFile(input=output) + a2e.analyseBeam() + return True if __name__ == '__main__': @@ -234,4 +246,5 @@ if __name__ == '__main__': sys.exit(0) status = modifyElegantDist(tmpfile,args.output,args.cut,args.nslice,args.nsize) + sys.exit(0)