diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index 1a1ebb3344d6a57b16d3d3cac83540c8c53ac6f5..e604d9082a8c8ac8b6de3b9ad5b4f2ec80e49c24 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -11,9 +11,12 @@ import warnings
import re
import pandas as pd
import numpy as np
+import matplotlib.pyplot as plt
+import collective_effects.tools as tools
from scipy.interpolate import interp1d
from tracking.element import Element
-
+from copy import deepcopy
+from scipy.integrate import trapz
class ComplexData:
"""
@@ -79,12 +82,14 @@ class ComplexData:
initial_data = pd.concat([self.data, data_to_concat], sort=True)
initial_data = initial_data[~initial_data.index.duplicated(keep='first')]
+ initial_data = initial_data.sort_index()
data_to_add = pd.concat(
[structure_to_add.data,
self.data.index.to_frame().set_index(index_name)],
sort=True)
data_to_add = data_to_add[~data_to_add.index.duplicated(keep='first')]
+ data_to_add = data_to_add.sort_index()
if method == 'common':
max_variable = min(structure_to_add.data.index.max(),
@@ -345,6 +350,10 @@ class Wakefield:
Add Impedance/WakeFunction to Wakefield.
list_to_attr(structure_list)
Add list of Impedance/WakeFunction to Wakefield.
+ add_wakefields(wake1, beta1, wake2, beta2)
+ Add two Wakefields taking into account beta functions.
+ add_several_wakefields(wakefields, beta)
+ Add a list of Wakefields taking into account beta functions.
"""
def __init__(self, structure_list=None, name=None):
@@ -396,6 +405,78 @@ class Wakefield:
"""
return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
+ @staticmethod
+ def add_wakefields(wake1, beta1, wake2, beta2):
+ """
+ Add two Wakefields taking into account beta functions.
+
+ Parameters
+ ----------
+ wake1 : Wakefield
+ Wakefield to add.
+ beta1 : array of shape (2,)
+ Beta functions at wake1 postion.
+ wake2 : Wakefield
+ Wakefield to add.
+ beta2 : array of shape (2,)
+ Beta functions at wake2 postion.
+
+ Returns
+ -------
+ wake_sum : Wakefield
+ Wakefield with summed components.
+
+ """
+ wake_sum = deepcopy(wake1)
+ for component_name1 in wake1.impedance_components:
+ impedance1 = getattr(wake_sum, component_name1)
+ weight1 = ((beta1[0] ** impedance1.power_x) *
+ (beta1[1] ** impedance1.power_y))
+ setattr(wake_sum, component_name1, weight1*impedance1)
+
+ for component_name2 in wake2.impedance_components:
+ impedance2 = getattr(wake2, component_name2)
+ weight2 = ((beta2[0] ** impedance2.power_x) *
+ (beta2[1] ** impedance2.power_y))
+ try:
+ impedance1 = getattr(wake_sum, component_name2)
+ setattr(wake_sum, component_name2, impedance1 +
+ weight2*impedance2)
+ except AttributeError:
+ setattr(wake_sum, component_name2, weight2*impedance2)
+
+ return wake_sum
+
+ @staticmethod
+ def add_several_wakefields(wakefields, beta):
+ """
+ Add a list of Wakefields taking into account beta functions.
+
+ Parameters
+ ----------
+ wakefields : list of Wakefield
+ Wakefields to add.
+ beta : array of shape (len(wakefields), 2)
+ Beta functions at the Wakefield postions.
+
+ Returns
+ -------
+ wake_sum : Wakefield
+ Wakefield with summed components..
+
+ """
+ if len(wakefields) == 1:
+ return wakefields[0]
+ elif len(wakefields) > 1:
+ wake_sum = Wakefield.add_wakefields(wakefields[0], beta[:,0],
+ wakefields[1], beta[:,1])
+ for i in range(len(wakefields) - 2):
+ wake_sum = Wakefield.add_wakefields(wake_sum, [1 ,1],
+ wakefields[i+2], beta[:,i+2])
+ return wake_sum
+ else:
+ raise ValueError("Error in input.")
+
class ImpedanceModel(Element):
"""
@@ -408,8 +489,40 @@ class ImpedanceModel(Element):
Wakefields to add to the model.
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfields.
-
+ Attributes
+ ----------
+ wakefields : list of Wakefield objects
+ Wakefields in the model.
+ positions : array
+ Wakefields positions.
+ names : array
+ Names of (unique) Wakefield objects.
+ sum : Wakefield
+ Sum of every Wakefield in the model.
+ sum_"name" : Wakefield
+ Sum of every Wakefield with the same "name".
+ sum_names : array
+ Names of attributes where the wakefields are summed by name.
+
+ Methods
+ -------
+ add(wakefield_list, wakefiled_postions)
+ Add elements to the model.
+ add_multiple_elements(wakefield, wakefiled_postions)
+ Add the same element at different locations to the model.
+ sum_elements()
+ Sum all wakefields into self.sum.
+ update_name_list()
+ Update self.names with uniques names of self.wakefields.
+ find_wakefield(name)
+ Return indexes of wakefields with the same name in self.wakefields.
+ sum_by_name(name)
+ Sum the elements with the same name in the model into sum_name.
+ sum_by_name_all(name)
+ Sum all the elements with the same name in the model into sum_name.
+ plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None)
+ Plot the contributions of different kind of wakefields.
"""
def __init__(self, ring, wakefield_list=None, wakefiled_postions=None):
@@ -417,6 +530,8 @@ class ImpedanceModel(Element):
self.optics = self.ring.optics
self.wakefields = []
self.positions = np.array([])
+ self.names = np.array([])
+ self.sum_names = np.array([])
self.add(wakefield_list, wakefiled_postions)
def track(self, beam):
@@ -430,27 +545,75 @@ class ImpedanceModel(Element):
raise NotImplementedError
def sum_elements(self):
+ """Sum all wakefields into self.sum"""
+ beta = self.optics.beta(self.positions)
+ self.sum = Wakefield.add_several_wakefields(self.wakefields, beta)
+ if "sum" not in self.sum_names:
+ self.sum_names = np.append(self.sum_names, "sum")
+
+ def update_name_list(self):
+ """Update self.names with uniques names of self.wakefields."""
+ for wakefield in self.wakefields:
+ if wakefield.name is None:
+ continue
+ if wakefield.name not in self.names:
+ self.names = np.append(self.names, wakefield.name)
+
+ def count_elements(self):
+ """Count number of each type of wakefield in the model."""
+ self.count = np.zeros(len(self.names))
+ for wakefield in self.wakefields:
+ if wakefield.name is not None:
+ self.count += (wakefield.name == self.names).astype(int)
+
+ def find_wakefield(self, name):
"""
- Sum all the elements in the model into sum_wakefield.
+ Return indexes of wakefields with the same name in self.wakefields.
+
+ Parameters
+ ----------
+ name : str
+ Wakefield name.
+
+ Returns
+ -------
+ index : list
+ Index of positions in self.wakefields.
+
"""
- beta = self.optics.beta(self.positions)
- list_impedance_components = ["Zlong","Zxdip","Zydip","Zxquad","Zyquad"]
- sum_wakefield = Wakefield()
- for component_name in list_impedance_components:
- sum_imp = Impedance(variable=np.array([0,1]),
- function=np.array([0, 0]),
- impedance_type=component_name[1:])
- for index, wakefield in enumerate(self.wakefields):
- try:
- impedance = wakefield.__getattribute__(component_name)
- weight = ((beta[0,index] ** impedance.power_x) *
- (beta[1,index] ** impedance.power_y))
- sum_imp += weight*impedance
- except AttributeError:
- pass
- sum_wakefield.append_to_model(sum_imp)
- self.sum_wakefield = sum_wakefield
+ index = []
+ for i, wakefield in enumerate(self.wakefields):
+ if wakefield.name == name:
+ index.append(i)
+ return index
+
+ def sum_by_name(self, name):
+ """
+ Sum the elements with the same name in the model into sum_name.
+ Parameters
+ ----------
+ name : str
+ Name of the wakefield to sum.
+ """
+ attribute_name = "sum_" + name
+ index = self.find_wakefield(name)
+ beta = self.optics.beta(self.positions[index])
+ wakes = []
+ for i in index:
+ wakes.append(self.wakefields[i])
+ wake_sum = Wakefield.add_several_wakefields(wakes, beta)
+ setattr(self, attribute_name, wake_sum)
+ if attribute_name not in self.sum_names:
+ self.sum_names = np.append(self.sum_names, attribute_name)
+
+ def sum_by_name_all(self):
+ """
+ Sum all the elements with the same name in the model into sum_name.
+ """
+ for name in self.names:
+ self.sum_by_name(name)
+
def add(self, wakefield_list, wakefiled_postions):
"""
Add elements to the model.
@@ -462,10 +625,115 @@ class ImpedanceModel(Element):
wakefiled_postions : list
Longitudinal positions corresponding to the added Wakfields.
"""
-
if (wakefield_list is not None) and (wakefiled_postions is not None):
for wakefield in wakefield_list:
self.wakefields.append(wakefield)
for position in wakefiled_postions:
- self.positions = np.append(self.positions, position)
\ No newline at end of file
+ self.positions = np.append(self.positions, position)
+
+ self.update_name_list()
+
+ def add_multiple_elements(self, wakefield, wakefiled_postions):
+ """
+ Add the same element at different locations to the model.
+
+ Parameters
+ ----------
+ wakefield : Wakefield object
+ Wakefield to add to the model.
+ wakefiled_postions : list
+ Longitudinal positions corresponding to the added Wakfield.
+ """
+ for position in wakefiled_postions:
+ self.positions = np.append(self.positions, position)
+ self.wakefields.append(wakefield)
+
+ self.update_name_list()
+
+ def plot_area(self, Z_type="Zlong", component="real", sigma=None,
+ attr_list=None):
+ """
+ Plot the contributions of different kind of wakefields.
+
+ Parameters
+ ----------
+ Z_type : str, optional
+ Type of impedance to plot.
+ component : str, optional
+ Component to plot, can be "real" or "imag".
+ sigma : float, optional
+ RMS bunch length in [s] to use for the spectral density. If equal
+ to None, the spectral density is not plotted.
+ attr_list : list or array of str, optional
+ Attributes to plot.
+
+ """
+ if attr_list is None:
+ attr_list = self.sum_names[self.sum_names != "sum"]
+
+ # 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)
+ sorted_index = area.argsort()[::-1]
+
+ # Init fig
+ fig = plt.figure()
+ ax = fig.gca()
+ zero_impedance = getattr(self.sum, Z_type)*0
+ total_imp = 0
+ legend = []
+
+ if sigma is not None:
+ legend.append("Spectral density for sigma = " + str(sigma) + " s")
+
+ # Main plot
+ 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, 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)
+
+ if sigma is not None:
+ spect = tools.spectral_density(zero_impedance.data.index, sigma)
+ spect = spect/spect.max()*total_imp.max()
+ ax.plot(sum_imp.data.index, spect, 'r', linewidth=2.5)
+
+ ax.legend(legend)
+ ax.set_xlabel("Frequency [Hz]")
+
+ if Z_type == "Zlong":
+ ax.set_ylabel(Z_type + " [Ohm] - " + component + " part")
+ else:
+ ax.set_ylabel(Z_type + " [Ohm/m] - " + component + " part")
+
+ 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
+
+ shape_array = (len(list_components), len(attr_list))
+ loss_array = np.zeros(shape_array)
+
+ 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] = tools.loss_factor(impedance, sigma)
+ except AttributeError:
+ pass
+
+ summary = pd.DataFrame(loss_array.T, index=attr_list,
+ columns=list_components)
+ return summary
+
\ No newline at end of file