Adds 3d field map writing using FieldMesh

docs/examples/fields/field_conversion.ipynb

@@ -447,7 +447,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.13"
+   "version": "3.9.18"
   }
  },
  "nbformat": 4,

pmd_beamphysics/fields/fieldmesh.py

@@ -10,7 +10,7 @@
 
 from pmd_beamphysics.interfaces.ansys import read_ansys_ascii_3d_fields
 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.gpt import write_gpt_fieldmap
 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
 
@@ -407,7 +407,7 @@ def write_gpt(self, filePath, asci2gdf_bin=None, verbose=True):
         Writes a GPT field file. 
         """
 
-        return write_gpt_fieldmesh(self, filePath, asci2gdf_bin=asci2gdf_bin, verbose=verbose)
+        return write_gpt_fieldmap(self, filePath, asci2gdf_bin=asci2gdf_bin, verbose=verbose)
 
     # Superfish
     @functools.wraps(write_superfish_t7)

pmd_beamphysics/interfaces/gpt.py

@@ -87,10 +87,96 @@ def run_asci2gdf(outfile, asci2gdf_bin, verbose=False):
 
 
 
-def write_gpt_fieldmesh(fm,           
-               outfile,
-               asci2gdf_bin=None,
-               verbose=False): 
+def write_gpt_fieldmap(fm,           
+                       outfile,
+                       asci2gdf_bin=None,
+                       verbose=False): 
+    """
+    Writes a GPT fieldmap file from a FieldMesh object.
+    
+    Requires cylindrical geometry for now.
+    """
+
+    if fm.geometry == 'cylindrical' and fm.coord_vec('r')[0] == fm.coord_vec('r')[-1]:
+        return write_gpt_1d_fieldmap(fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose)
+
+    elif fm.geometry == 'cylindrical':
+        return write_gpt_2d_fieldmap(fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose)
+
+    elif fm.geometry == 'rectangular':
+        return write_gpt_3d_fieldmap(fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose)
+
+
+    else:
+        raise ValueError(f'Unknown geometry {fm.geometry}')
+
+
+def write_gpt_1d_fieldmap(fm,           
+                          outfile,
+                          asci2gdf_bin=None,
+                          verbose=False):
+
+    """
+    Writes a GPT fieldmap file from a FieldMesh object.
+    
+    Requires cylindrical geometry for now.
+    """
+
+    assert fm.geometry == 'cylindrical', f'Geometry: {fm.geometry} not implemented'
+
+    assert fm.shape[1] == 1, 'Cylindrical symmetry required'
+
+    assert fm.coord_vec('r')[0]==0, 'r[0] must equal 0'
+
+    dat = {}
+    dat['Z'] = fm.coord_vec('z')
+
+    keys = ['Z']
+    if fm.is_static:
+        if fm.is_pure_magnetic:
+            keys = ['Z', 'Bz']
+            dat['Bz'] = np.real(fm['Bz'][0,0,:])
+        elif fm.is_pure_electric:
+            keys = ['Z', 'Ez']
+            dat['Er'] = np.real(fm['Er'][0,0,:])
+            dat['Ez'] = np.real(fm['Ez'][0,0,:])            
+        else:
+            raise ValueError('Mixed static field TODO')
+
+    else:
+        # Use internal Superfish routine 
+        keys = ['Z', 'Ez']
+        _, dat['Ez'], _, _ = fish_complex_to_real_fields(fm, verbose=verbose)     
+        dat['Ez'] = dat['Ez'][0,0,:]
+
+
+    # Flatten dat     
+    gptdata = np.array([dat[k].flatten() for k in keys]).T            
+
+    # Write file. 
+    # Hack to delete final newline
+    # https://stackoverflow.com/questions/28492954/numpy-savetxt-stop-newline-on-final-line
+    with open(outfile, 'w') as fout:
+        NEWLINE_SIZE_IN_BYTES = 1 # 2 on Windows?
+        np.savetxt(fout, gptdata, header=' '.join(keys), comments='')
+        fout.seek(0, os.SEEK_END) # Go to the end of the file.
+        # Go backwards one byte from the end of the file.
+        fout.seek(fout.tell() - NEWLINE_SIZE_IN_BYTES, os.SEEK_SET)
+        fout.truncate() # Truncate the file to this point.    
+
+
+    if asci2gdf_bin:
+        run_asci2gdf(outfile, asci2gdf_bin, verbose=verbose)
+    elif verbose: 
+        print(f'ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}')        
+
+    return outfile    
+
+
+def write_gpt_2d_fieldmap(fm,           
+                          outfile,
+                          asci2gdf_bin=None,
+                          verbose=False): 
     """
     Writes a GPT fieldmap file from a FieldMesh object.
     
@@ -143,7 +229,68 @@ def write_gpt_fieldmesh(fm,
     elif verbose: 
         print(f'ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}')        
 
-    return outfile    
+    return outfile  
+
+
+def write_gpt_3d_fieldmap(fm,           
+                          outfile,
+                          asci2gdf_bin=None,
+                          verbose=False): 
+    """
+    Writes a 3D GPT fieldmap file from a FieldMesh object.
+    """
+
+    assert fm.geometry == 'rectangular', f'Geometry: {fm.geometry} not implemented'
+
+    dat = {}
+    dat['X'], dat['Y'], dat['Z'] = np.meshgrid(fm.coord_vec('x'), fm.coord_vec('y'), fm.coord_vec('z'), indexing='ij')
+
+    keys = ['X', 'Y', 'Z']
+    if fm.is_static:
+        if fm.is_pure_magnetic:
+            keys = keys + ['Bx', 'By', 'Bz']
+            dat['Bx'] = np.real(fm['Bx'])
+            dat['By'] = np.real(fm['By'])
+            dat['Bz'] = np.real(fm['Bz'])
+
+        elif fm.is_pure_electric:
+            keys = keys + ['Ex', 'Ey', 'Ez']
+            dat['Ex'] = np.real(fm['Ex'])
+            dat['Ey'] = np.real(fm['Ey'])     
+            dat['Ez'] = np.real(fm['Ez'])            
+        else:
+            raise ValueError('Mixed static field TODO')
+
+    else:
+        # Use internal Superfish routine 
+        raise ValueError('Complex 3D Fields not implement yet!')
+        #keys = ['X', 'Y', 'Z', + ['Ex', 'Ey', 'Ez']'Bx', 'By', 'Bz']
+        #dat['Er'], dat['Ez'], dat['Bphi'], _ = fish_complex_to_real_fields(fm, verbose=verbose)            
+
+
+    # Flatten dat     
+    gptdata = np.array([dat[k].flatten() for k in keys]).T            
+
+    # Write file. 
+    # Hack to delete final newline
+    # https://stackoverflow.com/questions/28492954/numpy-savetxt-stop-newline-on-final-line
+    with open(outfile, 'w') as fout:
+        NEWLINE_SIZE_IN_BYTES = 1 # 2 on Windows?
+        np.savetxt(fout, gptdata, header=' '.join(keys), comments='')
+        fout.seek(0, os.SEEK_END) # Go to the end of the file.
+        # Go backwards one byte from the end of the file.
+        fout.seek(fout.tell() - NEWLINE_SIZE_IN_BYTES, os.SEEK_SET)
+        fout.truncate() # Truncate the file to this point.    
+
+
+    if asci2gdf_bin:
+        run_asci2gdf(outfile, asci2gdf_bin, verbose=verbose)
+    elif verbose: 
+        print(f'ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}')        
+
+    return outfile  
+
+