diff --git a/pmd_beamphysics/interfaces/astra.py b/pmd_beamphysics/interfaces/astra.py index 72338f3..c7e11e9 100644 --- a/pmd_beamphysics/interfaces/astra.py +++ b/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 diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index ed33b25..2e99816 100644 --- a/pmd_beamphysics/particles.py +++ b/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)