Add new methods to the WakeField class

WakeField.drop allows to delete a component or a list of component from the WakeField.
WakeField.save and WakeField.load add a shortcut to pickle WakeField objects.

mbtrack2/impedance/impedance_model.py

@@ -531,6 +531,9 @@ class ImpedanceModel(Element):
     def save(self, file):
         """
         Save impedance model to file.
+        
+        The same pandas version is needed on both saving and loading computer
+        for the pickle to work.
 
         Parameters
         ----------
@@ -550,6 +553,9 @@ class ImpedanceModel(Element):
     def load(self, file):
         """
         Load impedance model from file.
+        
+        The same pandas version is needed on both saving and loading computer
+        for the pickle to work.
 
         Parameters
         ----------

mbtrack2/impedance/wakefield.py

@@ -5,6 +5,7 @@ wake functions.
 """
 
 import warnings
+import pickle
 import pandas as pd
 import numpy as np
 import scipy as sc
@@ -667,6 +668,12 @@ class WakeField:
         Add two WakeFields taking into account beta functions.
     add_several_wakefields(wakefields, beta)
         Add a list of WakeFields taking into account beta functions.
+    drop(component)
+        Delete a component or a list of component from the WakeField.
+    save(file)
+        Save WakeField element to file.
+    load(file)
+        Load WakeField element from file.
     """
 
     def __init__(self, structure_list=None, name=None):
@@ -736,6 +743,79 @@ class WakeField:
                  "Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad"]
         return np.array([comp for comp in dir(self) if comp in valid])
     
+    def drop(self, component):
+        """
+        Delete a component or a list of component from the WakeField.
+
+        Parameters
+        ----------
+        component : str or list of str
+            Component or list of components to drop.
+            If "Z", drop all impedance components.
+            If "W"", drop all wake function components.
+
+        Returns
+        -------
+        None.
+
+        """
+        if component == "Z":
+            component = self.impedance_components
+        elif component == "W":
+            component = self.wake_components
+        
+        if isinstance(component, str):
+            delattr(self, component)
+        elif isinstance(component, list) or isinstance(component, np.ndarray):
+            for comp in component:
+                delattr(self, comp)
+        else:
+            raise TypeError("component should be a str or list of str.")
+            
+    def save(self, file):
+        """
+        Save WakeField element to file.
+        
+        The same pandas version is needed on both saving and loading computer
+        for the pickle to work.
+
+        Parameters
+        ----------
+        file : str
+            File where the WakeField element is saved.
+
+        Returns
+        -------
+        None.
+
+        """
+        with open(file,"wb") as f:
+            pickle.dump(self, f)
+            
+    @staticmethod
+    def load(file):
+        """
+        Load WakeField element from file.
+        
+        The same pandas version is needed on both saving and loading computer
+        for the pickle to work.
+
+        Parameters
+        ----------
+        file : str
+            File where the WakeField element is saved.
+
+        Returns
+        -------
+        wakefield : WakeField
+            Loaded wakefield.
+
+        """
+        with open(file, 'rb') as f:
+            wakefield = pickle.load(f)
+            
+        return wakefield
+    
     @staticmethod
     def add_wakefields(wake1, beta1, wake2, beta2):
         """