diff --git a/collective_effects/resonator.py b/collective_effects/resonator.py
new file mode 100644
index 0000000000000000000000000000000000000000..410a287a18d3a526b05b477429e9fbf385f73049
--- /dev/null
+++ b/collective_effects/resonator.py
@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines the impedances and wake functions from the resonator model
+based on the Wakefield class.
+
+@author: Alexis Gamelin, Watanyu Foosang
+@date: 25/09/2020
+"""
+
+import numpy as np
+from mbtrack2.collective_effects.wakefield import (Wakefield, Impedance,
+ WakeFunction)
+
+class Resonator(Wakefield):
+ def __init__(self, Rs, fr, Q, plane, n_wake=1e4, n_imp=1e4, imp_freq_lim=50e9):
+ """
+ Resonator model Wakefield element which computes the impedance and the
+ wake function in both longitudinal and transverse case.
+
+ Parameters
+ ----------
+ Rs : float
+ Shunt impedance in [ohm].
+ fr : float
+ Resonance frequency in [Hz].
+ Q : float
+ Quality factor.
+ plane : str
+ Plane on which the resonator is used: "long", "x" or "y".
+ n_wake : int or float, optional
+ Number of points used in the wake function. The default is 1e4.
+ n_imp : int or float, optional
+ Number of points used in the impedance. The default is 1e4.
+ imp_freq_lim : float, optional
+ Maximum frequency used in the impedance. The default is 50e9.
+
+ 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.
+
+ """
+ super().__init__()
+ self.Rs = Rs
+ self.fr = fr
+ self.wr = 2 * np.pi * self.fr
+ self.Q = Q
+ self.n_wake = int(n_wake)
+ self.n_imp = int(n_imp)
+ self.imp_freq_lim = imp_freq_lim
+ self.plane = plane
+
+ self.timestop = round(np.log(1000)/self.wr*2*self.Q, 12)
+
+ if self.Q >= 0.5:
+ self.Q_p = np.sqrt(self.Q**2 - 0.25)
+ else:
+ self.Q_p = np.sqrt(0.25 - self.Q**2)
+
+ self.wr_p = (self.wr*self.Q_p)/self.Q
+
+ if self.plane == "long":
+
+ freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
+ imp = Impedance(variable=freq,
+ function=self.long_impedance(freq),
+ impedance_type="long")
+ super().append_to_model(imp)
+
+ time = np.linspace(start=0, stop=self.timestop, num=self.n_wake)
+ wake = WakeFunction(variable=time,
+ function=self.long_wake_function(time),
+ wake_type="long")
+ super().append_to_model(wake)
+
+ elif self.plane == "x" or self.plane == "y" :
+
+ freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
+ imp = Impedance(variable=freq,
+ function=self.transverse_impedance(freq),
+ impedance_type=self.plane + "dip")
+ super().append_to_model(imp)
+
+ time = np.linspace(start=0, stop=self.timestop, num=self.n_wake)
+ wake = WakeFunction(variable=time,
+ function=self.transverse_wake_function(time),
+ wake_type=self.plane + "dip")
+ super().append_to_model(wake)
+ else:
+ raise ValueError("Plane must be: long, x or y")
+
+ def long_wake_function(self, t):
+ if self.Q >= 0.5:
+ return ( (self.wr * self.Rs / self.Q) *
+ np.exp(-1* self.wr * t / (2 * self.Q) ) *
+ (np.cos(self.wr_p * t) -
+ np.sin(self.wr_p * t) / (2 * self.Q_p) ) )
+
+ elif self.Q < 0.5:
+ return ( (self.wr * self.Rs / self.Q) *
+ np.exp(-1* self.wr * t / (2 * self.Q) ) *
+ (np.cosh(self.wr_p * t) -
+ np.sinh(self.wr_p * t) / (2 * self.Q_p) ) )
+
+ def long_impedance(self, f):
+ return self.Rs / (1 + 1j * self.Q * (f/self.fr - self.fr/f))
+
+ def transverse_impedance(self, f):
+ return self.Rs * self.fr / f / (
+ 1 + 1j * self.Q * (f / self.fr - self.fr / f) )
+
+ def transverse_wake_function(self, t):
+ if self.Q >= 0.5:
+ return (self.wr * self.Rs / self.Q_p *
+ np.exp(-1 * t * self.wr / 2 / self.Q_p) *
+ np.sin(self.wr_p * t) )
+ else:
+ return (self.wr * self.Rs / self.Q_p *
+ np.exp(-1 * t * self.wr / 2 / self.Q_p) *
+ np.sinh(self.wr_p * t) )
\ No newline at end of file
diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index 89b6cb3d48d0518eda37cc4ac4174f4cc89074b0..7854a48b66a4d130476650d607763c1ff2b1d4af 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -1,10 +1,10 @@
# -*- coding: utf-8 -*-
"""
This module defines general classes to describes wakefields, impedances and
-wake functions. Based on impedance library by David Amorim.
+wake functions.
-@author: David Amorin, Alexis Gamelin
-@date: 14/01/2020
+@author: David Amorin, Alexis Gamelin, Watanyu Foosang
+@date: 24/09/2020
"""
import warnings
@@ -12,14 +12,13 @@ import re
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
-from scipy.interpolate import interp1d
-from mbtrack2.tracking.element import Element
+import scipy as sc
from copy import deepcopy
-from scipy.integrate import trapz
from scipy.interpolate import interp1d
+from scipy.integrate import trapz
from scipy.constants import c
-import scipy as sc
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
+from mbtrack2.tracking.element import Element
class ComplexData:
"""
@@ -260,7 +259,7 @@ class WakeFunction(ComplexData):
def wake_type(self, value):
self._wake_type = value
- def wakefunction_from_imp(self, imp_obj, nout=None, trim_zero=True):
+ def wakefunction_from_imp(self, imp_obj, nout=None, trim=False):
"""
Compute a wake function from an Impedance object and load it to
the WakeFunction object.
@@ -273,13 +272,20 @@ class WakeFunction(ComplexData):
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
+ 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.
+ 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 = (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]
@@ -293,15 +299,15 @@ class WakeFunction(ComplexData):
time = sc.fft.fftshift(time_array)
Wlong = sc.fft.fftshift(Wlong_raw)
- if trim_zero is True:
- i_neg = np.where(time<0)[0]
+ 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=1e-3,
- bunch_charge = 1.44e-9, nout=None, freq_lim=200e9):
+ 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.
@@ -310,8 +316,10 @@ class WakeFunction(ComplexData):
----------
Wp_filename : str
Text file that contains wake potential data.
- bunch_length : float, optional
- Spatial bunch length [m].
+ 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
@@ -327,7 +335,7 @@ class WakeFunction(ComplexData):
freq_lim=freq_lim,
bunch_length=bunch_length,
bunch_charge=bunch_charge)
- self.wakefunction_from_imp(imp, nout=nout)
+ self.wakefunction_from_imp(imp, nout=nout, trim=trim)
class Impedance(ComplexData):
"""
@@ -634,6 +642,13 @@ class Wakefield:
"""
return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
+ @property
+ def wake_components(self):
+ """
+ Return an array of the wake function component names for the element.
+ """
+ return np.array([comp for comp in dir(self) if re.match(r'[W]', comp)])
+
@staticmethod
def add_wakefields(wake1, beta1, wake2, beta2):
"""
@@ -705,129 +720,7 @@ class Wakefield:
return wake_sum
else:
raise ValueError("Error in input.")
-
-class Resonator(Wakefield):
- """
- Compute an analytical wake function and an analytical impedance of a
- resonator [1] and load it to the Resonator object.
-
- Parameter
- ---------
- ring : Synchrotron object.
- Q_factor : float, optional
- Resonator quality factor. The default value is 1.
- f_res : float, optional
- 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,
- n_wake=1000, n_imp=1000, imp_freq_lim=50e9):
- super().__init__()
- self.ring = ring
- self.Q = Q_factor
- self.omega_res = 2*np.pi*f_res
- 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_p = np.sqrt(self.Q**2 - 0.25)
- else:
- 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, 12)
-
- def time_array(self):
- self.tau = np.linspace(start=0, stop=self.timestop, num=self.n_wake)
- def long_wakefunction(self):
- """
- 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))
-
- 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, the timestop is extended by 50 ps.
-
- """
-
- ratio = np.abs(max(self.wlong)) / np.abs(self.wlong[-6:-1])
- while any(ratio < 1000):
- self.timestop += 50e-12
- self.time_array()
- self.long_wakefunction()
- 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()
- self.check_wake_tail()
-
- 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):
"""
Define the impedance model of the machine.
diff --git a/tracking/element.py b/tracking/element.py
index bcc1940b11038a3c5fa0c405d1c2ea923bfb6f1b..9123e901224ce2d8e399ffa9ba9fbae603f2e9fe 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -9,13 +9,9 @@ included in the tracking.
"""
import numpy as np
-import pandas as pd
from abc import ABCMeta, abstractmethod
from functools import wraps
from mbtrack2.tracking.particles import Beam
-from scipy import signal
-import matplotlib.pyplot as plt
-from scipy.interpolate import interp1d
class Element(metaclass=ABCMeta):
"""
@@ -98,334 +94,7 @@ class LongitudinalMap(Element):
"""
bunch["delta"] -= self.ring.U0 / self.ring.E0
bunch["tau"] -= self.ring.ac * self.ring.T0 * bunch["delta"]
-
-class WakePotential(Element):
- """
- Resonator model based wake potential calculation for one turn.
-
- Parameters
- ----------
- ring : Synchrotron object.
- Q_factor : float, optional
- Resonator quality factor. The default value is 1.
- f_res : float, optional
- 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_bin : int, optional
- Number of bins for constructing the longitudinal bunch profile.
- The default is 65.
-
- Attributes
- ----------
- rho : array of shape (n_bin, )
- Bunch charge density profile.
- tau : array of shape (n_bin + time_extra, )
- Time array starting from the head of the bunch until the wake tail
- called timestop.
-
- The length of time_extra is determined by the last position of the
- bunch time_bunch[-1], timestop, and the mean bin width of the bunch
- profile mean_bin_size as
- len(time_extra) = (timestop - time_bunch[-1]) / mean_bin_size
- W_long : array of shape (n_bin + time_extra, )
- Wakefunction profile.
- W_p : array of shape (n_bin + time_extra, )
- Wake potential profile.
- wp : array of shape (mp_number, )
- Wake potential exerted on each macro-particle.
-
- Methods
- -------
- charge_density(bunch, n_bin=65)
- Calculate bunch charge density.
- plot(self, var, plot_rho=True)
- Plotting wakefunction or wake potential.
- track(bunch)
- Tracking method for the element.
-
- """
-
- def __init__(self, ring, Q_factor=1, f_res=10e9, R_shunt=100, n_bin=65):
- self.ring = ring
- self.n_bin = n_bin
-
- self.Q_factor = Q_factor
- self.omega_res = 2*np.pi*f_res
- self.R_shunt = R_shunt
-
- 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
- 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
-
- def charge_density(self, bunch, n_bin):
- self.bins = bunch.binning(n_bin=self.n_bin)
- self.bin_data = self.bins[2]
- self.bin_size = self.bins[0].length
-
- self.rho = bunch.charge_per_mp*self.bin_data/ \
- (self.bin_size*bunch.charge)
-
- def init_timestop(self):
- self.timestop = round(np.log(1000)/self.omega_res*2*self.Q_factor, 12)
-
- def time_array(self):
- time_bunch = self.bins[0].mid
- mean_bin_size = np.mean(self.bin_size)
- time_extra = np.arange(start = time_bunch[-1]+mean_bin_size,
- stop = self.timestop, step = mean_bin_size)
-
- self.tau = np.concatenate((time_bunch,time_extra))
-
- def long_wakefunction(self):
- w_list = []
- if self.Q_factor >= 0.5:
- for t in self.tau:
- if t >= 0:
- w_long = -(self.omega_res*self.R_shunt/self.Q_factor)*\
- np.exp(-self.omega_res*t/(2*self.Q_factor))*\
- (np.cos(self.omega_res_p*t)-\
- np.sin(self.omega_res_p*t)/(2*self.Q_factor_p))
- else:
- w_long = 0
-
- w_list.append(w_long)
-
- elif self.Q_factor < 0.5:
- for t in self.tau:
- if t >= 0:
- w_long = -(self.omega_res*self.R_shunt/self.Q_factor)*\
- np.exp(-self.omega_res*t/(2*self.Q_factor))*\
- (np.cosh(self.omega_res_p*t)-\
- np.sinh(self.omega_res_p*t)/(2*self.Q_factor_pp))
- else:
- w_long = 0
-
- w_list.append(w_long)
-
- self.W_long = np.array(w_list)
-
- def wake_potential(self):
- dtau = (self.tau[-1]-self.tau[0]) / len(self.tau)
- self.W_p = signal.convolve(self.W_long, self.rho, mode="same") * dtau
-
- def plot(self, var, plot_rho=True):
- """
- Plotting wakefunction or wake potential.
-
- Parameters
- ----------
- var : {'W_p', 'W_long' }
- If 'W_p', the wake potential is plotted.
- If 'W_long', the wakefunction is plotted.
- plot_rho : bool, optional
- Overlay the bunch charge density profile on the plot.
- The default is True.
- Returns
- -------
- fig : Figure
- Figure object with the plot on it.
-
- """
-
- fig, ax = plt.subplots()
-
- if var == "W_p":
- ax.plot(self.tau*1e12, self.W_p*1e-12)
-
- ax.set_xlabel("$\\tau$ (ps)")
- ax.set_ylabel("W$_p$ (V/pC)")
-
- elif var == "W_long":
- ax.plot(self.tau*1e12, self.W_long*1e-12)
- ax.set_xlabel("$\\tau$ (ps)")
- ax.set_ylabel("W$_{||}$ (V/ps)")
-
- if plot_rho is True:
- rho_array = np.array(self.rho)
- rho_rescaled = rho_array/max(rho_array)*max(self.W_p)
-
- ax.plot(self.bins[0].mid*1e12, rho_rescaled*1e-12)
-
- else:
- pass
-
- return fig
-
- def check_wake_tail(self):
- """
- Checking whether the full wakefunction is obtained by the calculated
- initial timestop.
-
- """
-
- ratio = np.abs(min(self.W_long) / self.W_long[-6:-1])
- while any(ratio < 1000):
- # Initial timestop is too short.
- # Extending timestop by 50 ps and recalculating."
- self.timestop += 50e-12
- self.time_array()
- self.long_wakefunction()
- ratio = np.abs(min(self.W_long) / self.W_long[-6:-1])
-
- @Element.parallel
- def track(self, bunch):
- """
- Tracking method for the element.
- No bunch to bunch interaction, so written for Bunch objects and
- @Element.parallel is used to handle Beam objects.
-
- Parameters
- ----------
- bunch : Bunch or Beam object.
-
- """
-
- self.charge_density(bunch, n_bin = self.n_bin)
- self.init_timestop()
- self.time_array()
- self.long_wakefunction()
- self.check_wake_tail()
- self.wake_potential()
-
- f = interp1d(self.tau, self.W_p, fill_value = 0, bounds_error = False)
- self.wp = f(bunch["tau"])
-
- bunch["delta"] += self.wp * bunch.charge / self.ring.E0
-
-class WakePotential2(Element):
- """
- Compute a wake potential from a wake function by performing a convolution
- with a bunch charge profile.
-
- Parameters
- ----------
- wake_obj : WakeFunction object
- ring : Synchrotron object
- n_bin : int, optional
- Number of bins for constructing the longitudinal bunch profile.
- bunch : Bunch or Beam object
-
- Attributes
- ----------
- rho : array of shape (n_bin, )
- Bunch charge density profile in the unit [1/s].
- Wlong : array
- Wake function profile.
- Wp : array
- Wake potential profile.
- Wp_interp : array of shape (mp_number, )
- Wake potential, obtained from interpolating Wp, exerted on each macro-particle.
-
- Methods
- -------
- charge_density(bunch)
- 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.
- track(bunch)
- Tracking method for the element.
-
- """
-
- def __init__(self, ring, wake_obj, n_bin=65):
- self.wake_obj = wake_obj
- self.tau0 = self.wake_obj.data.index
- self.Wlong0 = self.wake_obj.data['real']
- self.ring = ring
- self.n_bin = n_bin
-
- def charge_density(self, bunch):
- self.bins = bunch.binning(n_bin=self.n_bin)
- self.bin_data = self.bins[2]
- self.bin_size = self.bins[0].length
-
- 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)
- 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):
- """
- Plotting wakefunction or wake potential.
-
- Parameters
- ----------
- var : {'Wp', 'Wlong'}
- If 'Wp', the wake potential is plotted.
- If 'Wlong', the wakefunction is plotted.
- plot_rho : bool, optional
- Overlay the bunch charge density profile on the plot.
-
- Returns
- -------
- fig : Figure
- Figure object with the plot on it.
-
- """
-
- fig, ax = plt.subplots()
-
- if var == "Wp":
- ax.plot(self.tau*1e12, self.Wp*1e-12)
- ax.set_xlabel("$\\tau$ (ps)")
- ax.set_ylabel("W$_p$ (V/pC)")
-
- elif var == "Wlong":
- ax.plot(self.tau*1e12, self.Wlong*1e-12)
- ax.set_xlabel("$\\tau$ (ps)")
- ax.set_ylabel("W$_{||}$ (V/ps)")
-
- if plot_rho is True:
- rho_array = np.array(self.rho)
- 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
-
- @Element.parallel
- def track(self, bunch):
- """
- Tracking method for the element.
- No bunch to bunch interaction, so written for Bunch objects and
- @Element.parallel is used to handle Beam objects.
-
- Parameters
- ----------
- bunch : Bunch or Beam object.
-
- """
-
- self.charge_density(bunch)
- self.get_wakepotential()
-
- f = interp1d(self.tau, self.Wp, fill_value = 0, bounds_error = False)
- self.Wp_interp = f(bunch["tau"])
-
- bunch["delta"] += self.Wp_interp * bunch.charge / self.ring.E0
-
class SynchrotronRadiation(Element):
"""
Element to handle synchrotron radiation, radiation damping and quantum
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
new file mode 100644
index 0000000000000000000000000000000000000000..02dee0325a3603c3405cfdea53b6168a9825322d
--- /dev/null
+++ b/tracking/wakepotential.py
@@ -0,0 +1,253 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines the WakePotential class which deals with the single bunch
+effects, both longitudinal and transverse.
+
+@author: Alexis Gamelin, Watanyu Foosang
+@date: 25/09/2020
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+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.
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ wakefield : Wakefield object
+ Wakefield object which contains the wake functions to be used.
+ n_bin : int, optional
+ Number of bins for constructing the longitudinal bunch profile.
+
+ Attributes
+ ----------
+ rho : array of shape (n_bin, )
+ Bunch charge density profile in the unit [1/s].
+ Wp : array
+ Wake potential profile.
+ Wp_interp : array of shape (mp_number, )
+ Wake potential, obtained from interpolating Wp, exerted on each macro-particle.
+
+ Methods
+ -------
+ charge_density(bunch)
+ Compute bunch charge density profile in [1/s].
+ dipole_moment(bunch, plane)
+ Return the dipole moment of the bunch computed on the same time array
+ as the bunch profile.
+ get_wakepotential(bunch, wake_type)
+ Return the wake potential computed on the same time array as the bunch
+ profile.
+ track(bunch)
+ Tracking method for the element.
+ plot_last_wake(wake_type, plot_rho=True, plot_dipole=False)
+ Plot the last wake potential of a given type computed during the last
+ call of the track method.
+
+ """
+
+ def __init__(self, ring, wakefield, n_bin=65):
+ self.wakefield = wakefield
+ self.types = self.wakefield.wake_components
+ self.n_types = len(self.wakefield.wake_components)
+ self.ring = ring
+ self.n_bin = n_bin
+
+ def charge_density(self, bunch):
+ """
+ Compute bunch charge density profile in [1/s].
+
+ Parameters
+ ----------
+ bunch : Bunch object
+
+ """
+
+ # Get binning data
+ a, b, c = bunch.binning(n_bin=self.n_bin)
+ self.bins = a
+ self.sorted_index = b
+ self.profile = c
+ self.bin_size = self.bins.length
+
+ # Compute charge density
+ self.rho = bunch.charge_per_mp*self.profile/(self.bin_size*bunch.charge)
+ self.rho = np.array(self.rho)
+
+ # Compute time array
+ self.tau = np.array(self.bins.mid)
+ self.dtau = self.tau[1] - self.tau[0]
+
+ # Add N values before and after rho and tau
+ if self.n_bin % 2 == 0:
+ N = int(self.n_bin/2)
+ self.tau = np.arange(self.tau[0] - self.dtau*N,
+ self.tau[-1] + self.dtau*N,
+ self.dtau)
+ self.rho = np.append(self.rho, np.zeros(N))
+ self.rho = np.insert(self.rho, 0, np.zeros(N))
+ else:
+ N = int(np.floor(self.n_bin/2))
+ self.tau = np.arange(self.tau[0] - self.dtau*N,
+ self.tau[-1] + self.dtau*(N+1),
+ self.dtau)
+ self.rho = np.append(self.rho, np.zeros(N))
+ self.rho = np.insert(self.rho, 0, np.zeros(N+1))
+
+ if len(self.tau) != len(self.rho):
+ self.tau = np.append(self.tau, self.tau[-1] + self.dtau)
+
+ def dipole_moment(self, bunch, plane):
+ """
+ Return the dipole moment of the bunch computed on the same time array
+ as the bunch profile.
+
+ Parameters
+ ----------
+ bunch : Bunch object
+ plane : str
+ Plane on which the dipole moment is computed, "x" or "y".
+
+ Returns
+ -------
+ dipole : array
+ Dipole moment of the bunch.
+
+ """
+ dipole = np.zeros((self.n_bin,))
+ for i in range(self.n_bin):
+ dipole[i] = bunch[plane][self.sorted_index == self.bins[i]].sum()
+ dipole = dipole/self.profile
+ dipole[np.isnan(dipole)] = 0
+
+ # Add N values to get same size as tau/rho
+ if self.n_bin % 2 == 0:
+ N = int(self.n_bin/2)
+ dipole = np.append(dipole, np.zeros(N))
+ dipole = np.insert(dipole, 0, np.zeros(N))
+ else:
+ N = int(np.floor(self.n_bin/2))
+ dipole = np.append(dipole, np.zeros(N))
+ dipole = np.insert(dipole, 0, np.zeros(N+1))
+
+ setattr(self, "dipole_" + plane, dipole)
+ return dipole
+
+ def get_wakepotential(self, bunch, wake_type):
+ """
+ Return the wake potential computed on the same time array as the bunch
+ profile.
+
+ Parameters
+ ----------
+ bunch : Bunch object
+ wake_type : str
+ Wake function type: "Wlong", "Wxdip", ...
+
+ Returns
+ -------
+ Wp : array
+ Wake potential.
+
+ """
+
+ 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,
+ bounds_error=False)
+ W = interp_func_W0(self.tau)
+ if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
+ Wp = signal.convolve(self.rho, W*-1, mode='same')*self.dtau
+ elif wake_type == "Wxdip":
+ dipole = self.dipole_moment(bunch, "x")
+ Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+ elif wake_type == "Wydip":
+ dipole = self.dipole_moment(bunch, "y")
+ Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+ else:
+ raise ValueError("This type of wake is not taken into account.")
+
+ setattr(self, wake_type, Wp)
+ return Wp
+
+ @Element.parallel
+ def track(self, bunch):
+ """
+ Tracking method for the element.
+ No bunch to bunch interaction, so written for Bunch objects and
+ @Element.parallel is used to handle Beam objects.
+
+ Parameters
+ ----------
+ bunch : Bunch or Beam object.
+
+ """
+
+ 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)
+ Wp_interp = f(bunch["tau"])
+ if wake_type == "Wlong":
+ bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
+ elif wake_type == "Wxdip":
+ bunch["xp"] += Wp_interp * bunch.charge / self.ring.E0
+ elif wake_type == "Wydip":
+ bunch["yp"] += Wp_interp * bunch.charge / self.ring.E0
+ elif wake_type == "Wxquad":
+ bunch["xp"] += (bunch["x"] * Wp_interp * bunch.charge
+ / self.ring.E0)
+ elif wake_type == "Wyquad":
+ bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge
+ / self.ring.E0)
+
+ def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False):
+ """
+ Plot the last wake potential of a given type computed during the last
+ call of the track method.
+
+ Parameters
+ ----------
+ wake_type : str
+ Type of the wake to plot: "Wlong", "Wxdip", ...
+ plot_rho : bool, optional
+ Plot the normalised bunch profile. The default is True.
+ plot_dipole : bool, optional
+ Plot the normalised dipole moment. The default is False.
+
+ Returns
+ -------
+ fig : figure
+
+ """
+
+ labels = {"Wlong" : r"$W_{p,long}$ (V/pC)",
+ "Wxdip" : r"$W_{p,x}^{D} (V/pC)$",
+ "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])
+ ax.set_xlabel("$\\tau$ (ps)")
+ ax.set_ylabel(labels[wake_type])
+
+ if plot_rho is True:
+ 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.)")
+ 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.)")
+ plt.legend()
+
+ return fig
\ No newline at end of file