diff --git a/examples/000_resonator.py b/examples/000_resonator.py new file mode 100644 index 00000000..2c1c09ee --- /dev/null +++ b/examples/000_resonator.py @@ -0,0 +1,45 @@ +import numpy as np + +import xwakes.wit as wit + +from scipy.constants import c as clight + +import xtrack as xt +p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 1000)) +p.x[p.zeta > 0] += 1e-3 +p_ref = p.copy() + +res = wit.ComponentResonator( + r=1e8, q=1e7, f_r=1e9, + source_exponents=(1, 0), + test_exponents=(0, 0), + plane='x' +) + +import xfields as xf + + +wake = xf.Wakefield(components=[res], zeta_range=(-1, 1), + num_slices=100) + +xfwake = xf.Wakefield(components=[ + xf.ResonatorWake( + r_shunt=1e8, q_factor=1e7, frequency=1e9, + source_exponents=(1, 0), test_exponents=(0, 0), + kick='px')], + zeta_range=(-1, 1), num_slices=100) + + + +wake.track(p) +xfwake.track(p_ref) + + +import matplotlib.pyplot as plt +plt.close('all') +plt.plot(p.zeta, p.px, label='xwakes') +plt.plot(p_ref.zeta, p_ref.px, '--', label='xfields') + +plt.legend() + +plt.show() \ No newline at end of file diff --git a/examples/gr_resonator_sps.py b/examples/gr_resonator_sps.py new file mode 100644 index 00000000..89324274 --- /dev/null +++ b/examples/gr_resonator_sps.py @@ -0,0 +1,131 @@ +import numpy as np + +import xwakes.wit as wit + +from scipy.constants import c as clight + +import xtrack as xt +import xpart as xp +import xfields as xf + +import matplotlib.pyplot as plt +from scipy.signal import hilbert +from scipy.stats import linregress + +f_rev = 43347.648455970964 +fline = (9240*f_rev - (0.13)*f_rev) + +res = wit.ComponentResonator( + r=0.942e9*70, q=5e5, f_r=fline, + source_exponents=(1, 0), + test_exponents=(0, 0), + plane='x' +) + +# Simulation settings +n_turns = 2_048 + +circumference = 6911.5662 +#bucket_length_m = circumference / 35640 + +wake = xf.Wakefield(components=[res], zeta_range=(-1, 1), + num_slices=100, num_turns=100, circumference=circumference) + + + +one_turn_map = xt.LineSegmentMap( + length=circumference, betx=70., bety=80., + qx=20.13, qy=20.18, + longitudinal_mode='linear_fixed_qs', + dqx=2.41, dqy=2.41, + qs=0.017843254299369695, bets=731.27 +) + +# Generate line +line = xt.Line(elements=[one_turn_map, wake], + element_names=['one_turn_map', 'wake']) + + +line.particle_ref = xt.Particles(p0c=26e9) +line.build_tracker() + +# Generate particles +particles = xp.generate_matched_gaussian_bunch(line=line, + num_particles=100_000, total_intensity_particles=2.3e11, + nemitt_x=2e-6, nemitt_y=2e-6, sigma_z=0.075) + +# Apply a distortion to the bunch to trigger an instability +amplitude = 1e-3 +particles.x += amplitude +particles.y += amplitude + +flag_plot = True + +mean_x_xt = np.zeros(n_turns) +mean_y_xt = np.zeros(n_turns) + +plt.ion() + +fig1 = plt.figure(figsize=(6.4*1.7, 4.8)) +ax_x = fig1.add_subplot(121) +line1_x, = ax_x.plot(mean_x_xt, 'r-', label='average x-position') +line2_x, = ax_x.plot(mean_x_xt, 'm-', label='exponential fit') +ax_x.set_ylim(-3.5, -1) +ax_x.set_xlim(0, n_turns) +ax_y = fig1.add_subplot(122, sharex=ax_x) +line1_y, = ax_y.plot(mean_y_xt, 'b-', label='average y-position') +line2_y, = ax_y.plot(mean_y_xt, 'c-', label='exponential fit') +ax_y.set_ylim(-3.5, -1) +ax_y.set_xlim(0, n_turns) + +plt.xlabel('turn') +plt.ylabel('log10(average x-position)') +plt.legend() + +turns = np.linspace(0, n_turns - 1, n_turns) + + +for i_turn in range(n_turns): + line.track(particles, num_turns=1) + + mean_x_xt[i_turn] = np.mean(particles.x) + mean_y_xt[i_turn] = np.mean(particles.y) + + if i_turn % 50 == 0: + print(f'Turn: {i_turn}') + + if (i_turn % 50 == 0 and i_turn > 1) or i_turn == n_turns - 1: + i_fit_end = np.argmax(mean_x_xt) # i_turn + i_fit_start = int(i_fit_end * 0.9) + + # compute x instability growth rate + ampls_x_xt = np.abs(hilbert(mean_x_xt)) + fit_x_xt = linregress(turns[i_fit_start: i_fit_end], + np.log(ampls_x_xt[i_fit_start: i_fit_end])) + + # compute y instability growth rate + + ampls_y_xt = np.abs(hilbert(mean_y_xt)) + fit_y_xt = linregress(turns[i_fit_start: i_fit_end], + np.log(ampls_y_xt[i_fit_start: i_fit_end])) + + line1_x.set_xdata(turns[:i_turn]) + line1_x.set_ydata(np.log10(np.abs(mean_x_xt[:i_turn]))) + line2_x.set_xdata(turns[:i_turn]) + line2_x.set_ydata(np.log10(np.exp(fit_x_xt.intercept + + fit_x_xt.slope*turns[:i_turn]))) + + line1_y.set_xdata(turns[:i_turn]) + line1_y.set_ydata(np.log10(np.abs(mean_y_xt[:i_turn]))) + line2_y.set_xdata(turns[:i_turn]) + line2_y.set_ydata(np.log10(np.exp(fit_y_xt.intercept + + fit_y_xt.slope*turns[:i_turn]))) + print(f'xtrack h growth rate: {fit_x_xt.slope}') + print(f'xtrack v growth rate: {fit_y_xt.slope}') + + fig1.canvas.draw() + fig1.canvas.flush_events() + + out_folder = '.' + np.save(f'{out_folder}/mean_x.npy', mean_x_xt) + np.save(f'{out_folder}/mean_y.npy', mean_y_xt) diff --git a/setup.py b/setup.py index b8952221..22b008aa 100644 --- a/setup.py +++ b/setup.py @@ -44,6 +44,7 @@ 'numpy>=1.0', 'scipy', 'pyyaml', + 'pandas' ], extras_require={ 'tests': ['pytest'], diff --git a/tests/test_utilities.py b/tests/test_utilities.py index c32bbe20..29740e1f 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -4,13 +4,12 @@ create_resistive_wall_single_layer_approx_element, create_taper_RW_approx_component, create_taper_RW_approx_element) -from xwakes.wit.utilities import (_zlong_round_single_layer_approx, - _zdip_round_single_layer_approx) +from xwakes.wit.component import ComponentSingleLayerResistiveWall from test_common import relative_error from pywit.parameters import * from pywit.interface import (FlatIW2DInput, RoundIW2DInput, Sampling, component_names) -from xwakes.wit.interface import _IW2DInputBase +from xwakes.wit.interface_dataclasses import _IW2DInputBase from pywit.materials import layer_from_json_material_library, copper_at_temperature from typing import Dict @@ -450,15 +449,14 @@ def test_taper_RW_algorithm(freq, component_id, round_single_layer_input): for radius in radii: if component_id == 'zlong': - z_steps += _zlong_round_single_layer_approx(freq, + z_steps += ComponentSingleLayerResistiveWall._zlong_round_single_layer_approx(freq, round_single_layer_input.relativistic_gamma, round_single_layer_input.layers[0], radius, round_single_layer_input.length/float(npts)) else: - z_steps += _zdip_round_single_layer_approx(freq, + z_steps += ComponentSingleLayerResistiveWall._zdip_round_single_layer_approx(freq, round_single_layer_input.relativistic_gamma, round_single_layer_input.layers[0], radius, round_single_layer_input.length/float(npts)) npt.assert_allclose(comp_taper_trapz.impedance(freq), z_steps, rtol=1e-3) - diff --git a/xwakes/wit/__init__.py b/xwakes/wit/__init__.py index 94937423..3a263fb5 100644 --- a/xwakes/wit/__init__.py +++ b/xwakes/wit/__init__.py @@ -10,4 +10,9 @@ from . import plot from . import sacherer_formula from . import utilities -from . import utils \ No newline at end of file +from . import utils + +from .component import (Component, ComponentClassicThickWall, ComponentResonator, + ComponentTaperSingleLayerRestsistiveWall, + ComponentSingleLayerResistiveWall) +from .element import Element \ No newline at end of file diff --git a/xwakes/wit/component.py b/xwakes/wit/component.py index d91ea08f..6140303d 100644 --- a/xwakes/wit/component.py +++ b/xwakes/wit/component.py @@ -1,12 +1,21 @@ from __future__ import annotations +from .interface_dataclasses import Layer, FlatIW2DInput, RoundIW2DInput from .parameters import * from .utils import unique_sigfigs - from typing import Optional, Callable, Tuple, Union, List import numpy as np +from scipy.constants import c as c_light, mu_0 as mu0, epsilon_0 as eps0 +from numpy.typing import ArrayLike +from scipy import special as sp + +if hasattr(np, 'trapezoid'): + trapz = np.trapezoid # numpy 2.0 +else: + trapz = np.trapz +Z0 = mu0 * c_light def mix_fine_and_rough_sampling(start: float, stop: float, rough_points: int, fine_points: int, rois: List[Tuple[float, float]]): @@ -59,8 +68,8 @@ def __init__(self, impedance: Optional[Callable] = None, wake: Optional[Callable # Enforces that either impedance or wake is defined. assert impedance or wake, "The impedance- and wake functions cannot both be undefined." # The impedance- and wake functions as callable objects, e.g lambda functions - self.impedance = impedance - self.wake = wake + self._impedance = impedance + self._wake = wake self.name = name # The plane of the Component, either 'x', 'y' or 'z' @@ -76,8 +85,46 @@ def __init__(self, impedance: Optional[Callable] = None, wake: Optional[Callable self.test_exponents = test_exponents self.power_x = (source_exponents[0] + test_exponents[0] + (plane == 'x')) / 2 self.power_y = (source_exponents[1] + test_exponents[1] + (plane == 'y')) / 2 - self.f_rois = f_rois if f_rois else [] - self.t_rois = t_rois if t_rois else [] + self._f_rois = f_rois if f_rois else [] + self._t_rois = t_rois if t_rois else [] + + @property + def impedance(self) -> Optional[Callable]: + return self._impedance + + @property + def wake(self) -> Optional[Callable]: + return self._wake + + @property + def f_rois(self) -> List[Tuple[float, float]]: + return self._f_rois + + @property + def t_rois(self) -> List[Tuple[float, float]]: + return self._t_rois + + def function_vs_t(self, t, beta0): + out = -self.wake(-t) + return out + + def function_vs_zeta(self, zeta, beta0): + out = -self.wake(-zeta / beta0 / c_light) + return out + + @property + def kick(self) -> str: + return {'x': 'px', 'y': 'py', 'z': 'delta'}[self.plane] + + @property + def source_moments(self) -> List[str]: + out = ['num_particles'] + if self.source_exponents[0] != 0: + out.append('x') + if self.source_exponents[1] != 0: + out.append('y') + + return out def generate_wake_from_impedance(self) -> None: """ @@ -385,3 +432,732 @@ def get_shorthand_type(self) -> str: :return: A 5-character string indicating the plane as well as source- and test exponents of the component """ return self.plane + "".join(str(x) for x in (self.source_exponents + self.test_exponents)) + + +class ComponentResonator(Component): + def __init__(self, plane: str, + exponents: Tuple[int, int, int, int]| None = None, + source_exponents: Tuple[int, int] | None = None, + test_exponents: Tuple[int, int] | None = None, + r: float = None, q: float = None, f_r: float = None, + f_roi_level: float = 0.5) -> ComponentResonator: + """ + Creates a resonator component with the given parameters + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, + test_x, test_y) aka (a, b, c, d) + :param r: the shunt impedance of the given component of the resonator + :param q: the quality factor of the given component of the resonator + :param f_r: the resonance frequency of the given component of the + resonator + :param f_roi_level: fraction of the peak ok the resonator which is + covered by the ROI. I.e. the roi will cover + the frequencies for which the resonator impedance is larger than + f_roi_level*r + :return: A component object of a resonator, specified by the input + arguments + """ + self.r = r + self.q = q + self.f_r = f_r + self.f_roi_level = f_roi_level + + source_exponents, test_exponents = _handle_exponents_input( + exponents, source_exponents, test_exponents) + + # we set impedance and wake to a dummy callable because they will be + # overridden by methods + super().__init__(impedance=lambda x: 0, wake=lambda x: 0, plane=plane, + source_exponents=source_exponents, + test_exponents=test_exponents, + name="Resonator") + + def impedance(self, f): + f = np.atleast_1d(f) + f_r = self.f_r + q = self.q + r = self.r + plane = self.plane + if plane == 'z': + out = r / (1 - 1j * q * (f_r / f - f / f_r)) + else: + out = (f_r * r) / (f * (1 - 1j * q * (f_r / f - f / f_r))) + + return out + + + def wake(self, t): + t = np.atleast_1d(t) + f_r = self.f_r + q = self.q + r = self.r + plane = self.plane + omega_r = 2 * np.pi * f_r + root_term = np.sqrt(1 - 1 / (4 * q ** 2) + 0J) + if plane == 'z': + omega_bar = omega_r * root_term + alpha = omega_r / (2 * q) + out = np.zeros_like(t) + mask = t > 0 + out[mask] = (omega_r * r * np.exp(-alpha * t[mask]) * ( + np.cos(omega_bar * t[mask]) - + alpha * np.sin(omega_bar * t[mask]) / omega_bar) / q).real + else: + omega_bar = omega_r * root_term + out = np.zeros_like(t) + mask = t > 0 + out[mask] = (omega_r * r * np.exp(-omega_r * t[mask] / (2 * q)) * + np.sin(omega_r * root_term * t[mask]) / + (q * root_term)).real + return out + + @property + def f_rois(self): + if self.q > 1: + d = self._compute_resonator_f_roi_half_width( + q=self.q, f_r=self.f_r, f_roi_level=self.f_roi_level) + return [(self.f_r - d, self.f_r + d)] + else: + return [] + + @property + def t_rois(self): + return [] + + @staticmethod + def _compute_resonator_f_roi_half_width(q: float, f_r: float, + f_roi_level: float = 0.5): + aux = np.sqrt((1 - f_roi_level) / f_roi_level) + + return (aux + np.sqrt(aux**2 + 4*q**2))*f_r/(2*q) - f_r + + +class ComponentClassicThickWall(Component): + def __init__(self, plane: str, + exponents: Tuple[int, int, int, int] | None = None, + source_exponents: Tuple[int, int] | None = None, + test_exponents: Tuple[int, int] | None = None, + layer: Layer = None, radius: float = None) -> ComponentClassicThickWall: + """ + Creates a single component object modeling a resistive wall + impedance/wake, based on the "classic thick wall formula" (see e.g. + A. W. Chao, chap. 2 in "Physics of Collective Beams Instabilities in + High Energy Accelerators", John Wiley and Sons, 1993). + Only longitudinal and transverse dipolar impedances are supported here. + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, + test_x, test_y) aka (a, b, c, d) + :param layer: the chamber material, as a wit Layer object + :param radius: the chamber radius in m + :return: A component object + """ + self.layer = layer + self.radius = radius + self.plane = plane + + source_exponents, test_exponents = _handle_exponents_input( + exponents, source_exponents, test_exponents) + + # we set impedance and wake to a dummy callable because they will be + # overridden by methods + super().__init__(impedance=lambda x: 0, wake=lambda x: 0, plane=plane, + source_exponents=source_exponents, + test_exponents=test_exponents, + name="Classic Thick Wall") + + def impedance(self, f): + layer = self.layer + material_resistivity = layer.dc_resistivity + material_relative_permeability = 1. + layer.magnetic_susceptibility + material_permeability = material_relative_permeability * mu0 + radius = self.radius + plane = self.plane + exponents = self.exponents + + if plane == 'z' and exponents == (0, 0, 0, 0): + out = ((1/2) * + (1 + np.sign(f)*1j) * + material_resistivity / (np.pi * radius) * + (1 / self.delta_skin(f, material_resistivity, + material_permeability))) + # Transverse dipolar impedance + elif ((plane == 'x' and exponents == (1, 0, 0, 0)) or + (plane == 'y' and exponents == (0, 1, 0, 0))): + out = ((c_light/(2*np.pi*f)) * (1+np.sign(f)*1j) * + material_resistivity / (np.pi * radius**3) * + (1 / self.delta_skin(f, material_resistivity, + material_permeability))) + else: + print("Warning: resistive wall impedance not implemented for " + "component {}{}. Set to zero".format(plane, exponents)) + out = np.zeros_like(f) + + return out + + def wake(self, t): + layer = self.layer + material_resistivity = layer.dc_resistivity + radius = self.radius + plane = self.plane + exponents = self.exponents + + # Longitudinal impedance + if plane == 'z' and exponents == (0, 0, 0, 0): + out = (-c_light / (2*np.pi*radius) * + (Z0 * material_resistivity/np.pi)**(1/2) * + 1/(t**(1/2))) + # Transverse dipolar impedance + elif ((plane == 'x' and exponents == (1, 0, 0, 0)) or + (plane == 'y' and exponents == (0, 1, 0, 0))): + out = (-c_light / (2*np.pi*radius**3) * + (Z0 * material_resistivity/np.pi)**(1/2) * + 1/(t**(3/2))) + else: + print("Warning: resistive wall wake not implemented for " + "component {}{}. Set to zero".format(plane, exponents)) + out = np.zeros_like(t) + + return out + + @staticmethod + def delta_skin(f, material_resistivity, material_permeability): + return (material_resistivity / (2*np.pi*abs(f) * + material_permeability)) ** (1/2) + + @property + def f_rois(self): + return [] + + @property + def t_rois(self): + return [] + + +class ComponentSingleLayerResistiveWall(Component): + def __init__(self, plane: str, + exponents: Tuple[int, int, int, int] = None, + source_exponents: Tuple[int, int] | None = None, + test_exponents: Tuple[int, int] | None = None, + input_data: Union[FlatIW2DInput, RoundIW2DInput] = None): + """ + Creates a single component object modeling a resistive wall impedance, + based on the single-layer approximated formulas by E. Metral (see e.g. + Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, + https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and + Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, + 2009, https://doi.org/10.1103/PhysRevSTAB.12.084401) + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, + test_x, test_y) aka (a, b, c, d) + :param input_data: an IW2D input object (flat or round). If the input + is of type FlatIW2DInput and symmetric, we apply to the round formula + the Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, + KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, + https://cds.cern.ch/record/248630/files/p221.pdf), + while for a single plate we use those from A. Burov and V. Danilov, + PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other + kinds of asymmetric structure will raise an error. + If the input is of type RoundIW2DInput, the structure is in principle + round but the Yokoya factors put in the input will be used. + :return: A component object + """ + self.input_data = input_data + self.plane = plane + + source_exponents, test_exponents = _handle_exponents_input( + exponents, source_exponents, test_exponents) + + if isinstance(input_data, FlatIW2DInput): + if len(input_data.top_layers) > 1: + raise NotImplementedError("Input data can have only one layer") + self.yok_long = 1. + self.layer = input_data.top_layers[0] + self.radius = input_data.top_half_gap + if input_data.top_bottom_symmetry: + self.yok_dipx = np.pi**2/24. + self.yok_dipy = np.pi**2/12. + self.yok_quax = -np.pi**2/24. + self.yok_quay = np.pi**2/24. + elif input_data.bottom_half_gap == np.inf: + self.yok_dipx = 0.25 + self.yok_dipy = 0.25 + self.yok_quax = -0.25 + self.yok_quay = 0.25 + else: + raise NotImplementedError("For asymmetric structures, only the " + "case of a single plate is " + "implemented; hence the bottom " + "half-gap must be infinite") + elif isinstance(input_data, RoundIW2DInput): + self.radius = input_data.inner_layer_radius + if len(input_data.layers) > 1: + raise NotImplementedError("Input data can have only one layer") + self.layer = input_data.layers[0] + self.yok_long = input_data.yokoya_factors[0] + self.yok_dipx = input_data.yokoya_factors[1] + self.yok_dipy = input_data.yokoya_factors[2] + self.yok_quax = input_data.yokoya_factors[3] + self.yok_quay = input_data.yokoya_factors[4] + else: + raise NotImplementedError("Input of type neither FlatIW2DInput nor " + "RoundIW2DInput cannot be handled") + + + # we set impedance and wake to a dummy callable because they will be + # overridden by methods + super().__init__(impedance=lambda x: 0, wake=lambda x: 0, plane=plane, + source_exponents=source_exponents, + test_exponents=test_exponents, + name="Single layer resistive wall") + + def impedance(self, f): + # Longitudinal impedance + if (self.plane == 'z' and self.source_exponents == (0, 0) + and self.test_exponents == (0, 0)): + out = self.yok_long*self._zlong_round_single_layer_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius=self.radius, + length=self.input_data.length) + # Transverse impedances + elif (self.plane == 'x' and self.source_exponents == (1, 0) and + self.test_exponents == (0, 0)): + out = self.yok_dipx * self._zdip_round_single_layer_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius=self.radius, + length=self.input_data.length) + elif (self.plane == 'y' and self.source_exponents == (0, 1) and + self.test_exponents == (0, 0)): + out = self.yok_dipy * self._zdip_round_single_layer_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius=self.radius, + length=self.input_data.length) + elif (self.plane == 'x' and self.source_exponents == (0, 0) + and self.test_exponents == (1, 0)): + out = self.yok_quax * self._zdip_round_single_layer_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius=self.radius, + length=self.input_data.length) + elif (self.plane == 'y' and self.source_exponents == (0, 0) and + self.test_exponents == (0, 1)): + out = self.yok_quay * self._zdip_round_single_layer_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius=self.radius, + length=self.input_data.length) + else: + out = np.zeros_like(f) + + return out + + def wake(self, t): + raise NotImplementedError("Wake not implemented for single-layer " + "resistive wall impedance") + + @staticmethod + def _zlong_round_single_layer_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius: float, + length: float) -> ArrayLike: + """ + Function to compute the longitudinal resistive-wall impedance from + the single-layer, approximated formula for a cylindrical structure, + by E. Metral (see e.g. Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, + TUPD053, https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and + Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, + 2009, https://doi.org/10.1103/PhysRevSTAB.12.084401) + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius: the radius of the structure (in m) + :param length: the total length of the resistive object (in m) + :return: Zlong, the longitudinal impedance at these frequencies + """ + beta = np.sqrt(1 - 1 / gamma**2) + omega = 2 * np.pi * frequencies + k = omega / (beta * c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1 + layer.magnetic_susceptibility + eps1 = 1 - 1j / (eps0 * rho * omega * (1 + 1j * omega * tau)) + nu = k * np.sqrt(1 - beta**2 * eps1 * mu1) + + coef_long = 1j * omega * mu0 * length / (2 * np.pi * beta**2 * gamma**2) + + x1 = k * radius / gamma + x1sq = x1**2 + x2 = nu * radius + + zlong = coef_long * (sp.k0(x1) / sp.i0(x1) - + 1 / (x1sq * (1/2 + eps1 * sp.kve(1, x2)/ + (x2 * sp.kve(0, x2))))) + + return zlong + + @staticmethod + def _zdip_round_single_layer_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius: float, + length: float) -> ArrayLike: + """ + Function to compute the transverse dipolar resistive-wall impedance from + the single-layer, approximated formula for a cylindrical structure, + Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, + https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and + Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, + 2009, https://doi.org/10.1103/PhysRevSTAB.12.084401) + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius: the radius of the structure (in m) + :param length: the total length of the resistive object (in m) + :return: Zdip, the transverse dipolar impedance at these frequencies + """ + beta = np.sqrt(1 - 1 / gamma**2) + omega = 2 * np.pi * frequencies + k = omega / (beta * c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1 + layer.magnetic_susceptibility + eps1 = 1 - 1j/(eps0 * rho * omega*(1 + 1j * omega * tau)) + nu = k * np.sqrt(1 - beta**2 * eps1 * mu1) + + coef_dip = 1j * k**2 * Z0 * length/(4. * np.pi * beta * gamma**4) + + x1 = k * radius / gamma + x1sq = x1**2 + x2 = nu * radius + + zdip = coef_dip * (sp.k1(x1) / sp.i1(x1) + + 4 * beta**2 * gamma**2 / + (x1sq*(2 + x2 * sp.kve(0, x2) / + (mu1 * sp.kve(1, x2))))) + + return zdip + + +class ComponentTaperSingleLayerRestsistiveWall(Component): + def __init__(self, plane: str, + exponents: Tuple[int, int, int, int] = None, + source_exponents: Tuple[int, int] | None = None, + test_exponents: Tuple[int, int] | None = None, + input_data: Union[FlatIW2DInput, RoundIW2DInput] = None, + radius_small: float = None, radius_large: float = None, + step_size: float = 1e-3): + """ + Creates a single component object modeling a round or flat taper (flatness + along the horizontal direction, change of half-gap along the vertical one) + resistive-wall impedance, using the integration of the radius-dependent + approximated formula for a cylindrical structure (see + the above functions), over the length of the taper. + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param input_data: an IW2D input object (flat or round). If the input + is of type FlatIW2DInput and symmetric, we apply to the round formula the + Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, + KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, + https://cds.cern.ch/record/248630/files/p221.pdf), + while for a single plate we use those from A. Burov and V. Danilov, + PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other + kinds of asymmetric structure will raise an error. + If the input is of type RoundIW2DInput, the structure is in principle + round but the Yokoya factors put in the input will be used. + Note that the radius or half-gaps in input_data are not used (replaced + by the scan from radius_small to radius_large, for the integration). + :param radius_small: the smallest radius of the taper (in m) + :param radius_large: the largest radius of the taper (in m) + :param step_size: the step size (in the radial or vertical direction) + for the integration (in m) + :return: A component object + """ + self.input_data = input_data + self.radius_large = radius_large + self.radius_small = radius_small + self.step_size = step_size + self.plane = plane + + source_exponents, test_exponents = _handle_exponents_input( + exponents, source_exponents, test_exponents) + + if isinstance(input_data, FlatIW2DInput): + if len(input_data.top_layers) > 1: + raise NotImplementedError("Input data can have only one layer") + self.yok_long = 1. + self.layer = input_data.top_layers[0] + self.radius = input_data.top_half_gap + if input_data.top_bottom_symmetry: + self.yok_dipx = np.pi**2/24. + self.yok_dipy = np.pi**2/12. + self.yok_quax = -np.pi**2/24. + self.yok_quay = np.pi**2/24. + elif input_data.bottom_half_gap == np.inf: + self.yok_dipx = 0.25 + self.yok_dipy = 0.25 + self.yok_quax = -0.25 + self.yok_quay = 0.25 + else: + raise NotImplementedError("For asymmetric structures, only the case of a single plate is implemented; " + "hence the bottom half gap must be infinite") + elif isinstance(input_data, RoundIW2DInput): + self.radius = input_data.inner_layer_radius + if len(input_data.layers) > 1: + raise NotImplementedError("Input data can have only one layer") + self.layer = input_data.layers[0] + self.yok_long = input_data.yokoya_factors[0] + self.yok_dipx = input_data.yokoya_factors[1] + self.yok_dipy = input_data.yokoya_factors[2] + self.yok_quax = input_data.yokoya_factors[3] + self.yok_quay = input_data.yokoya_factors[4] + else: + raise NotImplementedError("Input of type neither FlatIW2DInput nor RoundIW2DInput cannot be handled") + + # we set impedance and wake to a dummy callable because they will be + # overridden by methods + super().__init__(impedance=lambda x: 0, wake=lambda x: 0, plane=plane, + source_exponents=source_exponents, + test_exponents=test_exponents, + name="Taper single layer resistive wall") + + def impedance(self, f): + # Longitudinal impedance + if (self.plane == 'z' and self.source_exponents == (0, 0) + and self.test_exponents == (0, 0)): + out = self.yok_long*self._zlong_round_taper_RW_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius_small=self.radius_small, + radius_large=self.radius_large, + length=self.input_data.length, + step_size=self.step_size) + # Transverse impedances + elif (self.plane == 'x' and self.source_exponents == (1, 0) + and self.test_exponents == (0, 0)): + out = self.yok_dipx*self._zdip_round_taper_RW_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius_small=self.radius_small, + radius_large=self.radius_large, + length=self.input_data.length, + step_size=self.step_size) + elif (self.plane == 'y' and self.source_exponents == (0, 1) + and self.test_exponents == (0, 0)): + out = self.yok_dipy*self._zdip_round_taper_RW_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius_small=self.radius_small, + radius_large=self.radius_large, + length=self.input_data.length, + step_size=self.step_size) + elif (self.plane == 'x' and self.source_exponents == (0, 0) + and self.test_exponents == (1, 0)): + out = self.yok_quax*self._zdip_round_taper_RW_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius_small=self.radius_small, + radius_large=self.radius_large, + length=self.input_data.length, + step_size=self.step_size) + elif (self.plane == 'y' and self.source_exponents == (0, 0) + and self.test_exponents == (0, 1)): + out = self.yok_quay*self._zdip_round_taper_RW_approx( + frequencies=f, + gamma=self.input_data.relativistic_gamma, + layer=self.layer, + radius_small=self.radius_small, + radius_large=self.radius_large, + length=self.input_data.length, + step_size=self.step_size) + else: + out = np.zeros_like(f) + + return out + + + def wake(self, t): + raise NotImplementedError("Wake not implemented for single-layer " + "resistive wall impedance") + + + @staticmethod + def _zlong_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius_small: float, + radius_large: float, length: float, + step_size: float = 1e-3) -> ArrayLike: + """ + Function to compute the longitudinal resistive-wall impedance for a + round taper, integrating the radius-dependent approximated formula + for a cylindrical structure (see_zlong_round_single_layer_approx above), + over the length of the taper. + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius_small: the smallest radius of the taper (in m) + :param radius_large: the largest radius of the taper (in m) + :param length: the total length of the taper (in m) + :param step_size: the step size (in the radial direction) for the + integration (in m) + :return: the longitudinal impedance at these frequencies + """ + if np.isscalar(frequencies): + frequencies = np.array(frequencies) + beta = np.sqrt(1.-1./gamma**2) + omega = 2*np.pi*frequencies.reshape((-1, 1)) + k = omega/(beta*c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1.+layer.magnetic_susceptibility + eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) + nu = k*np.sqrt(1.-beta**2*eps1*mu1) + + coef_long = 1j*omega*mu0/(2.*np.pi*beta**2*gamma**2) + + npts = int(np.floor(abs(radius_large-radius_small)/step_size)+1) + radii = np.linspace(radius_small, radius_large, npts).reshape((1, -1)) + one_array = np.ones(radii.shape) + + x1 = k.dot(radii)/gamma + x1sq = x1**2 + x2 = nu.dot(radii) + zlong = (coef_long.dot(length / float(npts) * one_array) * + (sp.k0(x1) / sp.i0(x1) - + 1. / (x1sq * (1. / 2. + eps1.dot(one_array) * + sp.kve(1, x2) / (x2 * sp.kve(0, x2))))) + ) + + return trapz(zlong, axis=1) + + + @staticmethod + def _zdip_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius_small: float, + radius_large: float, length: float, + step_size: float = 1e-3) -> ArrayLike: + """ + Function to compute the transverse dip. resistive-wall impedance for a + round taper, integrating the radius-dependent approximated formula + for a cylindrical structure (see_zdip_round_single_layer_approx above), + over the length of the taper. + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius_small: the smallest radius of the taper (in m) + :param radius_large: the largest radius of the taper (in m) + :param length: the total length of the taper (in m) + :param step_size: the step size (in the radial direction) for the + integration (in m) + :return: the transverse dipolar impedance at these frequencies + """ + if np.isscalar(frequencies): + frequencies = np.array(frequencies) + beta = np.sqrt(1.-1./gamma**2) + omega = 2*np.pi*frequencies.reshape((-1,1)) + k = omega/(beta*c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1.+layer.magnetic_susceptibility + eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) + nu = k*np.sqrt(1.-beta**2*eps1*mu1) + + coef_dip = 1j*k**2*Z0/(4.*np.pi*beta*gamma**4) + + npts = int(np.floor(abs(radius_large-radius_small)/step_size)+1) + radii = np.linspace(radius_small,radius_large,npts).reshape((1,-1)) + one_array = np.ones(radii.shape) + + x1 = k.dot(radii)/gamma + x1sq = x1**2 + x2 = nu.dot(radii) + zdip = ( + coef_dip.dot(length / float(npts) * one_array) * + (sp.k1(x1) / sp.i1(x1) + + 4 * beta**2 * gamma**2 / (x1sq * (2 + x2 * sp.kve(0, x2) / + (mu1 * sp.kve(1, x2))))) + ) + + return trapz(zdip, axis=1) + + +class ComponentInterpolated(Component): + def __init__(self, interpolation_frequencies: ArrayLike = None, + impedance_input: Optional[Callable] = None, + interpolation_times: ArrayLike = None, + wake_input: Optional[Callable] = None, plane: str = '', + source_exponents: Tuple[int, int] = (-1, -1), + test_exponents: Tuple[int, int] = (-1, -1), + name: str = "Interpolated Component", + f_rois: Optional[List[Tuple[float, float]]] = None, + t_rois: Optional[List[Tuple[float, float]]] = None): + + assert ((interpolation_frequencies is not None) == + (impedance_input is not None)), ("Either both or none of the " + "impedance and the interpolation frequencies must be given") + + self.interpolation_frequencies = interpolation_frequencies + self.impedance_input = impedance_input + + assert ((interpolation_times is not None) == + (wake_input is not None)), ("Either both or none of the wake " + "and the interpolation times must be given") + + self.interpolation_times = interpolation_times + self.wake_input = wake_input + + super().__init__(impedance=lambda x: 0, wake=lambda x: 0, plane=plane, + source_exponents=source_exponents, + test_exponents=test_exponents, + f_rois=f_rois, t_rois=t_rois, + name=name) + + def impedance(self, f): + if self.impedance_input is not None: + return np.interp(f, self.interpolation_frequencies, + self.impedance_input) + else: + return np.zeros_like(f) + + def wake(self, t): + if self.wake_input is not None: + return np.interp(t, self.interpolation_times, self.wake_input) + else: + return np.zeros_like(t) + + +def _handle_exponents_input(exponents, source_exponents, test_exponents): + + if exponents is not None: + assert source_exponents is None and test_exponents is None, ( + "If exponents is specified, source_exponents and test_exponents " + "should not be specified.") + source_exponents = exponents[0:2] + test_exponents = exponents[2:4] + elif source_exponents is not None or test_exponents is not None: + assert source_exponents is not None and test_exponents is not None, ( + "If source_exponents or test_exponents is specified, both should " + "be specified.") + assert exponents is None, ( + "If source_exponents or test_exponents is specified, exponents " + "should not be specified.") + + return source_exponents, test_exponents diff --git a/xwakes/wit/element.py b/xwakes/wit/element.py index f35309fc..96a9f942 100644 --- a/xwakes/wit/element.py +++ b/xwakes/wit/element.py @@ -4,9 +4,10 @@ from typing import List from collections import defaultdict - +from scipy.interpolate import interp1d from scipy.special import comb import numpy as np +import pandas as pd import copy class Element: @@ -308,3 +309,109 @@ def get_component(self, type_string: str): return comp raise KeyError(f"'{self.name}' has no component of the type '{type_string}'.") + + +class ElementFromTable(Element): + components_dict = {'longitudinal': ('z', (0, 0, 0, 0)), + 'constant_x': ('x', (0, 0, 0, 0)), + 'constant_y': ('y', (0, 0, 0, 0)), + 'dipole_x': ('x', (1, 0, 0, 0)), + 'dipole_y': ('y', (0, 1, 0, 0)), + 'dipole_xy': ('x', (0, 1, 0, 0)), + 'dipole_yx': ('y', (1, 0, 0, 0)), + 'quadrupole_x': ('x', (0, 0, 1, 0)), + 'quadrupole_y': ('y', (0, 0, 0, 1)), + 'quadrupole_xy': ('x', (0, 0, 0, 1)), + 'quadrupole_yx': ('y', (0, 0, 1, 0)), + } + ''' + An element whose components are defined point-by-point in a wake table + and/or an impedance table. The tables are given as pandas dataframes and + they must specify the time and frequency columns respectively and all the + components specified in the use_components list. + ''' + def __init__(self, wake_table: pd.dataframe = None, + impedance_table: pd.dataframe = None, + use_components: List['str'] = None, + length: float = 0, + beta_x: float = 0, beta_y: float = 0, + name: str = "Unnamed Element", + tag: str = "", description: str = ""): + ''' + Initializes an ElementFromTable object. + + Parameters + ---------- + wake_table : pd.dataframe + A pandas dataframe containing the wake function values at least for + all the components specified in use_components. The dataframe must + contain a column named 'time' which specifies the times at which the + components are evaluated. + impedance_table : pd.dataframe + A pandas dataframe containing the impedance function values at least + for all the components specified in use_components. The dataframe must + contain a column named 'frequency' which specifies the frequencies at + which the components are evaluated. + use_components : List[str] + A list of strings specifying the components to be used in the + ElementFromTable object. The following components are allowed: + longitudinal, constant_x, constant_y, dipole_x, dipole_y, + dipole_xy, dipole_yx, quadrupole_x, quadrupole_y, quadrupole_xy, + quadrupole_yx. + ''' + self.wake_table = wake_table + self.impedance_table = impedance_table + + for component in use_components: + if component not in self.component_names: + raise ValueError( + f"Invalid wake component: {component}. " + f"Valid wake components are: {self.component_names.keys()}") + + if 'time' not in list(wake_table.keys()): + raise ValueError("No wake_table column with name 'time'" + + "has been specified. \n") + + if 'frequency' not in list(impedance_table.keys()): + raise ValueError("No impedance_table column with name " + + "'frequency' has been specified. \n") + + components = [] + + for component_str in use_components: + + if wake_table is not None: + if not component_str in wake_table.keys(): + raise ValueError(f"{component_str} not in wake_table keys") + + wake_function = interp1d(wake_table['time'], + wake_table[component_str], + bounds_error=False, + fill_value=0.0) + else: + wake_function = lambda t: 0 + + if impedance_table is not None: + if not component_str in impedance_table.keys(): + raise ValueError(f"{component_str} not in impedance_table keys") + + impedance_function = interp1d(impedance_table['frequency'], + impedance_table[component_str], + bounds_error=False, + fill_value=0.0) + else: + impedance_function = lambda f: 0 + + plane, exponents = self.components_dict[component_str] + + source_exponents = exponents[:2] + test_exponents = exponents[2:] + + components.append(Component(impedance=impedance_function, + wake=wake_function, + plane=plane, + source_exponents=source_exponents, + test_exponents=test_exponents)) + + super().__init__(length, beta_x, beta_y, components, name, tag, + description) diff --git a/xwakes/wit/interface.py b/xwakes/wit/interface.py index 623c9556..b9ceafcb 100644 --- a/xwakes/wit/interface.py +++ b/xwakes/wit/interface.py @@ -2,6 +2,7 @@ from .component import Component from .element import Element +from .interface_dataclasses import Layer, Sampling, FlatIW2DInput, RoundIW2DInput, IW2DInput import subprocess from typing import Tuple, List, Optional, Dict, Any, Union @@ -148,96 +149,6 @@ def create_component_from_data(is_impedance: bool, plane: str, exponents: Tuple[ test_exponents=exponents[2:], ) -@dataclass(frozen=True, eq=True) -class Layer: - # The distance in mm of the inner surface of the layer from the reference orbit - thickness: float - dc_resistivity: float - resistivity_relaxation_time: float - re_dielectric_constant: float - magnetic_susceptibility: float - permeability_relaxation_frequency: float - - -@dataclass(frozen=True, eq=True) -class Sampling: - start: float - stop: float - # 0 = logarithmic, 1 = linear, 2 = both - scan_type: int - added: Tuple[float] - sampling_exponent: Optional[float] = None - points_per_decade: Optional[float] = None - min_refine: Optional[float] = None - max_refine: Optional[float] = None - n_refine: Optional[float] = None - - -# Define several dataclasses for IW2D input elements. We must split mandatory -# and optional arguments into private dataclasses to respect the resolution -# order. The public classes RoundIW2DInput and FlatIW2D input inherit from -# from the private classes. -# https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses - -@dataclass(frozen=True, eq=True) -class _IW2DInputBase: - machine: str - length: float - relativistic_gamma: float - calculate_wake: bool - f_params: Sampling - - -@dataclass(frozen=True, eq=True) -class _IW2DInputOptional: - z_params: Optional[Sampling] = None - long_factor: Optional[float] = None - wake_tol: Optional[float] = None - freq_lin_bisect: Optional[float] = None - comment: Optional[str] = None - - -@dataclass(frozen=True, eq=True) -class IW2DInput(_IW2DInputOptional, _IW2DInputBase): - pass - - -@dataclass(frozen=True, eq=True) -class _RoundIW2DInputBase(_IW2DInputBase): - layers: Tuple[Layer] - inner_layer_radius: float - # (long, xdip, ydip, xquad, yquad) - yokoya_factors: Tuple[float, float, float, float, float] - - -@dataclass(frozen=True, eq=True) -class _RoundIW2DInputOptional(_IW2DInputOptional): - pass - - -@dataclass(frozen=True, eq=True) -class RoundIW2DInput(_RoundIW2DInputOptional, _RoundIW2DInputBase): - pass - - -@dataclass(frozen=True, eq=True) -class _FlatIW2DInputBase(_IW2DInputBase): - top_bottom_symmetry: bool - top_layers: Tuple[Layer] - top_half_gap: float - - -@dataclass(frozen=True, eq=True) -class _FlatIW2DInputOptional(_IW2DInputOptional): - bottom_layers: Optional[Tuple[Layer]] = None - bottom_half_gap: Optional[float] = None - - -@dataclass(frozen=True, eq=True) -class FlatIW2DInput(_FlatIW2DInputOptional, _FlatIW2DInputBase): - pass - - def _iw2d_format_layer(layer: Layer, n: int) -> str: """ Formats the information describing a single layer into a string in accordance with IW2D standards. diff --git a/xwakes/wit/interface_dataclasses.py b/xwakes/wit/interface_dataclasses.py new file mode 100644 index 00000000..0840fc8c --- /dev/null +++ b/xwakes/wit/interface_dataclasses.py @@ -0,0 +1,91 @@ +from dataclasses import dataclass +from typing import Tuple, Optional + +@dataclass(frozen=True, eq=True) +class Sampling: + start: float + stop: float + # 0 = logarithmic, 1 = linear, 2 = both + scan_type: int + added: Tuple[float] + sampling_exponent: Optional[float] = None + points_per_decade: Optional[float] = None + min_refine: Optional[float] = None + max_refine: Optional[float] = None + n_refine: Optional[float] = None + + +@dataclass(frozen=True, eq=True) +class Layer: + # The distance in mm of the inner surface of the layer from the reference orbit + thickness: float + dc_resistivity: float + resistivity_relaxation_time: float + re_dielectric_constant: float + magnetic_susceptibility: float + permeability_relaxation_frequency: float + + +# Define several dataclasses for IW2D input elements. We must split mandatory +# and optional arguments into private dataclasses to respect the resolution +# order. The public classes RoundIW2DInput and FlatIW2D input inherit from +# from the private classes. +# https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses + +@dataclass(frozen=True, eq=True) +class _IW2DInputBase: + machine: str + length: float + relativistic_gamma: float + calculate_wake: bool + f_params: Sampling + + +@dataclass(frozen=True, eq=True) +class _IW2DInputOptional: + z_params: Optional[Sampling] = None + long_factor: Optional[float] = None + wake_tol: Optional[float] = None + freq_lin_bisect: Optional[float] = None + comment: Optional[str] = None + + +@dataclass(frozen=True, eq=True) +class IW2DInput(_IW2DInputOptional, _IW2DInputBase): + pass + + +@dataclass(frozen=True, eq=True) +class _RoundIW2DInputBase(_IW2DInputBase): + layers: Tuple[Layer] + inner_layer_radius: float + # (long, xdip, ydip, xquad, yquad) + yokoya_factors: Tuple[float, float, float, float, float] + + +@dataclass(frozen=True, eq=True) +class _RoundIW2DInputOptional(_IW2DInputOptional): + pass + + +@dataclass(frozen=True, eq=True) +class RoundIW2DInput(_RoundIW2DInputOptional, _RoundIW2DInputBase): + pass + + +@dataclass(frozen=True, eq=True) +class _FlatIW2DInputBase(_IW2DInputBase): + top_bottom_symmetry: bool + top_layers: Tuple[Layer] + top_half_gap: float + + +@dataclass(frozen=True, eq=True) +class _FlatIW2DInputOptional(_IW2DInputOptional): + bottom_layers: Optional[Tuple[Layer]] = None + bottom_half_gap: Optional[float] = None + + +@dataclass(frozen=True, eq=True) +class FlatIW2DInput(_FlatIW2DInputOptional, _FlatIW2DInputBase): + pass \ No newline at end of file diff --git a/xwakes/wit/materials.py b/xwakes/wit/materials.py index cdf73b2a..90ad952c 100644 --- a/xwakes/wit/materials.py +++ b/xwakes/wit/materials.py @@ -1,10 +1,12 @@ -from .interface import Layer from .utils import round_sigfigs from pathlib import Path import numpy as np import json from typing import Callable, Tuple +from .interface_dataclasses import Layer + + def layer_from_dict(thickness: float, material_dict: dict) -> Layer: diff --git a/xwakes/wit/utilities.py b/xwakes/wit/utilities.py index 148d5d41..299e1f40 100644 --- a/xwakes/wit/utilities.py +++ b/xwakes/wit/utilities.py @@ -1,7 +1,12 @@ -from .component import Component +from .component import (Component, ComponentResonator, + ComponentClassicThickWall, + ComponentSingleLayerResistiveWall, + ComponentTaperSingleLayerRestsistiveWall, + ComponentInterpolated) from .element import Element -from .interface import Layer, FlatIW2DInput, RoundIW2DInput +from .interface_dataclasses import FlatIW2DInput, RoundIW2DInput from .interface import component_names +from .materials import Layer from yaml import load, SafeLoader from typing import Tuple, Dict, List, Union, Sequence, Optional, Callable @@ -13,11 +18,7 @@ import scipy.constants from scipy import special as sp import numpy as np - -if hasattr(np, 'trapezoid'): - trapz = np.trapezoid # numpy 2.0 -else: - trapz = np.trapz +import pandas as pd c_light = scipy.constants.speed_of_light # m s-1 mu0 = scipy.constants.mu_0 @@ -83,48 +84,31 @@ def create_element_from_config(identifier: str) -> Element: return Element(length, beta_x, beta_y, components, name, tag) -def compute_resonator_f_roi_half_width(q: float, f_r: float, f_roi_level: float = 0.5): - aux = np.sqrt((1 - f_roi_level) / f_roi_level) - return (aux + np.sqrt(aux**2 + 4*q**2))*f_r/(2*q) - f_r def create_resonator_component(plane: str, exponents: Tuple[int, int, int, int], - r: float, q: float, f_r: float, f_roi_level: float = 0.5) -> Component: + r: float, q: float, f_r: float, + f_roi_level: float = 0.5) -> ComponentResonator: """ Creates a single component object belonging to a resonator :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param exponents: four integers corresponding to (source_x, source_y, + test_x, test_y) aka (a, b, c, d) :param r: the shunt impedance of the given component of the resonator :param q: the quality factor of the given component of the resonator :param f_r: the resonance frequency of the given component of the resonator - :param f_roi_level: fraction of the peak ok the resonator which is covered by the ROI. I.e. the roi will cover - the frequencies for which the resonator impedance is larger than f_roi_level*r + :param f_roi_level: fraction of the peak ok the resonator which is covered + by the ROI. I.e. the roi will cover the frequencies for which the resonator + impedance is larger thanf_roi_level*r :return: A component object of a resonator, specified by the input arguments """ - root_term = sqrt(1 - 1 / (4 * q ** 2) + 0J) - omega_r = 2 * pi * f_r - if plane == 'z': - impedance = lambda f: r / (1 - 1j * q * (f_r / f - f / f_r)) - omega_bar = omega_r * root_term - alpha = omega_r / (2 * q) - wake = lambda t: (omega_r * r * exp(-alpha * t) * ( - cos(omega_bar * t) - alpha * sin(omega_bar * t) / omega_bar) / q).real - else: - impedance = lambda f: (f_r * r) / (f * (1 - 1j * q * (f_r / f - f / f_r))) - wake = lambda t: (omega_r * r * exp(-omega_r * t / (2 * q)) * sin(omega_r * root_term * t) / - (q * root_term)).real - - if q > 1: - d = compute_resonator_f_roi_half_width(q=q, f_r=f_r, f_roi_level=f_roi_level) - f_rois = [(f_r - d, f_r + d)] - else: - f_rois = [] - # TODO: add ROI(s) for wake - - return Component(impedance, wake, plane, source_exponents=exponents[:2], - test_exponents=exponents[2:], - f_rois=f_rois) + return ComponentResonator(plane=plane, + exponents=exponents, + r=r, + q=q, + f_r=f_r, + f_roi_level=f_roi_level) def create_resonator_element(length: float, beta_x: float, beta_y: float, @@ -212,8 +196,10 @@ def create_many_resonators_element(length: float, beta_x: float, beta_y: float, return Element(length, beta_x, beta_y, components, tag=tag, description=description) -def create_classic_thick_wall_component(plane: str, exponents: Tuple[int, int, int, int], - layer: Layer, radius: float) -> Component: +def create_classic_thick_wall_component(plane: str, + exponents: Tuple[int, int, int, int], + layer: Layer, + radius: float)-> ComponentClassicThickWall: """ Creates a single component object modeling a resistive wall impedance/wake, based on the "classic thick wall formula" (see e.g. A. W. Chao, chap. 2 in @@ -221,120 +207,21 @@ def create_classic_thick_wall_component(plane: str, exponents: Tuple[int, int, i John Wiley and Sons, 1993). Only longitudinal and transverse dipolar impedances are supported here. :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param exponents: four integers corresponding to (source_x, source_y, + test_x, test_y) aka (a, b, c, d) :param layer: the chamber material, as a wit Layer object :param radius: the chamber radius in m :return: A component object """ + return ComponentClassicThickWall(plane=plane, + exponents=exponents, + layer=layer, + radius=radius) - # Material properties required for the skin depth computation are derived from the input Layer attributes - material_resistivity = layer.dc_resistivity - material_relative_permeability = 1. + layer.magnetic_susceptibility - material_permeability = material_relative_permeability * mu0 - - # Create the skin depth as a function offrequency and layer properties - delta_skin = lambda f: (material_resistivity / (2*pi*abs(f) * material_permeability)) ** (1/2) - - # Longitudinal impedance and wake - if plane == 'z' and exponents == (0, 0, 0, 0): - impedance = lambda f: (1/2) * (1+sign(f)*1j) * material_resistivity / (pi * radius) * (1 / delta_skin(f)) - wake = lambda t: - c_light / (2*pi*radius) * (Z0 * material_resistivity/pi)**(1/2) * 1/(t**(1/2)) - # Transverse dipolar impedance - elif (plane == 'x' and exponents == (1, 0, 0, 0)) or (plane == 'y' and exponents == (0, 1, 0, 0)): - impedance = lambda f: ((c_light/(2*pi*f)) * (1+sign(f)*1j) * - material_resistivity / (pi * radius**3) * - (1 / delta_skin(f))) - wake = lambda t: -c_light / (2*pi*radius**3) * (Z0 * material_resistivity/pi)**(1/2) * 1/(t**(3/2)) - else: - print("Warning: resistive wall impedance not implemented for component {}{}. Set to zero".format(plane, - exponents)) - impedance = lambda f: 0 - wake = lambda f: 0 - - return Component(vectorize(impedance), vectorize(wake), plane, source_exponents=exponents[:2], - test_exponents=exponents[2:]) - - -def _zlong_round_single_layer_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius: float, length: float) -> ArrayLike: - """ - Function to compute the longitudinal resistive-wall impedance from - the single-layer, approximated formula for a cylindrical structure, - by E. Metral (see e.g. Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, - https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and - Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, - https://doi.org/10.1103/PhysRevSTAB.12.084401) - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius: the radius of the structure (in m) - :param length: the total length of the resistive object (in m) - :return: Zlong, the longitudinal impedance at these frequencies - """ - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_long = 1j*omega*mu0*length/(2.*pi*beta**2*gamma**2) - - x1 = k*radius/gamma - x1sq = x1**2 - x2 = nu*radius - - zlong = coef_long * (sp.k0(x1)/sp.i0(x1) - 1./(x1sq*(1./2.+eps1*sp.kve(1, x2)/(x2*sp.kve(0, x2))))) - - return zlong - - -def _zdip_round_single_layer_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius: float, length: float) -> ArrayLike: - """ - Function to compute the transverse dipolar resistive-wall impedance from - the single-layer, approximated formula for a cylindrical structure, - Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, - https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and - Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, - https://doi.org/10.1103/PhysRevSTAB.12.084401) - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius: the radius of the structure (in m) - :param length: the total length of the resistive object (in m) - :return: Zdip, the transverse dipolar impedance at these frequencies - """ - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_dip = 1j*k**2*Z0*length/(4.*pi*beta*gamma**4) - x1 = k*radius/gamma - x1sq = x1**2 - x2 = nu*radius - - zdip = coef_dip * (sp.k1(x1)/sp.i1(x1) + 4.*beta**2*gamma**2/(x1sq*(2.+x2*sp.kve(0, x2)/(mu1*sp.kve(1, x2))))) - - return zdip - - -def create_resistive_wall_single_layer_approx_component(plane: str, exponents: Tuple[int, int, int, int], - input_data: Union[FlatIW2DInput, RoundIW2DInput]) -> Component: +def create_resistive_wall_single_layer_approx_component( + plane: str, exponents: Tuple[int, int, int, int], + input_data: Union[FlatIW2DInput, RoundIW2DInput]) -> Component: """ Creates a single component object modeling a resistive wall impedance, based on the single-layer approximated formulas by E. Metral (see e.g. @@ -356,63 +243,9 @@ def create_resistive_wall_single_layer_approx_component(plane: str, exponents: T but the Yokoya factors put in the input will be used. :return: A component object """ - gamma = input_data.relativistic_gamma - length = input_data.length - - if isinstance(input_data, FlatIW2DInput): - if len(input_data.top_layers) > 1: - raise NotImplementedError("Input data can have only one layer") - yok_long = 1. - layer = input_data.top_layers[0] - radius = input_data.top_half_gap - if input_data.top_bottom_symmetry: - yok_dipx = pi**2/24. - yok_dipy = pi**2/12. - yok_quax = -pi**2/24. - yok_quay = pi**2/24. - elif input_data.bottom_half_gap == inf: - yok_dipx = 0.25 - yok_dipy = 0.25 - yok_quax = -0.25 - yok_quay = 0.25 - else: - raise NotImplementedError("For asymmetric structures, only the case of a single plate is implemented; " - "hence the bottom half gap must be infinite") - elif isinstance(input_data, RoundIW2DInput): - radius = input_data.inner_layer_radius - if len(input_data.layers) > 1: - raise NotImplementedError("Input data can have only one layer") - layer = input_data.layers[0] - yok_long = input_data.yokoya_factors[0] - yok_dipx = input_data.yokoya_factors[1] - yok_dipy = input_data.yokoya_factors[2] - yok_quax = input_data.yokoya_factors[3] - yok_quay = input_data.yokoya_factors[4] - else: - raise NotImplementedError("Input of type neither FlatIW2DInput nor RoundIW2DInput cannot be handled") - - # Longitudinal impedance - if plane == 'z' and exponents == (0, 0, 0, 0): - impedance = lambda f: yok_long*_zlong_round_single_layer_approx( - f, gamma, layer, radius, length) - # Transverse impedances - elif plane == 'x' and exponents == (1, 0, 0, 0): - impedance = lambda f: yok_dipx*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - elif plane == 'y' and exponents == (0, 1, 0, 0): - impedance = lambda f: yok_dipy*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - elif plane == 'x' and exponents == (0, 0, 1, 0): - impedance = lambda f: yok_quax*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - elif plane == 'y' and exponents == (0, 0, 0, 1): - impedance = lambda f: yok_quay*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - else: - impedance = lambda f: 0 - - return Component(impedance=impedance, plane=plane, source_exponents=exponents[:2], - test_exponents=exponents[2:]) + return ComponentSingleLayerResistiveWall(plane=plane, + exponents=exponents, + input_data=input_data) def create_resistive_wall_single_layer_approx_element( @@ -456,106 +289,6 @@ def create_resistive_wall_single_layer_approx_element( description=description) -def _zlong_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius_small: float, - radius_large: float, length: float, - step_size: float = 1e-3) -> ArrayLike: - """ - Function to compute the longitudinal resistive-wall impedance for a - round taper, integrating the radius-dependent approximated formula - for a cylindrical structure (see_zlong_round_single_layer_approx above), - over the length of the taper. - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius_small: the smallest radius of the taper (in m) - :param radius_large: the largest radius of the taper (in m) - :param length: the total length of the taper (in m) - :param step_size: the step size (in the radial direction) for the - integration (in m) - :return: the longitudinal impedance at these frequencies - """ - if isscalar(frequencies): - frequencies = array(frequencies) - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies.reshape((-1, 1)) - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_long = 1j*omega*mu0/(2.*pi*beta**2*gamma**2) - - npts = int(floor(abs(radius_large-radius_small)/step_size)+1) - radii = linspace(radius_small, radius_large, npts).reshape((1, -1)) - one_array = ones(radii.shape) - - x1 = k.dot(radii)/gamma - x1sq = x1**2 - x2 = nu.dot(radii) - zlong = (coef_long.dot(length / float(npts) * one_array) * - (sp.k0(x1) / sp.i0(x1) - 1. / (x1sq * (1. / 2. + eps1.dot(one_array) * - sp.kve(1, x2) / (x2 * sp.kve(0, x2))))) - ) - - return trapz(zlong, axis=1) - - -def _zdip_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius_small: float, - radius_large: float, length: float, - step_size: float = 1e-3) -> ArrayLike: - """ - Function to compute the transverse dip. resistive-wall impedance for a - round taper, integrating the radius-dependent approximated formula - for a cylindrical structure (see_zdip_round_single_layer_approx above), - over the length of the taper. - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius_small: the smallest radius of the taper (in m) - :param radius_large: the largest radius of the taper (in m) - :param length: the total length of the taper (in m) - :param step_size: the step size (in the radial direction) for the - integration (in m) - :return: the transverse dipolar impedance at these frequencies - """ - if isscalar(frequencies): - frequencies = array(frequencies) - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies.reshape((-1,1)) - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_dip = 1j*k**2*Z0/(4.*pi*beta*gamma**4) - - npts = int(floor(abs(radius_large-radius_small)/step_size)+1) - radii = linspace(radius_small,radius_large,npts).reshape((1,-1)) - one_array = ones(radii.shape) - - x1 = k.dot(radii)/gamma - x1sq = x1**2 - x2 = nu.dot(radii) - zdip = ( - coef_dip.dot(length / float(npts) * one_array) * - (sp.k1(x1) / sp.i1(x1) + 4 * beta**2 * gamma**2 / (x1sq * (2 + x2 * sp.kve(0, x2) / (mu1 * sp.kve(1, x2))))) - ) - - return trapz(zdip, axis=1) - - def create_taper_RW_approx_component(plane: str, exponents: Tuple[int, int, int, int], input_data: Union[FlatIW2DInput, RoundIW2DInput], radius_small: float, radius_large: float, @@ -586,68 +319,12 @@ def create_taper_RW_approx_component(plane: str, exponents: Tuple[int, int, int, for the integration (in m) :return: A component object """ - gamma = input_data.relativistic_gamma - length = input_data.length - - if isinstance(input_data, FlatIW2DInput): - if len(input_data.top_layers) > 1: - raise NotImplementedError("Input data can have only one layer") - yok_long = 1. - layer = input_data.top_layers[0] - radius = input_data.top_half_gap - if input_data.top_bottom_symmetry: - yok_dipx = pi**2/24. - yok_dipy = pi**2/12. - yok_quax = -pi**2/24. - yok_quay = pi**2/24. - elif input_data.bottom_half_gap == inf: - yok_dipx = 0.25 - yok_dipy = 0.25 - yok_quax = -0.25 - yok_quay = 0.25 - else: - raise NotImplementedError("For asymmetric structures, only the case of a single plate is implemented; " - "hence the bottom half gap must be infinite") - elif isinstance(input_data, RoundIW2DInput): - radius = input_data.inner_layer_radius - if len(input_data.layers) > 1: - raise NotImplementedError("Input data can have only one layer") - layer = input_data.layers[0] - yok_long = input_data.yokoya_factors[0] - yok_dipx = input_data.yokoya_factors[1] - yok_dipy = input_data.yokoya_factors[2] - yok_quax = input_data.yokoya_factors[3] - yok_quay = input_data.yokoya_factors[4] - else: - raise NotImplementedError("Input of type neither FlatIW2DInput nor RoundIW2DInput cannot be handled") - - # Longitudinal impedance - if plane == 'z' and exponents == (0, 0, 0, 0): - impedance = lambda f: yok_long*_zlong_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - # Transverse impedances - elif plane == 'x' and exponents == (1, 0, 0, 0): - impedance = lambda f: yok_dipx*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - elif plane == 'y' and exponents == (0, 1, 0, 0): - impedance = lambda f: yok_dipy*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - elif plane == 'x' and exponents == (0, 0, 1, 0): - impedance = lambda f: yok_quax*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - elif plane == 'y' and exponents == (0, 0, 0, 1): - impedance = lambda f: yok_quay*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - else: - impedance = lambda f: 0 - - return Component(impedance=impedance, plane=plane, source_exponents=exponents[:2], - test_exponents=exponents[2:]) + return ComponentTaperSingleLayerRestsistiveWall(plane=plane, + exponents=exponents, + input_data=input_data, + radius_small=radius_small, + radius_large=radius_large, + step_size=step_size) def create_taper_RW_approx_element( @@ -700,33 +377,96 @@ def create_taper_RW_approx_element( def create_interpolated_impedance_component(interpolation_frequencies: ArrayLike, impedance: Optional[Callable] = None, - wake: Optional[Callable] = None, plane: str = '', + wake: Optional[Callable] = None, + interpolation_times: ArrayLike = None, + plane: str = '', source_exponents: Tuple[int, int] = (-1, -1), test_exponents: Tuple[int, int] = (-1, -1), name: str = "Unnamed Component", f_rois: Optional[List[Tuple[float, float]]] = None, t_rois: Optional[List[Tuple[float, float]]] = None): """ - Creates a component in which the impedance function is evaluated directly only on few points and it is interpolated - everywhere else. This helps when the impedance function is very slow to evaluate. - :param interpolation_frequencies: the frequencies where the function is evaluated for the interpolation - :param impedance: A callable function representing the impedance function of the Component. Can be undefined if + Creates a component in which the impedance function is evaluated directly + only on few points and it is interpolated + everywhere else. This helps when the impedance function is very slow to + evaluate. + :param interpolation_frequencies: the frequencies where the impedance + function is evaluated for the interpolation + :param impedance: A callable function representing the impedance function of + the Component. Can be undefined if the wake function is defined. - :param wake: A callable function representing the wake function of the Component. Can be undefined if - the impedance function is defined. - :param plane: The plane of the Component, either 'x', 'y' or 'z'. Must be specified for valid initialization - :param source_exponents: The exponents in the x and y planes experienced by the source particle. Also + :param interpolation_times: the times where the wake function is evaluated + for the interpolation + :param wake: A callable function representing the wake function of the + Component. Can be undefined if the impedance function is defined. + :param plane: The plane of the Component, either 'x', 'y' or 'z'. Must be + specified for valid initialization + :param source_exponents: The exponents in the x and y planes experienced by + the source particle. Also referred to as 'a' and 'b'. Must be specified for valid initialization - :param test_exponents: The exponents in the x and y planes experienced by the source particle. Also + :param test_exponents: The exponents in the x and y planes experienced by + the source particle. Also referred to as 'a' and 'b'. Must be specified for valid initialization :param name: An optional user-specified name of the component :param f_rois: a list of frequency regions of interest :param t_rois: a list of time regions of interest """ - def interpolated_impedance(x): - return np.interp(x, interpolation_frequencies, impedance(interpolation_frequencies)) - - return Component(impedance=interpolated_impedance, wake=wake, plane=plane, - source_exponents=source_exponents, test_exponents=test_exponents, - name=name, f_rois=f_rois, - t_rois=t_rois) + return ComponentInterpolated(impedance_input=impedance, wake_input=wake, + plane=plane, source_exponents=source_exponents, + test_exponents=test_exponents, name=name, + interpolation_frequencies=interpolation_frequencies, + f_rois=f_rois, t_rois=t_rois) + + +def read_pyheadtail_file(wake_file, wake_file_columns, + wake_pyheadtail_units=True) -> pd.DataFrame: + ''' + Read a wake or impedance file in headtail format and return a + pandas.DataFrame + + parameters: + ---------- + wake_file: str + Path to the file + wake_file_columns: list + List of strings with the names of the columns in the file + wake_headtail_units: bool + If True, assume that the units of the file are as in PyHEADTAIL + If False, SI units are used + ''' + valid_wake_components = ['constant_x', 'constant_y', 'dipole_x', + 'dipole_y', 'dipole_xy', 'dipole_yx', + 'quadrupole_x', 'quadrupole_y', + 'quadrupole_xy', 'quadrupole_yx', + 'longitudinal'] + + wake_data = np.loadtxt(wake_file) + if len(wake_file_columns) != wake_data.shape[1]: + raise ValueError("Length of wake_file_columns list does not" + + " correspond to the number of columns in the" + + " specified wake_file. \n") + if 'time' not in wake_file_columns: + raise ValueError("No wake_file_column with name 'time' has" + + " been specified. \n") + + dict_components = {} + + conversion_factor_time = -1E-9 if wake_pyheadtail_units else 1 + + itime = wake_file_columns.index('time') + dict_components['time'] = conversion_factor_time * wake_data[:, itime] + + for i_component, component in enumerate(wake_file_columns): + if component != 'time': + assert component in valid_wake_components + if component == 'longitudinal': + conversion_factor = -1E12 if wake_pyheadtail_units else 1 + else: + conversion_factor = -1E15 if wake_pyheadtail_units else 1 + + dict_components[component] = (wake_data[:, i_component] * + conversion_factor) + + df = pd.DataFrame(data=dict_components) + + return df