Merge pull request #45 from ChristopherMayes/astra_fixes

Add probe particle option

pmd_beamphysics/fields/fieldmesh.py

@@ -9,7 +9,7 @@
 from pmd_beamphysics.plot import plot_fieldmesh_cylindrical_2d, plot_fieldmesh_cylindrical_1d
 
 from pmd_beamphysics.interfaces.ansys import read_ansys_ascii_3d_fields
-from pmd_beamphysics.interfaces.astra import read_astra_3d_fieldmaps, write_astra_3d_fieldmaps
+from pmd_beamphysics.interfaces.astra import write_astra_1d_fieldmap, read_astra_3d_fieldmaps, write_astra_3d_fieldmaps, astra_1d_fieldmap_data
 from pmd_beamphysics.interfaces.gpt import write_gpt_fieldmesh
 from pmd_beamphysics.interfaces.impact import create_impact_solrf_fieldmap_fourier, create_impact_solrf_ele
 from pmd_beamphysics.interfaces.superfish import write_superfish_t7, read_superfish_t7
@@ -102,8 +102,10 @@ class FieldMesh:
     Writers
     
     - `.write`
+    - `.write_astra_1d`
     - `.write_astra_3d`
     - `.to_cylindrical`
+    - `.to_astra_1d`
     - `.to_impact_solrf`
     - `.write_gpt`
     - `.write_superfish`
@@ -345,7 +347,16 @@ def write(self, h5, name=None):
             g = h5
 
         write_pmd_field(g, self.data, name=name)   
-
+
+    @functools.wraps(write_astra_1d_fieldmap)
+    def write_astra_1d(self, filePath):      
+        return  write_astra_1d_fieldmap(self, filePath)
+
+    def to_astra_1d(self):
+        z, fz = astra_1d_fieldmap_data(self)   
+        dat = np.array([z, fz]).T 
+        return {'attrs': {'type': 'astra_1d'}, 'data': dat}
+
     def write_astra_3d(self, common_filePath, verbose=False):      
         return  write_astra_3d_fieldmaps(self, common_filePath)
 

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,13 +206,93 @@ 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
     np.savetxt(outfile, data, fmt = ' '.join(8*['%20.12e']+2*['%4i']))
 
 
 
+
+def write_astra_1d_fieldmap(fm, filePath): 
+    """
+    Writes an Astra fieldmap file from a FieldMesh object.
+    
+    Requires cylindrical geometry for now.
+    
+    Parameters
+    ----------
+    filePath: str
+        Filename to write to
+
+    """
+
+    z, fz = astra_1d_fieldmap_data(fm)
+
+    # Flatten dat   
+    dat = np.array([z, fz]).T    
+
+    np.savetxt(filePath, dat, comments='') 
+
+
+
+def astra_1d_fieldmap_data(fm): 
+    """
+    Astra fieldmap data.
+    
+    Requires cylindrical geometry for now.
+    
+    Returns
+    -------
+    z: array-like
+        z coordinate in meters
+    
+    fz: array-like
+        field amplitude corresponding to z
+    
+    """    
+
+
+    assert fm.geometry == 'cylindrical', f'Geometry: {fm.geometry} not implemented'
+
+    assert fm.shape[1] == 1, 'Cylindrical symmetry required'
+
+    dat = {}
+    z = fm.coord_vec('z')
+
+    if fm.is_static:
+        if fm.is_pure_magnetic:
+            fz = np.real(fm['Bz'][0,0,:])
+        elif fm.is_pure_electric:
+            fz = np.real(fm['Ez'][0,0,:])           
+        else:
+            raise ValueError('Mixed static field TODO')
+
+    else:
+        # Assume RF cavity, rotate onto the real axis
+        # Get complex fields
+        fz = fm.Ez[0,0,:]
+
+        # Get imaginary argument (angle)
+        iangle = np.unique(np.mod(np.angle(fz), np.pi))
+
+        # Make sure there is only one
+        assert len(iangle) == 1, print(iangle)
+        iangle = iangle[0]
+
+        if iangle != 0:
+            rot = np.exp(-1j*iangle)
+            # print(f'Rotating complex field by {-iangle}')
+        else:
+            rot = 1   
+        fz = np.real(fz*rot)
+
+    return z, fz  
+
+
+
+
 def vec_spacing(vec):
     """
     Returns the spacing and minimum of a coordinate.

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)