diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
index f9197bca6bf62f3d9633bb683e2a723bf7f469f7..41bd2d019270b976e91883c1c8f3d8722b93fa8a 100644
--- a/tracking/wakepotential.py
+++ b/tracking/wakepotential.py
@@ -13,17 +13,23 @@ import pandas as pd
from scipy import signal
from scipy.interpolate import interp1d
from mbtrack2.tracking.element import Element
-
+
class WakePotential(Element):
"""
- Compute a wake potential from wake functions by performing a convolution
- with a bunch charge profile.
+ Compute a wake potential from uniformly sampled wake functions by
+ performing a convolution with a bunch charge profile.
+
+ Two different time bases are used. The first one is controled by the n_bin
+ parameter and is used to compute the bunch profile. Then the bunch profile
+ is interpolated on the wake function time base which is used to perform the
+ convolution to get the wake potential.
Parameters
----------
ring : Synchrotron object
wakefield : Wakefield object
- Wakefield object which contains the wake functions to be used.
+ Wakefield object which contains the wake functions to be used. The wake
+ functions must be uniformly sampled!
n_bin : int, optional
Number of bins for constructing the longitudinal bunch profile.
@@ -40,17 +46,24 @@ class WakePotential(Element):
-------
charge_density(bunch)
Compute bunch charge density profile in [1/s].
- dipole_moment(bunch, plane)
+ dipole_moment(bunch, plane, tau0)
Return the dipole moment of the bunch computed on the same time array
- as the bunch profile.
+ as the wake function.
+ prepare_wakefunction(wake_type)
+ Prepare the wake function of a given wake_type to be used for the wake
+ potential computation.
get_wakepotential(bunch, wake_type)
- Return the wake potential computed on the same time array as the bunch
- profile.
+ Return the wake potential computed on the wake function time array
+ limited to the bunch profile.
track(bunch)
Tracking method for the element.
- plot_last_wake(wake_type, plot_rho=True, plot_dipole=False)
+ plot_last_wake(wake_type)
Plot the last wake potential of a given type computed during the last
call of the track method.
+ reference_loss(bunch)
+ Calculate the loss factor and kick factor from the wake potential and
+ compare it to a reference value assuming a Gaussian bunch computed in
+ the frequency domain.
"""
@@ -105,16 +118,21 @@ class WakePotential(Element):
if len(self.tau) != len(self.rho):
self.tau = np.append(self.tau, self.tau[-1] + self.dtau)
- def dipole_moment(self, bunch, plane):
+ self.tau_mean = np.mean(self.tau)
+ self.tau -= self.tau_mean
+
+ def dipole_moment(self, bunch, plane, tau0):
"""
Return the dipole moment of the bunch computed on the same time array
- as the bunch profile.
+ as the wake function.
Parameters
----------
bunch : Bunch object
plane : str
Plane on which the dipole moment is computed, "x" or "y".
+ tau0 : array
+ Time array on which the dipole moment will be interpolated, in [s].
Returns
-------
@@ -128,7 +146,7 @@ class WakePotential(Element):
dipole = dipole/self.profile
dipole[np.isnan(dipole)] = 0
- # Add N values to get same size as tau/rho
+ # Add N values to get same size as tau/profile
if self.n_bin % 2 == 0:
N = int(self.n_bin/2)
dipole = np.append(dipole, np.zeros(N))
@@ -138,13 +156,98 @@ class WakePotential(Element):
dipole = np.append(dipole, np.zeros(N))
dipole = np.insert(dipole, 0, np.zeros(N+1))
- setattr(self, "dipole_" + plane, dipole)
- return dipole
+ # Interpole on tau0 to get the same size as W0
+ interp_dipole = interp1d(self.tau, dipole, fill_value=0,
+ bounds_error=False)
+ dipole0 = interp_dipole(tau0)
+
+ setattr(self, "dipole_" + plane, dipole0)
+ return dipole0
+
+
+ def prepare_wakefunction(self, wake_type):
+ """
+ Prepare the wake function of a given wake_type to be used for the wake
+ potential computation.
+
+ The new time array keeps the same sampling time as given in the
+ WakeFunction definition but is restricted to the bunch profile time
+ array.
+
+ Parameters
+ ----------
+ wake_type : str
+ Type of the wake function to prepare: "Wlong", "Wxdip", ...
+
+ Returns
+ -------
+ tau0 : array
+ Time base of the wake function in [s].
+ dtau0 : float
+ Difference between two points of the wake function time base in
+ [s].
+ W0 : array
+ Wake function array in [V/C] or [V/C/m].
+
+ """
+
+ try:
+ tau0 = getattr(self, "tau0_" + wake_type)
+ dtau0 = getattr(self, "dtau0_" + wake_type)
+ W0 = getattr(self, "W0_" + wake_type)
+ except AttributeError:
+ tau0 = np.array(getattr(self.wakefield, wake_type).data.index)
+ dtau0 = tau0[1] - tau0[0]
+ W0 = np.array(getattr(self.wakefield, wake_type).data["real"])
+
+ # Keep only the wake function on the rho window
+ ind = np.all([min(self.tau[0], 0) < tau0, max(self.tau[-1], 0) > tau0],
+ axis=0)
+ tau0 = tau0[ind]
+ W0 = W0[ind]
+
+ # Check the wake function window for assymetry
+ assym = (np.abs(tau0[-1]) - np.abs(tau0[0])) / dtau0
+ n_assym = int(np.floor(assym))
+ if np.floor(assym) > 1:
+
+ # add at head
+ if np.abs(tau0[-1]) > np.abs(tau0[0]):
+ tau0 = np.arange(tau0[0] - dtau0*n_assym,
+ tau0[-1] + dtau0,
+ dtau0)
+ n_to_add = len(tau0) - len(W0)
+ W0 = np.insert(W0, 0, np.zeros(n_to_add))
+
+ # add at tail
+ elif np.abs(tau0[0]) > np.abs(tau0[-1]):
+ tau0 = np.arange(tau0[0],
+ tau0[-1] + dtau0*(n_assym+1),
+ dtau0)
+ n_to_add = len(tau0) - len(W0)
+ W0 = np.insert(W0, 0, np.zeros(n_to_add))
+
+ # Check is the wf is shorter than rho then add zeros
+ if (tau0[0] > self.tau[0]) or (tau0[-1] < self.tau[-1]):
+ n = max(int(np.ceil((tau0[0] - self.tau[0])/dtau0)),
+ int(np.ceil((self.tau[-1] - tau0[-1])/dtau0)))
+
+ tau0 = np.arange(tau0[0] - dtau0*n,
+ tau0[-1] + dtau0*(n+1),
+ dtau0)
+ W0 = np.insert(W0, 0, np.zeros(n))
+ n_to_add = len(tau0) - len(W0)
+ W0 = np.insert(W0, len(W0), np.zeros(n_to_add))
+
+ setattr(self, "tau0_" + wake_type, tau0)
+ setattr(self, "dtau0_" + wake_type, dtau0)
+ setattr(self, "W0_" + wake_type, W0)
+ return (tau0, dtau0, W0)
def get_wakepotential(self, bunch, wake_type):
"""
- Return the wake potential computed on the same time array as the bunch
- profile.
+ Return the wake potential computed on the wake function time array
+ limited to the bunch profile.
Parameters
----------
@@ -158,25 +261,28 @@ class WakePotential(Element):
Wake potential.
"""
+
+ (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type)
- tau0 = getattr(self.wakefield, wake_type).data.index
- W0 = getattr(self.wakefield, wake_type).data["real"]
- interp_func_W0 = interp1d(tau0, W0, fill_value=0,
+ interp_profile = interp1d(self.tau, self.rho, fill_value=0,
bounds_error=False)
- W = interp_func_W0(self.tau - np.mean(self.tau))
+
+ profile0 = interp_profile(tau0)
+
if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
- Wp = signal.convolve(self.rho, W*-1, mode='same')*self.dtau
+ Wp = signal.convolve(profile0, W0*-1, mode='same')*dtau0
elif wake_type == "Wxdip":
- dipole = self.dipole_moment(bunch, "x")
- Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+ dipole0 = self.dipole_moment(bunch, "x", tau0)
+ Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
elif wake_type == "Wydip":
- dipole = self.dipole_moment(bunch, "y")
- Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+ dipole0 = self.dipole_moment(bunch, "y", tau0)
+ Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
else:
raise ValueError("This type of wake is not taken into account.")
+ setattr(self,"profile0_" + wake_type, profile0)
setattr(self, wake_type, Wp)
- return Wp
+ return tau0, Wp
@Element.parallel
def track(self, bunch):
@@ -193,8 +299,8 @@ class WakePotential(Element):
self.charge_density(bunch)
for wake_type in self.types:
- Wp = self.get_wakepotential(bunch, wake_type)
- f = interp1d(self.tau, Wp, fill_value = 0, bounds_error = False)
+ tau0, Wp = self.get_wakepotential(bunch, wake_type)
+ f = interp1d(tau0 + self.tau_mean, Wp, fill_value = 0, bounds_error = False)
Wp_interp = f(bunch["tau"])
if wake_type == "Wlong":
bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
@@ -209,7 +315,8 @@ class WakePotential(Element):
bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge
/ self.ring.E0)
- def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False):
+ def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False,
+ plot_wake_function=True):
"""
Plot the last wake potential of a given type computed during the last
call of the track method.
@@ -222,6 +329,8 @@ class WakePotential(Element):
Plot the normalised bunch profile. The default is True.
plot_dipole : bool, optional
Plot the normalised dipole moment. The default is False.
+ plot_wake_function : bool, optional
+ Plot the normalised wake function. The default is True.
Returns
-------
@@ -234,84 +343,86 @@ class WakePotential(Element):
"Wydip" : r"$W_{p,y}^{D} (V/pC)$",
"Wxquad" : r"$W_{p,x}^{Q} (V/pC/m)$",
"Wyquad" : r"$W_{p,y}^{Q} (V/pC/m)$"}
- fig, ax = plt.subplots()
+
Wp = getattr(self, wake_type)
- ax.plot(self.tau*1e12, Wp*1e-12, label=labels[wake_type])
+ tau0 = getattr(self, "tau0_" + wake_type)
+
+ fig, ax = plt.subplots()
+ ax.plot(tau0*1e12, Wp*1e-12, label=labels[wake_type])
ax.set_xlabel("$\\tau$ (ps)")
ax.set_ylabel(labels[wake_type])
if plot_rho is True:
+ profile0 = getattr(self,"profile0_" + wake_type)
+ profile_rescaled = profile0/max(profile0)*max(np.abs(Wp))
rho_rescaled = self.rho/max(self.rho)*max(np.abs(Wp))
- ax.plot(self.tau*1e12, rho_rescaled*1e-12, label=r"$\rho$ (a.u.)")
+ ax.plot(tau0*1e12, profile_rescaled*1e-12, label=r"$\rho$ interpolated (a.u.)")
+ ax.plot((self.tau + self.tau_mean)*1e12, rho_rescaled*1e-12, label=r"$\rho$ (a.u.)", linestyle='dashed')
+ plt.legend()
+
+ if plot_wake_function is True:
+ W0 = getattr(self, "W0_" + wake_type)
+ W0_rescaled = W0/max(W0)*max(np.abs(Wp))
+ ax.plot(tau0*1e12, W0_rescaled*1e-12, label=r"$W_{function}$ (a.u.)")
plt.legend()
if plot_dipole is True:
dipole = getattr(self, "dipole_" + wake_type[1])
dipole_rescaled = dipole/max(dipole)*max(np.abs(Wp))
- ax.plot(self.tau*1e12, dipole_rescaled*1e-12, label=r"Dipole moment (a.u.)")
+ ax.plot(tau0*1e12, dipole_rescaled*1e-12, label=r"Dipole moment (a.u.)")
plt.legend()
return fig
- def energy_loss(self, bunch, compare_to='num'):
+ def reference_loss(self, bunch):
"""
- Calculate the energy loss from the wake potential and compare it to a
- reference value assuming a Gaussian bunch.
+ Calculate the loss factor and kick factor from the wake potential and
+ compare it to a reference value assuming a Gaussian bunch computed in
+ the frequency domain.
Parameters
----------
bunch : Bunch object
- compare_to : str, {"theory", "num"}, optional
- The method for calculating the reference energy loss. Use 'theory'
- for the analytical loss factor of a resonator [1], or use 'num'
- for a numerical result obtained from loss_factor method in
- Impedance class.
Returns
-------
- energy_loss_data : DataFrame
- An output showing the enegy loss compared to the reference value.
-
- Reference
- ---------
- [1] A. W. Chao and M. Tigner, "Handbook of Accelerator Physics and
- Engineering", 3rd printing, p.239
+ loss_data : DataFrame
+ An output showing the loss/kick factors compared to the reference
+ values.
"""
-
- if np.where(self.types=="Wlong")[0].size == 0:
- raise TypeError("No longitudinal wake function data.")
-
- else:
- Wp = self.get_wakepotential(bunch, "Wlong")
- Eloss = -1*np.trapz(Wp*self.rho, self.tau)*bunch.charge
-
- res = self.wakefield
- sigma = bunch['tau'].std()
- if compare_to == 'theory':
- Eloss_0 = 0.5*res.wr*res.Rs*1/res.Q *np.exp(-1*res.wr**2*sigma**2) \
- * bunch.charge
- elif compare_to == 'num':
- Eloss_0 = res.Zlong.loss_factor(sigma)*bunch.charge
- else:
- raise KeyError("Reference method not correct. Enter \"theory\" "
- "or \"num\" ")
+ loss = []
+ loss_0 = []
+ delta_loss = []
+ index = []
+ for wake_type in self.types:
+ tau0, Wp = self.get_wakepotential(bunch, wake_type)
+ profile0 = getattr(self, "profile0_" + wake_type)
+ factorTD = -1*np.trapz(Wp*profile0, tau0)
- delta_Eloss = (Eloss-Eloss_0) / Eloss_0 *100
+ if wake_type == "Wxdip":
+ factorTD = factorTD / bunch["x"].mean()
+ if wake_type == "Wydip":
+ factorTD = factorTD / bunch["y"].mean()
- column = ['Energy loss [eV]', 'Reference energy loss [eV]', 'Relative error [%]']
- energy_loss_data = pd.DataFrame(np.reshape([Eloss, Eloss_0, delta_Eloss], (1,3)),
- columns=column)
- return energy_loss_data
+ Z = getattr(self.wakefield, "Z" + wake_type[1:])
+ sigma = bunch['tau'].std()
+ factorFD = Z.loss_factor(sigma)
+ loss.append(factorTD)
+ loss_0.append(factorFD)
+ delta_loss.append( (factorTD - factorFD) / factorFD *100 )
+ if wake_type == "Wlong":
+ index.append("Wlong [V/C]")
+ else:
+ index.append(wake_type + " [V/C/m]")
-
-
-
+ column = ['TD factor', 'FD factor', 'Relative error [%]']
-
-
-
+ loss_data = pd.DataFrame(np.array([loss, loss_0, delta_loss]).T,
+ columns=column,
+ index=index)
+ return loss_data