From 8c764a24f991ee83c1ae2bd79ff9442429b06e41 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Wed, 30 Mar 2022 16:39:18 +0200
Subject: [PATCH] Rework code structure
Add utilities and instability folders
Change folder name from collective_effects to impedance
---
collective_effects/__init__.py | 21 -
collective_effects/utilities.py | 481 ------------------
impedance/__init__.py | 22 +
{collective_effects => impedance}/csr.py | 2 +-
.../impedance_model.py | 8 +-
.../resistive_wall.py | 2 +-
.../resonator.py | 2 +-
{collective_effects => impedance}/tapers.py | 2 +-
.../wakefield.py | 0
instability/__init__.py | 17 +
.../instabilities.py | 0
{collective_effects => instability}/ions.py | 0
tracking/__init__.py | 1 -
tracking/wakepotential.py | 2 +-
utilities/__init__.py | 21 +
{vlasov => utilities}/beamloading.py | 0
.../data/Yokoya_elliptic_from_Elias_USPAS.csv | 0
utilities/misc.py | 241 +++++++++
{tracking => utilities}/optics.py | 0
utilities/read_impedance.py | 135 +++++
utilities/spectrum.py | 121 +++++
vlasov/__init__.py | 8 -
22 files changed, 566 insertions(+), 520 deletions(-)
delete mode 100644 collective_effects/__init__.py
delete mode 100644 collective_effects/utilities.py
create mode 100644 impedance/__init__.py
rename {collective_effects => impedance}/csr.py (98%)
rename {collective_effects => impedance}/impedance_model.py (98%)
rename {collective_effects => impedance}/resistive_wall.py (99%)
rename {collective_effects => impedance}/resonator.py (98%)
rename {collective_effects => impedance}/tapers.py (98%)
rename {collective_effects => impedance}/wakefield.py (100%)
create mode 100644 instability/__init__.py
rename {collective_effects => instability}/instabilities.py (100%)
rename {collective_effects => instability}/ions.py (100%)
create mode 100644 utilities/__init__.py
rename {vlasov => utilities}/beamloading.py (100%)
rename {collective_effects => utilities}/data/Yokoya_elliptic_from_Elias_USPAS.csv (100%)
create mode 100644 utilities/misc.py
rename {tracking => utilities}/optics.py (100%)
create mode 100644 utilities/read_impedance.py
create mode 100644 utilities/spectrum.py
delete mode 100644 vlasov/__init__.py
diff --git a/collective_effects/__init__.py b/collective_effects/__init__.py
deleted file mode 100644
index a23651b..0000000
--- a/collective_effects/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-Created on Tue Jan 14 12:25:30 2020
-
-@author: gamelina
-"""
-
-from mbtrack2.collective_effects.instabilities import mbi_threshold, cbi_threshold, lcbi_growth_rate_mode, lcbi_growth_rate
-from mbtrack2.collective_effects.instabilities import rwmbi_growth_rate, rwmbi_threshold
-from mbtrack2.collective_effects.resistive_wall import (skin_depth, CircularResistiveWall, Coating)
-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, 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, read_IW2D_folder
-from mbtrack2.collective_effects.utilities import gaussian_bunch
-from mbtrack2.collective_effects.impedance_model import ImpedanceModel
-from mbtrack2.collective_effects.csr import (FreeSpaceCSR, ParallelPlatesCSR)
-from mbtrack2.collective_effects.ions import ion_cross_section, ion_frequency, fast_beam_ion, plot_critical_mass
\ No newline at end of file
diff --git a/collective_effects/utilities.py b/collective_effects/utilities.py
deleted file mode 100644
index a2d3858..0000000
--- a/collective_effects/utilities.py
+++ /dev/null
@@ -1,481 +0,0 @@
-# -*- 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 pathlib import Path
-from scipy.interpolate import interp1d
-from scipy.constants import c
-from mbtrack2.collective_effects.wakefield import Impedance, WakeFunction, WakeField
-
-
-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, file_type='Zlong'):
- """
- Read IW2D file format into an Impedance object or a WakeField object.
-
- Parameters
- ----------
- file : str
- Path to the file to read.
- file_type : str, optional
- Type of the Impedance or WakeField object.
-
- Returns
- -------
- result : Impedance or WakeField object
- Data from file.
- """
- if file_type[0] == "Z":
- df = pd.read_csv(file, sep = ' ', header = None,
- names = ["Frequency","Real","Imaginary"], skiprows=1)
- df.set_index("Frequency", inplace = True)
- df = df[df["Real"].notna()]
- df = df[df["Imaginary"].notna()]
- result = Impedance(variable = df.index,
- function = df["Real"] + 1j*df["Imaginary"],
- impedance_type=file_type[1:])
- elif file_type[0] == "W":
- df = pd.read_csv(file, sep = ' ', header = None,
- names = ["Distance","Wake"], skiprows=1)
- df["Time"] = df["Distance"] / c
- df.set_index("Time", inplace = True)
- if np.any(df.isna()):
- index = df.isna().values
- df = df.interpolate()
- print("Nan values have been interpolated to:")
- print(df[index])
- # if file_type == "Wlong":
- # df["Wake"] = df["Wake"]*-1
- result = WakeFunction(variable = df.index,
- function = df["Wake"],
- wake_type=file_type[1:])
- else:
- raise ValueError("file_type should begin by Z or W.")
- return result
-
-def read_IW2D_folder(folder, suffix, select="WZ"):
- """
- Read IW2D results into a WakeField object.
-
- Parameters
- ----------
- file : str
- Path to the file to read.
- suffix : str
- End of the name of each files. For example, in "Zlong_test.dat" the
- suffix should be "_test.dat".
- select : str, optional
- Select which object to load. "W" for WakeFunction, "Z" for Impedance
- and "WZ" or "ZW" for both.
-
- Returns
- -------
- result : WakeField object
- WakeField object with Impedance and WakeFunction objects from the
- different files.
- """
- if (select == "WZ") or (select == "ZW"):
- types = {"W" : WakeFunction,
- "Z" : Impedance}
- elif (select == "W"):
- types = {"W" : WakeFunction}
- elif (select == "Z"):
- types = {"Z" : Impedance}
- else:
- raise ValueError("select should be W, Z or WZ.")
-
- components = ["long", "xdip", "ydip", "xquad", "yquad"]
-
- data_folder = Path(folder)
-
- list_for_wakefield = []
- for key, item in types.items():
- for component in components:
- name = data_folder / (key + component + suffix)
- res = read_IW2D(file=name, file_type=key + component)
- list_for_wakefield.append(res)
-
- wake = WakeField(list_for_wakefield)
-
- return wake
-
-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 gaussian_bunch(time, sigma):
- """
- Compute a Gaussian bunch profile.
-
- Parameters
- ----------
- time : array
- sample points of the bunch profile in [s].
- sigma : float
- RMS bunch length in [s].
-
- Returns
- -------
- bunch_profile : array
- Bunch profile in [s**-1] sampled at points given in time.
- """
- return np.exp(-1/2*(time**2/sigma**2))/(sigma*np.sqrt(2*np.pi))
-
-
-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
-
- path_to_file = Path(__file__).parent
- file = path_to_file / "data" / "Yokoya_elliptic_from_Elias_USPAS.csv"
-
- # read Yokoya factors interpolation file
- # BEWARE: columns are ratio, dipy, dipx, quady, quadx
- yokoya_file = pd.read_csv(file)
- 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/impedance/__init__.py b/impedance/__init__.py
new file mode 100644
index 0000000..b58cd94
--- /dev/null
+++ b/impedance/__init__.py
@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 12:25:30 2020
+
+@author: gamelina
+"""
+
+from mbtrack2.impedance.resistive_wall import (skin_depth,
+ CircularResistiveWall,
+ Coating)
+from mbtrack2.impedance.resonator import (Resonator,
+ PureInductive,
+ PureResistive)
+from mbtrack2.impedance.tapers import (StupakovRectangularTaper,
+ StupakovCircularTaper)
+from mbtrack2.impedance.wakefield import (ComplexData,
+ Impedance,
+ WakeFunction,
+ WakeField)
+from mbtrack2.impedance.impedance_model import ImpedanceModel
+from mbtrack2.impedance.csr import (FreeSpaceCSR,
+ ParallelPlatesCSR)
\ No newline at end of file
diff --git a/collective_effects/csr.py b/impedance/csr.py
similarity index 98%
rename from collective_effects/csr.py
rename to impedance/csr.py
index b72681a..5f08765 100644
--- a/collective_effects/csr.py
+++ b/impedance/csr.py
@@ -9,7 +9,7 @@ import numpy as np
import mpmath as mp
from scipy.constants import c, mu_0
from scipy.special import gamma
-from mbtrack2.collective_effects.wakefield import WakeField, Impedance, WakeFunction
+from mbtrack2.impedance.wakefield import WakeField, Impedance, WakeFunction
class FreeSpaceCSR(WakeField):
"""
diff --git a/collective_effects/impedance_model.py b/impedance/impedance_model.py
similarity index 98%
rename from collective_effects/impedance_model.py
rename to impedance/impedance_model.py
index e2b1ef2..b398e9b 100644
--- a/collective_effects/impedance_model.py
+++ b/impedance/impedance_model.py
@@ -10,10 +10,10 @@ import matplotlib.pyplot as plt
from scipy.integrate import trapz
from scipy.interpolate import interp1d
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, double_sided_impedance)
-from mbtrack2.collective_effects.wakefield import WakeField
+from mbtrack2.utilities import (beam_spectrum, gaussian_bunch_spectrum,
+ beam_loss_factor, spectral_density,
+ effective_impedance, double_sided_impedance)
+from mbtrack2.impedance.wakefield import WakeField
from mbtrack2.tracking.element import Element
class ImpedanceModel(Element):
diff --git a/collective_effects/resistive_wall.py b/impedance/resistive_wall.py
similarity index 99%
rename from collective_effects/resistive_wall.py
rename to impedance/resistive_wall.py
index c0ae255..126f053 100644
--- a/collective_effects/resistive_wall.py
+++ b/impedance/resistive_wall.py
@@ -9,7 +9,7 @@ Define resistive wall elements based on the WakeField class.
import numpy as np
from scipy.constants import mu_0, epsilon_0, c
from scipy.integrate import quad
-from mbtrack2.collective_effects.wakefield import WakeField, Impedance, WakeFunction
+from mbtrack2.impedance.wakefield import WakeField, Impedance, WakeFunction
def skin_depth(frequency, rho, mu_r = 1, epsilon_r = 1):
"""
diff --git a/collective_effects/resonator.py b/impedance/resonator.py
similarity index 98%
rename from collective_effects/resonator.py
rename to impedance/resonator.py
index 41168b3..e41883e 100644
--- a/collective_effects/resonator.py
+++ b/impedance/resonator.py
@@ -8,7 +8,7 @@ based on the WakeField class.
"""
import numpy as np
-from mbtrack2.collective_effects.wakefield import (WakeField, Impedance,
+from mbtrack2.impedance.wakefield import (WakeField, Impedance,
WakeFunction)
class Resonator(WakeField):
diff --git a/collective_effects/tapers.py b/impedance/tapers.py
similarity index 98%
rename from collective_effects/tapers.py
rename to impedance/tapers.py
index 7703621..73d891e 100644
--- a/collective_effects/tapers.py
+++ b/impedance/tapers.py
@@ -8,7 +8,7 @@ Created on Tue Mar 31 13:15:36 2020
from scipy.constants import mu_0, c, pi
import numpy as np
from scipy.integrate import trapz
-from mbtrack2.collective_effects.wakefield import WakeField, Impedance
+from mbtrack2.impedance.wakefield import WakeField, Impedance
class StupakovRectangularTaper(WakeField):
"""
diff --git a/collective_effects/wakefield.py b/impedance/wakefield.py
similarity index 100%
rename from collective_effects/wakefield.py
rename to impedance/wakefield.py
diff --git a/instability/__init__.py b/instability/__init__.py
new file mode 100644
index 0000000..a8925c1
--- /dev/null
+++ b/instability/__init__.py
@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Mar 30 16:22:48 2022
+
+@author: gamelina
+"""
+
+from mbtrack2.instability.ions import (ion_cross_section,
+ ion_frequency,
+ fast_beam_ion,
+ plot_critical_mass)
+from mbtrack2.instability.instabilities import (mbi_threshold,
+ cbi_threshold,
+ lcbi_growth_rate_mode,
+ lcbi_growth_rate,
+ rwmbi_growth_rate,
+ rwmbi_threshold)
\ No newline at end of file
diff --git a/collective_effects/instabilities.py b/instability/instabilities.py
similarity index 100%
rename from collective_effects/instabilities.py
rename to instability/instabilities.py
diff --git a/collective_effects/ions.py b/instability/ions.py
similarity index 100%
rename from collective_effects/ions.py
rename to instability/ions.py
diff --git a/tracking/__init__.py b/tracking/__init__.py
index 3f5f79b..5ecae9b 100644
--- a/tracking/__init__.py
+++ b/tracking/__init__.py
@@ -10,7 +10,6 @@ from mbtrack2.tracking.particles import (Electron, Proton, Bunch, Beam,
from mbtrack2.tracking.synchrotron import Synchrotron
from mbtrack2.tracking.rf import RFCavity, CavityResonator
from mbtrack2.tracking.parallel import Mpi
-from mbtrack2.tracking.optics import Optics, PhysicalModel
from mbtrack2.tracking.element import (Element, LongitudinalMap,
TransverseMap, SynchrotronRadiation,
SkewQuadrupole)
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
index 061d4eb..655c266 100644
--- a/tracking/wakepotential.py
+++ b/tracking/wakepotential.py
@@ -14,7 +14,7 @@ from scipy import signal
from scipy.interpolate import interp1d
from scipy.constants import mu_0, c, pi
from mbtrack2.tracking.element import Element
-from mbtrack2.collective_effects.utilities import gaussian_bunch
+from mbtrack2.utilities.spectrum import gaussian_bunch
class WakePotential(Element):
"""
diff --git a/utilities/__init__.py b/utilities/__init__.py
new file mode 100644
index 0000000..0cee3f7
--- /dev/null
+++ b/utilities/__init__.py
@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+
+from mbtrack2.utilities.read_impedance import (read_CST,
+ read_IW2D,
+ read_IW2D_folder)
+from mbtrack2.utilities.misc import (effective_impedance,
+ yokoya_elliptic,
+ beam_loss_factor,
+ double_sided_impedance)
+from mbtrack2.utilities.spectrum import (spectral_density,
+ gaussian_bunch_spectrum,
+ gaussian_bunch,
+ beam_spectrum)
+from mbtrack2.utilities.optics import (Optics,
+ PhysicalModel)
+from mbtrack2.utilities.beamloading import BeamLoadingEquilibrium
\ No newline at end of file
diff --git a/vlasov/beamloading.py b/utilities/beamloading.py
similarity index 100%
rename from vlasov/beamloading.py
rename to utilities/beamloading.py
diff --git a/collective_effects/data/Yokoya_elliptic_from_Elias_USPAS.csv b/utilities/data/Yokoya_elliptic_from_Elias_USPAS.csv
similarity index 100%
rename from collective_effects/data/Yokoya_elliptic_from_Elias_USPAS.csv
rename to utilities/data/Yokoya_elliptic_from_Elias_USPAS.csv
diff --git a/utilities/misc.py b/utilities/misc.py
new file mode 100644
index 0000000..9dcac0d
--- /dev/null
+++ b/utilities/misc.py
@@ -0,0 +1,241 @@
+# -*- 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 pathlib import Path
+from mbtrack2.impedance.wakefield import Impedance
+
+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
+
+ path_to_file = Path(__file__).parent
+ file = path_to_file / "data" / "Yokoya_elliptic_from_Elias_USPAS.csv"
+
+ # read Yokoya factors interpolation file
+ # BEWARE: columns are ratio, dipy, dipx, quady, quadx
+ yokoya_file = pd.read_csv(file)
+ 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/tracking/optics.py b/utilities/optics.py
similarity index 100%
rename from tracking/optics.py
rename to utilities/optics.py
diff --git a/utilities/read_impedance.py b/utilities/read_impedance.py
new file mode 100644
index 0000000..cbaf301
--- /dev/null
+++ b/utilities/read_impedance.py
@@ -0,0 +1,135 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Mar 30 16:07:02 2022
+
+@author: gamelina
+"""
+
+import pandas as pd
+import numpy as np
+from pathlib import Path
+from scipy.constants import c
+from mbtrack2.impedance.wakefield import Impedance, WakeFunction, WakeField
+
+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, file_type='Zlong'):
+ """
+ Read IW2D file format into an Impedance object or a WakeField object.
+
+ Parameters
+ ----------
+ file : str
+ Path to the file to read.
+ file_type : str, optional
+ Type of the Impedance or WakeField object.
+
+ Returns
+ -------
+ result : Impedance or WakeField object
+ Data from file.
+ """
+ if file_type[0] == "Z":
+ df = pd.read_csv(file, sep = ' ', header = None,
+ names = ["Frequency","Real","Imaginary"], skiprows=1)
+ df.set_index("Frequency", inplace = True)
+ df = df[df["Real"].notna()]
+ df = df[df["Imaginary"].notna()]
+ result = Impedance(variable = df.index,
+ function = df["Real"] + 1j*df["Imaginary"],
+ impedance_type=file_type[1:])
+ elif file_type[0] == "W":
+ df = pd.read_csv(file, sep = ' ', header = None,
+ names = ["Distance","Wake"], skiprows=1)
+ df["Time"] = df["Distance"] / c
+ df.set_index("Time", inplace = True)
+ if np.any(df.isna()):
+ index = df.isna().values
+ df = df.interpolate()
+ print("Nan values have been interpolated to:")
+ print(df[index])
+ # if file_type == "Wlong":
+ # df["Wake"] = df["Wake"]*-1
+ result = WakeFunction(variable = df.index,
+ function = df["Wake"],
+ wake_type=file_type[1:])
+ else:
+ raise ValueError("file_type should begin by Z or W.")
+ return result
+
+def read_IW2D_folder(folder, suffix, select="WZ"):
+ """
+ Read IW2D results into a WakeField object.
+
+ Parameters
+ ----------
+ file : str
+ Path to the file to read.
+ suffix : str
+ End of the name of each files. For example, in "Zlong_test.dat" the
+ suffix should be "_test.dat".
+ select : str, optional
+ Select which object to load. "W" for WakeFunction, "Z" for Impedance
+ and "WZ" or "ZW" for both.
+
+ Returns
+ -------
+ result : WakeField object
+ WakeField object with Impedance and WakeFunction objects from the
+ different files.
+ """
+ if (select == "WZ") or (select == "ZW"):
+ types = {"W" : WakeFunction,
+ "Z" : Impedance}
+ elif (select == "W"):
+ types = {"W" : WakeFunction}
+ elif (select == "Z"):
+ types = {"Z" : Impedance}
+ else:
+ raise ValueError("select should be W, Z or WZ.")
+
+ components = ["long", "xdip", "ydip", "xquad", "yquad"]
+
+ data_folder = Path(folder)
+
+ list_for_wakefield = []
+ for key, item in types.items():
+ for component in components:
+ name = data_folder / (key + component + suffix)
+ res = read_IW2D(file=name, file_type=key + component)
+ list_for_wakefield.append(res)
+
+ wake = WakeField(list_for_wakefield)
+
+ return wake
\ No newline at end of file
diff --git a/utilities/spectrum.py b/utilities/spectrum.py
new file mode 100644
index 0000000..98e4512
--- /dev/null
+++ b/utilities/spectrum.py
@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Mar 30 16:09:05 2022
+
+@author: gamelina
+"""
+
+import numpy as np
+
+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 gaussian_bunch(time, sigma):
+ """
+ Compute a Gaussian bunch profile.
+
+ Parameters
+ ----------
+ time : array
+ sample points of the bunch profile in [s].
+ sigma : float
+ RMS bunch length in [s].
+
+ Returns
+ -------
+ bunch_profile : array
+ Bunch profile in [s**-1] sampled at points given in time.
+ """
+ return np.exp(-1/2*(time**2/sigma**2))/(sigma*np.sqrt(2*np.pi))
+
+
+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
\ No newline at end of file
diff --git a/vlasov/__init__.py b/vlasov/__init__.py
deleted file mode 100644
index 45a0308..0000000
--- a/vlasov/__init__.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-Created on Tue Jan 14 18:11:33 2020
-
-@author: gamelina
-"""
-
-from mbtrack2.vlasov.beamloading import BeamLoadingEquilibrium
\ No newline at end of file
--
GitLab