diff --git a/collective_effects/resistive_wall.py b/collective_effects/resistive_wall.py
index 764a59e61ae17515fcab23fe19a6460322f7fbbc..95ae59f68698a0d0bc8f51374bd0b33d66493714 100644
--- a/collective_effects/resistive_wall.py
+++ b/collective_effects/resistive_wall.py
@@ -8,7 +8,7 @@ Define resistive wall elements based on the Wakefield class.
import numpy as np
from scipy.constants import mu_0, epsilon_0, c
-from CollectiveEffect.wakefield import Wakefield, Impedance
+from collective_effects.wakefield import Wakefield, Impedance
def skin_depth(omega, rho, mu_r = 1, epsilon_r = 1):
""" General formula for the skin depth
diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index dea656d1b3bb4ac92c4c312f2bdb877a981a2d06..ea3fe51dc102b035efaa385d5c4c1ec40e0cfb42 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -12,6 +12,8 @@ import re
import pandas as pd
import numpy as np
from scipy.interpolate import interp1d
+from tracking.element import Element
+
class ComplexData:
"""
@@ -34,7 +36,8 @@ class ComplexData:
index=variable)
self.data.index.name = 'variable'
- def add(self, structure_to_add, method='zero', interp_kind='cubic'):
+ def add(self, structure_to_add, method='zero', interp_kind='cubic',
+ index_name="variable"):
"""
Method to add two structures. If the data don't have the same length,
different cases are handled.
@@ -55,6 +58,8 @@ class ComplexData:
interp_kind : str, optional
interpolation method which is passed to pandas and to
scipy.interpolate.interp1d.
+ index_name : str, optional
+ name of the dataframe index passed to the method
Returns
-------
@@ -72,13 +77,13 @@ class ComplexData:
initial_data = pd.concat(
[self.data,
- structure_to_add.data.index.to_frame().set_index('variable')],
+ structure_to_add.data.index.to_frame().set_index(index_name)],
sort=True)
initial_data = initial_data[~initial_data.index.duplicated(keep='first')]
data_to_add = pd.concat(
[structure_to_add.data,
- self.data.index.to_frame().set_index('variable')],
+ self.data.index.to_frame().set_index(index_name)],
sort=True)
data_to_add = data_to_add[~data_to_add.index.duplicated(keep='first')]
@@ -249,12 +254,15 @@ class Impedance(ComplexData):
compared so that the addition is self-consistent.
Beta functions can be precised as well.
"""
+
if isinstance(structure_to_add, (int, float, complex)):
- result = super().add(structure_to_add, method=method)
+ result = super().add(structure_to_add, method=method,
+ index_name="frequency [Hz]")
elif isinstance(structure_to_add, Impedance):
if (self.impedance_type == structure_to_add.impedance_type):
weight = (beta_x ** self.power_x) * (beta_y ** self.power_y)
- result = super().add(structure_to_add * weight, method=method)
+ result = super().add(structure_to_add * weight, method=method,
+ index_name="frequency [Hz]")
else:
warnings.warn(('The two Impedance objects do not have the '
'same coordinates or plane or type. '
@@ -322,10 +330,9 @@ class Wakefield:
and other informations: beta functions.
"""
- def __init__(self, betax=1, betay=1, structure_list=[]):
- self.betax = betax
- self.betay = betay
+ def __init__(self, structure_list=[], name=None):
self.list_to_attr(structure_list)
+ self.name = name
def append_to_model(self, structure_to_add):
list_of_attributes = dir(self)
@@ -358,68 +365,51 @@ class Wakefield:
"""
return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
- def add(self, element_to_add):
- """
- This function adds two Wakefield objects.
- The different impedance components are weighted and added.
- The result is a Wakefield object with all the components but the
- beta function equal to 1.
- """
-
- if not isinstance(element_to_add, Wakefield):
- raise TypeError("You can only add Wakefield objects to other"
- "Wakefields objects.")
-
- list_of_summed_impedances = []
- list_of_components_added = []
-
- betax1 = self.betax
- betay1 = self.betay
- betax2 = element_to_add.betax
- betay2 = element_to_add.betay
-
- # We first go through element1 impedance components and add them
- # to the element2 components. If the component is missing in element2,
- # we then simply weight the element1 component and add it to the model.
- for component_name in self.list_impedance_components:
+
+class ImpedanceModel(Element):
+
+ def __init__(self, ring, wakefields_to_add=None, wakefileds_postions=None):
+ self.ring = ring
+ self.optics = self.ring.optics
+ self.wakefields = []
+ self.positions = np.array([])
+ self.add(wakefields_to_add, wakefileds_postions)
- element1_impedance = self.__getattribute__(component_name)
+ def track(self, beam):
+ """
+ Track a beam object through this Element.
- element1_weight = ((betax1 ** element1_impedance.power_x)
- * (betay1 ** element1_impedance.power_y))
-
- try:
- element2_impedance = element_to_add.__getattribute__(component_name)
- list_of_components_added.append(component_name)
- element12_impedance_sum = element1_weight * element1_impedance
- element12_impedance_sum = element12_impedance_sum.add(element2_impedance, betax2, betay2)
- except:
- element12_impedance_sum = element1_weight * element1_impedance
+ Parameters
+ ----------
+ beam : Beam object
+ """
+ raise NotImplementedError
- list_of_summed_impedances.append(element12_impedance_sum)
+ def sum_elements(self):
- # We now go through the components which are unique to element2 and
- # add them to the impedance model, with the beta function weighting
- missing_components_list = list(set(element_to_add.list_impedance_components) -
- set(list_of_components_added))
-
- for component_name in missing_components_list:
-
- element2_impedance = element_to_add.__getattribute__(component_name)
- element2_weight = ((betax2 ** element2_impedance.power_x)
- * (betay2 ** element2_impedance.power_y))
- element12_impedance_sum = element2_weight * element2_impedance
-
- list_of_summed_impedances.append(element12_impedance_sum)
+ 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
- # Gather everything in an Wakefield object
- sum_wakefield = Wakefield(betax=1., betay=1.,
- structure_list=list_of_summed_impedances)
-
- return sum_wakefield
+ def add(self, wakefield_list, position_list):
- def __radd__(self, structure_to_add):
- return self.add(structure_to_add)
-
- def __add__(self, structure_to_add):
- return self.add(structure_to_add)
+ if (wakefield_list is not None) and (position_list is not None):
+ for wakefield in wakefield_list:
+ self.wakefields.append(wakefield)
+
+ for position in position_list:
+ self.positions = np.append(self.positions, position)
\ No newline at end of file