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

Merge branch 'numpy_bunch'

parents e32bfaf0 c670561f
Branches
Tags
No related merge requests found
Show whitespace changes
Inline Side-by-side
tracking/particles.py
+ 81
− 48
View file @ ac15e7e6
......@@ -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:
......
tracking/wakepotential.py
+ 4
− 3
View file @ ac15e7e6
......@@ -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
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment

The place to be for your source code