diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index 66a0683d07c2801848cdc6461c56c692065395f7..89b6cb3d48d0518eda37cc4ac4174f4cc89074b0 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -260,105 +260,74 @@ class WakeFunction(ComplexData):
def wake_type(self, value):
self._wake_type = value
- def deconvolve_wp(self, Wp_filename, bunchlength=1e-3, nout=None,
- stopfreq=200e9):
+ def wakefunction_from_imp(self, imp_obj, nout=None, trim_zero=True):
"""
- Compute wake function from wake potential by deconvolution method.
+ Compute a wake function from an Impedance object and load it to
+ the WakeFunction object.
Parameters
----------
- Wp_filename : str
- Text file that contains wake potential data.
- bunchlength : float, optional
- Spatial bunch length [m].
+ imp_obj : impedance object
nout : int, optional
- Length of the transformed axis of the output. For n output points,
- n//2+1 input points are necessary. If the input is longer than this,
- it is cropped. If it is shorter than this, it is padded with zeros.
- If n is not given, it is taken to be 2*(m-1) where m is the length
- of the input along the axis specified by axis.
- stopfreq : float, optional
- The maximum frequency for calculating the impedance.
-
- Returns
- -------
- Wake function object indexed with time.
-
+ 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_zero : bool, 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.
"""
- # FT of wake potential
- dist, wp0 = np.loadtxt(Wp_filename, unpack=True)
- tau_wp = dist*1e-3/c
- wp = wp0*1e12
- sampling_wp = tau_wp[1] - tau_wp[0]
- freq_wp = sc.fft.rfftfreq(len(tau_wp), sampling_wp)
- dtau_wp = (tau_wp[-1]-tau_wp[0])/len(tau_wp)
- dft_wp = sc.fft.rfft(wp) * dtau_wp
-
- # FT of analytical bunch profile and analytical impedance
- sigma = bunchlength/c
- mu = tau_wp[0]
- bunch_charge = 1.44e-9
-
- index_lim = freq_wp < stopfreq
- freq_wp_trun = freq_wp[index_lim]
- dft_wp_trun = dft_wp[index_lim]
+ Z0 = (imp_obj.data['real'] + imp_obj.data['imag']*1j)
+ Z = Z0[~np.isnan(Z0)]
+ freq = Z.index
+ fs = ( freq[-1] - freq[0] ) / len(freq)
+ sampling = freq[1] - freq[0]
- dft_rho_eq_trun = np.exp(-0.5*(sigma*2*np.pi*freq_wp_trun)**2 + \
- 1j*mu*2*np.pi*freq_wp_trun)*bunch_charge
- imp = dft_wp_trun/dft_rho_eq_trun*-1*bunch_charge
-
- # Wake function from impedance
if nout is None:
- nout = len(imp)
-
- fs = (freq_wp_trun[-1] - freq_wp_trun[0]) / len(freq_wp_trun)
- time_array = sc.fft.fftfreq(nout, fs)
- Wlong_raw = sc.fft.irfft(imp, n=nout) * nout * fs
+ 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_zero is True:
+ i_neg = np.where(time<0)[0]
+ Wlong[i_neg] = 0
+
super().__init__(variable=time, function=Wlong)
- self.data.index.name = "time [s]"
+ self.data.index.name = "time [s]"
- def deconvolve_imp(self, imp_obj, divided_by=1, nout=None):
+ def long_wakefunction(self, Wp_filename, bunch_length=1e-3,
+ bunch_charge = 1.44e-9, nout=None, freq_lim=200e9):
"""
- Compute wake function from impedance by deconvolution method.
+ Compute wake function from wake potential and load it to the WakeFunction
+ object.
Parameters
----------
- imp_obj : impedance object
- divided_by : int, optional
- Total number of elements that generate the impedance.
+ Wp_filename : str
+ Text file that contains wake potential data.
+ bunch_length : float, optional
+ Spatial bunch length [m].
nout : int, optional
- Length of the transformed axis of the output. For n output points,
- n//2+1 input points are necessary. If the input is longer than this,
- it is cropped. If it is shorter than this, it is padded with zeros.
- If n is not given, it is taken to be 2*(m-1) where m is the length of
- the input along the axis specified by axis.
-
- Returns
- -------
- Wake function object indexed with time.
+ 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].
"""
- Z = (imp_obj.data['real'] + imp_obj.data['imag']*1j) / divided_by
- fs = (Z.index[-1]-Z.index[0]) / len(Z.index)
- sampling = Z.index[1] - Z.index[0]
- if nout is not None:
- Wlong_raw = sc.fft.irfft(np.array(Z), n=nout, axis=0) * nout * fs
- time_array = sc.fft.fftfreq(nout, sampling)
- else:
- Wlong_raw = sc.fft.irfft(np.array(Z), n=Z.size, axis=0) *Z.size*fs
- time_array = sc.fft.fftfreq(Z.size, sampling)
-
- Wlong = sc.fft.fftshift(Wlong_raw)
- time = sc.fft.fftshift(time_array)
-
- super().__init__(variable=time, function=Wlong)
- self.data.index.name = "time [s]"
+ 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)
class Impedance(ComplexData):
"""
@@ -528,11 +497,12 @@ class Impedance(ComplexData):
return kloss
- def get_impedance(self, Wp_filename, axis0='dist', axis0_scale=1e-3,
- axis1_scale=1e-12, bunch_length=1e-3, bunch_charge=1.44e-9,
- stopfreq=200e9):
+ 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):
"""
- Compute impedance from wake potential data.
+ Compute impedance from wake potential data and load it to the Impedance
+ object.
Parameters
----------
@@ -543,21 +513,18 @@ class Impedance(ComplexData):
distance. Otherwise, 'time' for temporal distance.
axis0_scale : float, optional
Scale of the first column with respect to meter if axis0 = 'dist',
- or second if axis0 = 'time'.
+ 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].
- stopfreq : float, optional
- The maximum frequency for calculating the impedance.
-
- Returns
- -------
- Impedance object indexed with frequency.
"""
+
s, wp0 = np.loadtxt(Wp_filename, unpack=True)
if axis0 == 'dist':
tau = s*axis0_scale/c
@@ -566,25 +533,25 @@ class Impedance(ComplexData):
else:
raise TypeError('Type of axis 0 not recognized.')
- wp = wp0 * axis1_scale
+ wp = wp0 / axis1_scale
# FT of wake potential
sampling = tau[1] - tau[0]
freq = sc.fft.rfftfreq(len(tau), sampling)
dtau = (tau[-1]-tau[0])/len(tau)
- dft = sc.fft.rfft(wp) * dtau
+ dft_wake = sc.fft.rfft(wp) * dtau
# FT of analytical bunch profile and analytical impedance
sigma = bunch_length/c
mu = tau[0]
- index_lim = freq < stopfreq
- freq_trun = freq[index_lim]
- dft_trun = dft[index_lim]
+ i_limit = freq < freq_lim
+ freq_trun = freq[i_limit]
+ dft_wake_trun = dft_wake[i_limit]
- dft_rho_eq_trun = np.exp(-0.5*(sigma*2*np.pi*freq_trun)**2 + \
+ dft_rho_trun = np.exp(-0.5*(sigma*2*np.pi*freq_trun)**2 + \
1j*mu*2*np.pi*freq_trun)*bunch_charge
- imp = dft_trun/dft_rho_eq_trun*-1*bunch_charge
+ imp = dft_wake_trun/dft_rho_trun*-1*bunch_charge
super().__init__(variable=freq_trun, function=imp)
self.data.index.name = "frequency [Hz]"
@@ -741,7 +708,8 @@ class Wakefield:
class Resonator(Wakefield):
"""
- Compute analytical wake function of a resonator and return a Resonator object.
+ Compute an analytical wake function and an analytical impedance of a
+ resonator [1] and load it to the Resonator object.
Parameter
---------
@@ -752,56 +720,72 @@ class Resonator(Wakefield):
Resonator resonance frequency in [Hz]. The default value is 10e9 Hz.
R_shunt : float, optional
Resonator shunt impedance in [Ohm]. The default value is 100 Ohm.
+ n_wake : int, optional
+ Number of points for wake function calcuation.
+ n_imp : int, optional
+ Number of points for impedance calculation.
+ imp_freq_lim : float, optional
+ The frequency limit for impedance calculation [Hz].
+
+ References
+ ----------
+ [1] B. W. Zotter and S. A. Kheifets, "Impedances and Wakes in High-Energy
+ Particle Ac-celerators", Eq. (3.10) and (3.15), pp.51-53.
+
"""
- def __init__(self, ring, Q_factor=1, f_res=10e9, R_shunt=100):
+ def __init__(self, ring, Q_factor=1, f_res=10e9, R_shunt=100,
+ n_wake=1000, n_imp=1000, imp_freq_lim=50e9):
super().__init__()
self.ring = ring
- self.Q_factor = Q_factor
+ self.Q = Q_factor
self.omega_res = 2*np.pi*f_res
- self.R_shunt = R_shunt
+ self.R = R_shunt
+ self.n_wake = n_wake
+ self.n_imp = n_imp
+ self.f_stop = imp_freq_lim
if Q_factor >= 0.5:
- self.Q_factor_p = np.sqrt(self.Q_factor**2 - 0.25)
- self.omega_res_p = (self.omega_res*self.Q_factor_p)/self.Q_factor
+ self.Q_p = np.sqrt(self.Q**2 - 0.25)
else:
- self.Q_factor_pp = np.sqrt(0.25 - self.Q_factor**2)
- self.omega_res_p = (self.omega_res*self.Q_factor_pp)/self.Q_factor
+ self.Q_p = np.sqrt(0.25 - self.Q**2)
+
+ self.omega_res_p = (self.omega_res*self.Q_p)/self.Q
+
+ self.get_wakefunction()
+ self.long_impedance()
def init_timestop(self):
"""
Calculate an approximate position where the wake field depletes.
"""
- self.timestop = round(np.log(1000)/self.omega_res*2*self.Q_factor, 12)
+ self.timestop = round(np.log(1000)/self.omega_res*2*self.Q, 12)
def time_array(self):
- self.tau = np.linspace(start=-30e-12, stop=self.timestop, num=1000)
+ self.tau = np.linspace(start=0, stop=self.timestop, num=self.n_wake)
def long_wakefunction(self):
- tau_neg = self.tau[np.where(self.tau<0)[0]]
- tau_pos = self.tau[np.where(self.tau>=0)[0]]
-
- wlong_tau_neg = np.zeros((len(tau_neg),))
-
- if self.Q_factor >= 0.5:
- wlong_tau_pos = (self.omega_res*self.R_shunt/self.Q_factor)*\
- np.exp(-self.omega_res*tau_pos/(2*self.Q_factor))*\
- (np.cos(self.omega_res_p*tau_pos)-np.sin(self.omega_res_p*tau_pos)/ \
- (2*self.Q_factor_p))
-
- elif self.Q_factor < 0.5:
- wlong_tau_pos = (self.omega_res*self.R_shunt/self.Q_factor)*\
- np.exp(-self.omega_res*tau_pos/(2*self.Q_factor))*\
- (np.cosh(self.omega_res_p*tau_pos)-np.sinh(self.omega_res_p*tau_pos)/ \
- (2*self.Q_factor_pp))
+ """
+ The equations of resonator longitudinal wake function.
+
+ """
+ if self.Q >= 0.5:
+ self.wlong = (self.omega_res*self.R/self.Q)*\
+ np.exp(-self.omega_res*self.tau/(2*self.Q))*\
+ (np.cos(self.omega_res_p*self.tau)-\
+ np.sin(self.omega_res_p*self.tau)/(2*self.Q_p))
- self.wlong = np.concatenate((wlong_tau_neg, wlong_tau_pos))
+ elif self.Q < 0.5:
+ self.wlong = (self.omega_res*self.R/self.Q)*\
+ np.exp(-self.omega_res*self.tau/(2*self.Q))*\
+ (np.cosh(self.omega_res_p*self.tau)-\
+ np.sinh(self.omega_res_p*self.tau)/(2*self.Q_p))
def check_wake_tail(self):
"""
Checking whether the full wake function is obtained by the calculated
- initial timestop. If not, extend the timestop by 50 ps.
+ initial timestop. If not, the timestop is extended by 50 ps.
"""
@@ -813,6 +797,10 @@ class Resonator(Wakefield):
ratio = np.abs(max(self.wlong)) / np.abs(self.wlong[-6:-1])
def get_wakefunction(self):
+ """
+ Summary of all commands for wake function calculation.
+
+ """
self.init_timestop()
self.time_array()
self.long_wakefunction()
@@ -820,7 +808,25 @@ class Resonator(Wakefield):
wake_func = WakeFunction(variable=self.tau, function=self.wlong)
super().append_to_model(wake_func)
-
+
+ def long_impedance(self):
+ """
+ Calculate the analytical resonator longitudinal impedance.
+
+ """
+
+ freq = np.linspace(start=0, stop=self.f_stop, num=self.n_imp)
+ omega = 2*np.pi*freq
+
+ Re_Z = self.R * \
+ 1/(1 + self.Q**2 * (omega**2-self.omega_res**2)**2 / (omega*self.omega_res)**2 )
+ Im_Z = -1*self.R*self.Q * (omega**2-self.omega_res**2) * \
+ 1/(omega*self.omega_res) * \
+ 1/(1 + self.Q**2 * (omega**2-self.omega_res**2)**2 / (omega*self.omega_res)**2)
+
+ Z = Re_Z + 1j*Im_Z
+ imp = Impedance(variable=freq, function=Z)
+ super().append_to_model(imp)
class ImpedanceModel(Element):
"""
@@ -1608,3 +1614,16 @@ def double_sided_impedance(impedance):
all_data = impedance.data.append(negative_data)
all_data = all_data.sort_index()
impedance.data = all_data
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/tracking/element.py b/tracking/element.py
index 881e693d53e873fc7cfff3631f03dd8ddb7c48a4..bcc1940b11038a3c5fa0c405d1c2ea923bfb6f1b 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -313,7 +313,7 @@ class WakePotential2(Element):
Attributes
----------
rho : array of shape (n_bin, )
- Bunch charge density profile.
+ Bunch charge density profile in the unit [1/s].
Wlong : array
Wake function profile.
Wp : array
@@ -324,7 +324,10 @@ class WakePotential2(Element):
Methods
-------
charge_density(bunch)
- Calculate bunch charge density [A/C]
+ Calculate bunch charge density [1/s].
+ get_wakepotential
+ Compute a wake potential by convolving the bunch profile and the wake
+ function.
plot(self, var, plot_rho=True)
Plotting wakefunction or wake potential. If plot_rho=True, the bunch
charge density is overlayed.
@@ -333,11 +336,10 @@ class WakePotential2(Element):
"""
- def __init__(self, wake_obj, ring, n_bin=65):
+ def __init__(self, ring, wake_obj, n_bin=65):
self.wake_obj = wake_obj
- self.tau = self.wake_obj.data.index
- self.Wlong = self.wake_obj.data['real']
-
+ self.tau0 = self.wake_obj.data.index
+ self.Wlong0 = self.wake_obj.data['real']
self.ring = ring
self.n_bin = n_bin
@@ -349,13 +351,20 @@ class WakePotential2(Element):
self.rho = bunch.charge_per_mp*self.bin_data/(self.bin_size*bunch.charge)
def get_wakepotential(self):
+ dtau = (self.tau0[-1]-self.tau0[0]) / len(self.tau0)
+ first_tau = self.bins[0].mid[0]
+ last_tau = self.tau0[-1]
+ self.tau = np.arange(first_tau, last_tau+dtau, step=dtau)
+
+ self.tau_rho = np.arange(first_tau, self.bins[0].mid[-1]+dtau, step=dtau)
interp_func_rho = interp1d(self.bins[0].mid, self.rho, fill_value=0,
- bounds_error=False)
- rho_interp = interp_func_rho(self.tau)
- peak = np.where(rho_interp == max(rho_interp))[0]
- rho_interp_shift = np.roll(rho_interp, len(rho_interp)//2-peak)
- dtau = (self.tau[-1]-self.tau[0]) / len(self.tau)
- self.Wp = signal.convolve(self.Wlong*-1, rho_interp_shift, mode="same") * dtau
+ bounds_error=False)
+ self.rho_interp = interp_func_rho(self.tau_rho)
+
+ interp_func_wlong = interp1d(self.tau0, self.Wlong0, fill_value=0,
+ bounds_error=False)
+ self.Wlong = interp_func_wlong(self.tau)
+ self.Wp = signal.convolve(self.Wlong*-1, self.rho_interp, mode='same')*dtau
def plot(self, var, plot_rho=True):
"""
@@ -380,7 +389,6 @@ class WakePotential2(Element):
if var == "Wp":
ax.plot(self.tau*1e12, self.Wp*1e-12)
-
ax.set_xlabel("$\\tau$ (ps)")
ax.set_ylabel("W$_p$ (V/pC)")
@@ -391,8 +399,8 @@ class WakePotential2(Element):
if plot_rho is True:
rho_array = np.array(self.rho)
- rho_rescaled = rho_array/max(rho_array)*max(self.Wp)
-
+ dict_plot_rho = {"Wp":self.Wp, "Wlong":self.Wlong}
+ rho_rescaled = rho_array/max(rho_array)*max(dict_plot_rho[var])
ax.plot(self.bins[0].mid*1e12, rho_rescaled*1e-12)
return fig