REfined writing of ELegant distribution to include units in the distribution file.

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)