Improve read_IW2D_folder function

Allow to read both wake function and impedance with read_IW2D_folder
Fix yokoya_elliptic function

collective_effects/utilities.py

@@ -13,6 +13,8 @@ from pathlib import Path
 from scipy.interpolate import interp1d
 from scipy.constants import c
 from mbtrack2.collective_effects.wakefield import Impedance, WakeFunction, WakeField
+from pathlib import Path
+
 
 def read_CST(file, impedance_type='long', divide_by=None):
     """
@@ -91,7 +93,7 @@ def read_IW2D(file, file_type='Zlong'):
         raise ValueError("file_type should begin by Z or W.")
     return result
 
-def read_IW2D_folder(folder, suffix):
+def read_IW2D_folder(folder, suffix, select="WZ"):
     """
     Read IW2D results into a WakeField object.
     
@@ -99,6 +101,12 @@ def read_IW2D_folder(folder, suffix):
     ----------
     file : str
         Path to the file to read.
+    suffix : str
+        End of the name of each files. For example, in "Zlong_test.dat" the
+        suffix should be "_test.dat".
+    select : str, optional
+        Select which object to load. "W" for WakeFunction, "Z" for Impedance 
+        and "WZ" or "ZW" for both.
         
     Returns
     -------
@@ -106,8 +114,16 @@ def read_IW2D_folder(folder, suffix):
         WakeField object with Impedance and WakeFunction objects from the 
         different files.
     """
-    types = {"W" : WakeFunction,
-             "Z" : Impedance}
+    if (select == "WZ") or (select == "ZW"):
+        types = {"W" : WakeFunction,
+                 "Z" : Impedance}
+    elif (select == "W"):
+        types = {"W" : WakeFunction}
+    elif (select == "Z"):
+        types = {"Z" : Impedance}
+    else:
+        raise ValueError("select should be W, Z or WZ.")
+        
     components = ["long", "xdip", "ydip", "xquad", "yquad"]
     
     data_folder = Path(folder)
@@ -354,11 +370,13 @@ def yokoya_elliptic(x_radius , y_radius):
     else:
         small_semiaxis = x_radius
         large_semiaxis = y_radius
+        
+    path_to_file = Path(__file__).parent
+    file = path_to_file / "data" / "Yokoya_elliptic_from_Elias_USPAS.csv"
 
     # read Yokoya factors interpolation file
     # BEWARE: columns are ratio, dipy, dipx, quady, quadx
-    yokoya_file = pd.read_csv("collective_effects/data/" +
-                              "Yokoya_elliptic_from_Elias_USPAS.csv")
+    yokoya_file = pd.read_csv(file)
     ratio_col = yokoya_file["x"]
     # compute semi-axes ratio (first column of this file)
     ratio = (large_semiaxis - small_semiaxis)/(large_semiaxis + small_semiaxis)