Accept old Superfish T7 format, ignoring newlines in data block

pmd_beamphysics/fields/fieldmesh.py

@@ -10,6 +10,7 @@
 
 from pmd_beamphysics.interfaces.superfish import write_superfish_t7, read_superfish_t7
 from pmd_beamphysics.interfaces.gpt import write_gpt_fieldmesh
+from pmd_beamphysics.interfaces.astra import read_astra_3d_fieldmaps, write_astra_3d_fieldmaps
 
 from h5py import File
 import numpy as np
@@ -25,7 +26,7 @@
 
 
 axis_labels_from_geometry = {
-    'cartesian': ('x', 'y', 'z'),
+    'rectangular': ('x', 'y', 'z'),
     'cylindrical': ('r', 'theta', 'z')
 }
 
@@ -84,10 +85,12 @@ class FieldMesh:
     
     Writers
         .write
+        .write_astra_3d
         .write_gpt
         .write_superfish
         
     Readers (class methods):
+        .from_astra_3d
         .from_superfish
     
 
@@ -111,7 +114,6 @@ def __init__(self, h5=None, data=None):
             else:
                 pass
         else:
-            print('loading data')
             data = load_field_data_dict(data)
 
         # Internal data
@@ -318,7 +320,9 @@ def write(self, h5, name=None):
 
         write_pmd_field(g, self.data, name=name)   
 
-
+    def write_astra_3d(self, common_filePath, verbose=False):      
+        return  write_astra_3d_fieldmaps(self, common_filePath)
+
     def write_gpt(self, filePath, asci2gdf_bin=None, verbose=True):
         """
         Writes a GPT field file. 
@@ -349,6 +353,16 @@ def from_superfish(cls, filename, type='electric', geometry='cylindrical'):
         return c               
 
 
+    @classmethod
+    def from_astra_3d(cls, common_filename, frequency=0):
+        """
+        Class method to parse multiple 3D astra fieldmap files,
+        based on the common filename.
+        """
+
+        data = read_astra_3d_fieldmaps(common_filename, frequency=frequency)
+        return cls(data=data)
+
 
     def __eq__(self, other):
         """

pmd_beamphysics/interfaces/superfish.py

@@ -225,19 +225,31 @@ def read_superfish_t7(t7file,
 
     # ASCII parsing
 
-    # Read header
+    # Read header and data. 
     # zmin(cm), zmax(cm), nx-1
     # freq(MHz)
     # ymin(cm), ymax(cm), ny-1
     with open(t7file, 'r') as f:
         line1 = f.readline().split()
         line2 = f.readline().split()
-        line3 = f.readline().split()
-
+        # Read all lines and flatten the data. This accepts an old superfish format
+        dat = f.read().replace('\n' , ' ').replace('\t', ' ').split()
+        dat = np.loadtxt(dat)
+
+    # The length of the second line gives a clue about the problem type
     if len(line2) == 1:
         problem = 'fish'
+        line3 = dat[0:3] # Final header line
+        dat = dat[3:] # body data
+        n = len(dat)
+        assert n % 4 == 0, f'{n} should be divisible by 4'
+        dat = dat.reshape(n//4, 4)
+
     else:
         problem = 'poisson'
+        n = len(dat)
+        assert n % 2 == 0, f'{n} should be divisible by 2'
+        dat = dat.reshape(len(dat)//2, 2)
 
     components = {}
     if problem=='fish':
@@ -252,7 +264,7 @@ def read_superfish_t7(t7file,
         rmin, rmax, nr =  float(line3[0])*1e-2, float(line3[1])*1e-2, int(line3[2])+1  
 
         # Read and reshape
-        dat = np.loadtxt(t7file, skiprows=3)
+        # dat = np.loadtxt(t7file, skiprows=3)
         #labels=['Ez', 'Er', 'E', 'Hphi']
         dat = dat.reshape(nr, 1, nz, 4)
 
@@ -270,7 +282,7 @@ def read_superfish_t7(t7file,
         frequency=0
 
         # The structure here is different 
-        dat = np.loadtxt(t7file, skiprows=2)
+        # dat = np.loadtxt(t7file, skiprows=2)
         dat = dat.reshape(nz, 1, nr, 2)
 
         # type must be specified

pmd_beamphysics/tools.py

@@ -1,11 +1,5 @@
 import numpy as np
 
-
-
-
-
-
-
 def fstr(s):
     """
     Makes a fixed string for h5 files