diff --git a/tracking/particles.py b/tracking/particles.py
index 7a97fa35b64ca960442bcb856103b0e4c3ab8722..4378745f0bb80048aa21b266f3c674bf8df3c144 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,7 +103,8 @@ class Bunch:
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,
+ alive=True):
self.ring = ring
if not alive:
@@ -105,20 +112,22 @@ class Bunch:
current = 0
self._mp_number = int(mp_number)
- particles = {"x":np.zeros((self.mp_number,)),
- "xp":np.zeros((self.mp_number,)),
- "y":np.zeros((self.mp_number,)),
- "yp":np.zeros((self.mp_number,)),
- "tau":np.zeros((self.mp_number,)),
- "delta":np.zeros((self.mp_number,)),
- }
- self.particles = pd.DataFrame(particles)
- self.alive = pd.Series(np.ones((self.mp_number,),dtype=bool))
+ self.dtype = np.dtype([('x',np.float),
+ ('xp',np.float),
+ ('y',np.float),
+ ('yp',np.float),
+ ('tau',np.float),
+ ('delta',np.float)])
+
+ self.particles = np.zeros(self.mp_number, self.dtype)
+ self.track_alive = track_alive
+
+ self.alive = np.ones((self.mp_number,),dtype=bool)
self.current = current
if not alive:
- self.alive = pd.Series(np.zeros((self.mp_number,),dtype=bool))
+ self.alive = np.zeros((self.mp_number,),dtype=bool)
- self._energy_change = pd.Series(np.zeros((self.mp_number,),dtype=float))
+ self._energy_change = np.zeros((self.mp_number,),dtype=float)
def __len__(self):
"""Return the number of alive particles"""
@@ -126,19 +135,25 @@ class Bunch:
def __getitem__(self, label):
"""Return the columns label for alive particles"""
- return self.particles.loc[self.alive, label]
+ if self.track_alive is True:
+ return self.particles[label][self.alive]
+ else:
+ return self.particles[label]
def __setitem__(self, label, value):
"""Set value to the columns label for alive particles"""
- self.particles.loc[self.alive, label] = value
+ if self.track_alive is True:
+ self.particles[label][self.alive] = value
+ else:
+ self.particles[label] = value
def __iter__(self):
"""Iterate over labels"""
- return self[:].__iter__()
+ return self.dtype.names.__iter__()
def __repr__(self):
"""Return representation of alive particles"""
- return f'{self[:]!r}'
+ return f'{pd.DataFrame(self[:])!r}'
@property
def mp_number(self):
@@ -222,11 +237,17 @@ class Bunch:
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]
+ if self.track_alive is True:
+ return self._energy_change[self.alive]
+ else:
+ return self._energy_change
@energy_change.setter
def energy_change(self, value):
- self._energy_change.loc[self.alive] = value
+ if self.track_alive is True:
+ self._energy_change[self.alive] = value
+ else:
+ self._energy_change = value
def init_gaussian(self, cov=None, mean=None, **kwargs):
"""
@@ -289,22 +310,28 @@ 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.
"""
- bins = pd.interval_range(start=min(self[dimension]),
- end=max(self[dimension]),
- periods=n_bin,
- name=dimension)
- sorted_index = pd.cut(self[dimension], bins)
- sorted_index = sorted_index.reindex(self.alive.index)
- profile = sorted_index.value_counts().sort_index()
- return (bins, sorted_index, profile)
+ bin_min = self[dimension].min()
+ bin_min = min(bin_min*0.99, bin_min*1.01)
+ bin_max = self[dimension].max()
+ bin_max = max(bin_max*0.99, bin_max*1.01)
+
+ bins = np.linspace(bin_min, bin_max, n_bin)
+ center = (bins[1:] + bins[:-1])/2
+ sorted_index = np.searchsorted(bins, self[dimension], side='left')
+ sorted_index -= 1
+ profile = np.bincount(sorted_index, minlength=n_bin-1)
+
+ return (bins, sorted_index, profile, center)
def plot_profile(self, dimension="tau", n_bin=75):
"""
@@ -318,10 +345,10 @@ class Bunch:
Number of bins. The default is 75.
"""
- bins, sorted_index, profile = self.binning(dimension, n_bin)
+ bins, sorted_index, profile, center = self.binning(dimension, n_bin)
fig = plt.figure()
ax = fig.gca()
- ax.plot(bins.mid, profile)
+ ax.plot(center, profile)
def plot_phasespace(self, x_var="tau", y_var="delta", plot_type="j"):
"""
@@ -476,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
@@ -496,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):
@@ -507,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:
diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py
index 41bd2d019270b976e91883c1c8f3d8722b93fa8a..4a0c599d1cd9824fb566bbcf3eed8a068f01e324 100644
--- a/tracking/wakepotential.py
+++ b/tracking/wakepotential.py
@@ -85,18 +85,19 @@ class WakePotential(Element):
"""
# Get binning data
- a, b, c = bunch.binning(n_bin=self.n_bin)
+ a, b, c, d = bunch.binning(n_bin=self.n_bin)
self.bins = a
self.sorted_index = b
self.profile = c
- self.bin_size = self.bins.length
+ self.center = d
+ self.bin_size = self.bins[1] - self.bins[0]
# Compute charge density
self.rho = bunch.charge_per_mp*self.profile/(self.bin_size*bunch.charge)
self.rho = np.array(self.rho)
# Compute time array
- self.tau = np.array(self.bins.mid)
+ self.tau = np.array(self.center)
self.dtau = self.tau[1] - self.tau[0]
# Add N values before and after rho and tau