Add FieldMesh.to_astra_1d and related functions

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

@@ -214,6 +214,85 @@ def vprint(*a, **k):
 
 
 
+
+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.