ImpedanceModel object

Add a ImpedanceModel object which derives from Element and is used to sum up different Wakefield elements
Fix uppercase module in resistive_wall.py

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

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