diff --git a/collective_effects/utilities.py b/collective_effects/utilities.py
index 8f792b2ce17194998a0aea4c524dbcf6135b7d91..11803ea3092ffd52104506a5496a459579a992ee 100644
--- a/collective_effects/utilities.py
+++ b/collective_effects/utilities.py
@@ -9,7 +9,6 @@ collective_effects module.
import pandas as pd
import numpy as np
-import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from mbtrack2.collective_effects.wakefield import Impedance
diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index 13515477a3bb0db8ae77bb29b73f241efe7c4680..3895111b08909391a5a61888c1cc2e8d2d4be845 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -19,7 +19,9 @@ from scipy.integrate import trapz
from scipy.constants import c
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from mbtrack2.tracking.element import Element
-from mbtrack2.collective_effects.utilities import beam_spectrum, gaussian_bunch_spectrum, beam_loss_factor
+from mbtrack2.collective_effects.utilities import (beam_spectrum, gaussian_bunch_spectrum,
+ beam_loss_factor, spectral_density,
+ effective_impedance)
class ComplexData:
"""
@@ -193,6 +195,26 @@ class ComplexData:
class WakeFunction(ComplexData):
"""
Define a WakeFunction object based on a ComplexData object.
+
+ Parameters
+ ----------
+ variable : array-like
+ Time domain structure of the wake function in [s].
+ function : array-like
+ Wake function values in [V/C].
+ wake_type : str, optinal
+ Type of the wake function: "long", "xdip", "xquad", ...
+
+ Attributes
+ ----------
+ data : DataFrame
+ wake_type : str
+
+ Methods
+ -------
+ from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
+ Compute a wake function from a wake potential file and load it to the
+ WakeFunction object.
"""
def __init__(self,
@@ -260,14 +282,22 @@ class WakeFunction(ComplexData):
def wake_type(self, value):
self._wake_type = value
- def wakefunction_from_imp(self, imp_obj, nout=None, trim=False):
+ def from_wakepotential(self, file_name, bunch_length, bunch_charge,
+ freq_lim, nout=None, trim=False):
"""
- Compute a wake function from an Impedance object and load it to
- the WakeFunction object.
+ Compute a wake function from a wake potential file and load it to the
+ WakeFunction object.
Parameters
----------
- imp_obj : impedance object
+ file_name : str
+ Text file that contains wake potential data.
+ bunch_length : float
+ Spatial bunch length in [m].
+ bunch_charge : float
+ Bunch charge in [C].
+ freq_lim : float
+ The maximum frequency for calculating the impedance [Hz].
nout : int, optional
Length of the transformed axis of the output. If None, it is taken
to be 2*(n-1) where n is the length of the input. If nout > n, the
@@ -279,68 +309,42 @@ class WakeFunction(ComplexData):
If False, the original result is preserved.
If a float is given, the pseudo wake function is trimmed from
time <= trim to 0.
- """
-
- Z0 = (imp_obj.data['real'] + imp_obj.data['imag']*1j)
- Z = Z0[~np.isnan(Z0)]
-
- if imp_obj.impedance_type != "long":
- Z = Z / 1j
-
- freq = Z.index
- fs = ( freq[-1] - freq[0] ) / len(freq)
- sampling = freq[1] - freq[0]
-
- if nout is None:
- nout = len(Z)
-
- time_array = sc.fft.fftfreq(nout, sampling)
- Wlong_raw = sc.fft.irfft(np.array(Z), n=nout, axis=0) * nout * fs
-
- time = sc.fft.fftshift(time_array)
- Wlong = sc.fft.fftshift(Wlong_raw)
-
- if trim is not False:
- i_neg = np.where(time<trim)[0]
- Wlong[i_neg] = 0
-
- super().__init__(variable=time, function=Wlong)
- self.data.index.name = "time [s]"
-
- def long_wakefunction(self, Wp_filename, bunch_length, bunch_charge,
- nout=None, freq_lim=200e9, trim=False):
- """
- Compute wake function from wake potential and load it to the WakeFunction
- object.
-
- Parameters
- ----------
- Wp_filename : str
- Text file that contains wake potential data.
- bunch_length : float
- Spatial bunch length in [m].
- bunch_charge : float
- Bunch charge in [C].
- nout : int, optional
- Length of the transformed axis of the output. If None, it is taken
- to be 2*(n-1) where n is the length of the input. If nout > n, the
- input is padded with zeros. If nout < n, the inpu it cropped.
- Note that, for any nout value, nout//2+1 input points are required.
- freq_lim : float, optional
- The maximum frequency for calculating the impedance [Hz].
"""
imp = Impedance()
- imp.imp_from_wakepotential(Wp_filename=Wp_filename,
- freq_lim=freq_lim,
- bunch_length=bunch_length,
- bunch_charge=bunch_charge)
- self.wakefunction_from_imp(imp, nout=nout, trim=trim)
+ imp.from_wakepotential(file_name=file_name, bunch_length=bunch_length,
+ bunch_charge=bunch_charge, freq_lim=freq_lim)
+ wf = imp.to_wakefunction(nout=nout, trim=trim)
+ self.__init__(variable=wf.data.index, function=wf.data["real"])
class Impedance(ComplexData):
"""
Define an Impedance object based on a ComplexData object.
+
+ Parameters
+ ----------
+ variable : array-like
+ Frequency domain structure of the impedance in [Hz].
+ function : array-like
+ Impedance values in [Ohm].
+ impedance_type : str, optinal
+ Type of the impedance_type: "long", "xdip", "xquad", ...
+
+ Attributes
+ ----------
+ data : DataFrame
+ impedance_type : str
+
+ Methods
+ -------
+ from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
+ Compute impedance from wake potential data and load it to the Impedance
+ object.
+ loss_factor(sigma)
+ Compute the loss factor or the kick factor for a Gaussian bunch.
+ to_wakefunction()
+ Return a WakeFunction object from the impedance data.
"""
def __init__(self,
@@ -506,17 +510,23 @@ class Impedance(ComplexData):
return kloss
- def imp_from_wakepotential(self, Wp_filename, axis0='dist',
- axis0_scale=1e-3, axis1_scale=1e-12, freq_lim=200e9,
- bunch_length=1e-3, bunch_charge=1.44e-9):
+ def from_wakepotential(self, file_name, bunch_length, bunch_charge,
+ freq_lim, axis0='dist', axis0_scale=1e-3,
+ axis1_scale=1e-12):
"""
Compute impedance from wake potential data and load it to the Impedance
object.
Parameters
----------
- Wp_filename : str
+ file_name : str
Text file that contains wake potential data.
+ freq_lim : float
+ The maximum frequency for calculating the impedance [Hz].
+ bunch_length : float
+ Electron bunch lenth [m].
+ bunch_charge : float
+ Total bunch charge [C].
axis0 : {'dist', 'time'}, optional
Viariable of the data file's first column. Use 'dist' for spacial
distance. Otherwise, 'time' for temporal distance.
@@ -525,16 +535,10 @@ class Impedance(ComplexData):
or to second if axis0 = 'time'.
axis1_scale : float, optional
Scale of the wake potential (in second column) with respect to V/C.
- freq_lim : float, optional
- The maximum frequency for calculating the impedance [Hz].
- bunch_length : float, optional
- Electron bunch lenth [m].
- bunch_charge : float, optional
- Total bunch charge [C].
"""
- s, wp0 = np.loadtxt(Wp_filename, unpack=True)
+ s, wp0 = np.loadtxt(file_name, unpack=True)
if axis0 == 'dist':
tau = s*axis0_scale/c
elif axis0 == 'time':
@@ -564,6 +568,51 @@ class Impedance(ComplexData):
super().__init__(variable=freq_trun, function=imp)
self.data.index.name = "frequency [Hz]"
+
+ def to_wakefunction(self, nout=None, trim=False):
+ """
+ Return a WakeFunction object from the impedance data.
+
+ Parameters
+ ----------
+ nout : int, optional
+ Length of the transformed axis of the output. If None, it is taken
+ to be 2*(n-1) where n is the length of the input. If nout > n, the
+ input is padded with zeros. If nout < n, the inpu it cropped.
+ Note that, for any nout value, nout//2+1 input points are required.
+ trim : bool or float, optional
+ If True, the pseudo wake function is trimmed at time=0 and is forced
+ to zero where time<0.
+ If False, the original result is preserved.
+ If a float is given, the pseudo wake function is trimmed from
+ time <= trim to 0.
+ """
+
+ Z0 = (self.data['real'] + self.data['imag']*1j)
+ Z = Z0[~np.isnan(Z0)]
+
+ if self.impedance_type != "long":
+ Z = Z / 1j
+
+ freq = Z.index
+ fs = ( freq[-1] - freq[0] ) / len(freq)
+ sampling = freq[1] - freq[0]
+
+ if nout is None:
+ nout = len(Z)
+
+ time_array = sc.fft.fftfreq(nout, sampling)
+ Wlong_raw = sc.fft.irfft(np.array(Z), n=nout, axis=0) * nout * fs
+
+ time = sc.fft.fftshift(time_array)
+ Wlong = sc.fft.fftshift(Wlong_raw)
+
+ if trim is not False:
+ i_neg = np.where(time<trim)[0]
+ Wlong[i_neg] = 0
+
+ wf = WakeFunction(variable=time, function=Wlong)
+ return wf
class WakeField:
"""