Added initial Astra Distribution with 200000 particles

scripts/convertAstra2Elegant.py

@@ -86,8 +86,6 @@ class Astra2ElegantConverter:
         print('  Count: %d' % self.t.size)
         print('  Charge: %5.1f pC' % (self.Q*1e12))
         print('  Mean Energy: %5.1f MeV' % (self.meanP*0.511))
-        plt.scatter(self.t,self.p,s=0.5)
-        plt.show()
         return True
 
     def cutDistribution(self,sigma=0):
@@ -108,7 +106,7 @@ class Astra2ElegantConverter:
         print('  Count: %d' % self.t.size)
         print('  Charge: %5.1f pC' % (self.Q * 1e12))
 
-    def analyseBeam(self,N = 20):
+    def analyseBeam(self):
         betay, alphay,  emity = self.getTwiss(self.y, self.yp, self.p)
         betax, alphax, emitx = self.getTwiss(self.x, self.xp, self.p)
         print('Analysis of distribution:')
@@ -119,6 +117,46 @@ class Astra2ElegantConverter:
         print('  beta in y (m): %6.3f' % betay)
         print('  alpha in y (rad): %6.3f' % alphay)
 
+    def analyseBeamSlice(self,slice):
+        dt = (np.max(self.t)-np.min(self.t))/slice
+        tmin=np.min(self.t)
+        ts = np.linspace(0.5*dt+tmin,(slice-0.5)*dt+tmin,num=slice)
+        twiss=np.zeros((8,slice))
+        betay0, alphay0, emity0 = self.getTwiss(self.y, self.yp, self.p)
+        betax0, alphax0, emitx0 = self.getTwiss(self.x, self.xp, self.p)
+        gammax0 = (1+alphax0**2)/betax0
+        gammay0 = (1 + alphay0 ** 2) / betay0
+        for i,tt in enumerate(ts):
+            idx=np.argwhere(np.abs(self.t-tt)<0.5*dt)
+            betay, alphay, emity = self.getTwiss(self.y[idx], self.yp[idx], self.p[idx])
+            betax, alphax, emitx = self.getTwiss(self.x[idx], self.xp[idx], self.p[idx])
+            gammax = (1 + alphax ** 2) / betax
+            gammay = (1 + alphay ** 2) / betay
+            twiss[0, i] = betax
+            twiss[1, i] = alphax
+            twiss[2, i] = emitx
+            twiss[3, i] = betay
+            twiss[4, i] = alphay
+            twiss[5, i] = emity
+            twiss[6, i] = 0.5 * (betax * gammax0 + betax0 * gammax - 2 * alphax * alphax0)
+            twiss[7, i] = 0.5 * (betay * gammay0 + betay0 * gammay - 2 * alphay * alphay0)
+
+
+        plt.plot(ts*1e12, twiss[2, :]*1e9,label='x')
+        plt.plot(ts*1e12, twiss[5, :]*1e9,label='y')
+        plt.xlabel(r'$t$ (ps)')
+        plt.ylabel(r'$\epsilon_n$ (nm)')
+        plt.legend()
+        plt.title('Slice Emittance')
+        plt.show()
+        plt.plot(ts * 1e12, twiss[6, :], label='x')
+        plt.plot(ts * 1e12, twiss[7, :], label='y')
+        plt.xlabel(r'$t$ (ps)')
+        plt.ylabel(r'$\zeta$')
+        plt.legend()
+        plt.title(r'Slice Mismatch')
+        plt.show()
+
     def getTwiss(self,x,xp,p):
         x1 = np.mean(x)
         x2=  np.mean(x*x)
@@ -133,8 +171,6 @@ 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],}
         df_meas = pd.DataFrame.from_dict(meas_df)
@@ -144,28 +180,58 @@ class Astra2ElegantConverter:
         sdds.validate_data()
         pysdds.write(sdds, output)
 
+    def plotLPS(self):
+        plt.scatter(self.t*1e12, self.p*0.511, s=0.5)
+        plt.xlabel(r'$t$ (ps)')
+        plt.ylabel(r'$p$ (MeV)')
+        plt.show()
+
+    def cleanup(self,output='SF_input_dist.sdds'):
+        subprocess.run(['mv','applywakes.out',output])
+        subprocess.run(['rm','applywakes.ele'])
+        subprocess.run(['rm', 'applywakes.lat'])
+        subprocess.run(['rm','convert.sdds'])
 
 
+def convertAstra2Elegant(input,output):
+    a2e =Astra2ElegantConverter()
+    if not a2e.importAstraFile(input = input):
+        return False
+    if not a2e.importElegantFile('convert.sdds'):
+        return False
+    if not a2e.applyWakes():
+        return False
+    if not a2e.importElegantFile('applywakes.out'):
+        return False
+    a2e.analyseBeam()
+    a2e.cleanup(output=output)
+    return True
 
 
+def modifyElegantDist(input,output,cut,slice,scale):
+    a2e = Astra2ElegantConverter()
+    if not a2e.importElegantFile(input = input):
+        return False
+    a2e.cutDistribution(cut)
+    a2e.plotLPS()
+    if slice > 1:
+        a2e.analyseBeamSlice(slice)
+    a2e.exportElegantDistribution(output)
+    return True
 if __name__ == '__main__':
 
     parser = ArgumentParser()
-    parser.add_argument('-input', type=str, help='Astra Particle Distribution File', default='')
-    parser.add_argument('-output', type=str, help='SDDS Output Distribution', default='SF_input_dist.sdds')
+    parser.add_argument('-input', type=str, help='Astra Particle Distribution File Name', default='')
+    parser.add_argument('-output', type=str, help='SDDS Output Distribution File Name', default='SF_input_dist.sdds')
     parser.add_argument('-cut', type=float, help='Cut particles beyond the value, given in units of rms duration', default=2.2)
-    parser.add_argument('-nslice', type=int, help='number of slices for analysis', default=20)
+    parser.add_argument('-nslice', type=int, help='Number of slices for analysis', default=20)
+    parser.add_argument('-nsize', type=int, help='Number of particles after upscaling/downscaling', default=0)
+
     args = parser.parse_args()
-    a2e =Astra2ElegantConverter()
-    if not a2e.importAstraFile(input = args.input):
+    tmpfile = args.output+'.temp'
+    status = convertAstra2Elegant(args.input, tmpfile)
+    if not status:
         sys.exit(0)
-    if not a2e.importElegantFile('convert.sdds'):
-        sys.exit(0)
-    if not a2e.applyWakes():
-        sys.exit(0)
-    if not a2e.importElegantFile('applywakes.out'):
-        sys.exit(0)
-    a2e.cutDistribution(sigma = args.cut)
-    a2e.analyseBeam(N=args.nslice)
-    a2e.exportElegantDistribution(output = args.output)
+    status = modifyElegantDist(tmpfile,args.output,args.cut,args.nslice,args.nsize)
+
     sys.exit(0)