Add probe particle option

pmd_beamphysics/interfaces/astra.py

@@ -153,7 +153,10 @@ def vprint(*a, **k):
     sigma = {}
     for k in ['x', 'y', 'z', 'px', 'py', 'pz', 't']:
         ref_particle[k] = particle_group.avg(k)
-        sigma[k] =  particle_group.std(k)
+        std = particle_group.std(k)
+        if std == 0:
+            std = 1e-12 # Give some size
+        sigma[k] = std
     ref_particle['t'] *= 1e9 # s -> nS
 
     # Make structured array
@@ -203,6 +206,7 @@ def vprint(*a, **k):
         data[5]['x'] = 1.5*sigma['x'];data[5]['t'] =  1.5*sigma['t']
         data[6]['y'] = 1.5*sigma['y'];data[6]['t'] = -1.5*sigma['t']        
         data[1:7]['status'] = -3
+        data[1:7]['q'] = 0.5e-5 # ? Seems to be required 
         data[1:7]['pz'] = 0 #? This is what the Astra Generator does
 
     # Save in the 'high_res = T' format

pmd_beamphysics/particles.py

@@ -759,8 +759,9 @@ def drift_to_t(self, t=None):
         self.t = np.full(self.n_particle, t)
 
     # Writers
-    def write_astra(self, filePath, verbose=False):
-        write_astra(self, filePath, verbose=verbose)
+    @functools.wraps(write_astra)    
+    def write_astra(self, filePath, verbose=False, probe=False):
+        write_astra(self, filePath, verbose=verbose, probe=probe)
 
     def write_bmad(self, filePath, p0c=None, t_ref=0, verbose=False):
         write_bmad(self, filePath, p0c=p0c, t_ref=t_ref, verbose=verbose)