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