Add optional output to read_impedance functions

Add optional output for read_IW2D, read_IW2D_folder and read_ABCI
For read_ABCI, skip the longxdip and longydip cases as they are not implemented yet

mbtrack2/utilities/read_impedance.py

@@ -66,7 +66,7 @@ def read_CST(file, component_type='long', divide_by=None, imp=True):
     return result
 
 
-def read_IW2D(file, file_type='Zlong'):
+def read_IW2D(file, file_type='Zlong', output=False):
     """
     Read IW2D file format into an Impedance object or a WakeField object.
 
@@ -76,6 +76,9 @@ def read_IW2D(file, file_type='Zlong'):
         Path to the file to read.
     file_type : str, optional
         Type of the Impedance or WakeField object.
+    output : bool, optional
+        If True, print out the interpolated values.
+        Default is False.
 
     Returns
     -------
@@ -99,8 +102,9 @@ def read_IW2D(file, file_type='Zlong'):
         if np.any(df.isna()):
             index = df.isna().values
             df = df.interpolate()
-            print("Nan values have been interpolated to:")
-            print(df[index])
+            if output:
+                print("Nan values have been interpolated to:")
+                print(df[index])
         result = WakeFunction(variable=df.index,
                               function=df["Wake"],
                               component_type=file_type[1:])
@@ -109,7 +113,7 @@ def read_IW2D(file, file_type='Zlong'):
     return result
 
 
-def read_IW2D_folder(folder, suffix, select="WZ"):
+def read_IW2D_folder(folder, suffix, select="WZ", output=False):
     """
     Read IW2D results into a WakeField object.
 
@@ -123,6 +127,9 @@ def read_IW2D_folder(folder, suffix, select="WZ"):
     select : str, optional
         Select which object to load. "W" for WakeFunction, "Z" for Impedance
         and "WZ" or "ZW" for both.
+    output : bool, optional
+        If True, print out the interpolated values.
+        Default is False.
 
     Returns
     -------
@@ -148,7 +155,7 @@ def read_IW2D_folder(folder, suffix, select="WZ"):
     for key, item in types.items():
         for component in components:
             name = data_folder / (key + component + suffix)
-            res = read_IW2D(file=name, file_type=key + component)
+            res = read_IW2D(file=name, file_type=key + component, output=output)
             list_for_wakefield.append(res)
 
     wake = WakeField(list_for_wakefield)
@@ -156,7 +163,7 @@ def read_IW2D_folder(folder, suffix, select="WZ"):
     return wake
 
 
-def read_ABCI(file, azimuthal=False):
+def read_ABCI(file, azimuthal=False, output=False):
     """
     Read ABCI output files [1].
 
@@ -170,6 +177,9 @@ def read_ABCI(file, azimuthal=False):
         If False, it is loaded from the "TRANSVERSE" data. In that case, a -1
         factor is applied on the wake to agree with mbtrack2 sign convention.
         The default is False.
+    output : bool, optional
+        If True, print out the loaded components header.
+        Default is False.
 
     Returns
     -------
@@ -221,8 +231,7 @@ def read_ABCI(file, azimuthal=False):
                  f'  TITLE: {source} WAKE POTENTIAL               \n': 'Wxdip',
                  f'  TITLE: REAL PART OF {source} IMPEDANCE                        \n': 'Zxdip_re',
                  f'  TITLE: IMAGINARY PART OF {source} IMPEDANCE                   \n': 'Zxdip_im'}
-    impedance_type_specification_line = {
-        '   DATE:     TIME:      MROT:     NO. OF POINTS:\n'}
+
     wake_list = []
     start = True
     with open(file) as f:
@@ -238,7 +247,8 @@ def read_ABCI(file, azimuthal=False):
                 # read the header
                 header = [next(f) for _ in range(4)]
                 header.insert(0, line)
-            print(header)
+            if output:
+                print(header)
             # read the body until next TITLE field or end of file
             body = []
             while True:
@@ -263,8 +273,10 @@ def read_ABCI(file, azimuthal=False):
                         comp = _read_temp(tmp.name, abci_dict[header[0]])
                         wake_list.append(comp)
                     elif abci_dict[header[0]][1:] == "long" and impedance_type == 1:
-                        comp = _read_temp(tmp.name, abci_dict[header[0]]+'dip')
-                        wake_list.append(comp)
+                        pass
+                        # Wlongxdip & Wlongydip not implemented yet
+                        # comp = _read_temp(tmp.name, abci_dict[header[0]]+'dip')
+                        # wake_list.append(comp)
                     else:
                         comp_x = _read_temp(tmp.name, "Wxdip")
                         comp_y = _read_temp(tmp.name, "Wydip")
@@ -285,8 +297,10 @@ def read_ABCI(file, azimuthal=False):
                         comp = _read_temp(tmp1.name, "Zlong", tmp2.name)
                         wake_list.append(comp)
                     elif abci_dict[header[0]][1:-3] == "long" and impedance_type == 1:
-                        comp = _read_temp(tmp1.name, "Zlongdip", tmp2.name)
-                        wake_list.append(comp)
+                        pass
+                        # Zlongxdip & Zlongydip not implemented yet
+                        # comp = _read_temp(tmp1.name, "Zlongdip", tmp2.name)
+                        # wake_list.append(comp)
                     else:
                         comp_x = _read_temp(tmp1.name, "Zxdip", tmp2.name)
                         comp_y = _read_temp(tmp1.name, "Zydip", tmp2.name)