From 432853b81e5774d63da44bdb84bc97cde08986ed Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Fri, 21 Apr 2023 23:04:13 +0200
Subject: [PATCH] Various improvements for the wakefield module
Added class representation for ComplexData, Impedance, WakeFunction, WakeField and Beam classes.
Rework the deconvolution/from_wakepotential/to_impedance methods for them to work with any data loaded in the Impedance/WakeFunction classes and not just CST text files.
Add a plot method for the Impedance/WakeFunction classes.
Improved WakeField class by adding __iter__ and __add__ methods.
Add a read_ECHO2D function.
---
mbtrack2/impedance/wakefield.py | 247 ++++++++++++++-------------
mbtrack2/tracking/particles.py | 4 +
mbtrack2/utilities/__init__.py | 3 +-
mbtrack2/utilities/read_impedance.py | 34 +++-
4 files changed, 164 insertions(+), 124 deletions(-)
diff --git a/mbtrack2/impedance/wakefield.py b/mbtrack2/impedance/wakefield.py
index bc0c263..5de4a8f 100644
--- a/mbtrack2/impedance/wakefield.py
+++ b/mbtrack2/impedance/wakefield.py
@@ -183,6 +183,9 @@ class ComplexData:
y = self.data["imag"], kind=interp_kind)
return real_func(values) + 1j*imag_func(values)
+ def __repr__(self):
+ """Return representation of data"""
+ return f'{(self.data)!r}'
def initialize_coefficient(self):
"""
@@ -268,9 +271,12 @@ class WakeFunction(ComplexData):
Methods
-------
- from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
- Compute a wake function from a wake potential file and load it to the
- WakeFunction object.
+ to_impedance()
+ Return an Impedance object from the WakeFunction data.
+ deconvolution(freq_lim, sigma, mu)
+ Compute a wake function from wake potential data.
+ plot()
+ Plot the wake function data.
"""
def __init__(self,
@@ -281,6 +287,10 @@ class WakeFunction(ComplexData):
self.data.index.name = "time [s]"
self.initialize_coefficient()
+ def __repr__(self):
+ """Return representation of data"""
+ return f'WakeFunction of component_type {self.component_type}:\n {(self.data)!r}'
+
def add(self, structure_to_add, method='zero'):
"""
Method to add two WakeFunction objects. The two structures are
@@ -331,60 +341,100 @@ class WakeFunction(ComplexData):
def __rmul__(self, factor):
return self.multiply(factor)
- def from_wakepotential(self, file_name, bunch_length, bunch_charge,
- freq_lim, component_type="long", divide_by=None,
- nout=None, trim=False, axis0='dist',
- axis0_scale=1e-3, axis1_scale=1e-12):
+ def to_impedance(self, freq_lim, nout=None, sigma=None, mu=None):
"""
- Compute a wake function from a wake potential file and load it to the
- WakeFunction object.
-
+ Return an Impedance object from the WakeFunction data.
+ The WakeFunction data is assumed to be sampled equally.
+
+ If both sigma and mu are given, the impedance data is divided by the
+ Fourier transform of a Gaussian distribution. It then can be used to go
+ from wake potential data to impedance.
+
Parameters
----------
- file_name : str
- Text file that contains wake potential data.
- bunch_length : float
- Spatial bunch length in [m].
- bunch_charge : float
- Bunch charge in [C].
freq_lim : float
The maximum frequency for calculating the impedance [Hz].
- component_type : str, optional
- Type of the wake: "long", "xdip", "xquad", ...
- divide_by : float, optional
- Divide the wake potential by a value. Mainly used to normalize
- transverse wake function by displacement.
- nout : int, optional
+ nout : int or float, 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
- input is padded with zeros. If nout < n, the inpu it cropped.
+ input is padded with zeros. If nout < n, the input it cropped.
Note that, for any nout value, nout//2+1 input points are required.
- 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.
- If a float is given, the pseudo wake function is trimmed from
- time <= trim to 0.
- axis0 : {'dist', 'time'}, optional
- Viariable of the data file's first column. Use 'dist' for spacial
- distance. Otherwise, 'time' for temporal distance.
- axis0_scale : float, optional
- Scale of the first column with respect to meter if axis0 = 'dist',
- or to second if axis0 = 'time'.
- axis1_scale : float, optional
- Scale of the wake potential (in second column) with respect to V/C.
+ The default is None.
+ sigma : float
+ RMS of the Gaussian distribution in [s].
+ The default is None.
+ mu : float
+ Offset of the Gaussian distribution in [s].
+ The default is None.
+
+ """
+ tau = np.array(self.data.index)
+ wp = np.array(self.data["real"])
+
+ if nout is None:
+ nout = len(tau)
+ else:
+ nout = int(nout)
+
+ # FT of wake potential
+ sampling = tau[1] - tau[0]
+ freq = sc.fft.rfftfreq(nout, sampling)
+ dtau = (tau[-1]-tau[0])/len(tau)
+ dft_wake = sc.fft.rfft(wp, n=nout, axis=0) * dtau
+
+ i_limit = freq < freq_lim
+ freq_trun = freq[i_limit]
+ dft_wake_trun = dft_wake[i_limit]
+
+ if (sigma is not None) and (mu is not None):
+ dft_rho_trun = np.exp(-0.5*(sigma*2*np.pi*freq_trun)**2 + \
+ 1j*mu*2*np.pi*freq_trun)
+ else:
+ dft_rho_trun = -1
+
+ if self.component_type == "long":
+ imp = dft_wake_trun/dft_rho_trun*-1
+ elif (self.component_type == "xdip") or (self.component_type == "ydip"):
+ imp = dft_wake_trun/dft_rho_trun*-1j
+ else:
+ raise NotImplementedError(self.component_type + " is not correct.")
+
+ imp = Impedance(variable=freq_trun, function=imp,
+ component_type=self.component_type)
+
+ return imp
+ def deconvolution(self, freq_lim, sigma, mu, nout=None):
"""
+ Compute a wake function from wake potential data.
+
+ The present data loaded into the WakeFunction object is assumed to be
+ an uniformly sampled wake potential from a Gaussian bunch distribution
+ with RMS bunch length sigma and offset mu.
- imp = Impedance()
- imp.from_wakepotential(file_name=file_name, bunch_length=bunch_length,
- bunch_charge=bunch_charge, freq_lim=freq_lim,
- component_type=component_type, divide_by=divide_by,
- axis0=axis0, axis0_scale=axis0_scale,
- axis1_scale=axis1_scale)
- wf = imp.to_wakefunction(nout=nout, trim=trim)
- self.__init__(variable=wf.data.index, function=wf.data["real"],
- component_type=component_type)
+ The offset mu depends on the code used to compute the wake potential:
+ - for CST, it is the first time sample, so self.data.index[0].
+ - for ABCI, it is "-1 * ISIG * SIG / c".
+
+ Parameters
+ ----------
+ freq_lim : float
+ The maximum frequency for calculating the impedance [Hz].
+ sigma : float
+ RMS of the Gaussian distribution in [s].
+ mu : float
+ Offset of the Gaussian distribution in [s].
+ nout : int or float, 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
+ input is padded with zeros. If nout < n, the input it cropped.
+ Note that, for any nout value, nout//2+1 input points are required.
+ The default is None.
+
+ """
+ imp = self.to_impedance(freq_lim, sigma=sigma, mu=mu)
+ wf = imp.to_wakefunction(nout=nout)
+ return wf
def plot(self):
"""
@@ -425,13 +475,12 @@ class Impedance(ComplexData):
Methods
-------
- from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
- Compute impedance from wake potential data and load it to the Impedance
- object.
loss_factor(sigma)
Compute the loss factor or the kick factor for a Gaussian bunch.
to_wakefunction()
Return a WakeFunction object from the impedance data.
+ plot()
+ Plot the impedance data.
"""
def __init__(self,
@@ -441,6 +490,10 @@ class Impedance(ComplexData):
self._component_type = component_type
self.data.index.name = 'frequency [Hz]'
self.initialize_coefficient()
+
+ def __repr__(self):
+ """Return representation of data"""
+ return f'Impedance of component_type {self.component_type}:\n {(self.data)!r}'
def add(self, structure_to_add, beta_x=1, beta_y=1, method='zero'):
"""
@@ -535,78 +588,7 @@ class Impedance(ComplexData):
raise TypeError("Impedance type not recognized.")
return kloss
-
- def from_wakepotential(self, file_name, bunch_length, bunch_charge,
- freq_lim, component_type="long", divide_by=None,
- axis0='dist', axis0_scale=1e-3, axis1_scale=1e-12):
- """
- Compute impedance from wake potential data and load it to the Impedance
- object.
-
- Parameters
- ----------
- file_name : str
- Text file that contains wake potential data.
- bunch_length : float
- Electron bunch lenth [m].
- bunch_charge : float
- Total bunch charge [C].
- freq_lim : float
- The maximum frequency for calculating the impedance [Hz].
- component_type : str, optional
- Type of the impedance: "long", "xdip", "xquad", ...
- divide_by : float, optional
- Divide the impedance by a value. Mainly used to normalize transverse
- impedance by displacement.
- axis0 : {'dist', 'time'}, optional
- Viariable of the data file's first column. Use 'dist' for spacial
- distance. Otherwise, 'time' for temporal distance.
- axis0_scale : float, optional
- Scale of the first column with respect to meter if axis0 = 'dist',
- or to second if axis0 = 'time'.
- axis1_scale : float, optional
- Scale of the wake potential (in second column) with respect to V/C.
-
- """
-
- s, wp0 = np.loadtxt(file_name, unpack=True)
- if axis0 == 'dist':
- tau = s*axis0_scale/c
- elif axis0 == 'time':
- tau = s*axis0_scale
- else:
- raise TypeError('Type of axis 0 not recognized.')
-
- wp = wp0 / axis1_scale
- if divide_by is not None:
- wp = wp / divide_by
-
- # FT of wake potential
- sampling = tau[1] - tau[0]
- freq = sc.fft.rfftfreq(len(tau), sampling)
- dtau = (tau[-1]-tau[0])/len(tau)
- dft_wake = sc.fft.rfft(wp) * dtau
-
- # FT of analytical bunch profile and analytical impedance
- sigma = bunch_length/c
- mu = tau[0]
-
- i_limit = freq < freq_lim
- freq_trun = freq[i_limit]
- dft_wake_trun = dft_wake[i_limit]
-
- dft_rho_trun = np.exp(-0.5*(sigma*2*np.pi*freq_trun)**2 + \
- 1j*mu*2*np.pi*freq_trun)*bunch_charge
- if component_type == "long":
- imp = dft_wake_trun/dft_rho_trun*-1*bunch_charge
- elif (component_type == "xdip") or (component_type == "ydip"):
- imp = dft_wake_trun/dft_rho_trun*-1j*bunch_charge
- else:
- raise NotImplementedError(component_type + " is not correct.")
-
- self.__init__(variable=freq_trun, function=imp,
- component_type=component_type)
-
+
def to_wakefunction(self, nout=None, trim=False):
"""
Return a WakeFunction object from the impedance data.
@@ -717,6 +699,27 @@ class WakeField:
def __init__(self, structure_list=None, name=None):
self.list_to_attr(structure_list)
self.name = name
+
+ def __repr__(self):
+ """Return representation of WakeField components."""
+ return f'WakeField {self.name} with components:\n {(list(self.components))!r}'
+
+ def __radd__(self, structure_to_add):
+ return self._add(structure_to_add)
+
+ def __add__(self, structure_to_add):
+ return self._add(structure_to_add)
+
+ def __iter__(self):
+ """Iterate over components"""
+ comp = [getattr(self, comp) for comp in self.components]
+ return comp.__iter__()
+
+ def _add(self, structure_to_add):
+ """Allow to add two WakeField element with different components."""
+ for comp in structure_to_add.components:
+ self.append_to_model(getattr(structure_to_add, comp))
+ return self
def append_to_model(self, structure_to_add):
"""
@@ -746,7 +749,7 @@ class WakeField:
def list_to_attr(self, structure_list):
"""
- Add list of Impedance/WakeFunction components to WakeField.
+ Add list of Impedance/WakeFunction components to WakeField.
Parameters
----------
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index f086eaa..7a4c9e7 100644
--- a/mbtrack2/tracking/particles.py
+++ b/mbtrack2/tracking/particles.py
@@ -580,6 +580,10 @@ class Beam:
def __iter__(self):
"""Iterate over all bunches"""
return self.bunch_list.__iter__()
+
+ def __repr__(self):
+ """Return representation of the beam filling pattern"""
+ return f'{list((self.filling_pattern))!r}'
@property
def not_empty(self):
diff --git a/mbtrack2/utilities/__init__.py b/mbtrack2/utilities/__init__.py
index f62f2fb..1a0a7e2 100644
--- a/mbtrack2/utilities/__init__.py
+++ b/mbtrack2/utilities/__init__.py
@@ -2,7 +2,8 @@
from mbtrack2.utilities.read_impedance import (read_CST,
read_IW2D,
read_IW2D_folder,
- read_ABCI)
+ read_ABCI,
+ read_ECHO2D)
from mbtrack2.utilities.misc import (effective_impedance,
yokoya_elliptic,
beam_loss_factor,
diff --git a/mbtrack2/utilities/read_impedance.py b/mbtrack2/utilities/read_impedance.py
index ebe089a..77595e5 100644
--- a/mbtrack2/utilities/read_impedance.py
+++ b/mbtrack2/utilities/read_impedance.py
@@ -286,4 +286,36 @@ def read_ABCI(file, azimuthal=False):
wake = WakeField(wake_list)
- return wake
\ No newline at end of file
+ return wake
+
+def read_ECHO2D(file, component_type='long'):
+ """
+ Read ECHO2D text file format (after matlab post-processing) into a
+ WakeFunction object.
+
+ Parameters
+ ----------
+ file : str
+ Path to the text file to read.
+ component_type : str, optional
+ Type of the WakeFunction object to load.
+ Default is 'long'.
+
+ Returns
+ -------
+ result : WakeFunction object.
+ Data from file.
+ """
+
+ df = pd.read_csv(file, delim_whitespace=True,
+ header = None, names = ["Distance","Wake"])
+ df["Time"] = df["Distance"]/100/c
+ df["Wake"] = df["Wake"]*1e12
+ if component_type != 'long':
+ df["Wake"] = df["Wake"]*-1
+ df.set_index("Time", inplace = True)
+ result = WakeFunction(variable = df.index,
+ function = df["Wake"],
+ component_type=component_type)
+
+ return result
\ No newline at end of file
--
GitLab