Rework ImpedanceModel

Simplify and optimize adding element into ImpedanceModel.

mbtrack2/impedance/impedance_model.py

@@ -6,6 +6,7 @@ import pandas as pd
 import numpy as np
 import matplotlib.pyplot as plt
 import pickle
+from copy import deepcopy
 from scipy.integrate import trapz
 from scipy.interpolate import interp1d
 from mpl_toolkits.axes_grid1.inset_locator import inset_axes
@@ -14,52 +15,43 @@ from mbtrack2.utilities.misc import (beam_loss_factor, effective_impedance,
 from mbtrack2.utilities.spectrum import (beam_spectrum, 
                                          gaussian_bunch_spectrum, 
                                          spectral_density)
-from mbtrack2.impedance.wakefield import WakeField
-from mbtrack2.tracking.element import Element
 
-class ImpedanceModel(Element):
+class ImpedanceModel():
     """
     Define the impedance model of the machine.
     
+    The model must be completed with successive add(...) calls, then 
+    compute_sum() must be run.
+    
     Parameters
     ----------
     ring : Synchrotron object
-    wakefield_list : list of WakeField objects
-        WakeFields to add to the model.
-    wakefiled_positions : 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.
+    positions : list of arrays
+        Positions corresponding the different WakeFields.
+    names : list of str
+        Names of the WakeField objects.
     sum : WakeField
-        Sum of every WakeField in the model.
+        Sum of every WakeField in the model weigthed by beta functions.
     sum_"name" : WakeField
-        Sum of every Wakefield with the same "name".
+        Sum of the "name" Wakefield weigthed by beta functions.
     sum_names : array
         Names of attributes where the WakeFields are summed by name.    
     
     Methods
     -------
-    add(wakefield_list, wakefiled_positions)
-        Add elements to the model.
-    add_multiple_elements(wakefield, wakefiled_positions)
-        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.
+    add(wakefield, positions, name)
+        Add the same WakeField object at different locations to the model.
+    sum_beta(wake, beta)
+        Weight a WakeField object by an array of beta functions.
+    compute_sum_names()
+        Compute the weighted WakeField for each WakeField object type.
+    compute_sum()
+        Compute the sum of all weighted WakeField into self.sum.
     plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None)
         Plot the contributions of different kind of WakeFields.
     save(file)
@@ -68,132 +60,101 @@ class ImpedanceModel(Element):
         Load impedance model from file.
     """
     
-    def __init__(self, ring, wakefield_list=None, wakefiled_positions=None):
+    def __init__(self, ring):
         self.ring = ring
         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_positions)
-        
-    def track(self, beam):
+        self.positions = []
+        self.names = []
+        self.sum_names = []
+
+    def add(self, wakefield, positions, name=None):
         """
-        Track a beam object through this Element.
+        Add the same WakeField object at different locations to the model.
         
         Parameters
         ----------
-        beam : Beam object
+        wakefield : WakeField
+            WakeField object to add to the model.
+        positions : array, float or int
+            Array of longitudinal positions where the elements are loacted.
+        name : str, optional
+            Name of the element type. If None, the name of the WakeField object 
+            is used. The default is None.
+
+        Returns
+        -------
+        None.
+
         """
-        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):
+        self.wakefields.append(wakefield)
+        self.positions.append(positions)
+        if name is None:
+            name = wakefield.name
+        if name is None:
+            raise ValueError("Please give a valid name.")
+        if name not in self.names:
+            self.names.append(name)
+        else:
+            raise ValueError("This name is already taken.")
+        
+    @staticmethod
+    def sum_beta(wake, beta):
         """
-        Return indexes of WakeFields with the same name in self.wakefields.
+        Weight a WakeField object by an array of beta functions.
 
         Parameters
         ----------
-        name : str
-            WakeField name.
+        wake : WakeField
+            WakeField element object.
+        beta : array of shape (2, N)
+            Beta function at the locations of the elements.
 
         Returns
         -------
-        index : list
-            Index of positions in self.wakefields.
+        wake_sum : WakeField
+            WakeField object weighted by beta functions.
 
         """
-        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.
+        wake_sum = deepcopy(wake)
+        for component_name in wake.components:
+            comp = getattr(wake_sum, component_name)
+            weight = ((beta[0,:] ** comp.power_x) * 
+                      (beta[1,:] ** comp.power_y))
+            setattr(wake_sum, component_name, weight.sum()*comp)
+        return wake_sum
+    
+    def compute_sum_names(self):
         """
-        for name in self.names:
-            self.sum_by_name(name)
-                    
-    def add(self, wakefield_list, wakefiled_positions):
+        Compute the weighted WakeField for each WakeField object type.
+        The new summed WakeField object are set to into self.sum_name.
         """
-        Add elements to the model.
-
-        Parameters
-        ----------
-        wakefield_list : list of WakeField objects
-            WakeFields to add to the model.
-        wakefiled_positions : list
-            Longitudinal positions corresponding to the added Wakfields.
+        for idx, wake in enumerate(self.wakefields):
+            attribute_name = "sum_" + self.names[idx]
+            beta = self.optics.beta(self.positions[idx])
+            wake_sum = self.sum_beta(wake, beta)
+            setattr(self, attribute_name, wake_sum)
+            self.sum_names.append(attribute_name)
+        
+    def compute_sum(self):
         """
-        if (wakefield_list is not None) and (wakefiled_positions is not None):
-            for wakefield in wakefield_list:
-                self.wakefields.append(wakefield)
-                
-            for position in wakefiled_positions:
-                self.positions = np.append(self.positions, position)
-                
-        self.update_name_list()
-                
-    def add_multiple_elements(self, wakefield, wakefiled_positions):
+        Compute the sum of all weighted WakeField into self.sum.
+        self.compute_sum_names must be called before this.
         """
-        Add the same element at different locations to the model.
+        self.compute_sum_names()
+        for i, name in enumerate(self.sum_names):
+            if i==0:
+                self.sum = deepcopy(getattr(self, name))
+            else:
+                wake2 = getattr(self, name)
+                for component_name2 in wake2.components: 
+                    comp2 = getattr(wake2, component_name2)
+                    try:
+                        comp1 = getattr(self.sum, component_name2)
+                        setattr(self.sum, component_name2, comp1 + comp2)
+                    except AttributeError:
+                        setattr(self.sum, component_name2, comp2)
 
-        Parameters
-        ----------
-        WakeField : WakeField object
-            WakeField to add to the model.
-        wakefiled_positions : list
-            Longitudinal positions corresponding to the added Wakfield.
-        """
-        for position in wakefiled_positions:
-            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, zoom=False):
         """