diff --git a/collective_effects/__init__.py b/collective_effects/__init__.py
index 867d0d105cb7a1ffa8141c96eb0c45f72a3d7975..36d158e3ebce92f0ad91e7778a2186906b9dbc9e 100644
--- a/collective_effects/__init__.py
+++ b/collective_effects/__init__.py
@@ -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
diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index 095090b7c6509cf443b943494c1efd738086d0b0..971c759acf5dc1be131194c9033bd578a71dad19 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -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