diff --git a/collective_effects/__init__.py b/collective_effects/__init__.py
index 67c3d4f6011d9fc22f490119b06933d996d50655..bdea55f7291697a9d610b142d34d1f23d82ade5b 100644
--- a/collective_effects/__init__.py
+++ b/collective_effects/__init__.py
@@ -7,9 +7,11 @@ Created on Tue Jan 14 12:25:30 2020
 
 from mbtrack2.collective_effects.instabilities import mbi_threshold, cbi_threshold
 from mbtrack2.collective_effects.resistive_wall import skin_depth, CircularResistiveWall
+from mbtrack2.collective_effects.resonator import Resonator, PureInductive, PureResistive
 from mbtrack2.collective_effects.tapers import StupakovRectangularTaper, StupakovCircularTaper
-from mbtrack2.collective_effects.wakefield import ComplexData, Impedance, WakeFunction, ImpedanceModel, Wakefield
-from mbtrack2.collective_effects.wakefield import read_CST, read_IW2D, spectral_density
-from mbtrack2.collective_effects.wakefield import gaussian_bunch_spectrum, beam_spectrum, effective_impedance
-from mbtrack2.collective_effects.wakefield import yokoya_elliptic, beam_loss_factor
-from mbtrack2.collective_effects.wakefield import double_sided_impedance
\ No newline at end of file
+from mbtrack2.collective_effects.wakefield import ComplexData, Impedance, WakeFunction, WakeField
+from mbtrack2.collective_effects.utilities import read_CST, read_IW2D, spectral_density
+from mbtrack2.collective_effects.utilities import gaussian_bunch_spectrum, beam_spectrum, effective_impedance
+from mbtrack2.collective_effects.utilities import yokoya_elliptic, beam_loss_factor
+from mbtrack2.collective_effects.utilities import double_sided_impedance
+from mbtrack2.collective_effects.impedance_model import ImpedanceModel
\ No newline at end of file
diff --git a/collective_effects/impedance_model.py b/collective_effects/impedance_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e1db30d0094dafd625d10f9cc09d9633b51fb1a
--- /dev/null
+++ b/collective_effects/impedance_model.py
@@ -0,0 +1,424 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Sep 28 15:18:23 2020
+
+@author: gamelina
+"""
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.integrate import trapz
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes
+from mbtrack2.collective_effects.utilities import (beam_spectrum, gaussian_bunch_spectrum, 
+                                                   beam_loss_factor, spectral_density,
+                                                   effective_impedance)
+from mbtrack2.collective_effects.wakefield import WakeField
+from mbtrack2.tracking.element import Element
+
+class ImpedanceModel(Element):
+    """
+    Define the impedance model of the machine.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    wakefield_list : list of WakeField objects
+        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):
+        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_postions)
+        
+    def track(self, beam):
+        """
+        Track a beam object through this Element.
+        
+        Parameters
+        ----------
+        beam : Beam object
+        """
+        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):
+        """
+        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.
+
+        """
+        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.
+
+        Parameters
+        ----------
+        wakefield_list : list of WakeField objects
+            WakeFields to add to the model.
+        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)
+                
+        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, zoom=False):
+        """
+        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.
+        zoom : bool
+            If True, add a zoomed plot on top right corner.
+
+        """
+        if attr_list is None:
+            attr_list = self.sum_names[self.sum_names != "sum"]
+        
+        # manage legend
+        Ztype_dict = {"Zlong":0, "Zxdip":1, "Zydip":2, "Zxquad":3, "Zyquad":4}
+        scale = [1e-3, 1e-6, 1e-6, 1e-6, 1e-6]
+        label_list =  [r"$Z_{long} \; [k\Omega]$",
+                       r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Dip} \; [M\Omega]$",
+                       r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Dip} \; [M\Omega]$",
+                       r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Quad} \; [M\Omega]$",
+                       r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Quad} \; [M\Omega]$"]
+        leg = Ztype_dict[Z_type]
+        
+        # sort plot by decresing area size        
+        area = np.zeros((len(attr_list),))
+        for index, attr in enumerate(attr_list):
+            try:
+                sum_imp = getattr(getattr(self, attr), Z_type)
+                area[index] = trapz(sum_imp.data[component], sum_imp.data.index)
+            except AttributeError:
+                pass
+        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
+            try:
+                sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
+                ax.fill_between(sum_imp.data.index*1e-9, total_imp, 
+                                total_imp + sum_imp.data[component]*scale[leg])
+                total_imp += sum_imp.data[component]*scale[leg]
+                if attr[:4] == "sum_":
+                    legend.append(attr[4:])
+                else:
+                    legend.append(attr)
+            except AttributeError:
+                pass
+            
+        if sigma is not None:
+            spect = spectral_density(zero_impedance.data.index, sigma)
+            spect = spect/spect.max()*total_imp.max()
+            ax.plot(sum_imp.data.index*1e-9, spect, 'r', linewidth=2.5)
+        
+        ax.legend(legend, loc="upper left")            
+        ax.set_xlabel("Frequency [GHz]")
+        ax.set_ylabel(label_list[leg] + " - " + component + " part")
+        ax.set_title(label_list[leg] + " - " + component + " part")
+        
+        if zoom is True:
+            in_ax = inset_axes(ax,
+                            width="30%", # width = 30% of parent_bbox
+                            height=1.5, # height : 1 inch
+                            loc=1)
+            
+            total_imp = 0
+            for index in  sorted_index:
+                attr = attr_list[index]
+                # Set all impedances with common indexes using + zero_impedance
+                try:
+                    sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
+                    in_ax.fill_between(sum_imp.data.index*1e-3, total_imp, 
+                                    total_imp + sum_imp.data[component]*1e-9)
+                    total_imp += sum_imp.data[component]*1e-9
+                except AttributeError:
+                    pass
+            in_ax.set_xlim([0, 200])
+            in_ax.set_xlabel("Frequency [kHz]")
+            in_ax.set_ylabel(r"$[G\Omega]$")
+                
+        return fig
+            
+    def effective_impedance(self, m, mu, sigma, M, tuneS, xi=None, 
+                            mode="Hermite"):
+        """
+        Compute the longitudinal and transverse effective impedance.
+
+        Parameters
+        ----------
+        mu : int
+            coupled bunch mode number, goes from 0 to (M-1) where M is the
+            number of bunches
+        m : int
+            head-tail (or azimutal/synchrotron) mode number
+        sigma : float
+            RMS bunch length in [s]
+        M : int
+            Number of bunches.
+        tuneS : float
+            Synchrotron tune.
+        xi : float, optional
+            (non-normalized) chromaticity
+        mode: str, optional
+            type of the mode taken into account for the computation:
+            -"Hermite" modes for Gaussian bunches
+
+        Returns
+        -------
+        summary : DataFrame
+            Longitudinal and transverse effective impedance.
+
+        """
+
+        attr_list = self.sum_names
+        
+        eff_array = np.zeros((len(attr_list),3), dtype=complex)
+        
+        for i, attr in enumerate(attr_list):
+            try:
+                impedance = getattr(getattr(self, attr), "Zlong")
+                eff_array[i,0] = effective_impedance(self.ring, impedance, 
+                                                     m, mu, sigma, M, tuneS,
+                                                     xi, mode)
+            except AttributeError:
+                pass
+            
+            try:
+                impedance = getattr(getattr(self, attr), "Zxdip")
+                eff_array[i,1] = effective_impedance(self.ring, impedance, 
+                                                     m, mu, sigma, M, tuneS,
+                                                     xi, mode)
+            except AttributeError:
+                pass
+            
+            try:
+                impedance = getattr(getattr(self, attr), "Zydip")
+                eff_array[i,2] = effective_impedance(self.ring, impedance, 
+                                                     m, mu, sigma, M, tuneS,
+                                                     xi, mode)
+            except AttributeError:
+                pass
+            
+        eff_array[:,0] = eff_array[:,0]*self.ring.omega0*1e3
+        eff_array[:,1] = eff_array[:,1]*1e-3
+        eff_array[:,2] = eff_array[:,2]*1e-3
+        
+        summary = pd.DataFrame(eff_array, index=attr_list, 
+                               columns=["Z/n [mOhm]", 
+                                        "sum betax x Zxeff [kOhm]", 
+                                        "sum betay x Zyeff [kOhm]"])
+        
+        return summary
+
+
+    def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
+        """
+        Compute the beam and bunch loss factor and energy losses for each type 
+        of element in the model assuming Gaussian bunches and constant spacing
+        between bunches.
+
+        Parameters
+        ----------
+        sigma : float
+            RMS bunch length in [s].
+        M : int
+            Number of bunches in the beam.
+        bunch_spacing : float
+            Time between two bunches in [s].
+        I : float
+            Total beam current in [A].
+        n_points : float, optional
+            Number of points used in the frequency spectrums.
+
+        Returns
+        -------
+        summary : Dataframe
+            Contains the beam and bunch loss factor and energy loss for the 
+            full model and for each type of different component.
+
+        """
+        
+        fmax = self.sum.Zlong.data.index.max()
+        fmin = self.sum.Zlong.data.index.min()
+        
+        Q = I*self.ring.T0/M
+        
+        if fmin >= 0:
+            fmin = -1*fmax
+        f = np.linspace(fmin, fmax, int(n_points))
+
+        beam_spect = beam_spectrum(f, M, bunch_spacing, sigma= sigma)
+        
+        bunch_spect = gaussian_bunch_spectrum(f, sigma)
+        
+        attr_list = self.sum_names
+        
+        loss_array = np.zeros((len(attr_list),2))
+        
+        for i, attr in enumerate(attr_list):
+            try:
+                impedance = getattr(getattr(self, attr), "Zlong")
+                loss_array[i,0] = beam_loss_factor(impedance, f, beam_spect, self.ring)
+                loss_array[i,1] = beam_loss_factor(impedance, f, bunch_spect, self.ring)
+            except AttributeError:
+                pass
+
+        loss_array = loss_array*1e-12
+        summary = pd.DataFrame(loss_array, index=attr_list, 
+                               columns=["loss factor (beam) [V/pC]", "loss factor (bunch) [V/pC]"])
+        
+        summary["P (beam) [W]"] = summary["loss factor (beam) [V/pC]"]*1e12*Q**2/(self.ring.T0)
+        summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"]*1e12*Q**2/(self.ring.T0)*M
+                
+        return summary
diff --git a/collective_effects/resistive_wall.py b/collective_effects/resistive_wall.py
index 240e30e268dfa956dd1a2dbc005e88cfe6897dca..2df32f47833f7caba38c24dc733523e1097e44be 100644
--- a/collective_effects/resistive_wall.py
+++ b/collective_effects/resistive_wall.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """
-Define resistive wall elements based on the Wakefield class.
+Define resistive wall elements based on the WakeField class.
 
 @author: Alexis Gamelin
 @date: 14/01/2020
@@ -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 mbtrack2.collective_effects.wakefield import Wakefield, Impedance
+from mbtrack2.collective_effects.wakefield import WakeField, Impedance
 
 def skin_depth(omega, rho, mu_r = 1, epsilon_r = 1):
     """ General formula for the skin depth
@@ -32,9 +32,9 @@ def skin_depth(omega, rho, mu_r = 1, epsilon_r = 1):
     return delta
     
 
-class CircularResistiveWall(Wakefield):
+class CircularResistiveWall(WakeField):
     """
-    Resistive wall Wakefield element for a circular beam pipe, approximated
+    Resistive wall WakeField element for a circular beam pipe, approximated
     formulas from Eq. (2.77) of Chao book [1].
     
     Parameters
diff --git a/collective_effects/resonator.py b/collective_effects/resonator.py
new file mode 100644
index 0000000000000000000000000000000000000000..f11948ee3abe5b923d4d5f66bbcf2e483b64df09
--- /dev/null
+++ b/collective_effects/resonator.py
@@ -0,0 +1,192 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines the impedances and wake functions from the resonator model 
+based on the WakeField class.
+
+@author: Alexis Gamelin, Watanyu Foosang
+@date: 25/09/2020
+"""
+
+import numpy as np
+from mbtrack2.collective_effects.wakefield import (WakeField, Impedance, 
+                                                   WakeFunction)
+
+class Resonator(WakeField):
+    def __init__(self, Rs, fr, Q, plane, n_wake=1e6, n_imp=1e6, imp_freq_lim=100e9):
+        """
+        Resonator model WakeField element which computes the impedance and the 
+        wake function in both longitudinal and transverse case.
+
+        Parameters
+        ----------
+        Rs : float
+            Shunt impedance in [ohm].
+        fr : float
+            Resonance frequency in [Hz].
+        Q : float
+            Quality factor.
+        plane : str
+            Plane on which the resonator is used: "long", "x" or "y".
+        n_wake : int or float, optional
+            Number of points used in the wake function.
+        n_imp : int or float, optional
+            Number of points used in the impedance.
+        imp_freq_lim : float, optional
+            Maximum frequency used in the impedance.
+            
+        References
+        ----------
+        [1]  B. W. Zotter and S. A. Kheifets, "Impedances and Wakes in High-Energy 
+        Particle Ac-celerators", Eq. (3.10) and (3.15), pp.51-53.
+
+        """
+        super().__init__()
+        self.Rs = Rs
+        self.fr = fr
+        self.wr = 2 * np.pi * self.fr
+        self.Q = Q
+        self.n_wake = int(n_wake)
+        self.n_imp = int(n_imp)
+        self.imp_freq_lim = imp_freq_lim
+        self.plane = plane
+
+        self.timestop = round(np.log(1000)/self.wr*2*self.Q, 12)
+        
+        if self.Q >= 0.5:
+            self.Q_p = np.sqrt(self.Q**2 - 0.25)
+        else:
+            self.Q_p = np.sqrt(0.25 - self.Q**2)
+            
+        self.wr_p = (self.wr*self.Q_p)/self.Q
+        
+        if self.plane == "long":
+            
+            freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
+            imp = Impedance(variable=freq, 
+                            function=self.long_impedance(freq),
+                            impedance_type="long")
+            super().append_to_model(imp)
+            
+            time = np.linspace(start=0, stop=self.timestop, num=self.n_wake)
+            wake = WakeFunction(variable=time,
+                                function=self.long_wake_function(time),
+                                wake_type="long")
+            super().append_to_model(wake)
+            
+        elif self.plane == "x" or self.plane == "y" :
+            
+            freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
+            imp = Impedance(variable=freq, 
+                            function=self.transverse_impedance(freq),
+                            impedance_type=self.plane + "dip")
+            super().append_to_model(imp)
+            
+            time = np.linspace(start=0, stop=self.timestop, num=self.n_wake)
+            wake = WakeFunction(variable=time,
+                                function=self.transverse_wake_function(time),
+                                wake_type=self.plane + "dip")
+            super().append_to_model(wake)
+        else:
+            raise ValueError("Plane must be: long, x or y")
+        
+    def long_wake_function(self, t):
+        if self.Q >= 0.5:
+            return ( (self.wr * self.Rs / self.Q) * 
+                    np.exp(-1* self.wr * t / (2 * self.Q) ) *
+                     (np.cos(self.wr_p * t) - 
+                      np.sin(self.wr_p * t) / (2 * self.Q_p) ) )
+                        
+        elif self.Q < 0.5:
+            return ( (self.wr * self.Rs / self.Q) * 
+                    np.exp(-1* self.wr * t / (2 * self.Q) ) *
+                     (np.cosh(self.wr_p * t) - 
+                      np.sinh(self.wr_p * t) / (2 * self.Q_p) ) )
+                            
+    def long_impedance(self, f):
+        return self.Rs / (1 + 1j * self.Q * (f/self.fr - self.fr/f))
+    
+    def transverse_impedance(self, f):
+        return self.Rs * self.fr / f / (
+            1 + 1j * self.Q * (f / self.fr - self.fr / f) )
+    
+    def transverse_wake_function(self, t):
+        if self.Q >= 0.5:
+            return (self.wr * self.Rs / self.Q_p * 
+                    np.exp(-1 * t * self.wr / 2 / self.Q_p) *
+                    np.sin(self.wr_p * t) )
+        else:
+            return (self.wr * self.Rs / self.Q_p * 
+                    np.exp(-1 * t * self.wr / 2 / self.Q_p) *
+                    np.sinh(self.wr_p * t) )
+    
+class PureInductive(WakeField):
+    """
+    Pure inductive Wakefield element which computes associated longitudinal 
+    impedance and wake function.
+    
+    Parameters
+    ----------
+    L : float
+        Inductance value in [Ohm/Hz].
+    n_wake : int or float, optional
+        Number of points used in the wake function.
+    n_imp : int or float, optional
+        Number of points used in the impedance.
+    imp_freq_lim : float, optional
+        Maximum frequency used in the impedance. 
+    nout, trim : see Impedance.to_wakefunction
+    """
+    def __init__(self, L, n_wake=1e6, n_imp=1e6, imp_freq_lim=1e11, nout=None,
+                 trim=False):
+        self.L = L
+        self.n_wake = int(n_wake)
+        self.n_imp = int(n_imp)
+        self.imp_freq_lim = imp_freq_lim
+        
+        freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
+        imp = Impedance(variable=freq, 
+                        function=self.long_impedance(freq),
+                        impedance_type="long")
+        super().append_to_model(imp)
+        
+        wf = imp.to_wakefunction(nout=nout, trim=trim)
+        super().append_to_model(wf)
+        
+    def long_impedance(self, f):
+        return 1j*self.L*f
+    
+class PureResistive(WakeField):
+    """
+    Pure resistive Wakefield element which computes associated longitudinal 
+    impedance and wake function.
+    
+    Parameters
+    ----------
+    R : float
+        Resistance value in [Ohm].
+    n_wake : int or float, optional
+        Number of points used in the wake function.
+    n_imp : int or float, optional
+        Number of points used in the impedance.
+    imp_freq_lim : float, optional
+        Maximum frequency used in the impedance. 
+    nout, trim : see Impedance.to_wakefunction
+    """
+    def __init__(self, R, n_wake=1e6, n_imp=1e6, imp_freq_lim=1e11, nout=None,
+                 trim=False):
+        self.R = R
+        self.n_wake = int(n_wake)
+        self.n_imp = int(n_imp)
+        self.imp_freq_lim = imp_freq_lim
+        
+        freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
+        imp = Impedance(variable=freq, 
+                        function=self.long_impedance(freq),
+                        impedance_type="long")
+        super().append_to_model(imp)
+        
+        wf = imp.to_wakefunction(nout=nout, trim=trim)
+        super().append_to_model(wf)
+        
+    def long_impedance(self, f):
+        return self.R
\ No newline at end of file
diff --git a/collective_effects/tapers.py b/collective_effects/tapers.py
index 34dde1855a4ab3b3b2fcb0f220da80010286ab3b..8d944e2b84c772ceb99823fc57783c41daabf209 100644
--- a/collective_effects/tapers.py
+++ b/collective_effects/tapers.py
@@ -8,11 +8,11 @@ Created on Tue Mar 31 13:15:36 2020
 from scipy.constants import mu_0, epsilon_0, c, pi
 import numpy as np
 from scipy.integrate import trapz
-from mbtrack2.collective_effects.wakefield import Wakefield, Impedance
+from mbtrack2.collective_effects.wakefield import WakeField, Impedance
 
-class StupakovRectangularTaper(Wakefield):
+class StupakovRectangularTaper(WakeField):
     """
-    Rectangular vertical taper Wakefield element, using the low frequency 
+    Rectangular vertical taper WakeField element, using the low frequency 
     approxmiation. Assume constant taper angle. Formulas from [1].
     
     ! Valid for low w
@@ -136,9 +136,9 @@ class StupakovRectangularTaper(Wakefield):
         
         return -1j*pi*self.Z0/4*integral
     
-class StupakovCircularTaper(Wakefield):
+class StupakovCircularTaper(WakeField):
     """
-    Circular taper Wakefield element, using the low frequency 
+    Circular taper WakeField element, using the low frequency 
     approxmiation. Assume constant taper angle. Formulas from [1].
     
     ! Valid for low w
diff --git a/collective_effects/utilities.py b/collective_effects/utilities.py
new file mode 100644
index 0000000000000000000000000000000000000000..11803ea3092ffd52104506a5496a459579a992ee
--- /dev/null
+++ b/collective_effects/utilities.py
@@ -0,0 +1,392 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines utilities functions, helping to deals with the 
+collective_effects module.
+
+@author: Alexis Gamelin
+@date: 28/09/2020
+"""
+
+import pandas as pd
+import numpy as np
+from scipy.interpolate import interp1d
+from mbtrack2.collective_effects.wakefield import Impedance
+
+def read_CST(file, impedance_type='long', divide_by=None):
+    """
+    Read CST file format into an Impedance object.
+    
+    Parameters
+    ----------
+    file : str
+        Path to the file to read.
+    impedance_type : str, optional
+        Type of the Impedance object.
+    divide_by : float, optional
+        Divide the impedance by a value. Mainly used to normalize transverse 
+        impedance by displacement.
+        
+    Returns
+    -------
+    result : Impedance object
+        Data from file.
+    """
+    df = pd.read_csv(file, comment="#", header = None, sep = "\t", 
+                    names = ["Frequency","Real","Imaginary"])
+    df["Frequency"] = df["Frequency"]*1e9 
+    if divide_by is not None:
+        df["Real"] = df["Real"]/divide_by
+        df["Imaginary"] = df["Imaginary"]/divide_by
+    if impedance_type == "long":
+        df["Real"] = np.abs(df["Real"])
+    df.set_index("Frequency", inplace = True)
+    result = Impedance(variable = df.index,
+                       function = df["Real"] + 1j*df["Imaginary"],
+                       impedance_type=impedance_type)
+    return result
+
+def read_IW2D(file, impedance_type='long'):
+    """
+    Read IW2D file format into an Impedance object.
+    
+    Parameters
+    ----------
+    file : str
+        Path to the file to read.
+    impedance_type : str, optional
+        Type of the Impedance object.
+        
+    Returns
+    -------
+    result : Impedance object
+        Data from file.
+    """
+    df = pd.read_csv(file, sep = ' ', header = None, 
+                     names = ["Frequency","Real","Imaginary"], skiprows=1)
+    df.set_index("Frequency", inplace = True)
+    result = Impedance(variable = df.index,
+                       function = df["Real"] + 1j*df["Imaginary"],
+                       impedance_type=impedance_type)
+    return result
+
+def spectral_density(frequency, sigma, m = 1, mode="Hermite"):
+    """
+    Compute the spectral density of different modes for various values of the
+    head-tail mode number, based on Table 1 p238 of [1].
+    
+    Parameters
+    ----------
+    frequency : list or numpy array
+        sample points of the spectral density in [Hz]
+    sigma : float
+        RMS bunch length in [s]
+    m : int, optional
+        head-tail (or azimutal/synchrotron) mode number
+    mode: str, optional
+        type of the mode taken into account for the computation:
+        -"Hermite" modes for Gaussian bunches
+
+    Returns
+    -------
+    numpy array
+    
+    References
+    ----------
+    [1] : Handbook of accelerator physics and engineering, 3rd printing.
+    """
+    
+    if mode == "Hermite":
+        return (2*np.pi*frequency*sigma)**(2*m)*np.exp(
+                -1*(2*np.pi*frequency*sigma)**2)
+    else:
+        raise NotImplementedError("Not implemanted yet.")
+        
+def gaussian_bunch_spectrum(frequency, sigma): 
+    """
+    Compute a Gaussian bunch spectrum [1].
+
+    Parameters
+    ----------
+    frequency : array
+        sample points of the beam spectrum in [Hz].
+    sigma : float
+        RMS bunch length in [s].
+
+    Returns
+    -------
+    bunch_spectrum : array
+        Bunch spectrum sampled at points given in frequency.
+        
+    References
+    ----------
+    [1] : Gamelin, A. (2018). Collective effects in a transient microbunching 
+    regime and ion cloud mitigation in ThomX. p86, Eq. 4.19
+    """
+    return np.exp(-1/2*(2*np.pi*frequency)**2*sigma**2)
+        
+        
+def beam_spectrum(frequency, M, bunch_spacing, sigma=None, 
+                  bunch_spectrum=None): 
+    """
+    Compute the beam spectrum assuming constant spacing between bunches [1].
+
+    Parameters
+    ----------
+    frequency : list or numpy array
+        sample points of the beam spectrum in [Hz].
+    M : int
+        Number of bunches.
+    bunch_spacing : float
+        Time between two bunches in [s].
+    sigma : float, optional
+        If bunch_spectrum is None then a Gaussian bunch with sigma RMS bunch 
+        length in [s] is assumed.
+    bunch_spectrum : array, optional
+        Bunch spectrum sampled at points given in frequency.
+
+    Returns
+    -------
+    beam_spectrum : array
+
+    References
+    ----------
+    [1] Rumolo, G - Beam Instabilities - CAS - CERN Accelerator School: 
+        Advanced Accelerator Physics Course - 2014, Eq. 9
+    """
+    
+    if bunch_spectrum is None:
+        bunch_spectrum = gaussian_bunch_spectrum(frequency, sigma)
+        
+    beam_spectrum = (bunch_spectrum * np.exp(1j*np.pi*frequency * 
+                                             bunch_spacing*(M-1)) * 
+                     np.sin(M*np.pi*frequency*bunch_spacing) / 
+                     np.sin(np.pi*frequency*bunch_spacing))
+    
+    return beam_spectrum
+    
+    
+def effective_impedance(ring, imp, m, mu, sigma, M, tuneS, xi=None, 
+                        mode="Hermite"):
+    """
+    Compute the effective (longitudinal or transverse) impedance. 
+    Formulas from Eq. (1) and (2) p238 of [1].
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    imp : Impedance object
+    mu : int
+        coupled bunch mode number, goes from 0 to (M-1) where M is the
+        number of bunches
+    m : int
+        head-tail (or azimutal/synchrotron) mode number
+    sigma : float
+        RMS bunch length in [s]
+    M : int
+        Number of bunches.
+    tuneS : float
+        Synchrotron tune.
+    xi : float, optional
+        (non-normalized) chromaticity
+    mode: str, optional
+        type of the mode taken into account for the computation:
+        -"Hermite" modes for Gaussian bunches
+
+    Returns
+    -------
+    Zeff : float 
+        effective impedance in [ohm] or in [ohm/m] depanding on the impedance
+        type.
+        
+    References
+    ----------
+    [1] : Handbook of accelerator physics and engineering, 3rd printing.
+    """
+    
+    if not isinstance(imp, Impedance):
+        raise TypeError("{} should be an Impedance object.".format(imp))
+        
+    fmin = imp.data.index.min()
+    fmax = imp.data.index.max()
+    if fmin > 0:
+        double_sided_impedance(imp)
+        
+    if mode == "Hermite":
+        def h(f):
+            return spectral_density(frequency=f, sigma=sigma, m=m,
+                                    mode="Hermite")
+    else:
+        raise NotImplementedError("Not implemanted yet.")
+    
+    pmax = fmax/(ring.f0 * M) - 1
+    pmin = fmin/(ring.f0 * M) + 1
+    
+    p = np.arange(pmin,pmax+1)
+    
+    if imp.impedance_type == "long":
+        fp = ring.f0*(p*M + mu + m*tuneS)
+        fp = fp[np.nonzero(fp)] # Avoid division by 0
+        num = np.sum( imp(fp) * h(fp) / (fp*2*np.pi) )
+        den = np.sum( h(fp) )
+        Zeff = num/den
+        
+    elif imp.impedance_type == "xdip" or imp.impedance_type == "ydip":
+        if imp.impedance_type == "xdip":
+            tuneXY = ring.tune[0]
+            if xi is None :
+                xi = ring.chro[0]
+        elif imp.impedance_type == "ydip":
+            tuneXY = ring.tune[1]
+            if xi is None:
+                xi = ring.chro[1]
+        fp = ring.f0*(p*M + mu + tuneXY + m*tuneS)
+        f_xi = xi/ring.eta*ring.f0
+        num = np.sum( imp(fp) * h(fp - f_xi) )
+        den = np.sum( h(fp - f_xi) )
+        Zeff = num/den
+    else:
+        raise TypeError("Effective impedance is only defined for long, xdip"
+                        " and ydip impedance type.")
+        
+    return Zeff
+
+
+def yokoya_elliptic(x_radius , y_radius):
+    """
+    Compute Yokoya factors for an elliptic beam pipe.
+    Function adapted from N. Mounet IW2D.
+
+    Parameters
+    ----------
+    x_radius : float
+        Horizontal semi-axis of the ellipse in [m].
+    y_radius : float
+        Vertical semi-axis of the ellipse in [m].
+
+    Returns
+    -------
+    yoklong : float
+        Yokoya factor for the longitudinal impedance.
+    yokxdip : float
+        Yokoya factor for the dipolar horizontal impedance.
+    yokydip : float
+        Yokoya factor for the dipolar vertical impedance.
+    yokxquad : float
+        Yokoya factor for the quadrupolar horizontal impedance.
+    yokyquad : float
+        Yokoya factor for the quadrupolar vertical impedance.
+    """
+    if y_radius < x_radius:
+        small_semiaxis = y_radius
+        large_semiaxis = x_radius
+    else:
+        small_semiaxis = x_radius
+        large_semiaxis = y_radius
+
+    # read Yokoya factors interpolation file
+    # BEWARE: columns are ratio, dipy, dipx, quady, quadx
+    yokoya_file = pd.read_csv("collective_effects/data/" +
+                              "Yokoya_elliptic_from_Elias_USPAS.csv")
+    ratio_col = yokoya_file["x"]
+    # compute semi-axes ratio (first column of this file)
+    ratio = (large_semiaxis - small_semiaxis)/(large_semiaxis + small_semiaxis)
+
+    # interpolate Yokoya file at the correct ratio
+    yoklong = 1
+    
+    if y_radius < x_radius:
+        yokydip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
+        yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
+        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
+        yokxquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])
+    else:
+        yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
+        yokydip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
+        yokxquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
+        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])        
+
+    return (yoklong, yokxdip, yokydip, yokxquad, yokyquad)
+
+def beam_loss_factor(impedance, frequency, spectrum, ring):
+    """
+    Compute "beam" loss factor using the beam spectrum, uses a sum instead of 
+    integral compared to loss_factor [1].
+
+    Parameters
+    ----------
+    impedance : Impedance of type "long"
+    frequency : array
+        Sample points of spectrum.
+    spectrum : array
+        Beam spectrum to consider.
+    ring : Synchrotron object
+
+    Returns
+    -------
+    kloss_beam : float
+        Beam loss factor in [V/C].
+        
+    References
+    ----------
+    [1] : Handbook of accelerator physics and engineering, 3rd printing. 
+        Eq (3) p239.
+    """
+    pmax = np.floor(impedance.data.index.max()/ring.f0)
+    pmin = np.floor(impedance.data.index.min()/ring.f0)
+    
+    if pmin >= 0:
+        double_sided_impedance(impedance)
+        pmin = -1*pmax
+    
+    p = np.arange(pmin+1,pmax)    
+    pf0 = p*ring.f0
+    ReZ = np.real(impedance(pf0))
+    spectral_density = np.abs(spectrum)**2
+    # interpolation of the spectrum is needed to avoid problems liked to 
+    # division by 0
+    # computing the spectrum directly to the frequency points gives
+    # wrong results
+    spect = interp1d(frequency, spectral_density)
+    kloss_beam = ring.f0 * np.sum(ReZ*spect(pf0))
+    
+    return kloss_beam
+
+def double_sided_impedance(impedance):
+    """
+    Add negative frequency points to single sided impedance spectrum following
+    symetries depending on impedance type.
+
+    Parameters
+    ----------
+    impedance : Impedance object
+        Single sided impedance.
+    """
+    fmin = impedance.data.index.min()
+    
+    if fmin >= 0:
+        negative_index = impedance.data.index*-1
+        negative_data = impedance.data.set_index(negative_index)
+        
+        imp_type = impedance.impedance_type
+        
+        if imp_type == "long":
+            negative_data["imag"] = -1*negative_data["imag"]
+            
+        elif (imp_type == "xdip") or (imp_type == "ydip"):
+            negative_data["real"] = -1*negative_data["real"]
+        
+        elif (imp_type == "xquad") or (imp_type == "yquad"):
+            negative_data["real"] = -1*negative_data["real"]
+            
+        else:
+            raise ValueError("Wrong impedance type")
+            
+        try:
+            negative_data = negative_data.drop(0)
+        except KeyError:
+            pass
+            
+        all_data = impedance.data.append(negative_data)
+        all_data = all_data.sort_index()
+        impedance.data = all_data
+        
diff --git a/collective_effects/wakefield.py b/collective_effects/wakefield.py
index 323fdb92599b3676c7200c512e2f31047d6f846d..ee372cb5f06023070b4671336860de968e886eab 100644
--- a/collective_effects/wakefield.py
+++ b/collective_effects/wakefield.py
@@ -1,24 +1,21 @@
 # -*- coding: utf-8 -*-
 """
 This module defines general classes to describes wakefields, impedances and 
-wake functions. Based on impedance library by David Amorim.
+wake functions.
 
-@author: David Amorin, Alexis Gamelin
-@date: 14/01/2020
+@author: David Amorin, Alexis Gamelin, Watanyu Foosang
+@date: 24/09/2020
 """
 
 import warnings
 import re
 import pandas as pd
 import numpy as np
-import matplotlib.pyplot as plt
-from scipy.interpolate import interp1d
-from mbtrack2.tracking.element import Element
+import scipy as sc
 from copy import deepcopy
-from scipy.integrate import trapz
 from scipy.interpolate import interp1d
-from mpl_toolkits.axes_grid1.inset_locator import inset_axes
-
+from scipy.integrate import trapz
+from scipy.constants import c
 
 class ComplexData:
     """
@@ -192,17 +189,156 @@ class ComplexData:
 class WakeFunction(ComplexData):
     """
     Define a WakeFunction object based on a ComplexData object.
+    
+    Parameters
+    ----------
+    variable : array-like
+        Time domain structure of the wake function in [s].
+    function : array-like
+        Wake function values in [V/C].
+    wake_type : str, optinal
+        Type of the wake function: "long", "xdip", "xquad", ...
+    
+    Attributes
+    ----------
+    data : DataFrame
+    wake_type : str
+    
+    Methods
+    -------
+    from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
+        Compute a wake function from a wake potential file and load it to the 
+        WakeFunction object.
     """
 
     def __init__(self,
                  variable=np.array([-1e15, 1e15]),
                  function=np.array([0, 0]), wake_type='long'):
         super().__init__(variable, function)
-        pass
+        self._wake_type = wake_type
+        self.data.index.name = "time [s]"
+        
+    def add(self, structure_to_add, method='zero'):
+        """
+        Method to add two WakeFunction objects. The two structures are
+        compared so that the addition is self-consistent.
+        """
+ 
+        if isinstance(structure_to_add, (int, float, complex)):
+            result = super().add(structure_to_add, method=method,
+                                 index_name="time [s]")
+        elif isinstance(structure_to_add, WakeFunction):
+            if (self.wake_type == structure_to_add.wake_type):
+                result = super().add(structure_to_add, method=method,
+                                     index_name="time [s]")
+            else:
+                warnings.warn(('The two WakeFunction objects do not have the '
+                               'same coordinates or plane or type. '
+                               'Returning initial WakeFunction object.'),
+                              UserWarning)
+                result = self
+               
+        wake_to_return = WakeFunction(
+                                result.data.index,
+                                result.data.real.values,
+                                self.wake_type)   
+        return wake_to_return
+ 
+    def __radd__(self, structure_to_add):
+        return self.add(structure_to_add)
+ 
+    def __add__(self, structure_to_add):
+        return self.add(structure_to_add)
+ 
+    def multiply(self, factor):
+        """
+        Multiply a WakeFunction object with a float or an int.
+        If the multiplication is done with something else, throw a warning.
+        """
+        result = super().multiply(factor)
+        wake_to_return = WakeFunction(
+                                result.data.index,
+                                result.data.real.values,
+                                self.wake_type)   
+        return wake_to_return
+ 
+    def __mul__(self, factor):
+        return self.multiply(factor)
+ 
+    def __rmul__(self, factor):
+        return self.multiply(factor)
+       
+    @property
+    def wake_type(self):
+        return self._wake_type
+   
+    @wake_type.setter
+    def wake_type(self, value):
+        self._wake_type = value
+        
+    def from_wakepotential(self, file_name, bunch_length, bunch_charge, 
+                           freq_lim, nout=None, trim=False):
+        """
+        Compute a wake function from a wake potential file and load it to the 
+        WakeFunction object.
     
+        Parameters
+        ----------
+        file_name : str
+            Text file that contains wake potential data.
+        bunch_length : float
+            Spatial bunch length in [m].
+        bunch_charge : float
+            Bunch charge in [C].
+        freq_lim : float
+            The maximum frequency for calculating the impedance [Hz].
+        nout : int, optional
+            Length of the transformed axis of the output. If None, it is taken
+            to be 2*(n-1) where n is the length of the input. If nout > n, the
+            input is padded with zeros. If nout < n, the inpu it cropped.
+            Note that, for any nout value, nout//2+1 input points are required.
+        trim : bool or float, optional
+            If True, the pseudo wake function is trimmed at time=0 and is forced 
+            to zero where time<0. 
+            If False, the original result is preserved.
+            If a float is given, the pseudo wake function is trimmed from 
+            time <= trim to 0. 
+    
+        """
+        
+        imp = Impedance()
+        imp.from_wakepotential(file_name=file_name, bunch_length=bunch_length,
+                               bunch_charge=bunch_charge, freq_lim=freq_lim)
+        wf = imp.to_wakefunction(nout=nout, trim=trim)
+        self.__init__(variable=wf.data.index, function=wf.data["real"])
+        
 class Impedance(ComplexData):
     """
     Define an Impedance object based on a ComplexData object.
+    
+    Parameters
+    ----------
+    variable : array-like
+        Frequency domain structure of the impedance in [Hz].
+    function : array-like
+        Impedance values in [Ohm].
+    impedance_type : str, optinal
+        Type of the impedance_type: "long", "xdip", "xquad", ...
+    
+    Attributes
+    ----------
+    data : DataFrame
+    impedance_type : str
+    
+    Methods
+    -------
+    from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
+        Compute impedance from wake potential data and load it to the Impedance
+        object.
+    loss_factor(sigma)
+        Compute the loss factor or the kick factor for a Gaussian bunch.
+    to_wakefunction()
+        Return a WakeFunction object from the impedance data.
     """
 
     def __init__(self,
@@ -352,6 +488,9 @@ class Impedance(ComplexData):
         
         positive_index = self.data.index > 0
         frequency = self.data.index[positive_index]
+        
+        # Import here to avoid circular import
+        from mbtrack2.collective_effects.utilities import spectral_density
         sd = spectral_density(frequency, sigma, m=0)
         
         if(self.impedance_type == "long"):
@@ -368,16 +507,123 @@ class Impedance(ComplexData):
 
         return kloss
     
-class Wakefield:
+    def from_wakepotential(self, file_name, bunch_length, bunch_charge, 
+                           freq_lim, axis0='dist', axis0_scale=1e-3, 
+                           axis1_scale=1e-12):
+        """
+        Compute impedance from wake potential data and load it to the Impedance
+        object.
+
+        Parameters
+        ----------
+        file_name : str
+            Text file that contains wake potential data.
+        freq_lim : float
+            The maximum frequency for calculating the impedance [Hz].
+        bunch_length : float
+            Electron bunch lenth [m]. 
+        bunch_charge : float
+            Total bunch charge [C].
+        axis0 : {'dist', 'time'}, optional
+            Viariable of the data file's first column. Use 'dist' for spacial 
+            distance. Otherwise, 'time' for temporal distance. 
+        axis0_scale : float, optional
+            Scale of the first column with respect to meter if axis0 = 'dist',
+            or to second if axis0 = 'time'.
+        axis1_scale : float, optional
+            Scale of the wake potential (in second column) with respect to V/C.
+
+        """
+        
+        s, wp0 = np.loadtxt(file_name, unpack=True)
+        if axis0 == 'dist':
+            tau = s*axis0_scale/c
+        elif axis0 == 'time':
+            tau = s*axis0_scale
+        else:
+            raise TypeError('Type of axis 0 not recognized.')
+            
+        wp = wp0 / axis1_scale
+        
+        # FT of wake potential
+        sampling = tau[1] - tau[0]
+        freq = sc.fft.rfftfreq(len(tau), sampling)
+        dtau = (tau[-1]-tau[0])/len(tau)
+        dft_wake = sc.fft.rfft(wp) * dtau
+        
+        # FT of analytical bunch profile and analytical impedance
+        sigma = bunch_length/c
+        mu = tau[0]
+        
+        i_limit = freq < freq_lim
+        freq_trun = freq[i_limit]
+        dft_wake_trun = dft_wake[i_limit]
+        
+        dft_rho_trun = np.exp(-0.5*(sigma*2*np.pi*freq_trun)**2 + \
+                                  1j*mu*2*np.pi*freq_trun)*bunch_charge
+        imp = dft_wake_trun/dft_rho_trun*-1*bunch_charge
+        
+        super().__init__(variable=freq_trun, function=imp)
+        self.data.index.name = "frequency [Hz]"    
+        
+    def to_wakefunction(self, nout=None, trim=False):
+        """
+        Return a WakeFunction object from the impedance data.
+        The impedance data is assumed to be sampled equally.
+    
+        Parameters
+        ----------
+        nout : int or float, optional
+            Length of the transformed axis of the output. If None, it is taken
+            to be 2*(n-1) where n is the length of the input. If nout > n, the
+            input is padded with zeros. If nout < n, the input it cropped.
+            Note that, for any nout value, nout//2+1 input points are required.
+        trim : bool or float, optional
+            If True, the pseudo wake function is trimmed at time=0 and is forced 
+            to zero where time<0. 
+            If False, the original result is preserved.
+            If a float is given, the pseudo wake function is trimmed from 
+            time <= trim to 0. 
+        """
+        
+        Z0 = (self.data['real'] + self.data['imag']*1j)
+        Z = Z0[~np.isnan(Z0)]
+        
+        if self.impedance_type != "long":
+            Z = Z / 1j
+        
+        freq = Z.index
+        fs = ( freq[-1] - freq[0] ) / len(freq)
+        sampling = freq[1] - freq[0]
+        
+        if nout is None:
+            nout = len(Z)
+        else:
+            nout = int(nout)
+            
+        time_array = sc.fft.fftfreq(nout, sampling)
+        Wlong_raw = sc.fft.irfft(np.array(Z), n=nout, axis=0) * nout * fs
+        
+        time = sc.fft.fftshift(time_array)
+        Wlong = sc.fft.fftshift(Wlong_raw)
+        
+        if trim is not False:
+            i_neg = np.where(time<trim)[0]
+            Wlong[i_neg] = 0
+                    
+        wf = WakeFunction(variable=time, function=Wlong)
+        return wf
+    
+class WakeField:
     """
-    Defines a Wakefield which corresponds to a single physical element which 
+    Defines a WakeField which corresponds to a single physical element which 
     produces different types of wakes, represented by their Impedance or 
     WakeFunction objects.
     
     Parameters
     ----------
     structure_list : list, optional
-        list of Impedance/WakeFunction objects to add to the Wakefield.
+        list of Impedance/WakeFunction objects to add to the WakeField.
     name : str, optional
     
     Attributes
@@ -388,13 +634,13 @@ class Wakefield:
     Methods
     -------
     append_to_model(structure_to_add)
-        Add Impedance/WakeFunction to Wakefield.
+        Add Impedance/WakeFunction to WakeField.
     list_to_attr(structure_list)
-        Add list of Impedance/WakeFunction to Wakefield.
+        Add list of Impedance/WakeFunction to WakeField.
     add_wakefields(wake1, beta1, wake2, beta2)
-        Add two Wakefields taking into account beta functions.
+        Add two WakeFields taking into account beta functions.
     add_several_wakefields(wakefields, beta)
-        Add a list of Wakefields taking into account beta functions.
+        Add a list of WakeFields taking into account beta functions.
     """
 
     def __init__(self, structure_list=None, name=None):
@@ -403,7 +649,7 @@ class Wakefield:
 
     def append_to_model(self, structure_to_add):
         """
-        Add Impedance/WakeFunction component to Wakefield.
+        Add Impedance/WakeFunction component to WakeField.
 
         Parameters
         ----------
@@ -418,7 +664,7 @@ class Wakefield:
             else:
                 self.__setattr__(attribute_name, structure_to_add)
         elif isinstance(structure_to_add, WakeFunction):
-            attribute_name = "W" + structure_to_add.impedance_type
+            attribute_name = "W" + structure_to_add.wake_type
             if attribute_name in list_of_attributes:
                 raise ValueError("There is already a component of the type "
                                  "{} in this element.".format(attribute_name))
@@ -429,7 +675,7 @@ class Wakefield:
     
     def list_to_attr(self, structure_list):
         """
-         Add list of Impedance/WakeFunction components to Wakefield.
+         Add list of Impedance/WakeFunction components to WakeField.
 
         Parameters
         ----------
@@ -446,26 +692,33 @@ class Wakefield:
         """
         return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
     
+    @property
+    def wake_components(self):
+        """
+        Return an array of the wake function component names for the element.
+        """
+        return np.array([comp for comp in dir(self) if re.match(r'[W]', comp)])
+    
     @staticmethod
     def add_wakefields(wake1, beta1, wake2, beta2):
         """
-        Add two Wakefields taking into account beta functions.
+        Add two WakeFields taking into account beta functions.
 
         Parameters
         ----------
-        wake1 : Wakefield
-            Wakefield to add.
+        wake1 : WakeField
+            WakeField to add.
         beta1 : array of shape (2,)
             Beta functions at wake1 postion.
-        wake2 : Wakefield
-            Wakefield to add.
+        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 : WakeField
+            WakeField with summed components.
 
         """
         wake_sum = deepcopy(wake1)
@@ -491,818 +744,30 @@ class Wakefield:
     @staticmethod
     def add_several_wakefields(wakefields, beta):
         """
-        Add a list of Wakefields taking into account beta functions.
+        Add a list of WakeFields taking into account beta functions.
         
         Parameters
         ----------
-        wakefields : list of Wakefield
-            Wakefields to add.
+        wakefields : list of WakeField
+            WakeFields to add.
         beta : array of shape (len(wakefields), 2)
-            Beta functions at the Wakefield postions.
+            Beta functions at the WakeField postions.
 
         Returns
         -------
-        wake_sum : Wakefield
-            Wakefield with summed components..
+        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],
+            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], 
+                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):
-    """
-    Define the impedance model of the machine.
-    
-    Parameters
-    ----------
-    ring : Synchrotron object
-    wakefield_list : list of Wakefield objects
-        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):
-        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_postions)
-        
-    def track(self, beam):
-        """
-        Track a beam object through this Element.
-        
-        Parameters
-        ----------
-        beam : Beam object
-        """
-        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):
-        """
-        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.
-
-        """
-        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.
-
-        Parameters
-        ----------
-        wakefield_list : list of Wakefield objects
-            Wakefields to add to the model.
-        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)
-                
-        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, zoom=False):
-        """
-        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.
-        zoom : bool
-            If True, add a zoomed plot on top right corner.
-
-        """
-        if attr_list is None:
-            attr_list = self.sum_names[self.sum_names != "sum"]
-        
-        # manage legend
-        Ztype_dict = {"Zlong":0, "Zxdip":1, "Zydip":2, "Zxquad":3, "Zyquad":4}
-        scale = [1e-3, 1e-6, 1e-6, 1e-6, 1e-6]
-        label_list =  [r"$Z_{long} \; [k\Omega]$",
-                       r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Dip} \; [M\Omega]$",
-                       r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Dip} \; [M\Omega]$",
-                       r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Quad} \; [M\Omega]$",
-                       r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Quad} \; [M\Omega]$"]
-        leg = Ztype_dict[Z_type]
-        
-        # sort plot by decresing area size        
-        area = np.zeros((len(attr_list),))
-        for index, attr in enumerate(attr_list):
-            try:
-                sum_imp = getattr(getattr(self, attr), Z_type)
-                area[index] = trapz(sum_imp.data[component], sum_imp.data.index)
-            except AttributeError:
-                pass
-        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
-            try:
-                sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
-                ax.fill_between(sum_imp.data.index*1e-9, total_imp, 
-                                total_imp + sum_imp.data[component]*scale[leg])
-                total_imp += sum_imp.data[component]*scale[leg]
-                if attr[:4] == "sum_":
-                    legend.append(attr[4:])
-                else:
-                    legend.append(attr)
-            except AttributeError:
-                pass
-            
-        if sigma is not None:
-            spect = spectral_density(zero_impedance.data.index, sigma)
-            spect = spect/spect.max()*total_imp.max()
-            ax.plot(sum_imp.data.index*1e-9, spect, 'r', linewidth=2.5)
-        
-        ax.legend(legend, loc="upper left")            
-        ax.set_xlabel("Frequency [GHz]")
-        ax.set_ylabel(label_list[leg] + " - " + component + " part")
-        ax.set_title(label_list[leg] + " - " + component + " part")
-        
-        if zoom is True:
-            in_ax = inset_axes(ax,
-                            width="30%", # width = 30% of parent_bbox
-                            height=1.5, # height : 1 inch
-                            loc=1)
-            
-            total_imp = 0
-            for index in  sorted_index:
-                attr = attr_list[index]
-                # Set all impedances with common indexes using + zero_impedance
-                try:
-                    sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
-                    in_ax.fill_between(sum_imp.data.index*1e-3, total_imp, 
-                                    total_imp + sum_imp.data[component]*1e-9)
-                    total_imp += sum_imp.data[component]*1e-9
-                except AttributeError:
-                    pass
-            in_ax.set_xlim([0, 200])
-            in_ax.set_xlabel("Frequency [kHz]")
-            in_ax.set_ylabel(r"$[G\Omega]$")
-                
-        return fig
-            
-    def effective_impedance(self, m, mu, sigma, M, tuneS, xi=None, 
-                            mode="Hermite"):
-        """
-        Compute the longitudinal and transverse effective impedance.
-
-        Parameters
-        ----------
-        mu : int
-            coupled bunch mode number, goes from 0 to (M-1) where M is the
-            number of bunches
-        m : int
-            head-tail (or azimutal/synchrotron) mode number
-        sigma : float
-            RMS bunch length in [s]
-        M : int
-            Number of bunches.
-        tuneS : float
-            Synchrotron tune.
-        xi : float, optional
-            (non-normalized) chromaticity
-        mode: str, optional
-            type of the mode taken into account for the computation:
-            -"Hermite" modes for Gaussian bunches
-
-        Returns
-        -------
-        summary : DataFrame
-            Longitudinal and transverse effective impedance.
-
-        """
-
-        attr_list = self.sum_names
-        
-        eff_array = np.zeros((len(attr_list),3), dtype=complex)
-        
-        for i, attr in enumerate(attr_list):
-            try:
-                impedance = getattr(getattr(self, attr), "Zlong")
-                eff_array[i,0] = effective_impedance(self.ring, impedance, 
-                                                     m, mu, sigma, M, tuneS,
-                                                     xi, mode)
-            except AttributeError:
-                pass
-            
-            try:
-                impedance = getattr(getattr(self, attr), "Zxdip")
-                eff_array[i,1] = effective_impedance(self.ring, impedance, 
-                                                     m, mu, sigma, M, tuneS,
-                                                     xi, mode)
-            except AttributeError:
-                pass
-            
-            try:
-                impedance = getattr(getattr(self, attr), "Zydip")
-                eff_array[i,2] = effective_impedance(self.ring, impedance, 
-                                                     m, mu, sigma, M, tuneS,
-                                                     xi, mode)
-            except AttributeError:
-                pass
-            
-        eff_array[:,0] = eff_array[:,0]*self.ring.omega0*1e3
-        eff_array[:,1] = eff_array[:,1]*1e-3
-        eff_array[:,2] = eff_array[:,2]*1e-3
-        
-        summary = pd.DataFrame(eff_array, index=attr_list, 
-                               columns=["Z/n [mOhm]", 
-                                        "sum betax x Zxeff [kOhm]", 
-                                        "sum betay x Zyeff [kOhm]"])
-        
-        return summary
-
-
-    def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
-        """
-        Compute the beam and bunch loss factor and energy losses for each type 
-        of element in the model assuming Gaussian bunches and constant spacing
-        between bunches.
-
-        Parameters
-        ----------
-        sigma : float
-            RMS bunch length in [s].
-        M : int
-            Number of bunches in the beam.
-        bunch_spacing : float
-            Time between two bunches in [s].
-        I : float
-            Total beam current in [A].
-        n_points : float, optional
-            Number of points used in the frequency spectrums.
-
-        Returns
-        -------
-        summary : Dataframe
-            Contains the beam and bunch loss factor and energy loss for the 
-            full model and for each type of different component.
-
-        """
-        
-        fmax = self.sum.Zlong.data.index.max()
-        fmin = self.sum.Zlong.data.index.min()
-        
-        Q = I*self.ring.T0/M
-        
-        if fmin >= 0:
-            fmin = -1*fmax
-        f = np.linspace(fmin, fmax, int(n_points))
-
-        beam_spect = beam_spectrum(f, M, bunch_spacing, sigma= sigma)
-        
-        bunch_spect = gaussian_bunch_spectrum(f, sigma)
-        
-        attr_list = self.sum_names
-        
-        loss_array = np.zeros((len(attr_list),2))
-        
-        for i, attr in enumerate(attr_list):
-            try:
-                impedance = getattr(getattr(self, attr), "Zlong")
-                loss_array[i,0] = beam_loss_factor(impedance, f, beam_spect, self.ring)
-                loss_array[i,1] = beam_loss_factor(impedance, f, bunch_spect, self.ring)
-            except AttributeError:
-                pass
-
-        loss_array = loss_array*1e-12
-        summary = pd.DataFrame(loss_array, index=attr_list, 
-                               columns=["loss factor (beam) [V/pC]", "loss factor (bunch) [V/pC]"])
-        
-        summary["P (beam) [W]"] = summary["loss factor (beam) [V/pC]"]*1e12*Q**2/(self.ring.T0)
-        summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"]*1e12*Q**2/(self.ring.T0)*M
-                
-        return summary
-        
-    
-def read_CST(file, impedance_type='long', divide_by=None):
-    """
-    Read CST file format into an Impedance object.
-    
-    Parameters
-    ----------
-    file : str
-        Path to the file to read.
-    impedance_type : str, optional
-        Type of the Impedance object.
-    divide_by : float, optional
-        Divide the impedance by a value. Mainly used to normalize transverse 
-        impedance by displacement.
-        
-    Returns
-    -------
-    result : Impedance object
-        Data from file.
-    """
-    df = pd.read_csv(file, comment="#", header = None, sep = "\t", 
-                    names = ["Frequency","Real","Imaginary"])
-    df["Frequency"] = df["Frequency"]*1e9 
-    if divide_by is not None:
-        df["Real"] = df["Real"]/divide_by
-        df["Imaginary"] = df["Imaginary"]/divide_by
-    if impedance_type == "long":
-        df["Real"] = np.abs(df["Real"])
-    df.set_index("Frequency", inplace = True)
-    result = Impedance(variable = df.index,
-                       function = df["Real"] + 1j*df["Imaginary"],
-                       impedance_type=impedance_type)
-    return result
-
-def read_IW2D(file, impedance_type='long'):
-    """
-    Read IW2D file format into an Impedance object.
-    
-    Parameters
-    ----------
-    file : str
-        Path to the file to read.
-    impedance_type : str, optional
-        Type of the Impedance object.
-        
-    Returns
-    -------
-    result : Impedance object
-        Data from file.
-    """
-    df = pd.read_csv(file, sep = ' ', header = None, 
-                     names = ["Frequency","Real","Imaginary"], skiprows=1)
-    df.set_index("Frequency", inplace = True)
-    result = Impedance(variable = df.index,
-                       function = df["Real"] + 1j*df["Imaginary"],
-                       impedance_type=impedance_type)
-    return result
-
-def spectral_density(frequency, sigma, m = 1, mode="Hermite"):
-    """
-    Compute the spectral density of different modes for various values of the
-    head-tail mode number, based on Table 1 p238 of [1].
-    
-    Parameters
-    ----------
-    frequency : list or numpy array
-        sample points of the spectral density in [Hz]
-    sigma : float
-        RMS bunch length in [s]
-    m : int, optional
-        head-tail (or azimutal/synchrotron) mode number
-    mode: str, optional
-        type of the mode taken into account for the computation:
-        -"Hermite" modes for Gaussian bunches
-
-    Returns
-    -------
-    numpy array
-    
-    References
-    ----------
-    [1] : Handbook of accelerator physics and engineering, 3rd printing.
-    """
-    
-    if mode == "Hermite":
-        return (2*np.pi*frequency*sigma)**(2*m)*np.exp(
-                -1*(2*np.pi*frequency*sigma)**2)
-    else:
-        raise NotImplementedError("Not implemanted yet.")
-        
-def gaussian_bunch_spectrum(frequency, sigma): 
-    """
-    Compute a Gaussian bunch spectrum [1].
-
-    Parameters
-    ----------
-    frequency : array
-        sample points of the beam spectrum in [Hz].
-    sigma : float
-        RMS bunch length in [s].
-
-    Returns
-    -------
-    bunch_spectrum : array
-        Bunch spectrum sampled at points given in frequency.
-        
-    References
-    ----------
-    [1] : Gamelin, A. (2018). Collective effects in a transient microbunching 
-    regime and ion cloud mitigation in ThomX. p86, Eq. 4.19
-    """
-    return np.exp(-1/2*(2*np.pi*frequency)**2*sigma**2)
-        
-        
-def beam_spectrum(frequency, M, bunch_spacing, sigma=None, 
-                  bunch_spectrum=None): 
-    """
-    Compute the beam spectrum assuming constant spacing between bunches [1].
-
-    Parameters
-    ----------
-    frequency : list or numpy array
-        sample points of the beam spectrum in [Hz].
-    M : int
-        Number of bunches.
-    bunch_spacing : float
-        Time between two bunches in [s].
-    sigma : float, optional
-        If bunch_spectrum is None then a Gaussian bunch with sigma RMS bunch 
-        length in [s] is assumed.
-    bunch_spectrum : array, optional
-        Bunch spectrum sampled at points given in frequency.
-
-    Returns
-    -------
-    beam_spectrum : array
-
-    References
-    ----------
-    [1] Rumolo, G - Beam Instabilities - CAS - CERN Accelerator School: 
-        Advanced Accelerator Physics Course - 2014, Eq. 9
-    """
-    
-    if bunch_spectrum is None:
-        bunch_spectrum = gaussian_bunch_spectrum(frequency, sigma)
-        
-    beam_spectrum = (bunch_spectrum * np.exp(1j*np.pi*frequency * 
-                                             bunch_spacing*(M-1)) * 
-                     np.sin(M*np.pi*frequency*bunch_spacing) / 
-                     np.sin(np.pi*frequency*bunch_spacing))
-    
-    return beam_spectrum
-    
-    
-def effective_impedance(ring, imp, m, mu, sigma, M, tuneS, xi=None, 
-                        mode="Hermite"):
-    """
-    Compute the effective (longitudinal or transverse) impedance. 
-    Formulas from Eq. (1) and (2) p238 of [1].
-    
-    Parameters
-    ----------
-    ring : Synchrotron object
-    imp : Impedance object
-    mu : int
-        coupled bunch mode number, goes from 0 to (M-1) where M is the
-        number of bunches
-    m : int
-        head-tail (or azimutal/synchrotron) mode number
-    sigma : float
-        RMS bunch length in [s]
-    M : int
-        Number of bunches.
-    tuneS : float
-        Synchrotron tune.
-    xi : float, optional
-        (non-normalized) chromaticity
-    mode: str, optional
-        type of the mode taken into account for the computation:
-        -"Hermite" modes for Gaussian bunches
-
-    Returns
-    -------
-    Zeff : float 
-        effective impedance in [ohm] or in [ohm/m] depanding on the impedance
-        type.
-        
-    References
-    ----------
-    [1] : Handbook of accelerator physics and engineering, 3rd printing.
-    """
-    
-    if not isinstance(imp, Impedance):
-        raise TypeError("{} should be an Impedance object.".format(imp))
-        
-    fmin = imp.data.index.min()
-    fmax = imp.data.index.max()
-    if fmin > 0:
-        double_sided_impedance(imp)
-        
-    if mode == "Hermite":
-        def h(f):
-            return spectral_density(frequency=f, sigma=sigma, m=m,
-                                    mode="Hermite")
-    else:
-        raise NotImplementedError("Not implemanted yet.")
-    
-    pmax = fmax/(ring.f0 * M) - 1
-    pmin = fmin/(ring.f0 * M) + 1
-    
-    p = np.arange(pmin,pmax+1)
-    
-    if imp.impedance_type == "long":
-        fp = ring.f0*(p*M + mu + m*tuneS)
-        fp = fp[np.nonzero(fp)] # Avoid division by 0
-        num = np.sum( imp(fp) * h(fp) / (fp*2*np.pi) )
-        den = np.sum( h(fp) )
-        Zeff = num/den
-        
-    elif imp.impedance_type == "xdip" or imp.impedance_type == "ydip":
-        if imp.impedance_type == "xdip":
-            tuneXY = ring.tune[0]
-            if xi is None :
-                xi = ring.chro[0]
-        elif imp.impedance_type == "ydip":
-            tuneXY = ring.tune[1]
-            if xi is None:
-                xi = ring.chro[1]
-        fp = ring.f0*(p*M + mu + tuneXY + m*tuneS)
-        f_xi = xi/ring.eta*ring.f0
-        num = np.sum( imp(fp) * h(fp - f_xi) )
-        den = np.sum( h(fp - f_xi) )
-        Zeff = num/den
-    else:
-        raise TypeError("Effective impedance is only defined for long, xdip"
-                        " and ydip impedance type.")
-        
-    return Zeff
-
-
-def yokoya_elliptic(x_radius , y_radius):
-    """
-    Compute Yokoya factors for an elliptic beam pipe.
-    Function adapted from N. Mounet IW2D.
-
-    Parameters
-    ----------
-    x_radius : float
-        Horizontal semi-axis of the ellipse in [m].
-    y_radius : float
-        Vertical semi-axis of the ellipse in [m].
-
-    Returns
-    -------
-    yoklong : float
-        Yokoya factor for the longitudinal impedance.
-    yokxdip : float
-        Yokoya factor for the dipolar horizontal impedance.
-    yokydip : float
-        Yokoya factor for the dipolar vertical impedance.
-    yokxquad : float
-        Yokoya factor for the quadrupolar horizontal impedance.
-    yokyquad : float
-        Yokoya factor for the quadrupolar vertical impedance.
-    """
-    if y_radius < x_radius:
-        small_semiaxis = y_radius
-        large_semiaxis = x_radius
-    else:
-        small_semiaxis = x_radius
-        large_semiaxis = y_radius
-
-    # read Yokoya factors interpolation file
-    # BEWARE: columns are ratio, dipy, dipx, quady, quadx
-    yokoya_file = pd.read_csv("collective_effects/data/" +
-                              "Yokoya_elliptic_from_Elias_USPAS.csv")
-    ratio_col = yokoya_file["x"]
-    # compute semi-axes ratio (first column of this file)
-    ratio = (large_semiaxis - small_semiaxis)/(large_semiaxis + small_semiaxis)
-
-    # interpolate Yokoya file at the correct ratio
-    yoklong = 1
-    
-    if y_radius < x_radius:
-        yokydip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
-        yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
-        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
-        yokxquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])
-    else:
-        yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
-        yokydip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
-        yokxquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
-        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])        
-
-    return (yoklong, yokxdip, yokydip, yokxquad, yokyquad)
-
-def beam_loss_factor(impedance, frequency, spectrum, ring):
-    """
-    Compute "beam" loss factor using the beam spectrum, uses a sum instead of 
-    integral compared to loss_factor [1].
-
-    Parameters
-    ----------
-    impedance : Impedance of type "long"
-    frequency : array
-        Sample points of spectrum.
-    spectrum : array
-        Beam spectrum to consider.
-    ring : Synchrotron object
-
-    Returns
-    -------
-    kloss_beam : float
-        Beam loss factor in [V/C].
-        
-    References
-    ----------
-    [1] : Handbook of accelerator physics and engineering, 3rd printing. 
-        Eq (3) p239.
-    """
-    pmax = np.floor(impedance.data.index.max()/ring.f0)
-    pmin = np.floor(impedance.data.index.min()/ring.f0)
-    
-    if pmin >= 0:
-        double_sided_impedance(impedance)
-        pmin = -1*pmax
-    
-    p = np.arange(pmin+1,pmax)    
-    pf0 = p*ring.f0
-    ReZ = np.real(impedance(pf0))
-    spectral_density = np.abs(spectrum)**2
-    # interpolation of the spectrum is needed to avoid problems liked to 
-    # division by 0
-    # computing the spectrum directly to the frequency points gives
-    # wrong results
-    spect = interp1d(frequency, spectral_density)
-    kloss_beam = ring.f0 * np.sum(ReZ*spect(pf0))
-    
-    return kloss_beam
-
-def double_sided_impedance(impedance):
-    """
-    Add negative frequency points to single sided impedance spectrum following
-    symetries depending on impedance type.
-
-    Parameters
-    ----------
-    impedance : Impedance object
-        Single sided impedance.
-    """
-    fmin = impedance.data.index.min()
-    
-    if fmin >= 0:
-        negative_index = impedance.data.index*-1
-        negative_data = impedance.data.set_index(negative_index)
-        
-        imp_type = impedance.impedance_type
-        
-        if imp_type == "long":
-            negative_data["imag"] = -1*negative_data["imag"]
-            
-        elif (imp_type == "xdip") or (imp_type == "ydip"):
-            negative_data["real"] = -1*negative_data["real"]
-        
-        elif (imp_type == "xquad") or (imp_type == "yquad"):
-            negative_data["real"] = -1*negative_data["real"]
-            
-        else:
-            raise ValueError("Wrong impedance type")
-            
-        try:
-            negative_data = negative_data.drop(0)
-        except KeyError:
-            pass
-            
-        all_data = impedance.data.append(negative_data)
-        all_data = all_data.sort_index()
-        impedance.data = all_data
-    
\ No newline at end of file
+        
\ No newline at end of file
diff --git a/tracking/__init__.py b/tracking/__init__.py
index 45db00132ecab8571c819849c041e992ca8826fe..969d1bba6dd1244b8818cd1eb2824cdbbc3243f9 100644
--- a/tracking/__init__.py
+++ b/tracking/__init__.py
@@ -15,4 +15,5 @@ from mbtrack2.tracking.element import (Element, LongitudinalMap,
                                        TransverseMap, SynchrotronRadiation)
 from mbtrack2.tracking.aperture import (CircularAperture, ElipticalAperture,
                                         RectangularAperture)
+from mbtrack2.tracking.wakepotential import WakePotential
                                        
diff --git a/tracking/element.py b/tracking/element.py
index 81f506cac6478144ed78ee60a5c588fb972ce863..9123e901224ce2d8e399ffa9ba9fbae603f2e9fe 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -9,13 +9,9 @@ included in the tracking.
 """
 
 import numpy as np
-import pandas as pd
 from abc import ABCMeta, abstractmethod
 from functools import wraps
 from mbtrack2.tracking.particles import Beam
-from scipy import signal
-import matplotlib.pyplot as plt
-from scipy.interpolate import interp1d
 
 class Element(metaclass=ABCMeta):
     """
@@ -98,204 +94,7 @@ class LongitudinalMap(Element):
         """
         bunch["delta"] -= self.ring.U0 / self.ring.E0
         bunch["tau"] -= self.ring.ac * self.ring.T0 * bunch["delta"]
-        
-class WakePotential(Element):
-    """
-    Resonator model based wake potential calculation for one turn.
-    
-    Parameters
-    ----------
-    ring : Synchrotron object.
-    Q_factor : float, optional
-        Resonator quality factor. The default value is 1.
-    f_res : float, optional
-        Resonator resonance frequency in [Hz]. The default value is 10e9 Hz.
-    R_shunt : float, optional
-        Resonator shunt impedance in [Ohm]. The default value is 100 Ohm.
-    n_bin : int, optional
-        Number of bins for constructing the longitudinal bunch profile.
-        The default is 65.
-        
-    Attributes
-    ----------
-    rho : array of shape (n_bin, )
-        Bunch charge density profile.
-    tau : array of shape (n_bin + time_extra, )
-        Time array starting from the head of the bunch until the wake tail 
-        called timestop.
-        
-        The length of time_extra is determined by the last position of the
-        bunch time_bunch[-1], timestop, and the mean bin width of the bunch
-        profile mean_bin_size as
-            len(time_extra) = (timestop - time_bunch[-1]) / mean_bin_size
-    W_long : array of shape (n_bin + time_extra, )
-        Wakefunction profile.
-    W_p : array of shape (n_bin + time_extra, )
-        Wake potential profile.
-    wp : array of shape (mp_number, )
-        Wake potential exerted on each macro-particle.
-    
-    Methods
-    -------
-    charge_density(bunch, n_bin=65)
-        Calculate bunch charge density.
-    plot(self, var, plot_rho=True)
-        Plotting wakefunction or wake potential.
-    track(bunch)
-        Tracking method for the element.
-        
-    """
-    
-    def __init__(self, ring, Q_factor=1, f_res=10e9, R_shunt=100, n_bin=65):       
-        self.ring = ring
-        self.n_bin = n_bin
-        
-        self.Q_factor = Q_factor
-        self.omega_res = 2*np.pi*f_res
-        self.R_shunt = R_shunt
-        
-        if Q_factor >= 0.5:
-            self.Q_factor_p = np.sqrt(self.Q_factor**2 - 0.25)
-            self.omega_res_p = (self.omega_res*self.Q_factor_p)/self.Q_factor
-        else:
-            self.Q_factor_pp = np.sqrt(0.25 - self.Q_factor**2)
-            self.omega_res_p = (self.omega_res*self.Q_factor_pp)/self.Q_factor
-            
-    def charge_density(self, bunch, n_bin):
-        self.bins = bunch.binning(n_bin=self.n_bin)
-        self.bin_data = self.bins[2]
-        self.bin_size = self.bins[0].length
-        
-        self.rho = bunch.charge_per_mp*self.bin_data/ \
-            (self.bin_size*bunch.charge)
-    
-    def init_timestop(self):
-        self.timestop = round(np.log(1000)/self.omega_res*2*self.Q_factor, 12)
-    
-    def time_array(self):        
-        time_bunch = self.bins[0].mid
-        mean_bin_size = np.mean(self.bin_size)
-        time_extra = np.arange(start = time_bunch[-1]+mean_bin_size, 
-                               stop = self.timestop, step = mean_bin_size)
-        
-        self.tau = np.concatenate((time_bunch,time_extra))
-     
-    def long_wakefunction(self):
-        w_list = []
-        if self.Q_factor >= 0.5:
-            for t in self.tau: 
-                if t >= 0:
-                    w_long = -(self.omega_res*self.R_shunt/self.Q_factor)*\
-                        np.exp(-self.omega_res*t/(2*self.Q_factor))*\
-                            (np.cos(self.omega_res_p*t)-\
-                             np.sin(self.omega_res_p*t)/(2*self.Q_factor_p))
-                else: 
-                    w_long = 0
-                    
-                w_list.append(w_long)
-                        
-        elif self.Q_factor < 0.5:
-            for t in self.tau: 
-                if t >= 0:
-                    w_long = -(self.omega_res*self.R_shunt/self.Q_factor)*\
-                        np.exp(-self.omega_res*t/(2*self.Q_factor))*\
-                            (np.cosh(self.omega_res_p*t)-\
-                             np.sinh(self.omega_res_p*t)/(2*self.Q_factor_pp))
-                else:
-                    w_long = 0
-                    
-                w_list.append(w_long)
-                
-        self.W_long = np.array(w_list)
-                
-    def wake_potential(self): 
-        self.W_p = signal.convolve(self.W_long*1e-12, self.rho, mode="same") 
-    
-    def plot(self, var, plot_rho=True):
-        """
-        Plotting wakefunction or wake potential.
-
-        Parameters
-        ----------
-        var : {'W_p', 'W_long' }
-            If 'W_p', the wake potential is plotted. 
-            If 'W_long', the wakefunction is plotted.            
-        plot_rho : bool, optional
-            Overlay the bunch charge density profile on the plot.
-            The default is True.
-
-        Returns
-        -------
-        fig : Figure
-            Figure object with the plot on it.
 
-        """
-        
-        fig, ax = plt.subplots()
-            
-        if var == "W_p":
-            ax.plot(self.tau*1e12, self.W_p*1e-12)
-            
-            ax.set_xlabel("$\\tau$ (ps)")
-            ax.set_ylabel("W$_p$ (V/pC)")
-    
-        elif var == "W_long":
-            ax.plot(self.tau*1e12, self.W_long*1e-12)
-            ax.set_xlabel("$\\tau$ (ps)")
-            ax.set_ylabel("W$_{||}$ ($\\Omega$/ps)")
-            
-        if plot_rho is True:
-            rho_array = np.array(self.rho)
-            rho_rescaled = rho_array/max(rho_array)*max(self.W_p)
-            
-            ax.plot(self.bins[0].mid*1e12, rho_rescaled*1e-12)
-            
-        else:
-            pass
-            
-        return fig
-    
-    def check_wake_tail(self):
-        """
-        Checking whether the full wakefunction is obtained by the calculated
-        initial timestop.
-
-        """
-        
-        ratio = np.abs(min(self.W_long) / self.W_long[-6:-1])
-        while any(ratio < 1000):
-            # Initial timestop is too short. 
-            # Extending timestop by 50 ps and recalculating."
-            self.timestop += 50e-12      
-            self.time_array()
-            self.long_wakefunction()
-            ratio = np.abs(min(self.W_long) / self.W_long[-6:-1])
-        
-    @Element.parallel
-    def track(self, bunch):
-        """
-        Tracking method for the element.
-        No bunch to bunch interaction, so written for Bunch objects and
-        @Element.parallel is used to handle Beam objects.
-
-        Parameters
-        ----------
-        bunch : Bunch or Beam object.
-        
-        """
-
-        self.charge_density(bunch, n_bin = self.n_bin)
-        self.init_timestop()
-        self.time_array()
-        self.long_wakefunction()
-        self.check_wake_tail()
-        self.wake_potential()
-        
-        f = interp1d(self.tau, self.W_p, fill_value = 0, bounds_error = False)
-        self.wp = f(bunch["tau"])
-        
-        bunch["delta"] += self.wp * bunch.charge / self.ring.E0
-        
 class SynchrotronRadiation(Element):
     """
     Element to handle synchrotron radiation, radiation damping and quantum 
diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 028472c4c7fe270d4434391af894d9fcc56a35eb..2f94fe639ee6b2337461aa8155f57816b4c4e70b 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -8,6 +8,10 @@ Created on Tue Jan 14 18:11:33 2020
 from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor, 
                                                  PhaseSpaceMonitor,
                                                  BeamMonitor,
-                                                 ProfileMonitor)
+                                                 ProfileMonitor,
+                                                 WakePotentialMonitor)
 from mbtrack2.tracking.monitors.plotting import (plot_bunchdata, 
-                                                 plot_phasespacedata)
\ No newline at end of file
+                                                 plot_phasespacedata,
+                                                 plot_profiledata,
+                                                 plot_beamdata,
+                                                 plot_wakedata)
\ No newline at end of file
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index c7c6ac86173bda7087ac8e529879671b9ff32544..0382055537a4508fcd02f6ad24e30184c7224aa7 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -282,8 +282,7 @@ class BunchMonitor(Monitor):
         object_to_save : Bunch or Beam object
         """        
         self.track_bunch_data(object_to_save)
-
-            
+        
 class PhaseSpaceMonitor(Monitor):
     """
     Monitor a single bunch and save the full phase space.
@@ -626,4 +625,128 @@ class ProfileMonitor(Monitor):
         ----------
         object_to_save : Bunch or Beam object
         """        
-        self.track_bunch_data(object_to_save)
\ No newline at end of file
+        self.track_bunch_data(object_to_save)
+        
+class WakePotentialMonitor(Monitor):
+    """
+    Monitor the wake potential from a single bunch and save attributes (tau, 
+    ...).
+    
+    Parameters
+    ----------
+    bunch_number : int
+        Bunch to monitor.
+    wake_types : str or list of str
+        Wake types to save: "Wlong, "Wxdip", ...
+    n_bin : int
+        Number of bin used for the bunch slicing.
+    file_name : string, optional
+        Name of the HDF5 where the data will be stored. Must be specified
+        the first time a subclass of Monitor is instancied and must be None
+        the following times.
+    save_every : int or float, optional
+        Set the frequency of the save. The data is saved every save_every 
+        call of the montior.
+    buffer_size : int or float, optional
+        Size of the save buffer.
+    total_size : int or float, optional
+        Total size of the save. The following relationships between the 
+        parameters must exist: 
+            total_size % buffer_size == 0
+            number of call to track / save_every == total_size
+    mpi_mode : bool, optional
+        If True, open the HDF5 file in parallel mode, which is needed to
+        allow several cores to write in the same file at the same time.
+        If False, open the HDF5 file in standard mode.
+
+    Methods
+    -------
+    track(wake_potential_to_save)
+        Save data.
+    """
+    
+    def __init__(self, bunch_number, wake_types, n_bin, file_name=None, 
+                 save_every=5, buffer_size=500, total_size=2e4, mpi_mode=True):
+        
+        self.bunch_number = bunch_number
+        group_name = "WakePotentialData_" + str(self.bunch_number)
+        
+        if isinstance(wake_types, str):
+            self.wake_types = [wake_types]
+        else:
+            self.wake_types = wake_types
+            
+        self.n_bin = n_bin
+        
+        dict_buffer = {}
+        dict_file = {}
+        dict_buffer.update({"tau" : (self.n_bin, buffer_size)})
+        dict_file.update({"tau" : (self.n_bin, total_size)})
+        dict_buffer.update({"rho" : (self.n_bin, buffer_size)})
+        dict_file.update({"rho" : (self.n_bin, total_size)})
+        for index, dim in enumerate(self.wake_types):
+            dict_buffer.update({dim : (self.n_bin, buffer_size)})
+            dict_file.update({dim : (self.n_bin, total_size)})
+
+        self.monitor_init(group_name, save_every, buffer_size, total_size,
+                          dict_buffer, dict_file, file_name, mpi_mode)
+        
+        self.dict_buffer = dict_buffer
+        self.dict_file = dict_file
+        
+    def to_buffer(self, wp):
+        """
+        Save data to buffer.
+        
+        Parameters
+        ----------
+        wp : WakePotential object
+        """
+
+        self.time[self.buffer_count] = self.track_count
+        self.tau[:, self.buffer_count] = wp.tau
+        self.rho[:, self.buffer_count] = wp.rho
+        for index, dim in enumerate(self.wake_types):
+            self.__getattribute__(dim)[:, self.buffer_count] = wp.__getattribute__(dim)
+        
+        self.buffer_count += 1
+        
+        if self.buffer_count == self.buffer_size:
+            self.write()
+            self.buffer_count = 0
+            
+    def write(self):
+        """Write data from buffer to the HDF5 file."""
+        
+        self.file[self.group_name]["time"][self.write_count*self.buffer_size:(
+                    self.write_count+1)*self.buffer_size] = self.time
+
+        self.file[self.group_name]["tau"][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.tau
+        
+        self.file[self.group_name]["rho"][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.rho
+        
+        for dim in self.wake_types:
+            self.file[self.group_name][dim][:, 
+                    self.write_count * self.buffer_size:(self.write_count+1) * 
+                    self.buffer_size] = self.__getattribute__(dim)
+            
+        self.file.flush()
+        self.write_count += 1
+                    
+    def track(self, wake_potential_to_save):
+        """
+        Save data.
+        
+        Parameters
+        ----------
+        object_to_save : WakePotential object
+        """        
+        if self.track_count % self.save_every == 0:
+            self.to_buffer(wake_potential_to_save)
+        self.track_count += 1
+
+            
\ No newline at end of file
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 273832166adbda6625b9a9c9db891628943e876d..de90791b3d3b595a973b8398cfc691c24f909bcc 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -373,4 +373,112 @@ def plot_profiledata(filename, bunch_number, dimension="tau", start=0,
         return fig
     elif streak_plot is True:
         return fig2
+    
+def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
+                     stop=None, step=None, profile_plot=False, streak_plot=True):
+    """
+    Plot data recorded by WakePotentialMonitor
+
+    Parameters
+    ----------
+    filename : str
+        Name of the HDF5 file that contains the data.
+    bunch_number : int
+        Bunch to plot. This has to be identical to 'bunch_number' parameter in 
+        'WakePotentialMonitor' object.
+    wake_type : str, optional
+        Wake type to plot: "Wlong", "Wxdip", ... Can also be "rho" to show 
+        bunch profile.
+    start : int, optional
+        First turn to plot. The default is 0.
+    stop : int, optional
+        Last turn to plot. If None, the last turn of the record is selected.
+    step : int, optional
+        Plotting step. This has to be divisible by 'save_every' parameter in
+        'WakePotentialMonitor' object, i.e. step % save_every == 0. If None, 
+        step is equivalent to save_every.
+    profile_plot : bool, optional
+        If Ture, wake potential profile plot is plotted.
+    streak_plot : bool, optional
+        If True, strek plot is plotted.
+
+    Returns
+    -------
+    fig : Figure
+        Figure object with the plot on it.
+
+    """
+    
+    file = hp.File(filename, "r")
+    path = file['WakePotentialData_{0}'.format(bunch_number)]
+    
+    if stop is None:
+        stop = path['time'][-1]
+    elif stop not in path['time']:
+        raise ValueError("stop not found. Choose from {0}"
+                         .format(path['time'][:]))
+ 
+    if start not in path['time']:
+        raise ValueError("start not found. Choose from {0}"
+                         .format(path['time'][:]))
+    
+    save_every = path['time'][1] - path['time'][0]
+    
+    if step is None:
+        step = save_every
+    
+    if step % save_every != 0:
+        raise ValueError("step must be divisible by the recording step "
+                         "which is {0}.".format(save_every))
+    
+    dimension_dict = {"Wlong":0, "Wxdip":1, "Wydip":2, "Wxquad":3, "Wyquad":4, "rho":5}
+    scale = [1e-12, 1e-12, 1e-12, 1e-15, 1e-15, 1]
+    label = ["$W_p$ (V/pC)", "$W_{p,x}^D (V/pC)$", "$W_{p,y}^D (V/pC)$", "$W_{p,x}^Q (V/pC/mm)$",
+             "$W_{p,y}^Q (V/pC/mm)$", "$\\rho (a.u.)$"]
+    
+    num = int((stop - start)/step)
+    n_bin = len(path[wake_type][:,0])
+    
+    start_index = np.where(path['time'][:] == start)[0][0]
+    
+    x_var = np.zeros((num+1,n_bin))
+    turn_index_array = np.zeros((num+1,))
+    for i in range(num+1):
+        turn_index = start_index + i * step / save_every 
+        turn_index_array[i] = turn_index
+        # construct an array of bin mids
+        x_var[i,:] = path["tau"][:,turn_index]
+        
+    if profile_plot is True:
+        fig, ax = plt.subplots()
+        for i in range(num+1):
+            ax.plot(x_var[i]*1e12,
+                    path[wake_type][:,turn_index_array[i]]*scale[dimension_dict[wake_type]], 
+                    label="turn {0}".format(path['time'][turn_index_array[i]]))
+        ax.set_xlabel("$\\tau$ (ps)")
+        ax.set_ylabel(label[dimension_dict[wake_type]])         
+        ax.legend()
+            
+    if streak_plot is True:
+        turn = np.reshape(path['time'][turn_index_array], (num+1,1))
+        y_var = np.ones((num+1,n_bin)) * turn
+        z_var = np.transpose(path[wake_type][:,turn_index_array]*scale[dimension_dict[wake_type]])
+        fig2, ax2 = plt.subplots()
+        cmap = mpl.cm.cool
+        c = ax2.imshow(z_var, cmap=cmap, origin='lower' , aspect='auto',
+                       extent=[x_var.min()*1e12,
+                               x_var.max()*1e12,
+                               y_var.min(),y_var.max()])
+        ax2.set_xlabel("$\\tau$ (ps)")
+        ax2.set_ylabel("Number of turns")
+        cbar = fig2.colorbar(c, ax=ax2)
+        cbar.set_label(label[dimension_dict[wake_type]]) 
+
+    file.close()
+    if profile_plot is True and streak_plot is True:
+        return fig, fig2
+    elif profile_plot is True:
+        return fig
+    elif streak_plot is True:
+        return fig2
     
\ No newline at end of file
diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py
index d5f3107b2cdf0aa1097bf60d868911c7a360f413..f5e2d69a6f5f65e42229ff9f41e5c604a6a24b68 100644
--- a/tracking/synchrotron.py
+++ b/tracking/synchrotron.py
@@ -283,4 +283,4 @@ class Synchrotron:
         phase = np.pi - np.arcsin(self.U0 / Vrf)
         tuneS = np.sqrt( - (Vrf / self.E0) * (self.h * self.ac) / (2*np.pi) 
                         * np.cos(phase) )
-        return tuneS
\ No newline at end of file
+        return tuneS
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9197bca6bf62f3d9633bb683e2a723bf7f469f7
--- /dev/null
+++ b/tracking/wakepotential.py
@@ -0,0 +1,320 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines the WakePotential class which deals with the single bunch 
+effects, both longitudinal and transverse.
+
+@author: Alexis Gamelin, Watanyu Foosang
+@date: 25/09/2020
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import pandas as pd
+from scipy import signal
+from scipy.interpolate import interp1d
+from mbtrack2.tracking.element import Element
+
+class WakePotential(Element):
+    """
+    Compute a wake potential from wake functions by performing a convolution 
+    with a bunch charge profile.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    wakefield : Wakefield object
+        Wakefield object which contains the wake functions to be used.
+    n_bin : int, optional
+        Number of bins for constructing the longitudinal bunch profile.
+        
+    Attributes
+    ----------
+    rho : array of shape (n_bin, )
+        Bunch charge density profile in the unit [1/s].
+    Wp : array 
+        Wake potential profile.
+    Wp_interp : array of shape (mp_number, )
+        Wake potential, obtained from interpolating Wp, exerted on each macro-particle.
+    
+    Methods
+    -------
+    charge_density(bunch)
+        Compute bunch charge density profile in [1/s].
+    dipole_moment(bunch, plane)
+        Return the dipole moment of the bunch computed on the same time array 
+        as the bunch profile.
+    get_wakepotential(bunch, wake_type)
+        Return the wake potential computed on the same time array as the bunch 
+        profile.
+    track(bunch)
+        Tracking method for the element.
+    plot_last_wake(wake_type, plot_rho=True, plot_dipole=False)
+        Plot the last wake potential of a given type computed during the last
+        call of the track method.
+        
+    """
+    
+    def __init__(self, ring, wakefield, n_bin=65):
+        self.wakefield = wakefield
+        self.types = self.wakefield.wake_components
+        self.n_types = len(self.wakefield.wake_components)
+        self.ring = ring
+        self.n_bin = n_bin
+            
+    def charge_density(self, bunch):
+        """
+        Compute bunch charge density profile in [1/s].
+
+        Parameters
+        ----------
+        bunch : Bunch object
+
+        """
+        
+        # Get binning data
+        a, b, c = bunch.binning(n_bin=self.n_bin)
+        self.bins = a
+        self.sorted_index = b
+        self.profile = c
+        self.bin_size = self.bins.length
+        
+        # Compute charge density
+        self.rho = bunch.charge_per_mp*self.profile/(self.bin_size*bunch.charge)
+        self.rho = np.array(self.rho)
+        
+        # Compute time array
+        self.tau = np.array(self.bins.mid)
+        self.dtau = self.tau[1] - self.tau[0]
+        
+        # Add N values before and after rho and tau
+        if self.n_bin % 2 == 0:
+            N = int(self.n_bin/2)
+            self.tau = np.arange(self.tau[0] - self.dtau*N, 
+                                 self.tau[-1] + self.dtau*N,
+                                 self.dtau)
+            self.rho = np.append(self.rho, np.zeros(N))
+            self.rho = np.insert(self.rho, 0, np.zeros(N))
+        else:
+            N = int(np.floor(self.n_bin/2))
+            self.tau = np.arange(self.tau[0] - self.dtau*N, 
+                                 self.tau[-1] + self.dtau*(N+1),
+                                 self.dtau)
+            self.rho = np.append(self.rho, np.zeros(N))
+            self.rho = np.insert(self.rho, 0, np.zeros(N+1))
+            
+        if len(self.tau) != len(self.rho):
+            self.tau = np.append(self.tau, self.tau[-1] + self.dtau)
+            
+    def dipole_moment(self, bunch, plane):
+        """
+        Return the dipole moment of the bunch computed on the same time array 
+        as the bunch profile.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+        plane : str
+            Plane on which the dipole moment is computed, "x" or "y".
+
+        Returns
+        -------
+        dipole : array
+            Dipole moment of the bunch.
+
+        """
+        dipole = np.zeros((self.n_bin,))
+        for i in range(self.n_bin):
+            dipole[i] = bunch[plane][self.sorted_index == self.bins[i]].sum()
+        dipole = dipole/self.profile
+        dipole[np.isnan(dipole)] = 0
+        
+        # Add N values to get same size as tau/rho
+        if self.n_bin % 2 == 0:
+            N = int(self.n_bin/2)
+            dipole = np.append(dipole, np.zeros(N))
+            dipole = np.insert(dipole, 0, np.zeros(N))
+        else:
+            N = int(np.floor(self.n_bin/2))
+            dipole = np.append(dipole, np.zeros(N))
+            dipole = np.insert(dipole, 0, np.zeros(N+1))
+            
+        setattr(self, "dipole_" + plane, dipole)
+        return dipole
+        
+    def get_wakepotential(self, bunch, wake_type):
+        """
+        Return the wake potential computed on the same time array as the bunch 
+        profile.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+        wake_type : str
+            Wake function type: "Wlong", "Wxdip", ...
+
+        Returns
+        -------
+        Wp : array
+            Wake potential.
+
+        """
+        
+        tau0 = getattr(self.wakefield, wake_type).data.index
+        W0 = getattr(self.wakefield, wake_type).data["real"]
+        interp_func_W0 = interp1d(tau0, W0, fill_value=0, 
+                                     bounds_error=False)
+        W = interp_func_W0(self.tau - np.mean(self.tau))
+        if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
+            Wp = signal.convolve(self.rho, W*-1, mode='same')*self.dtau
+        elif wake_type == "Wxdip":
+            dipole = self.dipole_moment(bunch, "x")
+            Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+        elif wake_type == "Wydip":
+            dipole = self.dipole_moment(bunch, "y")
+            Wp = signal.convolve(self.rho*dipole, W*-1, mode='same')*self.dtau
+        else:
+            raise ValueError("This type of wake is not taken into account.")
+        
+        setattr(self, wake_type, Wp)
+        return Wp
+    
+    @Element.parallel
+    def track(self, bunch):
+        """
+        Tracking method for the element.
+        No bunch to bunch interaction, so written for Bunch objects and
+        @Element.parallel is used to handle Beam objects.
+
+        Parameters
+        ----------
+        bunch : Bunch or Beam object.
+        
+        """
+
+        self.charge_density(bunch)
+        for wake_type in self.types:
+            Wp = self.get_wakepotential(bunch, wake_type)
+            f = interp1d(self.tau, Wp, fill_value = 0, bounds_error = False)
+            Wp_interp = f(bunch["tau"])
+            if wake_type == "Wlong":
+                bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
+            elif wake_type == "Wxdip":
+                bunch["xp"] += Wp_interp * bunch.charge / self.ring.E0
+            elif wake_type == "Wydip":
+                bunch["yp"] += Wp_interp * bunch.charge / self.ring.E0
+            elif wake_type == "Wxquad":
+                bunch["xp"] += (bunch["x"] * Wp_interp * bunch.charge 
+                                / self.ring.E0)
+            elif wake_type == "Wyquad":
+                bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge 
+                                / self.ring.E0)
+                
+    def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False):
+        """
+        Plot the last wake potential of a given type computed during the last
+        call of the track method.
+
+        Parameters
+        ----------
+        wake_type : str
+            Type of the wake to plot: "Wlong", "Wxdip", ...
+        plot_rho : bool, optional
+            Plot the normalised bunch profile. The default is True.
+        plot_dipole : bool, optional
+            Plot the normalised dipole moment. The default is False.
+
+        Returns
+        -------
+        fig : figure
+
+        """
+        
+        labels = {"Wlong" : r"$W_{p,long}$ (V/pC)", 
+                  "Wxdip" : r"$W_{p,x}^{D} (V/pC)$",
+                  "Wydip" : r"$W_{p,y}^{D} (V/pC)$",
+                  "Wxquad" : r"$W_{p,x}^{Q} (V/pC/m)$",
+                  "Wyquad" : r"$W_{p,y}^{Q} (V/pC/m)$"}
+        fig, ax = plt.subplots()
+        Wp = getattr(self, wake_type)
+        ax.plot(self.tau*1e12, Wp*1e-12, label=labels[wake_type])
+        ax.set_xlabel("$\\tau$ (ps)")
+        ax.set_ylabel(labels[wake_type])
+
+        if plot_rho is True:
+            rho_rescaled = self.rho/max(self.rho)*max(np.abs(Wp))
+            ax.plot(self.tau*1e12, rho_rescaled*1e-12, label=r"$\rho$ (a.u.)")
+            plt.legend()
+            
+        if plot_dipole is True:
+            dipole = getattr(self, "dipole_" + wake_type[1])
+            dipole_rescaled = dipole/max(dipole)*max(np.abs(Wp))
+            ax.plot(self.tau*1e12, dipole_rescaled*1e-12, label=r"Dipole moment (a.u.)")
+            plt.legend()
+            
+        return fig
+    
+    def energy_loss(self, bunch, compare_to='num'):
+        """
+        Calculate the energy loss from the wake potential and compare it to a 
+        reference value assuming a Gaussian bunch.
+
+        Parameters
+        ----------
+        bunch : Bunch object
+        compare_to : str, {"theory", "num"}, optional
+            The method for calculating the reference energy loss. Use 'theory' 
+            for the analytical loss factor of a resonator [1], or use 'num' 
+            for a numerical result obtained from loss_factor method in 
+            Impedance class.
+
+        Returns
+        -------
+        energy_loss_data : DataFrame
+            An output showing the enegy loss compared to the reference value.
+            
+        Reference
+        ---------
+        [1] A. W. Chao and M. Tigner, "Handbook of Accelerator Physics and
+            Engineering", 3rd printing, p.239
+
+        """
+        
+        if np.where(self.types=="Wlong")[0].size == 0:
+            raise TypeError("No longitudinal wake function data.")
+                
+        else:
+            Wp = self.get_wakepotential(bunch, "Wlong")
+            Eloss = -1*np.trapz(Wp*self.rho, self.tau)*bunch.charge
+            
+            res = self.wakefield
+            sigma = bunch['tau'].std()
+            if compare_to == 'theory':
+                Eloss_0 = 0.5*res.wr*res.Rs*1/res.Q *np.exp(-1*res.wr**2*sigma**2) \
+                    * bunch.charge
+            elif compare_to == 'num':
+                Eloss_0 = res.Zlong.loss_factor(sigma)*bunch.charge
+            else:
+                raise KeyError("Reference method not correct. Enter \"theory\" "
+                               "or \"num\" ")
+            
+            delta_Eloss = (Eloss-Eloss_0) / Eloss_0 *100
+            
+            column = ['Energy loss [eV]', 'Reference energy loss [eV]', 'Relative error [%]']
+            energy_loss_data = pd.DataFrame(np.reshape([Eloss, Eloss_0, delta_Eloss], (1,3)), 
+                                            columns=column)
+        return energy_loss_data
+            
+            
+        
+                
+                           
+            
+    
+    
+    
+    
+    
+    
+    
+    
+    
\ No newline at end of file