Skip to content
Snippets Groups Projects
Commit aa1314e3 authored by Gamelin Alexis's avatar Gamelin Alexis
Browse files

Restructure code and docstrings 

To be PEP8 complient:
Lowercase folders
Extended docstrings
Rename and divide some files
parent 90ae8393
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 362 additions and 161 deletions
Machines/soleil.py
+ 9
− 7
View file @ aa1314e3
......@@ -7,8 +7,9 @@ SOLEIL synchrotron parameters script.
"""
import numpy as np
from Tracking.synchrotron import Synchrotron, Electron, Optics
from Tracking.beam import Beam
from tracking.synchrotron import Synchrotron
from tracking.optics import Optics
from tracking.particles import Electron, Beam
def soleil(mode = 'Uniform'):
"""
......@@ -19,13 +20,13 @@ def soleil(mode = 'Uniform'):
L = 3.540969742590899e+02
E0 = 2.75e9
particle = Electron()
ac = 4.16e-4
U0 = 1171e3
tau = np.array([0, 0, 3.27e-3])
ac = 1.47e-4
U0 = 0.310e6
tau = np.array([6.56e-3, 6.56e-3, 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
sigma_delta = 8.6e-4
chro = [2,3]
# mean values
......@@ -56,4 +57,5 @@ def soleil(mode = 'Uniform'):
raise ValueError("{} is not a correct operation mode.".format(mode))
return ring
\ No newline at end of file
return ring
Tracking/one_turn_matrix.py → Tracking/element.py
+ 108
− 62
View file @ aa1314e3
# -*- coding: utf-8 -*-
"""
One turn map matrix elements
This module defines the most basic elements for tracking, including Element,
an abstract base class which is to be used as mother class to every elements
included in the tracking.
@author: gamelina
@date: 04/03/2020
@date: 11/03/2020
"""
import numpy as np
from abc import ABCMeta, abstractmethod
from functools import wraps
import numpy as np
from tracking.particles import Beam
class Element(metaclass=ABCMeta):
"""
......@@ -21,101 +24,131 @@ class Element(metaclass=ABCMeta):
"""
Track a beam object through this Element.
This method needs to be overloaded in each Element subclass.
Parameters
----------
beam : Beam object
"""
raise NotImplementedError
@staticmethod
def all_bunches(function):
"""
"""
@wraps(function)
def wrapper(*args, **kwargs):
self = args[0]
beam = args[1]
if (beam.mpi_switch == True):
function(self, beam[beam.mpi.bunch_num], *args[2:], **kwargs)
else:
for bunch in beam:
function(self, bunch, *args[2:], **kwargs)
return wrapper
@staticmethod
def not_empty(function):
def parallel(track):
"""
Defines the decorator @parallel which handle the embarrassingly
parallel case which happens when there is no bunch to bunch
interaction in the tracking routine.
Adding @Element.parallel allows to write the track method of the
Element subclass for a Bunch object instead of a Beam object.
Parameters
----------
track : function, method of an Element subclass
track method of an Element subclass which takes a Bunch object as
input
Returns
-------
track_wrapper: function, method of an Element subclass
track method of an Element subclass which takes a Beam object or a
Bunch object as input
"""
@wraps(function)
def wrapper(*args, **kwargs):
self = args[0]
beam = args[1]
if (beam.mpi_switch == True):
function(self, beam[beam.mpi.bunch_num], *args[2:], **kwargs)
@wraps(track)
def track_wrapper(*args, **kwargs):
if isinstance(args[1], Beam):
self = args[0]
beam = args[1]
if (beam.mpi_switch == True):
track(self, beam[beam.mpi.bunch_num], *args[2:], **kwargs)
else:
for bunch in beam.not_empty:
track(self, bunch, *args[2:], **kwargs)
else:
for bunch in beam.not_empty:
function(self, bunch, *args[2:], **kwargs)
return wrapper
self = args[0]
bunch = args[1]
track(self, bunch, *args[2:], **kwargs)
return track_wrapper
class Long_one_turn(Element):
class LongitudinalMap(Element):
"""
Longitudinal transform for a single turn.
Longitudinal map for a single turn in the synchrotron.
Parameters
----------
ring : Synchrotron object
"""
def __init__(self, ring):
self.ring = ring
@Element.not_empty
@Element.parallel
def track(self, bunch):
"""track"""
bunch["delta"] -= self.ring.U0/self.ring.E0
bunch["tau"] -= self.ring.ac*self.ring.T0*bunch["delta"]
"""
Tracking method for the element.
No bunch to bunch interaction, so written for Bunch objects and
@Element.parallel is used to handle Beam objects.
Parameters
----------
bunch : Bunch or Beam object
"""
bunch["delta"] -= self.ring.U0 / self.ring.E0
bunch["tau"] -= self.ring.ac * self.ring.T0 * bunch["delta"]
class SynchrotronRadiation(Element):
"""SyncRad"""
"""
Element to handle synchrotron radiation, radiation damping and quantum
excitation, for a single turn in the synchrotron.
Parameters
----------
ring : Synchrotron object
switch : bool array of shape (3,), optional
allow to choose on which plane the synchrotron radiation is active
"""
def __init__(self, ring, switch = np.ones((3,), dtype=bool)):
self.ring = ring
self.switch = switch
@Element.not_empty
@Element.parallel
def track(self, bunch):
"""track"""
"""
Tracking method for the element.
No bunch to bunch interaction, so written for Bunch objects and
@Element.parallel is used to handle Beam objects.
Parameters
----------
bunch : Bunch or Beam object
"""
if (self.switch[0] == True):
rand = np.random.normal(size=len(bunch))
bunch["delta"] = ((1 - 2*self.ring.T0/self.ring.tau[2])*bunch["delta"] +
2*self.ring.sigma_delta*(self.ring.T0/self.ring.tau[2])**0.5*rand)
if (self.switch[1] == True):
rand = np.random.normal(size=(len(bunch),2))
bunch["x"] += self.ring.sigma[0]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,0]
bunch["xp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["xp"]
bunch["xp"] += self.ring.sigma[1]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,1]
if (self.switch[2] == True):
rand = np.random.normal(size=(len(bunch),2))
bunch["y"] += self.ring.sigma[2]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,0]
bunch["yp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["yp"]
bunch["yp"] += self.ring.sigma[3]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,1]
bunch.energy_change = 0
class RF_cavity(Element):
""" Perfect RF cavity class for main and harmonic RF cavities."""
def __init__(self, ring, m, Vc, theta):
self.ring = ring
self.m = m # Harmonic number of the cavity
self.Vc = Vc # Amplitude of Cavity voltage [V]
self.theta = theta # phase of Cavity voltage
@Element.not_empty
def track(self,bunch):
"""Track """
energy_change = self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*bunch["tau"] + self.theta)
bunch["delta"] += energy_change
bunch.energy_change += energy_change
def value(self, val):
return self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*val + self.theta)
# Reset energy change to 0 for next turn
bunch.energy_change = 0
class Trans_one_turn(Element):
class TransverseMap(Element):
"""
Transverse transformation for a single turn.
Transverse map for a single turn in the synchrotron.
Parameters
----------
ring : Synchrotron object
"""
def __init__(self, ring):
......@@ -127,13 +160,26 @@ class Trans_one_turn(Element):
self.dispp = self.ring.mean_optics.dispp
self.phase_advance = self.ring.tune[0:2]*2*np.pi
@Element.not_empty
@Element.parallel
def track(self, bunch):
"""track"""
"""
Tracking method for the element.
No bunch to bunch interaction, so written for Bunch objects and
@Element.parallel is used to handle Beam objects.
Parameters
----------
bunch : Bunch or Beam object
"""
# Compute phase adcence which depends on energy via chromaticity
phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"])
phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"])
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
matrix = np.zeros((6,6,len(bunch)))
# Horizontal
matrix[0,0,:] = np.cos(phase_advance_x) + self.alpha[0]*np.sin(phase_advance_x)
matrix[0,1,:] = self.beta[0]*np.sin(phase_advance_x)
matrix[0,2,:] = self.disp[0]
......@@ -142,6 +188,7 @@ class Trans_one_turn(Element):
matrix[1,2,:] = self.dispp[0]
matrix[2,2,:] = 1
# Vertical
matrix[3,3,:] = np.cos(phase_advance_y) + self.alpha[1]*np.sin(phase_advance_y)
matrix[3,4,:] = self.beta[1]*np.sin(phase_advance_y)
matrix[3,5,:] = self.disp[1]
......@@ -158,5 +205,4 @@ class Trans_one_turn(Element):
bunch["x"] = x
bunch["xp"] = xp
bunch["y"] = y
bunch["yp"] = yp
bunch["yp"] = yp
\ No newline at end of file
Tracking/optics.py 0 → 100644
+ 18
− 0
View file @ aa1314e3
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 11 18:00:28 2020
@author: gamelina
"""
class Optics:
"""Handle optic functions"""
def __init__(self, beta, alpha, disp, dispp):
self.beta = beta
self.alpha = alpha
self.disp = disp
self.dispp = dispp
@property
def gamma(self):
return (1 + self.alpha**2)/self.beta
\ No newline at end of file
Tracking/parallel.py
+ 68
− 3
View file @ aa1314e3
......@@ -10,8 +10,44 @@ Module to handle parallel computation
import numpy as np
class Mpi:
""" Class which handle mpi
"""
Class which handle parallel computation via the mpi4py module [1].
Parameters
----------
filling_pattern : bool array of shape (h,)
Filling pattern of the beam, like Beam.filling_pattern
Attributes
----------
comm : MPI.Intracomm object
MPI intra-comminicator of the processor group, used to manage
communication between processors.
rank : int
Rank of the processor which run the program
size : int
Number of processor within the processor group (in fact in the
intra-comminicator group)
table : int array of shape (size, 2)
Table of correspondance between the rank of the processor and its
associated bunch number
bunch_num : int
Return the bunch number corresponding to the current processor
Methods
-------
write_table(filling_pattern)
Write a table with the rank and the corresponding bunch number for each
bunch of the filling pattern
rank_to_bunch(rank)
Return the bunch number corresponding to rank
bunch_to_rank(bunch_num)
Return the rank corresponding to the bunch number bunch_num
References
----------
[1] L. Dalcin, P. Kler, R. Paz, and A. Cosimo, Parallel Distributed
Computing using Python, Advances in Water Resources, 34(9):1124-1139, 2011.
"""
def __init__(self, filling_pattern):
......@@ -29,6 +65,11 @@ class Mpi:
"""
Write a table with the rank and the corresponding bunch number for each
bunch of the filling pattern
Parameters
----------
filling_pattern : bool array of shape (h,)
Filling pattern of the beam, like Beam.filling_pattern
"""
if(filling_pattern.sum() != self.size):
raise ValueError("The number of processors must be equal to the"
......@@ -39,11 +80,35 @@ class Mpi:
self.table = table
def rank_to_bunch(self, rank):
"""Return the bunch number corresponding to rank"""
"""
Return the bunch number corresponding to rank
Parameters
----------
rank : int
Rank of a processor
Returns
-------
bunch_num : int
Bunch number corresponding to the input rank
"""
return self.table[rank,1]
def bunch_to_rank(self, bunch_num):
"""Return the rank corresponding to the bunch number bunch_num"""
"""
Return the rank corresponding to the bunch number bunch_num
Parameters
----------
bunch_num : int
Bunch number
Returns
-------
rank : int
Rank of the processor which tracks the input bunch number
"""
try:
rank = np.where(self.table[:,1] == bunch_num)[0][0]
except IndexError:
......
Tracking/beam.py → Tracking/particles.py
+ 105
− 47
View file @ aa1314e3
# -*- coding: utf-8 -*-
"""
Beam and bunch elements
Module where particles, bunches and beams are described as objects.
@author: Alexis Gamelin
@date: 17/01/2020
@date: 11/03/2020
"""
import numpy as np
import pandas as pd
from Tracking.parallel import Mpi
from tracking.parallel import Mpi
from scipy.constants import c, m_e, m_p, e
class Particle:
""" Define a particle object
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 Bunch:
"""
......@@ -36,11 +67,28 @@ class Bunch:
Number of particles in the bunch
current : float
Bunch current in [A]
mean : array of shape (6,)
Mean position of alive particles for each coordinates
std : array of shape (6,)
Standard deviation of the position of alive particles for each
coordinates
emit : array of shape (3,)
Bunch emittance for each plane [1]. !!! -> Correct for long ?
energy_change : series of shape (alive,)
Store the particle relative energy change in the last turn. Only the
values for alive particles are returned. Used by the
SynchrotronRadiation class to compute radiation damping. To be changed
by Element objects which change the energy, for example RF cavities.
Methods
-------
init_gaussian(cov=None, mean=None, **kwargs)
Initialize bunch particles with 6D gaussian phase space.
References
----------
[1] Wiedemann, H. (2015). Particle accelerator physics. 4th edition.
Springer, Eq.(8.39) of p224.
"""
def __init__(self, ring, mp_number=1e3, current=1e-3, alive=True):
......@@ -85,18 +133,6 @@ class Bunch:
def __repr__(self):
"""Return representation of alive particles"""
return f'{self[:]!r}'
@property
def energy_change(self):
"""Store the particle relative energy change in the last turn. Used by
the SynchrotronRadiation class to compute radiation damping. To be
changed by Element objects which change the energy, for example RF
cavities."""
return self._energy_change.loc[self.alive]
@energy_change.setter
def energy_change(self, value):
self._energy_change.loc[self.alive] = value
@property
def mp_number(self):
......@@ -147,46 +183,24 @@ class Bunch:
@property
def mean(self):
"""
Compute the mean position of alive particles for each
coordinates.
Returns
-------
mean : numpy array
mean position of alive particles
Return the mean position of alive particles for each coordinates.
"""
mean = [[self[name].mean()] for name in self]
return np.array(mean)
return np.squeeze(np.array(mean))
@property
def std(self):
"""
Compute the standard deviation of the position of alive
Return 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]
return np.array(std)
return np.squeeze(np.array(std))
@property
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.
Return the bunch emittance for each plane.
"""
emitX = (np.mean(self['x']**2)*np.mean(self['xp']**2) -
np.mean(self['x']*self['xp'])**2)**(0.5)
......@@ -195,12 +209,26 @@ class Bunch:
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])
@property
def energy_change(self):
"""Store the particle relative energy change in the last turn. Used by
the SynchrotronRadiation class to compute radiation damping. To be
changed by Element objects which change the energy, for example RF
cavities."""
return self._energy_change.loc[self.alive]
@energy_change.setter
def energy_change(self, value):
self._energy_change.loc[self.alive] = value
def init_gaussian(self, cov=None, mean=None, **kwargs):
"""
Initialize bunch particles with 6D gaussian phase space.
Covariance matrix is taken from [1].
!!! -> disp & dispp not included yet.
Parameters
----------
cov : (6,6) array, optional
......@@ -259,8 +287,25 @@ class Beam:
Total bunch charge in [C]
particle_number : int
Total number of particle in the beam
filling_pattern : array of bool
filling_pattern : bool array of shape (ring.h,)
Filling pattern of the beam
bunch_current : array of shape (ring.h,)
Current in each bunch in [A]
bunch_charge : array of shape (ring.h,)
Charge in each bunch in [C]
bunch_particle : array of shape (ring.h,)
Particle number in each bunch
bunch_mean : array of shape (6, ring.h)
Mean position of alive particles for each bunch
bunch_std : array of shape (6, ring.h)
Standard deviation of the position of alive particles for each bunch
bunch_emit : array of shape (6, ring.h)
Bunch emittance of alive particles for each bunch
mpi : Mpi object
mpi_switch : bool
Status of MPI parallelisation, should not be changed directly but with
mpi_init() and mpi_close()
Methods
------
......@@ -268,6 +313,13 @@ class Beam:
Initialize beam with a given filling pattern and marco-particle number
per bunch. Then initialize the different bunches with a 6D gaussian
phase space.
mpi_init()
Switch on MPI parallelisation and initialise a Mpi object
mpi_gather()
Gather beam, all bunches of the different processors are sent to
all processors. Rather slow
mpi_close()
Call mpi_gather and switch off MPI parallelisation
"""
def __init__(self, ring, bunch_list=None):
......@@ -408,7 +460,7 @@ class Beam:
bunches"""
bunch_mean = np.zeros((6,self.ring.h))
for index, bunch in enumerate(self):
bunch_mean[:,index] = np.squeeze(bunch.mean)
bunch_mean[:,index] = bunch.mean
return bunch_mean
@property
......@@ -417,7 +469,7 @@ class Beam:
particles for each bunches"""
bunch_std = np.zeros((6,self.ring.h))
for index, bunch in enumerate(self):
bunch_std[:,index] = np.squeeze(bunch.std)
bunch_std[:,index] = bunch.std
return bunch_std
@property
......@@ -426,11 +478,11 @@ class Beam:
bunches and each plane"""
bunch_emit = np.zeros((3,self.ring.h))
for index, bunch in enumerate(self):
bunch_emit[:,index] = np.squeeze(bunch.emit)
bunch_emit[:,index] = bunch.emit
return bunch_emit
def mpi_init(self):
""" Initialise mpi """
"""Switch on MPI parallelisation and initialise a Mpi object"""
self.mpi = Mpi(self.filling_pattern)
self.mpi_switch = True
......@@ -445,6 +497,12 @@ class Beam:
bunches = self.mpi.comm.allgather(bunch)
for rank in range(self.mpi.size):
self[self.mpi.rank_to_bunch(rank)] = bunches[rank]
def mpi_close(self):
"""Call mpi_gather and switch off MPI parallelisation"""
self.mpi_gather()
self.mpi_switch = False
self.mpi = None
......
Tracking/rf.py 0 → 100644
+ 53
− 0
View file @ aa1314e3
# -*- coding: utf-8 -*-
"""
This module handles radio-frequency (RF) cavitiy elements.
@author: gamelina
@date: 11/03/2020
"""
import numpy as np
from tracking.element import Element
class RFCavity(Element):
"""
Perfect RF cavity class for main and harmonic RF cavities.
Use cosine definition.
Parameters
----------
ring : Synchrotron object
m : int
Harmonic number of the cavity
Vc : float
Amplitude of cavity voltage [V]
theta : float
Phase of Cavity voltage
"""
def __init__(self, ring, m, Vc, theta):
self.ring = ring
self.m = m
self.Vc = Vc
self.theta = theta
@Element.parallel
def track(self,bunch):
"""
Tracking method for the element.
No bunch to bunch interaction, so written for Bunch objects and
@Element.parallel is used to handle Beam objects.
Parameters
----------
bunch : Bunch or Beam object
"""
energy_change = self.Vc / self.ring.E0 * np.cos(
self.m * self.ring.omega1 * bunch["tau"] + self.theta )
bunch["delta"] += energy_change
# energy_change is used by SynchrotronRadiation to compute radiation
# damping
bunch.energy_change += energy_change
def value(self, val):
return self.Vc / self.ring.E0 * np.cos(
self.m * self.ring.omega1 * val + self.theta )
\ No newline at end of file
Tracking/synchrotron.py
+ 1
− 42
View file @ aa1314e3
......@@ -7,49 +7,8 @@ General elements
"""
import numpy as np
from scipy.constants import c, m_e, m_p, e
from scipy.constants import c, 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, disp, dispp):
self.beta = beta
self.alpha = alpha
self.disp = disp
self.dispp = dispp
@property
def gamma(self):
return (1 + self.alpha**2)/self.beta
class Synchrotron:
""" Synchrotron class to store main properties. """
......
CollectiveEffect/__init__.py → collective_effects/__init__.py
+ 0
− 0
View file @ aa1314e3
File moved
CollectiveEffect/tools.py → collective_effects/tools.py
+ 0
− 0
View file @ aa1314e3
File moved
CollectiveEffect/wakefield.py → collective_effects/wakefield.py
+ 0
− 0
View file @ aa1314e3
File moved
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

The place to be for your source code