From 2470c02f66ed759834f481586ecbbe5d3ccbfe09 Mon Sep 17 00:00:00 2001
From: Watanyu Foosang <watanyu.f@gmail.com>
Date: Fri, 10 Dec 2021 17:59:13 +0100
Subject: [PATCH] Add WakePotential_LR class
The long-range wake potential tracking class has been added.
---
tracking/particles.py | 8 ++
tracking/wakepotential.py | 287 ++++++++++++++++++++++++++++++++++++++
2 files changed, 295 insertions(+)
diff --git a/tracking/particles.py b/tracking/particles.py
index 068fba4..107d374 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -362,6 +362,14 @@ class Bunch:
Dimension to plot. The default is "tau".
n_bin : int, optional
Number of bins. The default is 75.
+ bin_min, bin_max : float, optional
+ The lowest and the hightest bound of the bins. The range between
+ this two bounds must cover every particle in the bunch, i.e.
+ they must satisfy bin_min <= Bunch[dimension].min() and
+ Bunch[dimension].max() <= bin_max.
+ bin_size : float, optional
+ Define the bin size. If both n_bin and bin_size are provided,
+ n_bin value will be ignored.
"""
bins, sorted_index, profile, center = self.binning(dimension, n_bin,
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
index e85ff90..be05538 100644
--- a/tracking/wakepotential.py
+++ b/tracking/wakepotential.py
@@ -12,8 +12,12 @@ import matplotlib.pyplot as plt
import pandas as pd
from scipy import signal
from scipy.interpolate import interp1d
+from scipy.signal import convolve
from mbtrack2.tracking.element import Element
from mbtrack2.collective_effects.utilities import gaussian_bunch
+from mbtrack2.collective_effects import WakeField
+from mpi4py import MPI
+
class WakePotential(Element):
"""
@@ -526,6 +530,289 @@ class WakePotential(Element):
columns=column,
index=index)
return loss_data
+
+class WakePotential_LR(Element):
+ """
+ Long-range wake potential tracking class for multibunch and multiturn
+ effects tracking.
+
+ Wake potentials are first calculated individually from each bunch, then
+ they are shifted longitudinally depending on the location of the bunch
+ before being summed up. Each bunch is then exerted by a apart of the
+ total wake that corresponds to its location.
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ beam : Beam object
+ wakefield : Wakefield object
+ Wakefield object which contains the wake functions to be used. The wake
+ functions must be uniformly sampled!
+ m : int, optional
+ Unitless parameter defining the bin size relative to the revolution
+ period T1 as :
+ T1 = m * bin_size.
+ window_length : float, optional
+ Defines the length of a fixed size window for bunch binning in [s].
+ The window must cover the whole bunch.
+ wake_length : int, optional
+ Length of the wake potential to be stored in the unit of number of
+ revolutions.
+
+ """
+
+ def __init__(self, ring, beam, wakefield, m=3000, window_length=1e-10,
+ wake_length=10):
+
+ self.ring = ring
+ self.wakefield = wakefield
+ self.type = self.wakefield.wake_components
+
+ self.m = int(m)
+ self.slice_size = self.ring.T1/self.m
+ self.bunch_head = window_length/2
+ self.bunch_tail = -1*window_length/2
+ self.n_bins = int(window_length/self.slice_size)+1
+
+ self.wake_length = int(wake_length)
+
+ def charge_density(self, bunch, wake_length):
+ """
+ Computes a time array and a charge density array of an individual bunch
+ with the length specified in wake_length.
+
+ Parameters
+ ----------
+ bunch : Bunch object
+ wake_lengths : int
+ Length of the multiturn effect in the unit of number of turns.
+
+ """
+
+ _ , b, c, d = bunch.binning(n_bin=self.n_bins, bin_min=self.bunch_tail,
+ bin_max=self.bunch_head, bin_size=self.slice_size)
+
+ self.sorted_index = b
+ self.profile = c
+ self.center = d
+ self.bin_size = self.center[1]-self.center[0]
+
+ self.rho = bunch.charge_per_mp*self.profile/(self.bin_size*bunch.charge) # [mp_number/tau]
+ self.rho = np.array(self.rho)
+
+ # Compute time array relative to t0
+ self.tau = np.array(self.center)
+ self.dtau = self.tau[1] - self.tau[0]
+
+ if self.n_bins % 2 == 0:
+ self.n_bins_total = int(wake_length*self.ring.T0/self.bin_size)
+ self.N_front = int(self.n_bins_total/2)
+ self.N_back = int(self.n_bins_total/2)
+ self.tau = np.arange(self.tau[0] - self.dtau*self.N_front,
+ self.tau[-1] + self.dtau*self.N_back,
+ self.dtau)
+ self.rho = np.append(self.rho, np.zeros(self.N_back))
+ self.rho = np.insert(self.rho, 0, np.zeros(self.N_front))
+ else:
+ self.n_bins_total = int(np.floor(wake_length*self.ring.T0/self.bin_size))
+ self.N_front = int(np.floor(self.n_bins_total/2))
+ self.N_back = int(np.floor(self.n_bins_total/2))
+ self.tau = np.arange(self.tau[0] - self.dtau*self.N_front,
+ self.tau[-1] + self.dtau*(self.N_back+1),
+ self.dtau)
+ self.rho = np.append(self.rho, np.zeros(self.N_back))
+ self.rho = np.insert(self.rho, 0, np.zeros(self.N_front+1))
+
+ if len(self.tau) != len(self.rho):
+ self.tau = np.append(self.tau, self.tau[-1]+self.dtau)
+
+ def prepare_wakefunction(self, tau, wake_type):
+ """
+ Interpolates the wake function from the wakefield object on the time
+ array given by charge_density method. The parts outside the orginal
+ wake function are filled with zeros.
+
+ Parameters
+ ----------
+ tau : array
+ Time array on which the wake fucntion will be interpolated, in [s].
+ wake_type : str
+ Type of the wake function to prepare: "Wlong", "Wxdip", ...
+
+ Returns
+ -------
+ wf : array
+ Interpolated wake function array.
+
+ """
+
+ wf0 = getattr(self.wakefield, wake_type)
+ wf_interp = interp1d(wf0.data.index, wf0.data['real'],
+ bounds_error=False, fill_value=0)
+ wf = wf_interp(tau)
+
+ return wf
+
+ def dipole_moment(self, bunch, plane, tau):
+ """
+ Return the dipole moment of an individual bunch computed on the same
+ time array as the wake function.
+
+ Parameters
+ ----------
+ bunch : Bunch object
+ plane : str
+ Plane on which the dipole moment is computed, "x" or "y".
+ tau : array
+ Time array on which the dipole moment will be interpolated, in [s].
+
+ Returns
+ -------
+ dipole : array
+ Dipole moment of the bunch.
+
+ """
+
+ dipole0 = np.zeros((self.n_bins - 1,))
+
+ for i in range(self.n_bins - 1):
+ dipole0[i] = bunch[plane][self.sorted_index == i].sum()
+ dipole0 = dipole0/self.profile
+ dipole0[np.isnan(dipole0)] = 0
+
+ dipole0 = np.append(dipole0, np.zeros(self.N_back))
+ dipole0 = np.insert(dipole0, 0, np.zeros(self.N_front))
+
+ if len(tau) != len(dipole0):
+ if len(tau)-len(dipole0) == 1:
+ dipole0 = np.append(dipole0, 0)
+ elif len(tau)-len(dipole0) == 2:
+ dipole0 = np.append(dipole0, 0)
+ dipole0 = np.insert(dipole0, 0, 0)
+ elif len(tau)-len(dipole0) == 3:
+ dipole0 = np.append(dipole0, np.array([0,0]))
+ dipole0 = np.insert(dipole0, 0, 0)
+
+ # Interpole on tau to get the same size as wf
+ dipole_func = interp1d(tau, dipole0, fill_value=0, bounds_error=False)
+ dipole = dipole_func(tau)
+
+ return dipole
+
+ def get_wakepotential(self, bunch, wake_type):
+ """
+ Computes the wake potential of an individual bunch. The zero part of
+ the wake before the bunch is cut off to reduce memory usage. The time
+ array is also cut accordingly to the wake potential.
+
+ Parameters
+ ----------
+ bunch : Bunch object
+ wake_type : str
+ Wake function type: "Wlong", "Wxdip", ...
+
+ Returns
+ -------
+ tau0 : array
+ Reduced time array.
+ wp0 : array
+ Reduced wake potential.
+
+ """
+
+ wf = self.prepare_wakefunction(self.tau, wake_type)
+
+ if wake_type == "Wxdip" :
+ dipole = self.dipole_moment(bunch, plane="x", tau=self.tau)
+ wp0 = convolve(self.rho*dipole, wf, mode='same')*self.dtau
+ elif wake_type == "Wydip" :
+ dipole = self.dipole_moment(bunch, plane="y", tau=self.tau)
+ wp0 = convolve(self.rho*dipole, wf, mode='same')*self.dtau
+ elif wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
+ wp0 = convolve(self.rho, wf*-1, mode='same')*self.dtau
+ else:
+ raise ValueError("Incorrect wake field type.")
+
+ # Cut off the zero part before the bunch
+ cut_off = (self.n_bins_total-self.n_bins)//2
+ wp0 = wp0[cut_off:]
+ tau0 = self.tau[cut_off:]
+
+ return tau0, wp0
+
+ def track(self, beam, k, save_data=False):
+ """
+ Tracking method for the element.
+
+ Parameters
+ ----------
+ beam : Beam object
+ k : int
+ Temporary argument used for naming output files indicating the turn
+ number of the data.
+ save_data : bool, optional
+ Temporary argument used for saving wake potentials. This function
+ should be disabled when track for large number of turns or large
+ number of bunches to prevent output file flooding.
+
+ """
+
+ bunch_num = beam.mpi.bunch_num
+ bunch = beam[bunch_num]
+
+ beam.mpi.share_distributions(beam, n_bin=self.n_bins,
+ bin_min=self.bunch_tail, bin_max=self.bunch_head,
+ bin_size=self.slice_size)
+
+
+ self.charge_density(bunch, self.wake_length)
+ for wake_type in self.type:
+ tau, wp = self.get_wakepotential(bunch, wake_type)
+
+ shift = bunch_num * self.m # shift in number of bins unit
+ wp = np.roll(wp, shift)
+ wp[:shift] = 0
+
+ if save_data is True:
+ np.savez('testclass_turn{0}_bunch{1}'.format(k, bunch_num), tau=tau, wp=wp)
+
+ beam.mpi.__setattr__("wp_all", np.empty((len(tau),), dtype=np.float64))
+ beam.mpi.comm.Allreduce([wp, MPI.DOUBLE], [beam.mpi.__getattribute__("wp_all"), MPI.DOUBLE], op=MPI.SUM)
+
+ del wp
+
+ if save_data is True and bunch_num == 0:
+ np.savez('testclass_turn{0}_sum'.format(k), tau=tau, wp_all=beam.mpi.wp_all)
+
+ del tau
+
+ if wake_type == "Wxdip":
+ bunch["xp"] += beam.mpi.wp_all[self.sorted_index+shift] * \
+ bunch.charge / self.ring.E0
+ elif wake_type == "Wydip":
+ bunch["yp"] += beam.mpi.wp_all[self.sorted_index+shift] * \
+ bunch.charge / self.ring.E0
+ elif wake_type == "Wxquad":
+ bunch["xp"] += (bunch["x"] * beam.mpi.wp_all[self.sorted_index+shift] * \
+ bunch.charge / self.ring.E0)
+ elif wake_type == "Wyquad":
+ bunch["yp"] += (bunch["y"] * beam.mpi.wp_all[self.sorted_index+shift] * \
+ bunch.charge / self.ring.E0)
+ elif wake_type == "Wlong":
+ bunch["delta"] += beam.mpi.wp_all[self.sorted_index+shift] * \
+ bunch.charge / self.ring.E0
+
+
+
+
+
+
+
+
+
+
+
+
--
GitLab