wakeSolver module

class wakis.wakeSolver.WakeSolver(wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, xsource=0.0, ysource=0.0, xtest=0.0, ytest=0.0, chargedist=None, ti=None, fmax=None, compute_plane='both', skip_cells=0, add_space=None, Ez_file='Ez.h5', save=True, results_folder='results/', verbose=0, counter_moving=False)[source]

Bases: object

Class for wake potential and impedance calculation from 3D time domain E fields

assign_logs()[source]

Assign the parameters of the wake to the logger.

static calc_impedance_from_wake(wake, s=None, t=None, fmax=None, samples=None, verbose=True)[source]

Calculate impedance from wake function using FFT.

Parameters

  • wake (array_like or list) – Wake function or [t, wake] list.

  • s (ndarray, optional) – Distance array [m].

  • t (ndarray, optional) – Time array [s].

  • fmax (float, optional) – Maximum frequency.

  • samples (int, optional) – Number of FFT samples.

  • verbose (bool, optional) – If True, print information.

Returns

  • f (ndarray) – Frequency array.

  • Z (ndarray) – Impedance array.

calc_lambdas(**kwargs)[source]

Obtain normalized charge distribution in terms of s Ξ»(s) for impedance calculation.

Parameters

  • s (ndarray) – Wakelength vector s=c*t-z.

  • chargedist (ndarray, optional) – Charge distribution Ξ»(z).

  • q (float, optional) – Total beam charge in [C].

  • z (ndarray, optional) – Vector containing z-coordinates of the domain [m].

  • zf (ndarray, optional) – Vector containing z-coordinates of the field monitor [m].

  • **kwargs – Additional parameters to set as attributes.

calc_lambdas_analytic(**kwargs)[source]

Obtain normalized charge distribution in s Ξ»(z) as an analytical gaussian.

Parameters

  • s (ndarray) – Wakelength vector s=c*t-z.

  • sigmaz (float) – Beam sigma in the longitudinal direction [m].

  • **kwargs – Additional parameters to set as attributes.

calc_long_WP(Ezt=None, **kwargs)[source]

Obtain the wake potential from the pre-computed longitudinal Ez(z,t) field.

Parameters

  • t (ndarray) – Vector containing time values [s].

  • z (ndarray) – Vector containing z-coordinates [m].

  • sigmaz (float) – Beam longitudinal sigma, to calculate injection time [m].

  • q (float) – Beam charge, to normalize wake potential.

  • ti (float, optional) – Injection time needed to set the negative part of s vector and wakelength.

  • Ezt (ndarray, optional) – Matrix (nz x nt) containing Ez(x_test, y_test, z, t).

  • Ez_file (str, optional) – HDF5 file containing the Ez(x, y, z) field data for every timestep.

  • **kwargs – Additional parameters to set as attributes.

calc_long_WP_3d(**kwargs)[source]

Obtain the 3D wake potential from the pre-computed Ez(x,y,z) field.

Parameters

  • Ez_file (str, optional) – HDF5 file containing the Ez(x,y,z) field data for every timestep.

  • t (ndarray) – Vector containing time values [s].

  • z (ndarray) – Vector containing z-coordinates [m].

  • q (float) – Total beam charge in [C].

  • n_transverse_cells (int, optional) – Number of transverse cells used for the 3d calculation: 2*n+1.

  • **kwargs – Additional parameters to set as attributes.

calc_long_Z(samples=1001, fmax=None, **kwargs)[source]

Obtain the longitudinal impedance from the longitudinal wake potential and the beam charge distribution using a single-sided DFT.

Parameters

  • WP (ndarray) – Longitudinal wake potential WP(s).

  • s (ndarray) – Wakelength vector s=c*t-z.

  • lambdas (ndarray) – Charge distribution Ξ»(s) interpolated to s axis, normalized by the beam charge.

  • chargedist (ndarray, optional) – Charge distribution Ξ»(z).

  • q (float, optional) – Total beam charge in [C].

  • z (ndarray) – Vector containing z-coordinates [m].

  • sigmaz (float) – Beam sigma in the longitudinal direction [m].

  • samples (int, optional) – Number of DFT samples. Default is 1001.

  • fmax (float, optional) – Maximum frequency of interest.

  • **kwargs – Additional parameters to set as attributes.

calc_trans_WP(**kwargs)[source]

Obtain the transverse wake potential from the longitudinal wake potential in 3d using the Panofsky-Wenzel theorem.

Parameters

  • WP_3d (ndarray) – Longitudinal wake potential in 3d WP(x,y,s).

  • s (ndarray) – Wakelength vector s=c*t-z.

  • dx (float) – Ez field mesh step in transverse plane, x-dir [m].

  • dy (float) – Ez field mesh step in transverse plane, y-dir [m].

  • x (ndarray, optional) – Vector containing x-coordinates [m].

  • y (ndarray, optional) – Vector containing y-coordinates [m].

  • n_transverse_cells (int, optional) – Number of transverse cells used for the 3d calculation: 2*n+1.

  • **kwargs – Additional parameters to set as attributes.

calc_trans_Z(samples=1001, fmax=None)[source]

Obtain the transverse impedance from the transverse wake potential and the beam charge distribution using a single-sided DFT.

Parameters

  • samples (int, optional) – Number of DFT samples. Default is 1001.

  • fmax (float, optional) – Maximum frequency of interest.

static calc_wake_from_impedance(impedance, f=None, tmax=None, samples=None, pad=0, verbose=True)[source]

Calculate wake function from impedance using inverse FFT.

Parameters

  • impedance (array_like or list) – Impedance or [f, Z] list.

  • f (ndarray, optional) – Frequency array.

  • tmax (float, optional) – Maximum time.

  • samples (int, optional) – Number of FFT samples.

  • pad (int, optional) – Padding for FFT.

  • verbose (bool, optional) – If True, print information.

Returns

  • t (ndarray) – Time array [s].

  • wake (ndarray) – Wake function.

copy()[source]

Return a shallow copy of the WakeSolver object.

Returns

obj – Shallow copy of the WakeSolver.

Return type

WakeSolver

get_DEmodel_fitting(freq_data=None, impedance_data=None, plane='longitudinal', dim='z', parameterBounds=None, N_resonators=None, DE_kernel='DE', maxiter=1e5, cmaes_sigma=0.01, popsize=150, tol=1e-3, use_minimization=True, minimization_margin=[0.3, 0.2, 0.01], minimum_peak_height=1.0, distance=3, inspect_bounds=False, Rs_bounds=[0.8, 10], Q_bounds=[0.5, 5], fres_bounds=[- 0.01e9, + 0.01e9], verbose=False)[source]

Fit the impedance using Differential Evolution or CMA-ES.

Parameters

  • freq_data (ndarray, optional) – Frequency data.

  • impedance_data (ndarray, optional) – Impedance data.

  • plane (str, optional) – β€˜longitudinal’ or β€˜transverse’.

  • dim (str, optional) – β€˜z’, β€˜x’, or β€˜y’.

  • parameterBounds (object, optional) – Parameter bounds for fitting.

  • N_resonators (int, optional) – Number of resonators.

  • DE_kernel (str, optional) – β€˜DE’ or β€˜CMAES’.

  • maxiter (int, optional) – Maximum number of iterations.

  • cmaes_sigma (float, optional) – Sigma for CMA-ES.

  • popsize (int, optional) – Population size.

  • tol (float, optional) – Tolerance for convergence.

  • use_minimization (bool, optional) – Whether to run minimization after fitting.

  • minimization_margin (list, optional) – Margin for minimization.

  • minimum_peak_height (float, optional) – Minimum peak height for SmartBounds.

  • distance (int, optional) – Minimum distance between peaks for SmartBounds.

  • inspect_bounds (bool, optional) – If True, inspect and display bounds.

  • Rs_bounds (list, optional) – Bounds for Rs, Q, and fres.

  • Q_bounds (list, optional) – Bounds for Rs, Q, and fres.

  • fres_bounds (list, optional) – Bounds for Rs, Q, and fres.

Returns

DE_model – Fitted DE model.

Return type

object

get_SmartBounds(freq_data=None, impedance_data=None, minimum_peak_height=1.0, distance=3, inspect_bounds=True, Rs_bounds=[0.8, 10], Q_bounds=[0.5, 5], fres_bounds=[- 0.01e9, + 0.01e9])[source]

Use Smart Bound Determination to estimate bounds for resonator fitting.

Parameters

  • freq_data (ndarray) – Frequency data.

  • impedance_data (ndarray) – Impedance data.

  • minimum_peak_height (float, optional) – Minimum peak height for peak finding.

  • distance (int, optional) – Minimum distance between peaks.

  • inspect_bounds (bool, optional) – If True, inspect and display bounds.

  • Rs_bounds (list, optional) – Bounds for Rs, Q, and fres.

  • Q_bounds (list, optional) – Bounds for Rs, Q, and fres.

  • fres_bounds (list, optional) – Bounds for Rs, Q, and fres.

Returns

bounds – SmartBoundDetermination result.

Return type

object

get_extrapolated_impedance(f=None, use_minimization=True, wakelength=None)[source]

Get the extrapolated impedance [Ohm] from the DE model.

Parameters

  • f (ndarray, optional) – Frequency array.

  • use_minimization (bool, optional) – Whether to use minimization in the DE model.

  • wakelength (float, optional) – Wakelength for the DE model.

Returns

  • f (ndarray) – Frequency array.

  • impedance (ndarray) – Extrapolated impedance [Ohm].

get_extrapolated_wake(wakelength, sigma=None, use_minimization=True)[source]

Get the extrapolated wake potential [V/pC] from the DE model.

Parameters

  • wakelength (float) – Wakelength to extrapolate to.

  • sigma (float, optional) – Sigma for the Gaussian. Default is self.sigmaz / c_light.

  • use_minimization (bool, optional) – Whether to use minimization in the DE model.

Returns

  • s (ndarray) – Distance array [m].

  • wake_potential (ndarray) – Extrapolated wake potential [V/pC].

get_extrapolated_wake_function(wakelength, use_minimization=True)[source]

Get the extrapolated wake function (Green function) from the DE model.

Parameters

  • wakelength (float) – Wakelength to extrapolate to.

  • use_minimization (bool, optional) – Whether to use minimization in the DE model.

Returns

  • t (ndarray) – Time array [s].

  • wake_function (ndarray) – Extrapolated wake function.

load_results(folder=None)[source]

Load all txt results from a given folder.

Parameters

folder (str, optional) – Folder path containing result txt files. If None, uses self.folder.

Notes

The txt files are generated when the attribute save = True is used.

log(txt)[source]

Print a log message if verbose is enabled.

Parameters

txt (str) – Message to print.

read_Ez(filename=None, return_value=False)[source]

Read the Ez.h5 file containing the Ez field information.

Parameters

  • filename (str, optional) – Path to the Ez.h5 file. If None, uses self.Ez_file.

  • return_value (bool, optional) – If True, return the h5py.File object.

Returns

hf – Returned if return_value is True.

Return type

h5py.File, optional

read_cst_3d(path=None, folder='3d', filename='Ez.h5', units=1e-3)[source]

Read CST 3D exports folder and store the Ez field information into a matrix Ez(x,y,z) for every timestep into a single .h5 file compatible with wakis.

Parameters

  • path (str, optional) – Path to the field data. Default is None.

  • folder (str, optional) – Folder containing the CST field data .txt files. Default is β€˜3d’.

  • filename (str, optional) – Name of the h5 file that will be generated. Default is β€˜Ez.h5’.

  • units (float, optional) – Unit conversion factor for coordinates. Default is 1e-3.

read_txt(txt, skiprows=2, delimiter=None, usecols=None)[source]

Read variables from ASCII .txt files and return data in a dictionary.

Parameters

  • txt (str) – Path to the .txt file.

  • skiprows (int, optional) – Number of rows to skip. Default is 2.

  • delimiter (str, optional) – Delimiter for columns.

  • usecols (list, optional) – Columns to read.

Returns

d – Dictionary with keys as header names or integer indices.

Return type

dict

save_txt(f_name, x_data=None, y_data=None, x_name='X [-]', y_name='Y [-]')[source]

Save x and y data to a text file in a two-column format.

Parameters

  • f_name (str) – Name of the output file (without the .txt extension).

  • x_data (ndarray, optional) – Array containing x-axis data.

  • y_data (ndarray, optional) – Array containing y-axis data.

  • x_name (str, optional) – Label for the x-axis column in the output file.

  • y_name (str, optional) – Label for the y-axis column in the output file.

Notes

  • The data is saved in a two-column format where x_data and y_data

are combined column-wise. - If x_data or y_data is missing, the function prints a warning and does not save a file.

Examples

Save two NumPy arrays to data.txt:

>>> x = np.linspace(0, 10, 5)
>>> y = np.sin(x)
>>> save_txt("data", x, y, x_name="Time [s]", y_name="Amplitude")

The saved file will look like:

0.00 0.00 2.50 0.59 5.00 -0.99 7.50 0.94 10.00 -0.54

solve(compute_plane=None, **kwargs)[source]

Perform the wake potential and impedance calculation for longitudinal and transverse planes.

Parameters

  • compute_plane (str, optional) – Which planes to compute: β€˜both’, β€˜longitudinal’, or β€˜transverse’. Default is None (uses self.compute_plane).

  • **kwargs – Additional parameters to set as attributes.

update_long_WP(t)[source]

Calculate the wake potential on the fly for a given time step.

Parameters

t (float) – Current simulation time [s].

Notes

Work in progress. Calculation of wake potential on the fly.