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

Use pandas DF to hold data in Bunch class

Use pandas DF to hold data in Bunch class instead of numpy masked array -> better behaviour and allows to set to alive particles only
parent c611185b
Branches
Tags
No related merge requests found
......@@ -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):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment