diff --git a/tracking/element.py b/tracking/element.py
index bee031ce02323679e53832f4745196ef33dc1666..d7f7e44bdca7c13dbfa521f7cb288dd365709cc3 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -9,9 +9,13 @@ included in the tracking.
"""
import numpy as np
+import pandas as pd
from abc import ABCMeta, abstractmethod
from functools import wraps
from tracking.particles import Beam
+from scipy import signal
+import matplotlib.pyplot as plt
+from scipy.interpolate import interp1d
class Element(metaclass=ABCMeta):
"""
@@ -95,6 +99,203 @@ 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):
+ self.W_p = signal.convolve(self.W_long*1e-12, self.rho, mode="same")
+
+ 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$_{||}$ ($\\Omega$/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 SynchrotronRadiation(Element):
"""
Element to handle synchrotron radiation, radiation damping and quantum