From c670561f9a1811ce9fdf41f67842912f65b90902 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Thu, 12 Nov 2020 17:14:33 +0100
Subject: [PATCH] Add track_alive option
Update docstrings
Use pd.DataFrame representation for bunch __repr__
Add the track_alive option for Bunch and Beam.init_beam
---
tracking/particles.py | 60 ++++++++++++++++++++++++++-----------------
1 file changed, 37 insertions(+), 23 deletions(-)
diff --git a/tracking/particles.py b/tracking/particles.py
index b327da5..4378745 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -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:
--
GitLab