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

Add track_alive option

Update docstrings
Use pd.DataFrame representation for bunch __repr__
Add the track_alive option for Bunch and Beam.init_beam
parent ca905880
No related branches found
No related tags found
No related merge requests found
......@@ -8,15 +8,15 @@ Module where particles, bunches and beams are described as objects.
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mbtrack2.tracking.parallel import Mpi
import seaborn as sns
import pandas as pd
from mbtrack2.tracking.parallel import Mpi
from scipy.constants import c, m_e, m_p, e
class Particle:
"""
Define a particle object
Define a particle object.
Attributes
----------
......@@ -47,35 +47,41 @@ class Proton(Particle):
class Bunch:
"""
Define a bunch
Define a bunch object.
Parameters
----------
ring : Synchrotron object
mp_number : float, optional
Macro-particle number
Macro-particle number.
current : float, optional
Bunch current in [A]
Bunch current in [A].
track_alive : bool, optional
If False, the code no longer take into account alive/dead particles.
Should be set to True if element such as apertures are used.
Can be set to False to gain a speed increase.
alive : bool, optional
If False, the bunch is defined as empty
If False, the bunch is defined as empty.
Attributes
----------
mp_number : int
Macro-particle number
Macro-particle number.
alive : array of bool of shape (mp_number,)
Array used to monitor which particle is dead or alive.
charge : float
Bunch charge in [C]
Bunch charge in [C].
charge_per_mp : float
Charge per macro-particle in [C]
Charge per macro-particle in [C].
particle_number : int
Number of particles in the bunch
Number of particles in the bunch.
current : float
Bunch current in [A]
Bunch current in [A].
mean : array of shape (6,)
Mean position of alive particles for each coordinates
Mean position of alive particles for each coordinates.
std : array of shape (6,)
Standard deviation of the position of alive particles for each
coordinates
coordinates.
emit : array of shape (3,)
Bunch emittance for each plane [1]. !!! -> Correct for long ?
energy_change : series of shape (alive,)
......@@ -97,8 +103,8 @@ class Bunch:
Springer, Eq.(8.39) of p224.
"""
def __init__(self, ring, mp_number=1e3, current=1e-3, alive=True,
track_alive=False):
def __init__(self, ring, mp_number=1e3, current=1e-3, track_alive=True,
alive=True):
self.ring = ring
if not alive:
......@@ -147,7 +153,7 @@ class Bunch:
def __repr__(self):
"""Return representation of alive particles"""
return f'{self[:]!r}'
return f'{pd.DataFrame(self[:])!r}'
@property
def mp_number(self):
......@@ -304,12 +310,14 @@ class Bunch:
Returns
-------
bins : IntervalIndex of shape (n_bin,)
bins : array of shape (n_bin,)
Bins where the particles are sorted.
sorted_index : Series of shape (self.mp_number,)
sorted_index : array of shape (self.mp_number,)
Bin number of each macro-particles.
profile : Series of shape (n_bin,)
profile : array of shape (n_bin - 1,)
Number of marco-particles in each bin.
center : array of shape (n_bin - 1,)
Center of each bin.
"""
bin_min = self[dimension].min()
......@@ -495,7 +503,7 @@ class Beam:
return distance
def init_beam(self, filling_pattern, current_per_bunch=1e-3,
mp_per_bunch=1e3):
mp_per_bunch=1e3, track_alive=True):
"""
Initialize beam with a given filling pattern and marco-particle number
per bunch. Then initialize the different bunches with a 6D gaussian
......@@ -515,6 +523,10 @@ class Beam:
Current per bunch in [A]
mp_per_bunch : float, optional
Macro-particle number per bunch
track_alive : bool, optional
If False, the code no longer take into account alive/dead particles.
Should be set to True if element such as apertures are used.
Can be set to False to gain a speed increase.
"""
if (len(filling_pattern) != self.ring.h):
......@@ -526,13 +538,15 @@ class Beam:
if filling_pattern.dtype == np.dtype("bool"):
for value in filling_pattern:
if value == True:
bunch_list.append(Bunch(self.ring, mp_per_bunch, current_per_bunch))
bunch_list.append(Bunch(self.ring, mp_per_bunch,
current_per_bunch, track_alive))
elif value == False:
bunch_list.append(Bunch(self.ring, alive=False))
elif filling_pattern.dtype == np.dtype("float64"):
for current in filling_pattern:
if current != 0:
bunch_list.append(Bunch(self.ring, mp_per_bunch, current))
bunch_list.append(Bunch(self.ring, mp_per_bunch,
current_per_bunch, track_alive))
elif current == 0:
bunch_list.append(Bunch(self.ring, alive=False))
else:
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment