From c611185bfde52e81da9bcdfb87f50b715722146c Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Fri, 7 Feb 2020 17:28:08 +0100
Subject: [PATCH] Initial commit
---
.gitignore | 1 +
CollectiveEffect/__init__.py | 7 +
CollectiveEffect/instabilities.py | 44 +++
CollectiveEffect/resistive_wall.py | 65 +++++
CollectiveEffect/tools.py | 185 +++++++++++++
CollectiveEffect/wakefield.py | 425 +++++++++++++++++++++++++++++
Machines/__init__.py | 7 +
Machines/soleil.py | 55 ++++
Tracking/__init__.py | 7 +
Tracking/beam.py | 341 +++++++++++++++++++++++
Tracking/synchrotron.py | 181 ++++++++++++
__init__.py | 7 +
12 files changed, 1325 insertions(+)
create mode 100644 .gitignore
create mode 100644 CollectiveEffect/__init__.py
create mode 100644 CollectiveEffect/instabilities.py
create mode 100644 CollectiveEffect/resistive_wall.py
create mode 100644 CollectiveEffect/tools.py
create mode 100644 CollectiveEffect/wakefield.py
create mode 100644 Machines/__init__.py
create mode 100644 Machines/soleil.py
create mode 100644 Tracking/__init__.py
create mode 100644 Tracking/beam.py
create mode 100644 Tracking/synchrotron.py
create mode 100644 __init__.py
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..5c86bbc
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+*__pycache__*
diff --git a/CollectiveEffect/__init__.py b/CollectiveEffect/__init__.py
new file mode 100644
index 0000000..27c993f
--- /dev/null
+++ b/CollectiveEffect/__init__.py
@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 12:25:30 2020
+
+@author: gamelina
+"""
+
diff --git a/CollectiveEffect/instabilities.py b/CollectiveEffect/instabilities.py
new file mode 100644
index 0000000..34bd523
--- /dev/null
+++ b/CollectiveEffect/instabilities.py
@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+"""
+General calculations about instabilities
+
+@author: Alexis Gamelin
+@date: 15/01/2020
+"""
+
+from scipy.constants import c, m_e, e, pi, epsilon_0
+
+def MBI_threshold(ring, sigma, R, b):
+ """
+ Compute the microbunching instability (MBI) threshold for a bunched beam
+ considering the steady-state parallel plate model [1][2].
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ sigma : float
+ RMS bunch length in [s]
+ R : float
+ dipole bending radius in [m]
+ b : float
+ vertical distance between the conducting parallel plates in [m]
+
+ Returns
+ -------
+ I : float
+ MBI current threshold
+
+ [1] : Y. Cai, "Theory of microwave instability and coherent synchrotron
+ radiation in electron storage rings", SLAC-PUB-14561
+ [2] : D. Zhou, "Coherent synchrotron radiation and microwave instability in
+ electron storage rings", PhD thesis, p112
+ """
+
+ Ia = 4*pi*epsilon_0*m_e*c**3/e # Alfven current
+ chi = sigma*(R/b**3)**(1/2) # Shielding paramter
+ xi = 0.5 + 0.34*chi
+ N = (ring.L0 * Ia * ring.ac * ring.gamma * ring.sigma_delta**2 * xi *
+ sigma**(1/3) / ( c * e * R**(1/3) ))
+ I = N*e/ring.T0
+
+ return I
\ No newline at end of file
diff --git a/CollectiveEffect/resistive_wall.py b/CollectiveEffect/resistive_wall.py
new file mode 100644
index 0000000..764a59e
--- /dev/null
+++ b/CollectiveEffect/resistive_wall.py
@@ -0,0 +1,65 @@
+# -*- coding: utf-8 -*-
+"""
+Define resistive wall elements based on the Wakefield class.
+
+@author: Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import numpy as np
+from scipy.constants import mu_0, epsilon_0, c
+from CollectiveEffect.wakefield import Wakefield, Impedance
+
+def skin_depth(omega, rho, mu_r = 1, epsilon_r = 1):
+ """ General formula for the skin depth
+
+ Parameters
+ ----------
+ omega: angular frequency in [Hz]
+ rho: resistivity in [ohm.m]
+ mu_r : relative magnetic permeability, optional
+ epsilon_r : relative electric permittivity, optional
+
+ Returns
+ -------
+ delta : skin depth in [m]
+
+ """
+
+ delta = np.sqrt(2*rho/(np.abs(omega)*mu_r*mu_0))*np.sqrt(np.sqrt(1 +
+ (rho*np.abs(omega)*epsilon_r*epsilon_0)**2 )
+ + rho*np.abs(omega)*epsilon_r*epsilon_0)
+ return delta
+
+
+class CircularResistiveWall(Wakefield):
+ """
+ Resistive wall Wakefield element for a circular beam pipe, approximated
+ formulas from Eq. (2.77) of Chao book [1].
+
+ Parameters
+ ----------
+ frequency: frequency points where the impedance will be evaluated in [Hz]
+ length: beam pipe length in [m]
+ rho: resistivity in [ohm.m]
+ radius: beam pipe radius in [m]
+ mu_r : relative magnetic permeability, optional
+ epsilon_r : relative electric permittivity, optional
+ """
+
+ def __init__(self, frequency, length, rho, radius, mu_r = 1, epsilon_r = 1):
+ super().__init__()
+
+ omega = 2*np.pi*frequency
+ Z1 = length*(1 + np.sign(omega)*1j)*rho/(
+ 2*np.pi*radius*skin_depth(omega,rho))
+ Z2 = c/omega*length*(1 + np.sign(omega)*1j)*rho/(
+ np.pi*radius**3*skin_depth(omega,rho))
+
+ Zlong = Impedance(variable = frequency, function = Z1, impedance_type='long')
+ Zxdip = Impedance(variable = frequency, function = Z2, impedance_type='xdip')
+ Zydip = Impedance(variable = frequency, function = Z2, impedance_type='ydip')
+
+ super().append_to_model(Zlong)
+ super().append_to_model(Zxdip)
+ super().append_to_model(Zydip)
\ No newline at end of file
diff --git a/CollectiveEffect/tools.py b/CollectiveEffect/tools.py
new file mode 100644
index 0000000..3f3f633
--- /dev/null
+++ b/CollectiveEffect/tools.py
@@ -0,0 +1,185 @@
+# -*- coding: utf-8 -*-
+"""
+Tools and utilities functions for the Wakefield and impedances classes.
+
+@author: Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import pandas as pd
+import numpy as np
+from scipy.integrate import trapz
+from CollectiveEffect.wakefield import Wakefield, Impedance, WakeFunction
+#from ..Tracking.synchrotron import Synchrotron
+
+def read_CST(file, impedance_type='long'):
+ """Read CST file format into an Impedance object."""
+ df = pd.read_csv(file, comment="#", header = None, sep = "\t",
+ names = ["Frequency","Real","Imaginary"])
+ df["Frequency"] = df["Frequency"]*1e9
+ 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."""
+ 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 loss_factor(wake, sigma):
+ """
+ Compute the loss factor or the kick factor of a Wakefield, Impedance or
+ WakeFunction for a Gaussian bunch. Formulas from Eq. (2) p239 and
+ Eq.(7) p240 of [1].
+
+ Parameters
+ ----------
+ wake: Wakefield, Impedance or WakeFunction object.
+ sigma: RMS bunch length in [s]
+
+ Returns
+ -------
+ kloss: loss factor in [V/C] or kick factor in [V/C/m] depanding on the
+ impedance type.
+
+ [1] Handbook of accelerator physics and engineering, 3rd printing.
+ """
+
+ if isinstance(wake, Impedance):
+
+ positive_index = wake.data.index > 0
+ omega = 2*np.pi*wake.data.index[positive_index]
+ gaussian_spectrum = np.exp(-1*omega**2*sigma**2)
+
+ if(wake.impedance_type == "long"):
+ kloss = trapz(x = omega,
+ y = 1/np.pi*wake.data["real"][positive_index]*gaussian_spectrum)
+ if(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"):
+ kloss = trapz(x = omega,
+ y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum)
+ if isinstance(wake, WakeFunction):
+ # To be implemented.
+ pass
+
+ if isinstance(wake, Wakefield):
+ # To be implemented. Create field in wakefield with values ?
+ pass
+
+ return kloss
+
+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
+
+ [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 effective_impedance(ring, imp, m, mu, sigma, 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]
+ 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.
+
+ [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:
+ raise ValueError("A double sided impedance spectrum, with positive and"
+ " negative frequencies, is needed to compute the "
+ "effective impedance.")
+
+ if mode == "Hermite":
+ def h(f):
+ return spectral_density(frequency=f, sigma=sigma, m=m,
+ mode="Hermite")
+ else:
+ raise NotImplementedError("Not implemanted yet.")
+
+ M = 416
+
+ pmax = fmax/(ring.f0 * M) - 50
+ pmin = fmin/(ring.f0 * M) + 50
+
+ p = np.arange(pmin,pmax+1)
+
+ if imp.impedance_type == "long":
+ fp = ring.f0*(p*M + mu + m*ring.tune[2])
+ 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]
+ xi = ring.chro[0]
+ elif imp.impedance_type == "ydip":
+ tuneXY = ring.tune[1]
+ xi = ring.chro[1]
+ fp = ring.f0*(p*M + mu + tuneXY + m*ring.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
diff --git a/CollectiveEffect/wakefield.py b/CollectiveEffect/wakefield.py
new file mode 100644
index 0000000..dea656d
--- /dev/null
+++ b/CollectiveEffect/wakefield.py
@@ -0,0 +1,425 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines general classes to describes wakefields, impedances and
+wake functions. Based on impedance library by David Amorim.
+
+@author: David Amorin, Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import warnings
+import re
+import pandas as pd
+import numpy as np
+from scipy.interpolate import interp1d
+
+class ComplexData:
+ """
+ Define a general data structure for a complex function based on a pandas
+ DataFrame.
+
+ Parameters
+ ----------
+ variable : list, numpy array
+ contains the function variable values
+
+ function : list or numpy array of complex numbers
+ contains the values taken by the complex function
+ """
+
+ def __init__(self, variable=np.array([-1e15, 1e15]),
+ function=np.array([0, 0])):
+ self.data = pd.DataFrame({'real': np.real(function),
+ 'imag': np.imag(function)},
+ index=variable)
+ self.data.index.name = 'variable'
+
+ def add(self, structure_to_add, method='zero', interp_kind='cubic'):
+ """
+ Method to add two structures. If the data don't have the same length,
+ different cases are handled.
+
+ Parameters
+ ----------
+ structure_to_add : ComplexData object, int, float or complex.
+ structure to be added.
+ method : str, optional
+ manage how the addition is done, possibilties are:
+ -common: the common range of the index (variable) from the two
+ structures is used. The input data are cross-interpolated
+ so that the result data contains the points at all initial
+ points in the common domain.
+ -extrapolate: extrapolate the value of both ComplexData.
+ -zero: outside of the common range, the missing values are zero. In
+ the common range, behave as "common" method.
+ interp_kind : str, optional
+ interpolation method which is passed to pandas and to
+ scipy.interpolate.interp1d.
+
+ Returns
+ -------
+ ComplexData
+ Contains the sum of the two inputs.
+ """
+
+ # Create first two new DataFrames merging the variable
+ # from the two input data
+
+ if isinstance(structure_to_add, (int, float, complex)):
+ structure_to_add = ComplexData(
+ variable=self.data.index,
+ function=structure_to_add * np.ones(len(self.data.index)))
+
+ initial_data = pd.concat(
+ [self.data,
+ structure_to_add.data.index.to_frame().set_index('variable')],
+ sort=True)
+ initial_data = initial_data[~initial_data.index.duplicated(keep='first')]
+
+ data_to_add = pd.concat(
+ [structure_to_add.data,
+ self.data.index.to_frame().set_index('variable')],
+ sort=True)
+ data_to_add = data_to_add[~data_to_add.index.duplicated(keep='first')]
+
+ if method == 'common':
+ max_variable = min(structure_to_add.data.index.max(),
+ self.data.index.max())
+
+ min_variable = max(structure_to_add.data.index.min(),
+ self.data.index.min())
+
+ initial_data = initial_data.interpolate(method=interp_kind)
+ mask = ((initial_data.index <= max_variable)
+ & (initial_data.index >= min_variable))
+ initial_data = initial_data[mask]
+
+ data_to_add = data_to_add.interpolate(method=interp_kind)
+ mask = ((data_to_add.index <= max_variable)
+ & (data_to_add.index >= min_variable))
+ data_to_add = data_to_add[mask]
+
+ result_structure = ComplexData()
+ result_structure.data = initial_data + data_to_add
+ return result_structure
+
+ if method == 'extrapolate':
+ print('Not there yet')
+ return self
+
+ if method == 'zero':
+ max_variable = min(structure_to_add.data.index.max(),
+ self.data.index.max())
+
+ min_variable = max(structure_to_add.data.index.min(),
+ self.data.index.min())
+
+ mask = ((initial_data.index <= max_variable)
+ & (initial_data.index >= min_variable))
+ initial_data[mask] = initial_data[mask].interpolate(method=interp_kind)
+
+ mask = ((data_to_add.index <= max_variable)
+ & (data_to_add.index >= min_variable))
+ data_to_add[mask] = data_to_add[mask].interpolate(method=interp_kind)
+
+ initial_data.replace(to_replace=np.nan, value=0, inplace=True)
+ data_to_add.replace(to_replace=np.nan, value=0, inplace=True)
+
+ result_structure = ComplexData()
+ result_structure.data = initial_data + data_to_add
+ return result_structure
+
+ def __radd__(self, structure_to_add):
+ return self.add(structure_to_add, method='zero')
+
+ def __add__(self, structure_to_add):
+ return self.add(structure_to_add, method='zero')
+
+ def multiply(self, factor):
+ """
+ Multiply a data strucure with a float or an int.
+ If the multiplication is done with something else, throw a warning.
+ """
+ if isinstance(factor, (int, float)):
+ result_structure = ComplexData()
+ result_structure.data = self.data * factor
+ return result_structure
+ else:
+ warnings.warn(('The multiplication factor is not a float '
+ 'or an int.'), UserWarning)
+ return self
+
+ def __mul__(self, factor):
+ return self.multiply(factor)
+
+ def __rmul__(self, factor):
+ return self.multiply(factor)
+
+ def __call__(self, values, interp_kind="cubic"):
+ """
+ Interpolation of the data by calling the class to have a function-like
+ behaviour.
+
+ Parameters
+ ----------
+ values : list or numpy array of complex, int or float
+ values to be interpolated.
+ interp_kind : str, optional
+ interpolation method which is passed to scipy.interpolate.interp1d.
+
+ Returns
+ -------
+ numpy array
+ Contains the interpolated data.
+ """
+ real_func = interp1d(x = self.data.index,
+ y = self.data["real"], kind=interp_kind)
+ imag_func = interp1d(x = self.data.index,
+ y = self.data["imag"], kind=interp_kind)
+ return real_func(values) + 1j*imag_func(values)
+
+class WakeFunction(ComplexData):
+ """
+ Define a WakeFunction object based on a ComplexData object.
+ """
+
+ def __init__(self,
+ variable=np.array([-1e15, 1e15]),
+ function=np.array([0, 0]), wake_type='long'):
+ super().__init__(variable, function)
+ pass
+
+class Impedance(ComplexData):
+ """
+ Define an Impedance object based on a ComplexData object.
+ """
+
+ def __init__(self,
+ variable=np.array([-1e15, 1e15]),
+ function=np.array([0, 0]), impedance_type='long'):
+ super().__init__(variable, function)
+ self._impedance_type = impedance_type
+ self.data.index.name = 'frequency [Hz]'
+ self.initialize_impedance_coefficient()
+
+
+ def initialize_impedance_coefficient(self):
+ """
+ Define the impedance coefficients and the plane of the impedance that
+ are presents in attributes of the class.
+ """
+ table = self.impedance_name_and_coefficients_table()
+
+ try:
+ component_coefficients = table[self.impedance_type].to_dict()
+ except KeyError:
+ print('Impedance type {} does not exist'.format(self.impedance_type))
+
+ self.a = component_coefficients["a"]
+ self.b = component_coefficients["b"]
+ self.c = component_coefficients["c"]
+ self.d = component_coefficients["d"]
+ self.plane = component_coefficients["plane"]
+
+ def impedance_name_and_coefficients_table(self):
+ """
+ Return a table associating the human readbale names of an impedance
+ component and its associated coefficients and plane.
+ """
+
+ component_dict = {
+ 'long': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'z'},
+ 'xcst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
+ 'ycst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
+ 'xdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
+ 'ydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'y'},
+ 'xydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'x'},
+ 'yxdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
+ 'xquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'x'},
+ 'yquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'y'},
+ 'xyquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'x'},
+ 'yxquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'y'},
+ }
+
+ return pd.DataFrame(component_dict)
+
+ def add(self, structure_to_add, beta_x=1, beta_y=1, method='zero'):
+ """
+ Method to add two Impedance objects. The two structures are
+ compared so that the addition is self-consistent.
+ Beta functions can be precised as well.
+ """
+ if isinstance(structure_to_add, (int, float, complex)):
+ result = super().add(structure_to_add, method=method)
+ elif isinstance(structure_to_add, Impedance):
+ if (self.impedance_type == structure_to_add.impedance_type):
+ weight = (beta_x ** self.power_x) * (beta_y ** self.power_y)
+ result = super().add(structure_to_add * weight, method=method)
+ else:
+ warnings.warn(('The two Impedance objects do not have the '
+ 'same coordinates or plane or type. '
+ 'Returning initial Impedance object.'),
+ UserWarning)
+ result = self
+
+ impedance_to_return = Impedance(
+ result.data.index,
+ result.data.real.values + 1j*result.data.imag.values,
+ self.impedance_type)
+ return impedance_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 Impedance object with a float or an int.
+ If the multiplication is done with something else, throw a warning.
+ """
+ result = super().multiply(factor)
+ impedance_to_return = Impedance(
+ result.data.index,
+ result.data.real.values + 1j*result.data.imag.values,
+ self.impedance_type)
+ return impedance_to_return
+
+ def __mul__(self, factor):
+ return self.multiply(factor)
+
+ def __rmul__(self, factor):
+ return self.multiply(factor)
+
+ @property
+ def impedance_type(self):
+ return self._impedance_type
+
+ @impedance_type.setter
+ def impedance_type(self, value):
+ self._impedance_type = value
+ self.initialize_impedance_coefficient()
+
+ @property
+ def power_x(self):
+ power_x = self.a/2 + self.c/2.
+ if self.plane == 'x':
+ power_x += 1./2.
+ return power_x
+
+ @property
+ def power_y(self):
+ power_y = self.b/2. + self.d/2.
+ if self.plane == 'y':
+ power_y += 1./2.
+ return power_y
+
+class Wakefield:
+ """
+ Define a general Wakefield machine element.
+ It contains Imepdance or WakeFunction objects which defines the wakefield
+ and other informations: beta functions.
+ """
+
+ def __init__(self, betax=1, betay=1, structure_list=[]):
+ self.betax = betax
+ self.betay = betay
+ self.list_to_attr(structure_list)
+
+ def append_to_model(self, structure_to_add):
+ list_of_attributes = dir(self)
+ if isinstance(structure_to_add, Impedance):
+ attribute_name = "Z" + structure_to_add.impedance_type
+ if attribute_name in list_of_attributes:
+ raise ValueError("There is already a component of the type "
+ "{} in this element.".format(attribute_name))
+ else:
+ self.__setattr__(attribute_name, structure_to_add)
+ elif isinstance(structure_to_add, WakeFunction):
+ attribute_name = "W" + structure_to_add.impedance_type
+ if attribute_name in list_of_attributes:
+ raise ValueError("There is already a component of the type "
+ "{} in this element.".format(attribute_name))
+ else:
+ self.__setattr__(attribute_name, structure_to_add)
+ else:
+ raise ValueError("{} is not an Impedance nor a WakeFunction.".format(structure_to_add))
+
+ def list_to_attr(self, structure_list):
+ for component in structure_list:
+ self.append_to_model(component)
+
+ @property
+ def list_impedance_components(self):
+ """
+ Return the list of impedance components for the element,
+ based on the attributes names.
+ """
+ return np.array([comp for comp in dir(self) if re.match(r'[Z]', comp)])
+
+ def add(self, element_to_add):
+ """
+ This function adds two Wakefield objects.
+ The different impedance components are weighted and added.
+ The result is a Wakefield object with all the components but the
+ beta function equal to 1.
+ """
+
+ if not isinstance(element_to_add, Wakefield):
+ raise TypeError("You can only add Wakefield objects to other"
+ "Wakefields objects.")
+
+ list_of_summed_impedances = []
+ list_of_components_added = []
+
+ betax1 = self.betax
+ betay1 = self.betay
+ betax2 = element_to_add.betax
+ betay2 = element_to_add.betay
+
+ # We first go through element1 impedance components and add them
+ # to the element2 components. If the component is missing in element2,
+ # we then simply weight the element1 component and add it to the model.
+ for component_name in self.list_impedance_components:
+
+ element1_impedance = self.__getattribute__(component_name)
+
+ element1_weight = ((betax1 ** element1_impedance.power_x)
+ * (betay1 ** element1_impedance.power_y))
+
+ try:
+ element2_impedance = element_to_add.__getattribute__(component_name)
+ list_of_components_added.append(component_name)
+ element12_impedance_sum = element1_weight * element1_impedance
+ element12_impedance_sum = element12_impedance_sum.add(element2_impedance, betax2, betay2)
+ except:
+ element12_impedance_sum = element1_weight * element1_impedance
+
+ list_of_summed_impedances.append(element12_impedance_sum)
+
+ # We now go through the components which are unique to element2 and
+ # add them to the impedance model, with the beta function weighting
+ missing_components_list = list(set(element_to_add.list_impedance_components) -
+ set(list_of_components_added))
+
+ for component_name in missing_components_list:
+
+ element2_impedance = element_to_add.__getattribute__(component_name)
+ element2_weight = ((betax2 ** element2_impedance.power_x)
+ * (betay2 ** element2_impedance.power_y))
+ element12_impedance_sum = element2_weight * element2_impedance
+
+ list_of_summed_impedances.append(element12_impedance_sum)
+
+ # Gather everything in an Wakefield object
+ sum_wakefield = Wakefield(betax=1., betay=1.,
+ structure_list=list_of_summed_impedances)
+
+ return sum_wakefield
+
+ def __radd__(self, structure_to_add):
+ return self.add(structure_to_add)
+
+ def __add__(self, structure_to_add):
+ return self.add(structure_to_add)
diff --git a/Machines/__init__.py b/Machines/__init__.py
new file mode 100644
index 0000000..c67145e
--- /dev/null
+++ b/Machines/__init__.py
@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+
diff --git a/Machines/soleil.py b/Machines/soleil.py
new file mode 100644
index 0000000..b0ef691
--- /dev/null
+++ b/Machines/soleil.py
@@ -0,0 +1,55 @@
+# -*- coding: utf-8 -*-
+"""
+SOLEIL synchrotron parameters script.
+
+@author: Alexis Gamelin
+@date: 14/01/2020
+"""
+
+import numpy as np
+from Tracking.synchrotron import Synchrotron, Electron, Optics
+from Tracking.beam import Beam
+
+def soleil(mode = 'Uniform'):
+ """
+
+ """
+
+ h = 416
+ L = 3.540969742590899e+02
+ E0 = 2.75e9
+ particle = Electron()
+ ac = 4.16e-4
+ U0 = 1171e3
+ tau = np.array([0, 0, 3.27e-3])
+ tune = np.array([18.15687, 10.22824, 0.00502])
+ emit = np.array([4.5e-9, 4.5e-9*0.01])
+ sigma_0 = 8e-12
+ sigma_delta = 9e-4
+
+ # mean values
+ beta = np.array([3, 1.3])
+ alpha = np.array([0, 0])
+ mean_val = Optics(beta, alpha)
+
+ ring = Synchrotron(h, L, E0, particle, ac=ac, U0=U0, tau=tau,
+ mean_optics=mean_val, emit=emit, tune=tune,
+ sigma_delta=sigma_delta, sigma_0=sigma_0)
+
+ if mode == 'Uniform':
+ beam = Beam(ring, )
+ elif mode == 'Hybrid':
+ pass
+ elif mode == 'Low alpha/25':
+ pass
+ elif mode == 'Low alpha/25':
+ pass
+ elif mode == '8 bunches':
+ pass
+ elif mode == 'Single':
+ pass
+ else:
+ raise ValueError("{} is not a correct operation mode.".format(mode))
+
+
+ return ring
\ No newline at end of file
diff --git a/Tracking/__init__.py b/Tracking/__init__.py
new file mode 100644
index 0000000..c67145e
--- /dev/null
+++ b/Tracking/__init__.py
@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+
diff --git a/Tracking/beam.py b/Tracking/beam.py
new file mode 100644
index 0000000..e3c4f1b
--- /dev/null
+++ b/Tracking/beam.py
@@ -0,0 +1,341 @@
+# -*- coding: utf-8 -*-
+"""
+Beam and bunch elements
+
+@author: Alexis Gamelin
+@date: 17/01/2020
+"""
+
+import numpy as np
+import numpy.ma as ma
+import warnings
+
+class Bunch:
+ """
+ Define a bunch
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ mp_number : float, optional
+ Macro-particle number
+ current : float, optional
+ Bunch current in [A]
+
+ Attributes
+ ----------
+ mp_number : int
+ Macro-particle number
+ charge : float
+ Bunch charge in [C]
+ charge_per_mp : float
+ Charge per macro-particle in [C]
+ particle_number : int
+ Number of particles in the bunch
+ current : float
+ Bunch current in [A]
+
+ Methods
+ ------
+ stats()
+ Compute the mean and the standard deviation of alive particles for each
+ coordinates.
+ init_gaussian(cov=None, mean=None, **kwargs)
+ Initialize bunch particles with 6D gaussian phase space.
+ """
+
+ def __init__(self, ring, mp_number=1e3, current=1e-3):
+ dpart = np.dtype([("x",np.float64),("xp",np.float64),
+ ("y",np.float64),("yp",np.float64),
+ ("tau",np.float64),("delta",np.float64)])
+ self.dtype = dpart
+ particles = np.zeros((int(mp_number),6), dtype=np.float64)
+ self.particles = ma.masked_array(particles)
+
+ self.ring = ring
+ self._mp_number = int(mp_number)
+ self.current = current
+
+ def __len__(self):
+ """Return the number of alive particles"""
+ #mask = self.particles["alive"] == True
+ return len(self.particles)
+
+ def __getitem__(self, i):
+ """Return the alive particle number i"""
+ return self.particles[i]
+
+ def __setitem__(self, i, value):
+ """Set value to the alive particle number i"""
+ self.particles[i] = value
+
+ def __iter__(self):
+ """Iterate over alive particles"""
+ return self.particles.__iter__()
+
+ @property
+ def mp_number(self):
+ """Macro-particle number"""
+ return self._mp_number
+
+ @mp_number.setter
+ def mp_number(self, value):
+ self._mp_number = int(value)
+ self.__init__(self.ring, value, self.charge)
+
+ @property
+ def charge_per_mp(self):
+ """Charge per macro-particle [C]"""
+ return self._charge_per_mp
+
+ @charge_per_mp.setter
+ def charge_per_mp(self, value):
+ self._charge_per_mp = value
+
+ @property
+ def charge(self):
+ """Bunch charge in [C]"""
+ return self.__len__()*self.charge_per_mp
+
+ @charge.setter
+ def charge(self, value):
+ self.charge_per_mp = value / self.__len__()
+
+ @property
+ def particle_number(self):
+ """Particle number"""
+ return int(self.charge / np.abs(self.ring.particle.charge))
+
+ @particle_number.setter
+ def particle_number(self, value):
+ self.charge_per_mp = value * self.ring.particle.charge / self.__len__()
+
+ @property
+ def current(self):
+ """Bunch current [A]"""
+ return self.charge / self.ring.T0
+
+ @current.setter
+ def current(self, value):
+ self.charge_per_mp = value * self.ring.T0 / self.__len__()
+
+ def mean(self):
+ """
+ Compute the mean position of alive particles for each
+ coordinates.
+
+ Returns
+ -------
+ mean : numpy array
+ mean position of alive particles
+ """
+ mean = [[self[name].mean()] for name in self.dtype.names]
+ return np.array(mean)
+
+ def std(self):
+ """
+ Compute the standard deviation of the position of alive
+ particles for each coordinates.
+
+ Returns
+ -------
+ std : numpy array
+ standard deviation of the position of alive particles
+ """
+ std = [[self[name].std()] for name in self.dtype.names]
+ return np.array(std)
+
+ def emit(self):
+ """
+ Compute the bunch emittance for each plane [1].
+ Correct definition of long emit ?
+
+ Returns
+ -------
+ emit : numpy array
+ bunch emittance
+
+ References
+ ----------
+ [1] Wiedemann, H. (2015). Particle accelerator physics. 4th
+ edition. Springer, Eq.(8.39) of p224.
+ """
+ emitX = (np.mean(self['x']**2)*np.mean(self['xp']**2) -
+ np.mean(self['x']*self['xp'])**2)**(0.5)
+ emitY = (np.mean(self['y']**2)*np.mean(self['yp']**2) -
+ np.mean(self['y']*self['yp'])**2)**(0.5)
+ emitS = (np.mean(self['tau']**2)*np.mean(self['delta']**2) -
+ np.mean(self['tau']*self['delta'])**2)**(0.5)
+ return np.array([emitX, emitY, emitS])
+
+ def init_gaussian(self, cov=None, mean=None, **kwargs):
+ """
+ Initialize bunch particles with 6D gaussian phase space.
+ Covariance matrix is taken from [1].
+
+ Parameters
+ ----------
+ cov : (6,6) array, optional
+ Covariance matrix of the bunch distribution
+ mean : (6,) array, optional
+ Mean of the bunch distribution
+
+ References
+ ----------
+ [1] Wiedemann, H. (2015). Particle accelerator physics. 4th
+ edition. Springer, Eq.(8.38) of p223.
+
+ """
+ if mean is None:
+ mean = np.zeros((6,))
+
+ if cov is None:
+ sigma_0 = kwargs.get("sigma_0", self.ring.sigma_0)
+ sigma_delta = kwargs.get("sigma_delta", self.ring.sigma_delta)
+ optics = kwargs.get("optics", self.ring.mean_optics)
+
+ cov = np.zeros((6,6))
+ cov[0,0] = self.ring.emit[0]*optics.beta[0]
+ cov[1,1] = self.ring.emit[0]*optics.gamma[0]
+ cov[0,1] = -1*self.ring.emit[0]*optics.alpha[0]
+ cov[1,0] = -1*self.ring.emit[0]*optics.alpha[0]
+ cov[2,2] = self.ring.emit[1]*optics.beta[1]
+ cov[3,3] = self.ring.emit[1]*optics.gamma[1]
+ cov[2,3] = -1*self.ring.emit[1]*optics.alpha[1]
+ cov[3,2] = -1*self.ring.emit[1]*optics.alpha[1]
+ cov[4,4] = sigma_0**2
+ cov[5,5] = sigma_delta**2
+
+ values = np.random.multivariate_normal(mean, cov, size=self.mp_number)
+ self.particles["x"] = values[:,0]
+ self.particles["xp"] = values[:,1]
+ self.particles["y"] = values[:,2]
+ self.particles["yp"] = values[:,3]
+ self.particles["tau"] = values[:,4]
+ self.particles["delta"] = values[:,5]
+
+class Beam:
+ """
+ Define a Beam object composed of several Bunch objects.
+
+ Parameters
+ ----------
+ ring : Synchrotron object
+ bunch_list : list of Bunch object, optional
+
+ Attributes
+ ----------
+ current : float
+ Total beam current in [A]
+ charge : float
+ Total bunch charge in [C]
+ particle_number : int
+ Total number of particle in the beam
+ filling_pattern : list of bool
+ Filling pattern of the beam
+
+ Methods
+ ------
+ init_beam()
+ """
+
+ def __init__(self, ring, bunch_list=None):
+ self.ring = ring
+
+ if bunch_list is None:
+ self.init_beam(np.zeros((self.ring.h,1),dtype=bool))
+ else:
+ if (len(bunch_list) != self.ring.h):
+ raise ValueError(("The length of the bunch list is {} ".format(len(bunch_list)) +
+ "but should be {}".format(self.ring.h)))
+ self.bunch_list = bunch_list
+
+ def __len__(self):
+ """Return the number of bunches"""
+ length = 0
+ for bunch in self:
+ if bunch is not None:
+ length += 1
+ return length
+
+ def __getitem__(self, i):
+ """Return the bunch number i"""
+ return self.bunch_list.__getitem__(i)
+
+ def __setitem__(self, i, value):
+ """Set value to the bunch number i"""
+ self.bunch_list.__setitem__(i, value)
+
+ def __iter__(self):
+ """Iterate over the bunches"""
+ return self.bunch_list.__iter__()
+
+ def init_beam(self, filling_pattern, current_per_bunch=1e-3,
+ mp_per_bunch=1e3):
+ """
+ Initialize beam with a given filling pattern but with uniform current
+ per bunch and marco-particle number per bunch.
+
+ Parameters
+ ----------
+ filling_pattern : numpy array or list of length ring.h
+ Filling pattern of the beam,
+ charge_per_bunch : float, optional
+ Charge per bunch in [C]
+ mp_per_bunch : float, optional
+ Macro-particle number per bunch
+ """
+
+ if (len(filling_pattern) != self.ring.h):
+ raise ValueError(("The length of filling pattern is {} ".format(len(filling_pattern)) +
+ "but should be {}".format(self.ring.h)))
+
+ bunch_list = []
+ for value in filling_pattern:
+ if value == True:
+ bunch_list.append(Bunch(self.ring, mp_per_bunch, current_per_bunch))
+ elif value == False:
+ bunch_list.append(None)
+ else:
+ raise ValueError("{} should be True or False".format(value))
+ self.bunch_list = bunch_list
+ # init bunches => gauss
+
+ @property
+ def current(self):
+ """Total beam current in [A]"""
+ current = 0
+ for bunch in self:
+ if bunch is not None:
+ current += bunch.current
+ return current
+
+ @property
+ def charge(self):
+ """Total beam charge in [C]"""
+ charge = 0
+ for bunch in self:
+ if bunch is not None:
+ charge += bunch.charge
+ return charge
+
+ @property
+ def particle_number(self):
+ """Total number of particle in the beam"""
+ particle_number = 0
+ for bunch in self:
+ if bunch is not None:
+ particle_number += bunch.particle_number
+ return particle_number
+
+ @property
+ def filling_pattern(self):
+ """Filling pattern of the beam"""
+ filling_pattern = []
+ for bunch in self:
+ if isinstance(bunch, Bunch):
+ filling_pattern.append(True)
+ else:
+ filling_pattern.append(False)
+ return filling_pattern
+
\ No newline at end of file
diff --git a/Tracking/synchrotron.py b/Tracking/synchrotron.py
new file mode 100644
index 0000000..a7eb44b
--- /dev/null
+++ b/Tracking/synchrotron.py
@@ -0,0 +1,181 @@
+# -*- coding: utf-8 -*-
+"""
+General elements
+
+@author: Alexis Gamelin
+@date: 15/01/2020
+"""
+
+import numpy as np
+from scipy.constants import c, m_e, m_p, e
+
+class Particle:
+ """ Define a particle
+
+ Attributes
+ ----------
+ mass : float
+ total particle mass in [kg]
+ charge : float
+ electrical charge in [C]
+ E_rest : float
+ particle rest energy in [eV]
+ """
+ def __init__(self, mass, charge):
+ self.mass = mass
+ self.charge = charge
+
+ @property
+ def E_rest(self):
+ return self.mass * c ** 2 / e
+
+class Electron(Particle):
+ """ Define an electron"""
+ def __init__(self):
+ super().__init__(m_e, -1*e)
+
+class Proton(Particle):
+ """ Define a proton"""
+ def __init__(self):
+ super().__init__(m_p, e)
+
+class Optics:
+ """Handle optic functions"""
+ def __init__(self, beta, alpha):
+ self.beta = beta
+ self.alpha = alpha
+
+ @property
+ def gamma(self):
+ return (1 + self.alpha**2)/self.beta
+
+class Synchrotron:
+ """ Synchrotron class to store main properties. """
+ def __init__(self, h, L, E0, particle, **kwargs):
+ self.particle = particle
+ self._h = h
+ self.L = L
+ self.E0 = E0 # Nominal (total) energy of the ring [eV]
+
+ self.ac = kwargs.get('ac') # Momentum compaction factor
+ self.U0 = kwargs.get('U0') # Energy loss per turn [eV]
+ self.tau = kwargs.get('tau') # X/Y/S damping times [s]
+ self.sigma_delta = kwargs.get('sigma_delta') # Equilibrium energy spread
+ self.sigma_0 = kwargs.get('sigma_0') # Natural bunch length [s]
+ self.tune = kwargs.get('tune') # X/Y/S tunes
+ self.emit = kwargs.get('emit') # X/Y emittances in [m.rad]
+ self.chro = kwargs.get('chro') # X/Y (non-normalized) chromaticities
+ self.mean_optics = kwargs.get('mean_optics') # Optics object with mean values
+
+ @property
+ def h(self):
+ """Harmonic number"""
+ return self._h
+
+ @h.setter
+ def h(self,value):
+ self._h = value
+ self.L = self.L # call setter
+
+ @property
+ def L(self):
+ """Ring circumference [m]"""
+ return self._L
+
+ @L.setter
+ def L(self,value):
+ self._L = value
+ self._T0 = self.L/c
+ self._f0 = 1/self.T0
+ self._omega0 = 2*np.pi*self.f0
+ self._f1 = self.h*self.f0
+ self._omega1 = 2*np.pi*self.f1
+ self._k1 = self.omega1/c
+
+ @property
+ def T0(self):
+ """Revolution time [s]"""
+ return self._T0
+
+ @T0.setter
+ def T0(self, value):
+ self.L = c*value
+
+ @property
+ def f0(self):
+ """Revolution frequency [Hz]"""
+ return self._f0
+
+ @f0.setter
+ def f0(self,value):
+ self.L = c/value
+
+ @property
+ def omega0(self):
+ """Angular revolution frequency [Hz rad]"""
+ return self._omega0
+
+ @omega0.setter
+ def omega0(self,value):
+ self.L = 2*np.pi*c/value
+
+ @property
+ def f1(self):
+ """Fundamental RF frequency [Hz]"""
+ return self._f1
+
+ @f1.setter
+ def f1(self,value):
+ self.L = self.h*c/value
+
+ @property
+ def omega1(self):
+ """Fundamental RF angular frequency[Hz rad]"""
+ return self._omega1
+
+ @omega1.setter
+ def omega1(self,value):
+ self.L = 2*np.pi*self.h*c/value
+
+ @property
+ def k1(self):
+ """Fundamental RF wave number [m**-1]"""
+ return self._k1
+
+ @k1.setter
+ def k1(self,value):
+ self.L = 2*np.pi*self.h/value
+
+ @property
+ def gamma(self):
+ """Relativistic gamma"""
+ return self._gamma
+
+ @gamma.setter
+ def gamma(self, value):
+ self._gamma = value
+ self._beta = np.sqrt(1 - self.gamma**-2)
+ self._E0 = self.gamma*self.particle.mass*c**2/e
+
+ @property
+ def beta(self):
+ """Relativistic beta"""
+ return self._beta
+
+ @beta.setter
+ def beta(self, value):
+ self.gamma = 1/np.sqrt(1-value**2)
+
+ @property
+ def E0(self):
+ """Nominal (total) energy of the ring [eV]"""
+ return self._E0
+
+ @E0.setter
+ def E0(self, value):
+ self.gamma = value/(self.particle.mass*c**2/e)
+
+ @property
+ def eta(self):
+ """Momentum compaction"""
+ return self.ac - 1/(self.gamma**2)
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..c67145e
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+
--
GitLab