diff --git a/collective_effects/utilities.py b/collective_effects/utilities.py index f6c6eec27bc87770dfa644d969bdc5dc6db9a9e8..5a57ca97c595bf88e4fb9d5386e9d86490dc1b09 100644 --- a/collective_effects/utilities.py +++ b/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)