Improve imports, move functions to utilities and rename Wakefield to WakeField

collective_effects/__init__.py

@@ -7,9 +7,10 @@ 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
 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, ImpedanceModel, 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
\ No newline at end of file

collective_effects/resistive_wall.py

 # -*- 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

collective_effects/resonator.py

 # -*- coding: utf-8 -*-
 """
 This module defines the impedances and wake functions from the resonator model 
-based on the Wakefield class.
+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, 
+from mbtrack2.collective_effects.wakefield import (WakeField, Impedance, 
                                                    WakeFunction)
 
-class Resonator(Wakefield):
+class Resonator(WakeField):
     def __init__(self, Rs, fr, Q, plane, n_wake=1e4, n_imp=1e4, imp_freq_lim=50e9):
         """
-        Resonator model Wakefield element which computes the impedance and the 
+        Resonator model WakeField element which computes the impedance and the 
         wake function in both longitudinal and transverse case.
 
         Parameters

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

collective_effects/utilities.py

+# -*- 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
+import matplotlib.pyplot as plt
+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
+