Skip to content
Snippets Groups Projects
Commit 9ca25437 authored by Gamelin Alexis's avatar Gamelin Alexis
Browse files

Misc improvements on wakefield.py 

Add effective impedance calculation to ImpedanceModel
Rework plot_area to support transverse impedance and add better legend and scaling
double_sided_impedance : new function to add negative frequencies to impedance spectrums
effective_impedance : tuneS as input, use double_sided_impedance and xi input is used
beam_loss_factor : uses double_sided_impedance
Correct a few bugs in wakefield.py
parent cb8b05d8
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
collective_effects/__init__.py
+ 2
− 1
View file @ 9ca25437
......@@ -11,4 +11,5 @@ from mbtrack2.collective_effects.tapers import StupakovRectangularTaper, Stupako
from mbtrack2.collective_effects.wakefield import ComplexData, Impedance, WakeFunction, ImpedanceModel, Wakefield
from mbtrack2.collective_effects.wakefield import read_CST, read_IW2D, spectral_density
from mbtrack2.collective_effects.wakefield import gaussian_bunch_spectrum, beam_spectrum, effective_impedance
from mbtrack2.collective_effects.wakefield import yokoya_elliptic, beam_loss_factor
\ No newline at end of file
from mbtrack2.collective_effects.wakefield import yokoya_elliptic, beam_loss_factor
from mbtrack2.collective_effects.wakefield import double_sided_impedance
\ No newline at end of file
collective_effects/wakefield.py
+ 158
− 64
View file @ 9ca25437
......@@ -711,11 +711,24 @@ class ImpedanceModel(Element):
if attr_list is None:
attr_list = self.sum_names[self.sum_names != "sum"]
# manage legend
Ztype_dict = {"Zlong":0, "Zxdip":1, "Zydip":2, "Zxquad":3, "Zyquad":4}
scale = [1e-3, 1e-6, 1e-6, 1e-6, 1e-6]
label_list = [r"$Z_{long} \; [k\Omega]$",
r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Dip} \; [M\Omega]$",
r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Dip} \; [M\Omega]$",
r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Quad} \; [M\Omega]$",
r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Quad} \; [M\Omega]$"]
leg = Ztype_dict[Z_type]
# sort plot by decresing area size
area = np.zeros((len(attr_list),))
for index, attr in enumerate(attr_list):
sum_imp = getattr(getattr(self, attr), Z_type)
area[index] = trapz(sum_imp.data[component], sum_imp.data.index)
try:
sum_imp = getattr(getattr(self, attr), Z_type)
area[index] = trapz(sum_imp.data[component], sum_imp.data.index)
except AttributeError:
pass
sorted_index = area.argsort()[::-1]
# Init fig
......@@ -732,14 +745,17 @@ class ImpedanceModel(Element):
for index in sorted_index:
attr = attr_list[index]
# Set all impedances with common indexes using + zero_impedance
sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
ax.fill_between(sum_imp.data.index*1e-9, total_imp,
total_imp + sum_imp.data[component])
total_imp += sum_imp.data[component]
if attr_list is None:
legend.append(attr[4:])
else:
legend.append(attr)
try:
sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
ax.fill_between(sum_imp.data.index*1e-9, total_imp,
total_imp + sum_imp.data[component]*scale[leg])
total_imp += sum_imp.data[component]*scale[leg]
if attr[:4] == "sum_":
legend.append(attr[4:])
else:
legend.append(attr)
except AttributeError:
pass
if sigma is not None:
spect = spectral_density(zero_impedance.data.index, sigma)
......@@ -748,36 +764,83 @@ class ImpedanceModel(Element):
ax.legend(legend, loc="upper left")
ax.set_xlabel("Frequency [GHz]")
if Z_type == "Zlong":
ax.set_ylabel(Z_type + " [Ohm] - " + component + " part")
else:
ax.set_ylabel(Z_type + " [Ohm/m] - " + component + " part")
ax.set_ylabel(label_list[leg] + " - " + component + " part")
ax.set_title(label_list[leg] + " - " + component + " part")
return fig
def summary(self, sigma, attr_list=None, list_components=None):
if attr_list is None:
attr_list = self.sum_names
if list_components is None:
list_components = self.sum.impedance_components
def effective_impedance(self, m, mu, sigma, M, tuneS, xi=None,
mode="Hermite"):
"""
Compute the longitudinal and transverse effective impedance.
Parameters
----------
mu : int
coupled bunch mode number, goes from 0 to (M-1) where M is the
number of bunches
m : int
head-tail (or azimutal/synchrotron) mode number
sigma : float
RMS bunch length in [s]
M : int
Number of bunches.
tuneS : float
Synchrotron tune.
xi : float, optional
(non-normalized) chromaticity
mode: str, optional
type of the mode taken into account for the computation:
-"Hermite" modes for Gaussian bunches
Returns
-------
summary : DataFrame
Longitudinal and transverse effective impedance.
"""
attr_list = self.sum_names
shape_array = (len(list_components), len(attr_list))
loss_array = np.zeros(shape_array)
eff_array = np.zeros((len(attr_list),3), dtype=complex)
for i, attr in enumerate(attr_list):
for j, component_name in enumerate(list_components):
try:
impedance = getattr(getattr(self, attr), component_name)
loss_array[j,i] = impedance.loss_factor(sigma)
except AttributeError:
pass
summary = pd.DataFrame(loss_array.T, index=attr_list,
columns=list_components)
try:
impedance = getattr(getattr(self, attr), "Zlong")
eff_array[i,0] = effective_impedance(self.ring, impedance,
m, mu, sigma, M, tuneS,
xi, mode)
except AttributeError:
pass
try:
impedance = getattr(getattr(self, attr), "Zxdip")
eff_array[i,1] = effective_impedance(self.ring, impedance,
m, mu, sigma, M, tuneS,
xi, mode)
except AttributeError:
pass
try:
impedance = getattr(getattr(self, attr), "Zydip")
eff_array[i,2] = effective_impedance(self.ring, impedance,
m, mu, sigma, M, tuneS,
xi, mode)
except AttributeError:
pass
eff_array[:,0] = eff_array[:,0]*self.ring.omega0*1e3
eff_array[:,1] = eff_array[:,1]*1e-3
eff_array[:,2] = eff_array[:,2]*1e-3
summary = pd.DataFrame(eff_array, index=attr_list,
columns=["Z/n [mOhm]",
"sum betax x Zxeff [kOhm]",
"sum betay x Zyeff [kOhm]"])
return summary
def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
"""
Compute the beam and bunch loss factor and energy losses for each type
......@@ -814,9 +877,9 @@ class ImpedanceModel(Element):
fmin = -1*fmax
f = np.linspace(fmin, fmax, int(n_points))
beam_spect = tools.beam_spectrum(f, M, bunch_spacing, sigma= sigma)
beam_spect = beam_spectrum(f, M, bunch_spacing, sigma= sigma)
bunch_spect = tools.Gaussian_bunch_spectrum(f, sigma)
bunch_spect = gaussian_bunch_spectrum(f, sigma)
attr_list = self.sum_names
......@@ -825,8 +888,8 @@ class ImpedanceModel(Element):
for i, attr in enumerate(attr_list):
try:
impedance = getattr(getattr(self, attr), "Zlong")
loss_array[i,0] = tools.beam_loss_factor(impedance, f, beam_spect, self.ring)
loss_array[i,1] = tools.beam_loss_factor(impedance, f, bunch_spect, self.ring)
loss_array[i,0] = beam_loss_factor(impedance, f, beam_spect, self.ring)
loss_array[i,1] = beam_loss_factor(impedance, f, bunch_spect, self.ring)
except AttributeError:
pass
......@@ -983,7 +1046,7 @@ def beam_spectrum(frequency, M, bunch_spacing, sigma=None,
"""
if bunch_spectrum is None:
bunch_spectrum = Gaussian_bunch_spectrum(frequency, sigma)
bunch_spectrum = gaussian_bunch_spectrum(frequency, sigma)
beam_spectrum = (bunch_spectrum * np.exp(1j*np.pi*frequency *
bunch_spacing*(M-1)) *
......@@ -993,7 +1056,8 @@ def beam_spectrum(frequency, M, bunch_spacing, sigma=None,
return beam_spectrum
def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
def effective_impedance(ring, imp, m, mu, sigma, M, tuneS, xi=None,
mode="Hermite"):
"""
Compute the effective (longitudinal or transverse) impedance.
Formulas from Eq. (1) and (2) p238 of [1].
......@@ -1009,6 +1073,10 @@ def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
head-tail (or azimutal/synchrotron) mode number
sigma : float
RMS bunch length in [s]
M : int
Number of bunches.
tuneS : float
Synchrotron tune.
xi : float, optional
(non-normalized) chromaticity
mode: str, optional
......@@ -1021,7 +1089,9 @@ def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
effective impedance in [ohm] or in [ohm/m] depanding on the impedance
type.
[1] Handbook of accelerator physics and engineering, 3rd printing.
References
----------
[1] : Handbook of accelerator physics and engineering, 3rd printing.
"""
if not isinstance(imp, Impedance):
......@@ -1030,9 +1100,7 @@ def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
fmin = imp.data.index.min()
fmax = imp.data.index.max()
if fmin > 0:
raise ValueError("A double sided impedance spectrum, with positive and"
" negative frequencies, is needed to compute the "
"effective impedance.")
double_sided_impedance(imp)
if mode == "Hermite":
def h(f):
......@@ -1041,15 +1109,13 @@ def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
else:
raise NotImplementedError("Not implemanted yet.")
M = 416
pmax = fmax/(ring.f0 * M) - 50
pmin = fmin/(ring.f0 * M) + 50
pmax = fmax/(ring.f0 * M) - 1
pmin = fmin/(ring.f0 * M) + 1
p = np.arange(pmin,pmax+1)
if imp.impedance_type == "long":
fp = ring.f0*(p*M + mu + m*ring.tune[2])
fp = ring.f0*(p*M + mu + m*tuneS)
fp = fp[np.nonzero(fp)] # Avoid division by 0
num = np.sum( imp(fp) * h(fp) / (fp*2*np.pi) )
den = np.sum( h(fp) )
......@@ -1058,11 +1124,13 @@ def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
elif imp.impedance_type == "xdip" or imp.impedance_type == "ydip":
if imp.impedance_type == "xdip":
tuneXY = ring.tune[0]
xi = ring.chro[0]
if xi is None :
xi = ring.chro[0]
elif imp.impedance_type == "ydip":
tuneXY = ring.tune[1]
xi = ring.chro[1]
fp = ring.f0*(p*M + mu + tuneXY + m*ring.tuneS)
if xi is None:
xi = ring.chro[1]
fp = ring.f0*(p*M + mu + tuneXY + m*tuneS)
f_xi = xi/ring.eta*ring.f0
num = np.sum( imp(fp) * h(fp - f_xi) )
den = np.sum( h(fp - f_xi) )
......@@ -1158,19 +1226,7 @@ def beam_loss_factor(impedance, frequency, spectrum, ring):
pmin = np.floor(impedance.data.index.min()/ring.f0)
if pmin >= 0:
# Add negative part
negative_index = impedance.data.index*-1
negative_data = impedance.data.set_index(negative_index)
negative_data["imag"] = -1*negative_data["imag"]
try:
negative_data = negative_data.drop(0)
except KeyError:
pass
all_data = impedance.data.append(negative_data)
all_data = all_data.sort_index()
impedance.data = all_data
double_sided_impedance(impedance)
pmin = -1*pmax
p = np.arange(pmin+1,pmax)
......@@ -1185,5 +1241,43 @@ def beam_loss_factor(impedance, frequency, spectrum, ring):
kloss_beam = ring.f0 * np.sum(ReZ*spect(pf0))
return kloss_beam
def double_sided_impedance(impedance):
"""
Add negative frequency points to single sided impedance spectrum following
symetries depending on impedance type.
Parameters
----------
impedance : Impedance object
Single sided impedance.
"""
fmin = impedance.data.index.min()
if fmin >= 0:
negative_index = impedance.data.index*-1
negative_data = impedance.data.set_index(negative_index)
imp_type = impedance.impedance_type
if imp_type == "long":
negative_data["imag"] = -1*negative_data["imag"]
elif (imp_type == "xdip") or (imp_type == "ydip"):
negative_data["real"] = -1*negative_data["real"]
elif (imp_type == "xquad") or (imp_type == "yquad"):
negative_data["real"] = -1*negative_data["real"]
else:
raise ValueError("Wrong impedance type")
try:
negative_data = negative_data.drop(0)
except KeyError:
pass
all_data = impedance.data.append(negative_data)
all_data = all_data.sort_index()
impedance.data = all_data
\ No newline at end of file
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

The place to be for your source code