From a44dbbf152bf2c8a65fe6d16bb3644f63b3638d6 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Fri, 20 Nov 2020 15:58:08 +0100
Subject: [PATCH] Various improvements
WakePotential -> add method to plot wake potential for a perfect gaussian bunch
utilities -> new function gaussian_bunch which gives the gaussian bunch profile
utilities -> read_IW2D now also reads wake function from IW2D
utilities -> new function read_IW2D_folder which reads the full results folder of IW2D into a WakeField element.
resistive_wall -> change the sign of the wake functions to agree with the code convention
---
collective_effects/__init__.py | 3 +-
collective_effects/resistive_wall.py | 4 +-
collective_effects/utilities.py | 96 +++++++++++++++++---
collective_effects/wakefield.py | 1 +
tracking/wakepotential.py | 125 +++++++++++++++++++++++++--
5 files changed, 205 insertions(+), 24 deletions(-)
diff --git a/collective_effects/__init__.py b/collective_effects/__init__.py
index a2b820a..39ac402 100644
--- a/collective_effects/__init__.py
+++ b/collective_effects/__init__.py
@@ -13,5 +13,6 @@ from mbtrack2.collective_effects.wakefield import ComplexData, Impedance, WakeFu
from mbtrack2.collective_effects.utilities import read_CST, read_IW2D, spectral_density
from mbtrack2.collective_effects.utilities import gaussian_bunch_spectrum, beam_spectrum, effective_impedance
from mbtrack2.collective_effects.utilities import yokoya_elliptic, beam_loss_factor
-from mbtrack2.collective_effects.utilities import double_sided_impedance
+from mbtrack2.collective_effects.utilities import double_sided_impedance, read_IW2D_folder
+from mbtrack2.collective_effects.utilities import gaussian_bunch
from mbtrack2.collective_effects.impedance_model import ImpedanceModel
\ No newline at end of file
diff --git a/collective_effects/resistive_wall.py b/collective_effects/resistive_wall.py
index 922cc94..d3ae124 100644
--- a/collective_effects/resistive_wall.py
+++ b/collective_effects/resistive_wall.py
@@ -155,7 +155,7 @@ class CircularResistiveWall(WakeField):
ind = np.isnan(wl)
wl[ind] = 0
- return wl
+ return -1*wl
def TransverseWakeFunction(self, time, exact=False):
"""
@@ -206,7 +206,7 @@ class CircularResistiveWall(WakeField):
/ time**(1/2) * self.length * -1)
ind = np.isnan(wt)
wt[ind] = 0
- return wt
+ return -1*wt
def function(self, t, x):
return ( (x**2 * np.exp(-1* (x**2) * t / self.t0) ) / (x**6 + 8) )
diff --git a/collective_effects/utilities.py b/collective_effects/utilities.py
index 11803ea..f6c6eec 100644
--- a/collective_effects/utilities.py
+++ b/collective_effects/utilities.py
@@ -9,8 +9,10 @@ collective_effects module.
import pandas as pd
import numpy as np
+from pathlib import Path
from scipy.interpolate import interp1d
-from mbtrack2.collective_effects.wakefield import Impedance
+from scipy.constants import c
+from mbtrack2.collective_effects.wakefield import Impedance, WakeFunction, WakeField
def read_CST(file, impedance_type='long', divide_by=None):
"""
@@ -45,30 +47,82 @@ def read_CST(file, impedance_type='long', divide_by=None):
impedance_type=impedance_type)
return result
-def read_IW2D(file, impedance_type='long'):
+def read_IW2D(file, file_type='Zlong'):
"""
- Read IW2D file format into an Impedance object.
+ Read IW2D file format into an Impedance object or a WakeField object.
Parameters
----------
file : str
Path to the file to read.
- impedance_type : str, optional
- Type of the Impedance object.
+ file_type : str, optional
+ Type of the Impedance or WakeField object.
Returns
-------
- result : Impedance object
+ result : Impedance or WakeField object
Data from file.
"""
- df = pd.read_csv(file, sep = ' ', header = None,
- names = ["Frequency","Real","Imaginary"], skiprows=1)
- df.set_index("Frequency", inplace = True)
- result = Impedance(variable = df.index,
- function = df["Real"] + 1j*df["Imaginary"],
- impedance_type=impedance_type)
+ if file_type[0] == "Z":
+ df = pd.read_csv(file, sep = ' ', header = None,
+ names = ["Frequency","Real","Imaginary"], skiprows=1)
+ df.set_index("Frequency", inplace = True)
+ df = df[df["Real"].notna()]
+ df = df[df["Imaginary"].notna()]
+ result = Impedance(variable = df.index,
+ function = df["Real"] + 1j*df["Imaginary"],
+ impedance_type=file_type[1:])
+ elif file_type[0] == "W":
+ df = pd.read_csv(file, sep = ' ', header = None,
+ names = ["Distance","Wake"], skiprows=1)
+ df["Time"] = df["Distance"] / c
+ df.set_index("Time", inplace = True)
+ if np.any(df.isna()):
+ index = df.isna().values
+ df = df.interpolate()
+ print("Nan values have been interpolated to:")
+ print(df[index])
+ # if file_type == "Wlong":
+ # df["Wake"] = df["Wake"]*-1
+ result = WakeFunction(variable = df.index,
+ function = df["Wake"],
+ wake_type=file_type[1:])
+ else:
+ raise ValueError("file_type should begin by Z or W.")
return result
+def read_IW2D_folder(folder, suffix):
+ """
+ Read IW2D results into a WakeField object.
+
+ Parameters
+ ----------
+ file : str
+ Path to the file to read.
+
+ Returns
+ -------
+ result : WakeField object
+ WakeField object with Impedance and WakeFunction objects from the
+ different files.
+ """
+ types = {"W" : WakeFunction,
+ "Z" : Impedance}
+ components = ["long", "xdip", "ydip", "xquad", "yquad"]
+
+ data_folder = Path(folder)
+
+ list_for_wakefield = []
+ for key, item in types.items():
+ for component in components:
+ name = data_folder / (key + component + suffix)
+ res = read_IW2D(file=name, file_type=key + component)
+ list_for_wakefield.append(res)
+
+ wake = WakeField(list_for_wakefield)
+
+ return wake
+
def spectral_density(frequency, sigma, m = 1, mode="Hermite"):
"""
Compute the spectral density of different modes for various values of the
@@ -123,6 +177,24 @@ def gaussian_bunch_spectrum(frequency, sigma):
regime and ion cloud mitigation in ThomX. p86, Eq. 4.19
"""
return np.exp(-1/2*(2*np.pi*frequency)**2*sigma**2)
+
+def gaussian_bunch(time, sigma):
+ """
+ Compute a Gaussian bunch profile.
+
+ Parameters
+ ----------
+ time : array
+ sample points of the bunch profile in [s].
+ sigma : float
+ RMS bunch length in [s].
+
+ Returns
+ -------
+ bunch_profile : array
+ Bunch profile in [s**-1] sampled at points given in time.
+ """
+ return np.exp(-1/2*(time**2/sigma**2))/(sigma*np.sqrt(2*np.pi))
def beam_spectrum(frequency, M, bunch_spacing, sigma=None,
diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index ee372cb..36b37f2 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -361,6 +361,7 @@ class Impedance(ComplexData):
component_coefficients = table[self.impedance_type].to_dict()
except KeyError:
print('Impedance type {} does not exist'.format(self.impedance_type))
+ raise
self.a = component_coefficients["a"]
self.b = component_coefficients["b"]
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
index 41e5644..444309b 100644
--- a/tracking/wakepotential.py
+++ b/tracking/wakepotential.py
@@ -13,6 +13,7 @@ import pandas as pd
from scipy import signal
from scipy.interpolate import interp1d
from mbtrack2.tracking.element import Element
+from mbtrack2.collective_effects.utilities import gaussian_bunch
class WakePotential(Element):
"""
@@ -166,7 +167,7 @@ class WakePotential(Element):
return dipole0
- def prepare_wakefunction(self, wake_type):
+ def prepare_wakefunction(self, wake_type, tau, save_data=True):
"""
Prepare the wake function of a given wake_type to be used for the wake
potential computation.
@@ -179,6 +180,10 @@ class WakePotential(Element):
----------
wake_type : str
Type of the wake function to prepare: "Wlong", "Wxdip", ...
+ tau : array
+ Time domain array of the bunch profile in [s].
+ save_data : bool, optional
+ If True, the results are saved as atributes.
Returns
-------
@@ -202,7 +207,7 @@ class WakePotential(Element):
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],
+ ind = np.all([min(tau[0], 0) < tau0, max(tau[-1], 0) > tau0],
axis=0)
tau0 = tau0[ind]
W0 = W0[ind]
@@ -229,9 +234,9 @@ class WakePotential(Element):
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)))
+ if (tau0[0] > tau[0]) or (tau0[-1] < tau[-1]):
+ n = max(int(np.ceil((tau0[0] - tau[0])/dtau0)),
+ int(np.ceil((tau[-1] - tau0[-1])/dtau0)))
tau0 = np.arange(tau0[0] - dtau0*n,
tau0[-1] + dtau0*(n+1),
@@ -239,10 +244,12 @@ class WakePotential(Element):
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))
+
+ if save_data:
+ setattr(self, "tau0_" + wake_type, tau0)
+ setattr(self, "dtau0_" + wake_type, dtau0)
+ setattr(self, "W0_" + wake_type, W0)
- 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):
@@ -263,7 +270,7 @@ class WakePotential(Element):
"""
- (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type)
+ (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type, self.tau)
interp_profile = interp1d(self.tau, self.rho, fill_value=0,
bounds_error=False)
@@ -375,6 +382,106 @@ class WakePotential(Element):
return fig
+ def get_gaussian_wakepotential(self, sigma, wake_type, dipole=1e-3):
+ """
+ Return the wake potential computed using a perfect gaussian profile.
+
+ Parameters
+ ----------
+ sigma : float
+ RMS bunch length in [s].
+ wake_type : str
+ Wake function type: "Wlong", "Wxdip", ...
+ dipole : float, optional
+ Dipole moment to consider in [m], (uniform dipole moment).
+ Returns
+ -------
+ tau0 : array
+ Time base in [s].
+ W0 : array
+ Wake function.
+ Wp : array
+ Wake potential.
+ profile0 : array
+ Gaussian bunch profile.
+ dipole0 : array
+ Dipole moment.
+
+ """
+
+ tau = np.linspace(-10*sigma,10*sigma, int(1e3))
+ (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type, tau, False)
+
+ profile0 = gaussian_bunch(tau0, sigma)
+ dipole0 = np.ones_like(profile0)*dipole
+
+ if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
+ Wp = signal.convolve(profile0, W0*-1, mode='same')*dtau0
+ elif wake_type == "Wxdip":
+ Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
+ elif wake_type == "Wydip":
+ Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
+ else:
+ raise ValueError("This type of wake is not taken into account.")
+
+ return tau0, W0, Wp, profile0, dipole0
+
+ def plot_gaussian_wake(self, sigma, wake_type, dipole=1e-3, plot_rho=True,
+ plot_dipole=False, plot_wake_function=True):
+ """
+ Plot the wake potential of a given type for a perfect gaussian bunch.
+
+ Parameters
+ ----------
+ sigma : float
+ RMS bunch length in [s].
+ wake_type : str
+ Type of the wake to plot: "Wlong", "Wxdip", ...
+ dipole : float, optional
+ Dipole moment to consider in [m], (uniform dipole moment).
+ 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.
+ plot_wake_function : bool, optional
+ Plot the normalised wake function. The default is True.
+
+ 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)$"}
+
+ tau0, W0, Wp, profile0, dipole0 = self.get_gaussian_wakepotential(sigma, wake_type, dipole)
+
+ 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:
+ profile_rescaled = profile0/max(profile0)*max(np.abs(Wp))
+ ax.plot(tau0*1e12, profile_rescaled*1e-12, label=r"$\rho$ (a.u.)")
+ plt.legend()
+
+ if plot_wake_function is True:
+ 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_rescaled = dipole0/max(dipole0)*max(np.abs(Wp))
+ ax.plot(tau0*1e12, dipole_rescaled*1e-12, label=r"Dipole moment (a.u.)")
+ plt.legend()
+
+ return fig
+
def reference_loss(self, bunch):
"""
Calculate the loss factor and kick factor from the wake potential and
--
GitLab