minor improvement of ibsimu2tw script

scripts/ibsimu2tw

@@ -17,9 +17,14 @@
     R. Miyamoto (2018.06.25)
 '''
 
-import argparse
+#-- Libraries
+
+import numpy as np
+import sys
 import os
-import subprocess
+import argparse
+from ess import TraceWin
+print(TraceWin.__file__)
 
 parser = argparse.ArgumentParser(description="Converts IBSimu output to TraceWin dst format")
 
@@ -60,14 +65,6 @@ if args.h3:
 
 #W_cut = 0.01  # Filter for W [MeV]. Only for experts.
 
-##################################################
-# No need to change below
-##################################################
-
-#-- Libraries
-
-import numpy as np
-import sys
 
 #-- Function for saving
 
@@ -77,19 +74,9 @@ def save_dst(file_name, x6, Ibeam, mc2=938.272, freq=352.21):
         - Coordinates: (  x,  x',  y,  y', phi,   W)
         - Units      : ( cm, rad, cm, rad, rad, MeV)
     '''
-    import numpy, struct
-
-    x6 = numpy.reshape(x6, (numpy.size(x6)))    
-    with open(file_name, 'wb') as file:
-        out  = struct.pack('b'                  , 125        )
-        out += struct.pack('b'                  , 100        )
-        out += struct.pack('i'                  , len(x6)//6 )
-        out += struct.pack('d'                  , Ibeam      )
-        out += struct.pack('d'                  , freq       )
-        out += struct.pack('b'                  , 125        )  # Typo in the manual !!!!
-        out += struct.pack('{}d'.format(len(x6)), *x6        )
-        out += struct.pack('d'                  , mc2        )
-        file.write(out)
+    dst=TraceWin.dst(Ib=Ibeam,freq=freq,mass=mc2)
+    dst.append_many(x6)
+    dst.save(file_name)
 
 #-- Parameters from the header
 
@@ -105,47 +92,36 @@ data = np.loadtxt(file_name_in, skiprows=8)
 
 #-- Adjust coordinates and units
 
-data[:,1] *= 1e2  # x  : m => cm
-data[:,3] *= 1e2  # y  : m => cm
 data[:,9] *= 1e3  # mc2: GeV => MeV
 data[:,6]  = np.sqrt(data[:,9]**2+(1.0+data[:,2]**2+data[:,4]**2)*(cp*(1.0+data[:,6]))**2)-data[:,9]  # dpz => W
 
 try   : data = data[np.where(np.abs(data[:,6]-0.075) < W_cut)]
 except: pass
 
-#-- Separate species
+#-- Separate species, save outputs
 
 mc2_lst = sorted(list(set(list(data[:,9]))))
 
 x6_H1    = data[np.where(data[:,9] == mc2_lst[0])]
 x6_H1    = x6_H1[:,1:7]
-Ibeam_H1 = Ibeam*len(x6_H1)/len(data)
+Ibeam_H1 = Ibeam*sum(x6_H1[8])/len(data)
 
-try:
+save_dst(        file_name_out_H1, x6_H1, Ibeam_H1, 1.0*mc2_H1, freq)
+
+
+if args.h2:
     x6_H2    = data[np.where(data[:,9] == mc2_lst[1])]
     x6_H2    = x6_H2[:,1:7]
     Ibeam_H2 = Ibeam*len(x6_H2)/len(data)
-except:
-    pass
 
-try:
+    save_dst(file_name_out_H2, x6_H2, Ibeam_H2, 2.0*mc2_H1, freq)
+
+if args.h3:
     x6_H3    = data[np.where(data[:,9] == mc2_lst[2])]
     x6_H3    = x6_H3[:,1:7]
     Ibeam_H3 = Ibeam*len(x6_H3)/len(data)
-except:
-    pass
 
-#-- Save outputs
-
-save_dst(        file_name_out_H1, x6_H1, Ibeam_H1, 1.0*mc2_H1, freq)
-
-try   : save_dst(file_name_out_H2, x6_H2, Ibeam_H2, 2.0*mc2_H1, freq)
-except: pass
-try   : save_dst(file_name_out_H3, x6_H3, Ibeam_H3, 3.0*mc2_H1, freq)
-except: pass
+    save_dst(file_name_out_H3, x6_H3, Ibeam_H3, 3.0*mc2_H1, freq)
 
 #-- Ending
 
-sys.exit()
-
-#--