diff --git a/.gitignore b/.gitignore index 2ec607fe88cb09dace70c5df9b331a9a54f82aff..1004a589d6dfe0933e63e33a7f96e320a4c79e62 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ *.ipynb_checkpoints* test_*.py *.hdf5 +*.pyc diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py index 04f26b8dc64acc0e795ea4fda9fdf45e647d6f97..f334be881a7fd0a559eadb6826fb30e9b8835b4b 100644 --- a/mbtrack2/tracking/particles.py +++ b/mbtrack2/tracking/particles.py @@ -129,12 +129,12 @@ class Bunch: self._mp_number = int(mp_number) self.particles = { - "x": np.empty(self.mp_number, dtype=np.float64), - "xp": np.empty(self.mp_number, dtype=np.float64), - "y": np.empty(self.mp_number, dtype=np.float64), - "yp": np.empty(self.mp_number, dtype=np.float64), - "tau": np.empty(self.mp_number, dtype=np.float64), - "delta": np.empty(self.mp_number, dtype=np.float64), + "x": np.zeros(self.mp_number, dtype=np.float64), + "xp": np.zeros(self.mp_number, dtype=np.float64), + "y": np.zeros(self.mp_number, dtype=np.float64), + "yp": np.zeros(self.mp_number, dtype=np.float64), + "tau": np.zeros(self.mp_number, dtype=np.float64), + "delta": np.zeros(self.mp_number, dtype=np.float64), } self.track_alive = track_alive self.alive = np.ones((self.mp_number, ), dtype=bool) @@ -165,11 +165,13 @@ class Bunch: def __iter__(self): """Iterate over labels""" - return self.dtype.names.__iter__() + return self.particles.keys().__iter__() def __repr__(self): """Return representation of alive particles""" - return f'Bunch with macro-particles: \n {pd.DataFrame(self[:])!r}' + rep = pd.DataFrame(np.array([self[l] for l in self.__iter__()]).T, + columns=list(self.__iter__())) + return f'Bunch with macro-particles: \n {rep!r}' @property def mp_number(self): diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py index eb6a105f192deb9e9a65ad3fe77d99b507a57ec1..d64865caba6f83842729db9f9104365c1e8d17c1 100644 --- a/mbtrack2/tracking/wakepotential.py +++ b/mbtrack2/tracking/wakepotential.py @@ -231,7 +231,7 @@ class WakePotential(Element): dtau0) W0 = np.insert(W0, 0, np.zeros(n)) n_to_add = len(tau0) - len(W0) - W0 = np.pad(W0, (n, n_to_add), + W0 = np.pad(W0, (0, n_to_add), mode="constant", constant_values=(0, 0))