diff --git a/Tracking/beam.py b/Tracking/beam.py index e3c4f1bcbe98dc2c778ffb89cd10ad53b67742a6..a28bfd4ee8283c5d152079cca07136d854c20f1b 100644 --- a/Tracking/beam.py +++ b/Tracking/beam.py @@ -7,7 +7,7 @@ Beam and bunch elements """ import numpy as np -import numpy.ma as ma +import pandas as pd import warnings class Bunch: @@ -45,33 +45,38 @@ class Bunch: """ def __init__(self, ring, mp_number=1e3, current=1e-3): - dpart = np.dtype([("x",np.float64),("xp",np.float64), - ("y",np.float64),("yp",np.float64), - ("tau",np.float64),("delta",np.float64)]) - self.dtype = dpart - particles = np.zeros((int(mp_number),6), dtype=np.float64) - self.particles = ma.masked_array(particles) - + particles = {"x":np.zeros((int(mp_number),)), + "xp":np.zeros((int(mp_number),)), + "y":np.zeros((int(mp_number),)), + "yp":np.zeros((int(mp_number),)), + "tau":np.zeros((int(mp_number),)), + "delta":np.zeros((int(mp_number),)), + } + self.particles = pd.DataFrame(particles) + self.alive = pd.Series(np.ones((int(mp_number),),dtype=bool)) self.ring = ring self._mp_number = int(mp_number) self.current = current def __len__(self): """Return the number of alive particles""" - #mask = self.particles["alive"] == True - return len(self.particles) + return len(self[:]) - def __getitem__(self, i): - """Return the alive particle number i""" - return self.particles[i] + def __getitem__(self, label): + """Return the columns label for alive particles""" + return self.particles.loc[self.alive, label] - def __setitem__(self, i, value): - """Set value to the alive particle number i""" - self.particles[i] = value + def __setitem__(self, label, value): + """Set value to the columns label for alive particles""" + self.particles.loc[self.alive, label] = value def __iter__(self): - """Iterate over alive particles""" - return self.particles.__iter__() + """Iterate over labels""" + return self[:].__iter__() + + def __repr__(self): + """Return representation of alive particles""" + return f'{self[:]!r}' @property def mp_number(self): @@ -129,7 +134,7 @@ class Bunch: mean : numpy array mean position of alive particles """ - mean = [[self[name].mean()] for name in self.dtype.names] + mean = [[self[name].mean()] for name in self] return np.array(mean) def std(self): @@ -142,7 +147,7 @@ class Bunch: std : numpy array standard deviation of the position of alive particles """ - std = [[self[name].std()] for name in self.dtype.names] + std = [[self[name].std()] for name in self] return np.array(std) def emit(self):