#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Mon Mar 23 13:20:11 2020
@author: sjoly
@modified by: MaltheRaschke, edelafue
"""
from functools import partial
from typing import Any, Callable, Sequence
import numpy as np
import numpy.typing as npt
from scipy.optimize import minimize
from .objectiveFunctions import ObjectiveFunctions as obj
from .resonatorFormulas import Impedances as imp
from .resonatorFormulas import Wakes as wak
from .solvers import Solvers
from .uncertainties import get_uncertainties
from .utils import compute_fft
ArrayLike = npt.ArrayLike
ParameterBounds = list[tuple[float, float]]
ObjectiveCallable = Callable[..., float]
[docs]class EvolutionaryAlgorithm:
def __init__(
self,
x_data: ArrayLike,
y_data: ArrayLike,
N_resonators: int,
parameterBounds: ParameterBounds,
plane: str = "longitudinal",
fitFunction: str = "impedance",
objectiveFunction: ObjectiveCallable | str | None = None,
wake_length: float | None = None,
sigma: float | None = None,
uncertainty_warning: float = 0.2,
) -> None:
"""
Implements an evolutionary algorithm for fitting impedance models to data.
This class optimizes a resonator-based impedance model to fit measured or
simulated impedance data using an evolutionary algorithm. It supports both
longitudinal and transverse impedance models with a variable number of resonators.
Parameters
----------
x_data : numpy.ndarray
Stores the input x data: frequencies for impedance [Hz], or
times for wake function and wake potential [s]
y_data : numpy.ndarray
Stores the input y data: impedance data [Ohm] or Wake function/potential [V/C/s]
N_resonators : int
Number of resonators in the model.
parameterBounds : list of tuple
List of parameter bounds for the optimization. Each resonator has three
parameters: Rs (shunt resistance), Q (quality factor), and fr (resonant frequency).
fitFunction: str, optional
Specify which fitFunction to use: ["impedance", "wake", "wake_potential"]
plane : str, optional
Type of impedance model, either `"longitudinal"` or `"transverse"`.
Default is `"longitudinal"`.
objectiveFunction : callable or str, optional
The objective function to minimize. If str, it should be in
['Real', 'Complex', 'Abs'] (no case distinction). Default if
`y_data` is complex: ``obj.sumOfSquaredError``; otherwise it is
``obj.sumOfSquaredErrorReal``.
wake_length : float, optional
Length of the wake function in meters. Required for wake computations.
sigma : float, optional
Standard deviation for the bunch's Gaussian distribution used in
wake potential getters. Default is ``None``. When ``None`` and
wake-potential getter is used, a default value of 1e-10 s is applied internally.
uncertainty_warning : float, optional
Threshold for relative uncertainty warning. Default is 0.2 -->20%
If any parameter has a relative uncertainty (sigma/value) above this
threshold, a warning is printed after optimization.
Attributes
----------
fitFunction : callable
Partial function used to compute impedance based on the chosen model
(`imp.Resonator_longitudinal_imp`, `imp.n_Resonator_longitudinal_imp`, etc.).
evolutionParameters : dict or None
Stores parameters of the evolutionary optimization algorithm.
minimizationParameters : numpy.ndarray or None
Stores the best-fit parameters obtained from the optimization.
time_data : numpy.ndarray
Stores the input x data: times for wake function and wake potential [s]
wake_data : numpy.ndarray
Stores the input Wake function [V/C/s]
wake_potential_data : numpy.ndarray
Stores the input Wake potential [V/C/s]
frequency_data : numpy.ndarray
Stores the input x data: frequencies for impedance [Hz]
impedance_data : numpy.ndarray
Stores the input impedance data [Ohm]
Notes
-----
- The ``fitFunction`` is assigned based on the ``plane`` type, the
fit mode, and the number of resonators.
- The impedance and wake model is based on resonators and can be
used for both single-resonator and multi-resonator systems.
- The optimization is performed using an evolutionary algorithm, with results
stored in `minimizationParameters`.
Examples
--------
>>> import numpy as np
>>> import iddefix
>>> freq = np.linspace(1e9, 5e9, 100) # Frequency range from 1 GHz to 5 GHz
>>> Z = iddefix.Impedances.Resonator_longitudinal_imp(freq, 400, 30, 0.2e9) # Rs, Q, fr
>>> bounds = [(10, 1000), (1, 100), (1e9, 5e9)] # Example bounds for Rs, Q, fr
>>> model = iddefix.EvolutionaryAlgorithm(
... x_data=freq,
... y_data=Z,
... N_resonators=1,
... parameterBounds=bounds,
... )
>>> print(model.plane)
'longitudinal'
>>> model.run_differential_evolution(maxiter=100, popsize=10)
"""
self.x_data = x_data
self.y_data = y_data
self.N_resonators = N_resonators
self.parameterBounds = parameterBounds
self.objectiveFunction = objectiveFunction
self.wake_length = wake_length
self.plane = plane
self.sigma = sigma
self.time_data = None
self.wake_data = None
self.wake_potential_data = None
self.frequency_data = None
self.impedance_data = None
self.evolutionParameters = None
self.evolutionParametersUncertainties = None
self.minimizationParameters = None
self.minimizationParametersUncertainties = None
self.flagged_params = None
# Threshold on relative uncertainty (sigma/|value|) above which a
# warning is printed after evolution/minimization
self.uncertainty_warning = uncertainty_warning
if self.objectiveFunction is None:
if np.iscomplex(y_data).any():
self.objectiveFunction = obj.sumOfSquaredError
print(
"[!] Objective function set to default \
`iddefix.objectiveFunctions.sumOfSquaredError`"
)
else:
self.objectiveFunction = obj.sumOfSquaredErrorReal
print(
"[!] Objective function set to \
`iddefix.objectiveFunctions.sumOfSquaredErrorReal` \
for real-valued only data"
)
elif type(self.objectiveFunction) is str:
if self.objectiveFunction.lower() == "complex":
self.objectiveFunction = obj.sumOfSquaredError
if self.objectiveFunction.lower() == "real":
self.objectiveFunction = obj.sumOfSquaredErrorReal
elif self.objectiveFunction.lower() == "abs":
self.objectiveFunction = obj.sumOfSquaredErrorAbs
else:
print(
"[!] Objective function set to default \
`iddefix.objectiveFunctions.sumOfSquaredError`"
)
self.objectiveFunction = obj.sumOfSquaredError
if fitFunction == "wake" or fitFunction == "wake function":
if plane == "longitudinal":
self.fitFunction = wak.n_Resonator_longitudinal_wake
elif plane == "transverse":
self.fitFunction = wak.n_Resonator_transverse_wake
else:
raise Exception("Algorithm needs N_resonartors >= 1")
self.time_data = x_data
self.wake_data = y_data
elif fitFunction == "wake potential":
if self.sigma is None:
print("[!] sigma not specified, using the default sigma=1e-10 s")
self.sigma = 1e-10
if plane == "longitudinal":
self.fitFunction = partial(
wak.n_Resonator_longitudinal_wake_potential,
sigma=self.sigma,
)
elif plane == "transverse":
self.fitFunction = partial(
wak.n_Resonator_transverse_wake_potential, sigma=self.sigma
)
else:
raise Exception("Algorithm needs N_resonartors >= 1")
self.time_data = x_data
self.wake_potential_data = y_data
else: # Default to "impedance"
if wake_length is not None:
print(
"[!] Using the partially decayed resonator formalism for impedance"
)
else:
print("[!] Using the fully decayed resonator formalism for impedance")
if plane == "longitudinal":
self.fitFunction = partial(
imp.n_Resonator_longitudinal_imp, wake_length=wake_length
)
elif plane == "transverse":
self.fitFunction = partial(
imp.n_Resonator_transverse_imp, wake_length=wake_length
)
else:
raise Exception("Algorithm needs N_resonartors >= 1")
self.frequency_data = x_data
self.impedance_data = y_data
[docs] def check_y_data(self) -> None:
"""
Small function to avoid 0 frequency leading to zero division when
using resonators.
"""
mask = np.where(self.x_data > 0.0)[0]
self.x_data = self.x_data[mask]
self.y_data = self.y_data[mask]
[docs] def generate_Initial_Parameters(
self,
parameterBounds: ParameterBounds,
objectiveFunction: ObjectiveCallable,
fitFunction: Callable[..., np.ndarray],
x_values_data: ArrayLike,
y_values_data: ArrayLike,
maxiter: int = 2000,
popsize: int = 150,
mutation: tuple[float, float] = (0.1, 0.5),
crossover_rate: float = 0.8,
tol: float = 0.01,
solver: str = "scipy",
) -> tuple[np.ndarray, str]:
"""
Generates initial parameter estimates using a
Differential Evolution (DE) solver.
This function applies a DE optimization method to identify
suitable initial parameters for resonance fitting.
These parameters can be further refined using local minimization.
Parameters
----------
parameterBounds : list of tuple
A list of (min, max) bounds for each parameter.
objectiveFunction : callable
The objective function to minimize. It should accept parameters,
a fitting function, x-data, and y-data.
fitFunction : callable
The fitting function that models the impedance response.
x_values_data : numpy.ndarray
Array containing x-values of the data (frequency points).
y_values_data : numpy.ndarray
Array containing y-values of the data (impedance magnitudes).
maxiter : int, optional
Maximum number of iterations for the DE solver. Default is 2000.
popsize : int, optional
Population size for the DE algorithm. Default is 150.
mutation : tuple of float, optional
Range of mutation factors that control parameter variation.
Default is (0.1, 0.5).
crossover_rate : float, optional
Probability of recombining individuals in the DE algorithm.
Default is 0.8.
tol : float, optional
Convergence tolerance for stopping criteria. Default is 0.01.
solver : str, optional
The solver to use for differential evolution. Available options:
- `"scipy"`: Uses SciPy's built-in DE solver.
- `"pyfde"`: Uses `pyfde`, an alternative DE implementation.
- `"pyfde_jade"`: Uses JADE, a self-adaptive DE variant
(automatically adjusts `mutation` and `crossover_rate`).
Default is `"scipy"`.
Returns
-------
tuple
- **solution** : numpy.ndarray
Optimized parameter estimates found by the DE solver.
- **message** : str
Solver status message.
Notes
-----
- Calls the appropriate solver function based on the `solver` argument.
- If `solver='pyfde_jade'`, mutation and crossover rates are automatically adjusted.
- The result can be used as an initial guess for further optimization.
"""
objective_function = partial(
objectiveFunction,
fitFunction=fitFunction,
x=x_values_data,
y=y_values_data,
)
# Map solver names to functions
solver_functions = {
"scipy": Solvers.run_scipy_solver,
"pyfde": Solvers.run_pyfde_solver,
"pyfde_jade": Solvers.run_pyfde_jade_solver,
}
solver_function = solver_functions.get(solver)
if solver == "pyfde_jade":
mutation, crossover_rate = None, None
if not solver_function:
raise ValueError(f"Invalid solver name: {solver}")
solution, message = solver_function(
parameterBounds,
objective_function,
maxiter=maxiter,
popsize=popsize,
mutation=mutation,
crossover_rate=crossover_rate,
tol=tol,
)
return solution, message
[docs] def run_cmaes(
self,
maxiter: int = 1000,
popsize: int = 50,
sigma: float = 0.6,
verbose: bool = False,
**kwargs: Any,
) -> Any:
"""
Runs the CMA-ES (Covariance Matrix Adaptation Evolution Strategy) algorithm
from `pymoo` to optimize resonance parameters.
This function applies the CMA-ES global optimization method to minimize the
objective function based on the given impedance data and parameter bounds.
The resulting optimized parameters are stored for further analysis or refinement.
Parameters
----------
maxiter : int, optional
Maximum number of iterations for the CMA-ES solver. Default is 1000.
popsize : int, optional
Population size for the CMA-ES algorithm. Default is 50.
sigma : float, optional
Initial standard deviation for the sampling distribution. Default is 0.1.
**kwargs : dict, optional
Additional arguments passed to the `pymoo.CMAES` solver.
Returns
-------
res : pymoo.optimize.Result
The optimization result object containing the solution and metadata.
Notes
-----
- Uses `Solvers.run_pymoo_cmaes_solver()` to execute the optimization.
- The optimized parameters are stored in `self.evolutionParameters`.
- Calls `self.display_resonator_parameters()` to present the results.
"""
objective_function = partial(
self.objectiveFunction,
fitFunction=self.fitFunction,
x=self.x_data,
y=self.y_data,
)
solution, message, res = Solvers.run_pymoo_cmaes_solver(
self.parameterBounds,
objective_function,
sigma=sigma,
maxiter=maxiter,
popsize=popsize,
verbose=verbose,
**kwargs,
)
self.evolutionParameters = solution
self.evolutionParametersUncertainties = get_uncertainties(
self.evolutionParameters,
self.fitFunction,
self.x_data,
self.y_data,
)
self.warning = message
self._warn_large_uncertainties(
self.evolutionParameters,
self.evolutionParametersUncertainties,
)
self.display_resonator_parameters(
params=self.evolutionParameters,
to_markdown=False,
uncertainties=self.evolutionParametersUncertainties,
)
return res
[docs] def run_differential_evolution(
self,
maxiter: int = 2000,
popsize: int = 15,
mutation: tuple[float, float] = (0.1, 0.5),
crossover_rate: float = 0.8,
tol: float = 0.01,
solver: str = "scipy",
) -> None:
"""
Runs the differential evolution (DE) algorithm to estimate optimal
resonance parameters.
This function applies a global optimization technique using a DE solver
to determine the best-fitting parameters for the given impedance data.
The results can later be refined using a local minimization algorithm.
Parameters
----------
maxiter : int, optional
Maximum number of iterations for the DE solver. Default is 2000.
popsize : int, optional
Population size for the DE algorithm. Default is 15.
mutation : tuple of float, optional
Range of mutation factors controlling parameter variation.
Default is (0.1, 0.5).
crossover_rate : float, optional
Probability of recombining individuals in the DE algorithm.
Default is 0.8.
tol : float, optional
Convergence tolerance for stopping criteria. Default is 0.01.
solver : str, optional
Specifies the DE solver to use. Valid options are:
- `"scipy"`: Uses SciPy's built-in DE solver.
- `"pyfde"`: Uses `pyfde`, an alternative DE implementation.
- `"pyfde_jade"`: Uses JADE, a self-adaptive DE variant.
Default is `"scipy"`.
Notes
-----
- Uses `generate_Initial_Parameters()` to perform the differential
evolution process.
- The optimized parameters are stored in `self.evolutionParameters`.
- Calls `self.display_resonator_parameters()` to present the estimated
parameters.
Returns
-------
None
The optimized parameters are stored in `self.evolutionParameters`.
"""
evolutionParameters, warning = self.generate_Initial_Parameters(
self.parameterBounds,
self.objectiveFunction,
self.fitFunction,
self.x_data,
self.y_data,
maxiter=maxiter,
popsize=popsize,
mutation=mutation,
crossover_rate=crossover_rate,
tol=tol,
solver=solver,
# workers=workers,
# vectorized=vectorized,
# iteration_convergence=iteration_convergence,
)
self.evolutionParameters = evolutionParameters
self.evolutionParametersUncertainties = get_uncertainties(
self.evolutionParameters,
self.fitFunction,
self.x_data,
self.y_data,
)
self.warning = warning
self._warn_large_uncertainties(
self.evolutionParameters,
self.evolutionParametersUncertainties,
)
self.display_resonator_parameters(
params=self.evolutionParameters,
to_markdown=False,
uncertainties=self.evolutionParametersUncertainties,
)
[docs] def run_minimization_algorithm(
self,
margin: float | Sequence[float] = [0.1, 0.1, 0.1],
method: str = "Nelder-Mead",
) -> None:
"""
Runs a minimization algorithm to refine resonance parameters.
This function refines the parameters obtained from the Differential
Evolution (DE) algorithm by using a local optimization method.
If the DE algorithm has not been run, it directly minimizes the
objective function using initial parameter bounds.
Each parameter is allowed to vary within a specified margin, where:
- Rs values use `margin[0]`
- Q values use `margin[1]`
- fres values use `margin[2]`
Parameters
----------
margin : float or list of float, optional
A list of three values specifying the relative margins for Rs, Q, and fres.
Each parameter is allowed to vary by ±(margin * value).
Default is [0.1, 0.1, 0.1].
method : str, optional
Optimization method for `scipy.optimize.minimize`. Default is 'Nelder-Mead'.
Notes
-----
- The optimization is constrained within `minimizationBounds`, which
are computed using `margin` and the current `evolutionParameters`.
- If the DE algorithm has not been run, the function initializes parameters
using `self.parameterBounds` and minimizes the objective function.
- The minimization results are stored in `self.minimizationParameters`.
- Calls `self.display_resonator_parameters()` to display the refined parameters.
Returns
-------
None
The refined parameters are stored in `self.minimizationParameters`.
"""
print("Method for minimization : " + method)
objective_function = partial(
self.objectiveFunction,
fitFunction=self.fitFunction,
x=self.x_data,
y=self.y_data,
)
if type(margin) is float:
margin = [margin] * 3
if self.evolutionParameters is not None:
# Apply different margins based on parameter type (Rs, Q, fres)
minimizationBounds = [
sorted(((1 - margin[i % 3]) * p, (1 + margin[i % 3]) * p))
for i, p in enumerate(self.evolutionParameters)
]
minimizationParameters = minimize(
objective_function,
x0=self.evolutionParameters,
bounds=minimizationBounds,
tol=1, # empiric value, documentation is cryptic
method=method,
options={
"maxiter": self.N_resonators * 1000,
"maxfev": self.N_resonators * 1000,
"disp": False,
"adaptive": True,
},
)
else:
print("Differential Evolution algorithm not run, minimization only")
minimizationParameters = minimize(
objective_function,
x0=np.mean(self.parameterBounds, axis=1),
bounds=self.parameterBounds,
method=method,
tol=1,
options={
"maxiter": self.N_resonators * 5000,
"maxfev": self.N_resonators * 5000,
"disp": False,
"adaptive": True,
},
)
self.minimizationParameters = minimizationParameters.x
self.minimizationParametersUncertainties = get_uncertainties(
self.minimizationParameters,
self.fitFunction,
self.x_data,
self.y_data,
)
self._warn_large_uncertainties(
self.minimizationParameters,
self.minimizationParametersUncertainties,
)
self.display_resonator_parameters(
params=self.minimizationParameters,
to_markdown=False,
uncertainties=self.minimizationParametersUncertainties,
)
[docs] def display_resonator_parameters(
self,
params: ArrayLike | None = None,
to_markdown: bool = False,
uncertainties: ArrayLike | None = None,
display_uncertainties: bool = True,
) -> None:
"""
Displays resonance parameters in a formatted table using ASCII characters.
Args:
solution: A NumPy array of resonator parameters, typically shaped
(n_resonators, 3).
"""
# Use self parameters if not provided as argument.
if params is None:
if self.minimizationParameters is not None:
params = self.minimizationParameters
print("[*] Displaying parameters after minimization")
if uncertainties is None:
uncertainties = self.minimizationParametersUncertainties
elif self.evolutionParameters is not None:
params = self.evolutionParameters
if uncertainties is None:
uncertainties = self.evolutionParametersUncertainties
else:
print("[!] No parameters to display.")
return
flagged_mask = self._compute_flagged_params_mask(params, uncertainties)
if not display_uncertainties:
uncertainties = None
flagged_mask = None
if to_markdown:
print("\n")
print("| Resonator | Rs [Ohm/m or Ohm] | Q | fres [Hz] |")
print("|-----------|------------------|---|-----------|")
params_reshaped = params.reshape(-1, 3)
if uncertainties is not None:
uncert_reshaped = uncertainties.reshape(-1, 3)
for i, (p_row, u_row) in enumerate(
zip(params_reshaped, uncert_reshaped)
):
rs, q, fres = p_row
urs, uq, ufres = u_row
print(
f"| {i + 1} | {rs:.6g} ± {urs:.2g} | {q:.6g} ± {uq:.2g} | {fres:.6g} ± {ufres:.2g} |"
)
else:
for i, parameters in enumerate(params_reshaped):
print(
f"| {i + 1} | {parameters[0]:.6g} | {parameters[1]:.6g} | {parameters[2]:.6g} |"
)
else:
print("\n")
print("-" * 76)
# Print header
header_format = "{:^10}|{:^24}|{:^18}|{:^24}"
print(
header_format.format("Resonator", "Rs [Ohm/m or Ohm]", "Q", "fres [Hz]")
)
print("-" * 76)
# Print data
params_reshaped = params.reshape(-1, 3)
if uncertainties is not None:
uncert_reshaped = uncertainties.reshape(-1, 3)
if flagged_mask is not None:
flagged_reshaped = flagged_mask.reshape(-1, 3)
else:
flagged_reshaped = np.zeros_like(params_reshaped, dtype=bool)
for i, (p_row, u_row) in enumerate(
zip(params_reshaped, uncert_reshaped)
):
rs, q, fres = p_row
urs, uq, ufres = u_row
rs_str = f"{rs:.2e} ± {urs:.1e}"
q_str = f"{q:.2f} ± {uq:.2f}"
fres_str = f"{fres:.2e} ± {ufres:.1e}"
rs_cell = f"{rs_str:^24}"
q_cell = f"{q_str:^18}"
fres_cell = f"{fres_str:^24}"
# Highlight cells whose relative uncertainty exceeds threshold.
if flagged_reshaped[i, 0]:
rs_cell = f"\033[31m{rs_cell}\033[0m"
if flagged_reshaped[i, 1]:
q_cell = f"\033[31m{q_cell}\033[0m"
if flagged_reshaped[i, 2]:
fres_cell = f"\033[31m{fres_cell}\033[0m"
print(f"{i + 1:^10d}|{rs_cell}|{q_cell}|{fres_cell}")
else:
for i, parameters in enumerate(params_reshaped):
print(
f"{i + 1:^10d}|{parameters[0]:^24.2e}|{parameters[1]:^18.2f}|{parameters[2]:^24.3e}"
)
print("-" * 76)
def _compute_flagged_params_mask(
self,
params: ArrayLike | None,
uncertainties: ArrayLike | None,
) -> np.ndarray | None:
"""Compute boolean mask for parameters above relative uncertainty threshold."""
if params is None or uncertainties is None:
return None
threshold = getattr(self, "uncertainty_warning", None)
if threshold is None:
return None
params = np.asarray(params, dtype=float)
uncertainties = np.asarray(uncertainties, dtype=float)
if params.shape != uncertainties.shape:
return None
with np.errstate(divide="ignore", invalid="ignore"):
rel_unc = np.abs(uncertainties / params)
return np.isfinite(rel_unc) & (rel_unc >= threshold)
def _warn_large_uncertainties(
self,
params: ArrayLike | None,
uncertainties: ArrayLike | None,
) -> None:
"""Print a warning if any parameter has a large relative uncertainty.
The relative uncertainty is defined as ``abs(sigma / value)``. If this
exceeds ``self.uncertainty_warning`` for any
parameter, a concise warning is printed.
"""
mask = self._compute_flagged_params_mask(params, uncertainties)
self.flagged_params = mask
if mask is None:
return
threshold = getattr(self, "uncertainty_warning", None)
n_flagged = int(mask.sum())
if n_flagged == 0:
return
params = np.asarray(params, dtype=float)
uncertainties = np.asarray(uncertainties, dtype=float)
with np.errstate(divide="ignore", invalid="ignore"):
rel_unc = np.abs(uncertainties / params)
max_rel = float(np.nanmax(rel_unc[mask]))
print(
f"[!] Warning: {n_flagged} parameter(s) have relative uncertainty >= {threshold:.2f} (max {max_rel:.2f})."
)
[docs] def get_wake(
self,
time_data: ArrayLike | None = None,
use_minimization: bool = True,
) -> np.ndarray:
# Check for time data
if time_data is None:
if self.time_data is None:
raise AttributeError("Provide time data array")
time_data = self.time_data
else:
if self.time_data is None:
self.time_data = time_data
# Which pars to use
if use_minimization and self.minimizationParameters is not None:
pars = self.minimizationParameters
else:
pars = self.evolutionParameters
# Which plane and formula
if self.plane == "longitudinal" and self.N_resonators > 1:
wake_data = wak.n_Resonator_longitudinal_wake(time_data, pars)
elif self.plane == "transverse" and self.N_resonators > 1:
wake_data = wak.n_Resonator_transverse_wake(time_data, pars)
elif self.plane == "longitudinal" and self.N_resonators == 1:
wake_data = wak.Resonator_longitudinal_wake(time_data, pars)
elif self.plane == "transverse" and self.N_resonators == 1:
wake_data = wak.Resonator_transverse_wake(time_data, pars)
if self.wake_data is None:
self.wake_data = wake_data
return wake_data
[docs] def get_wake_potential(
self,
time_data: ArrayLike | None = None,
sigma: float | None = None,
use_minimization: bool = True,
) -> np.ndarray:
# Check for time data
if time_data is None:
if self.time_data is None:
raise AttributeError("Provide time data array")
time_data = self.time_data
else:
if self.time_data is None:
self.time_data = time_data
# Check for sigma
if sigma is None:
if self.sigma is None:
self.sigma = 1e-10
sigma = self.sigma
print(f"[!] sigma not specified, using sigma = {sigma:.2e} s")
# Which pars to use
if use_minimization and self.minimizationParameters is not None:
pars = self.minimizationParameters
else:
pars = self.evolutionParameters
# Which plane and formula - TODO check normalization
if self.plane == "longitudinal" and self.N_resonators > 1:
wake_potential_data = wak.n_Resonator_longitudinal_wake_potential(
time_data, pars, sigma=sigma
)
elif self.plane == "transverse" and self.N_resonators > 1:
wake_potential_data = wak.n_Resonator_transverse_wake_potential(
time_data, pars, sigma=sigma
)
elif self.plane == "longitudinal" and self.N_resonators == 1:
wake_potential_data = wak.Resonator_longitudinal_wake_potential(
time_data, pars, sigma=sigma
)
elif self.plane == "transverse" and self.N_resonators == 1:
wake_potential_data = wak.Resonator_transverse_wake_potential(
time_data, pars, sigma=sigma
)
return wake_potential_data
[docs] def get_impedance_from_fitFunction(
self,
frequency_data: ArrayLike | None = None,
use_minimization: bool = True,
) -> np.ndarray:
# Check for frequency data
if frequency_data is None:
if self.frequency_data is None:
raise AttributeError("Provide frequency data array")
frequency_data = self.frequency_data
else:
if self.frequency_data is None:
self.frequency_data = frequency_data
# Which pars to use
if use_minimization and self.minimizationParameters is not None:
pars = self.minimizationParameters
else:
pars = self.evolutionParameters
impedance_data = self.fitFunction(frequency_data, pars)
return impedance_data
[docs] def get_impedance(
self,
frequency_data: ArrayLike | None = None,
use_minimization: bool = True,
wakelength: float | None = None,
) -> np.ndarray:
# Check for frequency data
if frequency_data is None:
if self.frequency_data is None:
raise AttributeError("Provide frequency data array")
frequency_data = self.frequency_data
else:
if self.frequency_data is None:
self.frequency_data = frequency_data
# Which pars to use
if use_minimization and self.minimizationParameters is not None:
pars = self.minimizationParameters
else:
pars = self.evolutionParameters
# Which plane and formula
if self.plane == "longitudinal" and self.N_resonators > 1:
impedance_data = imp.n_Resonator_longitudinal_imp(
frequency_data, pars, wakelength
)
elif self.plane == "transverse" and self.N_resonators > 1:
impedance_data = imp.n_Resonator_transverse_imp(
frequency_data, pars, wakelength
)
elif self.plane == "longitudinal" and self.N_resonators == 1:
impedance_data = imp.Resonator_longitudinal_imp(
frequency_data, pars, wakelength
)
elif self.plane == "transverse" and self.N_resonators == 1:
impedance_data = imp.Resonator_transverse_imp(
frequency_data, pars, wakelength
)
return impedance_data
[docs] def get_impedance_from_fft(
self,
time_data: ArrayLike | None = None,
wake_data: ArrayLike | None = None,
fmax: float = 3e9,
samples: int = 1001,
) -> tuple[np.ndarray, np.ndarray]:
# Check for time data
if time_data is None:
if self.time_data is None:
raise AttributeError("Provide time data array")
time_data = self.time_data
else:
if self.time_data is None:
self.time_data = time_data
wake_data = self.get_wake(self.time_data)
f, Z = compute_fft(
data_time=time_data,
data_wake=wake_data,
fmax=fmax,
samples=samples,
)
# Apply convention
if self.plane == "transverse":
Z *= -1j
elif self.plane == "longitudinal":
Z *= -1.0
return f, Z
[docs] def save_txt(
self,
f_name: str,
x_data: ArrayLike | None = None,
y_data: ArrayLike | None = None,
x_name: str = "X [-]",
y_name: str = "Y [-]",
) -> None:
"""
Saves x and y data to a text file in a two-column format.
This function exports the provided `x_data` and `y_data` to a `.txt` file,
formatting the output with a header that includes custom column names.
Parameters
----------
f_name : str
Name of the output file (with or without the `.txt` extension).
x_data : numpy.ndarray, optional
Array containing x-axis data. If None, the file is not saved.
y_data : numpy.ndarray, optional
Array containing y-axis data. If None, the file is not saved.
x_name : str, optional
Label for the x-axis column in the output file. Default is `"X [-]"`.
y_name : str, optional
Label for the y-axis column in the output file. Default is `"Y [-]"`.
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 [a.u.]")
The saved file will look like:
Time [s] Amplitude
--------------------------------
0.00 0.00
2.50 0.59
5.00 -0.99
7.50 0.94
10.00 -0.54
"""
if not f_name.endswith(".txt"):
f_name += ".txt"
if x_data is not None and y_data is not None:
np.savetxt(
f_name + ".txt",
np.c_[x_data, y_data],
header=" " + x_name + " " * 20 + y_name + "\n" + "-" * 48,
)
else:
print("txt not saved, please provide x_data and y_data")
[docs] def read_txt(
self,
txt: str,
skiprows: int = 2,
delimiter: str | None = None,
usecols: Sequence[int] | None = None,
as_dict: bool = False,
) -> dict[Any, np.ndarray] | tuple[np.ndarray, np.ndarray]:
"""
Reads data from an ASCII text file and returns it as a dictionary or tuple.
This function reads a structured text file containing numerical data,
where the first line is expected to contain column headers. It attempts
to parse the headers and assign them as dictionary keys. If headers are
not properly formatted, integer indices are used instead.
Parameters
----------
txt : str
Path to the text file to read.
skiprows : int, optional
Number of initial rows to skip before reading the data. Default is 2.
delimiter : str, optional
Character used to separate values in the file. If None, whitespace is used.
usecols : list of int, optional
Indices of columns to read from the file. If None, all columns are read.
as_dict : bool, optional
If True, returns a dictionary where keys are the column headers
(if available) or integers (if headers are missing).
If False, returns `x_data` and `y_data`as separate arrays. Default is False.
Returns
-------
dict or tuple
- If `as_dict=True`, returns a dictionary `{header: column_data}`.
- If `as_dict=False`, returns `(x_data, y_data)`, where:
- `x_data` is the first column of data.
- `y_data` is the second column of data.
Notes
-----
- If an error occurs while reading the file, the function attempts to reload
the data assuming complex numbers (`dtype=complex`).
- If column headers are missing or unreadable, integer indices `[0, 1, ...]`
are assigned as dictionary keys.
- The first line of the file is expected to contain column headers.
Examples
--------
Read a file and return as a dictionary:
>>> data = read_txt("data.txt", as_dict=True)
>>> print(data.keys())
# Example output: {'Time[s]': array([...]), 'Amplitude': array([...])}
Read a file and return x and y data separately:
>>> x, y = read_txt("data.txt")
>>> print(x.shape, y.shape)
Example of an expected file format:
```
# Time[s] Amplitude
------------------------
0.00 0.00
2.50 0.59
5.00 -0.99
```
"""
try:
load = np.loadtxt(
txt, skiprows=skiprows, delimiter=delimiter, usecols=usecols
)
except Exception:
load = np.loadtxt(
txt,
skiprows=skiprows,
delimiter=delimiter,
usecols=usecols,
dtype=complex,
)
try: # keys == header names
with open(txt) as f:
header = f.readline()
header = header.replace(" ", "")
header = header.replace("#", "")
header = header.replace("\n", "")
header = header.split("]")
d = {}
for i in range(len(load[0, :])):
d[header[i] + "]"] = load[:, i]
except Exception: # keys == int 0, 1, ...
d = {}
for i in range(len(load[0, :])):
d[i] = load[:, i]
if as_dict:
return d
else:
x_data = list(d.values())[0]
y_data = list(d.values())[1]
return x_data, y_data