From 1694aa5904e9cdd341ca0d3c3d400459baff40b3 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Thu, 20 May 2021 16:30:10 +0200
Subject: [PATCH] Add a bunch_index property to Beam class
bunch_index returns the array with the positions (index) of the non-empty bunches, to be used with Beam.not_empty iterator.
bunch_mean, bunch_std and bunch_emit now iterate only over non empty bunches.
---
tracking/particles.py | 16 +++++++++++++---
tracking/rf.py | 2 +-
2 files changed, 14 insertions(+), 4 deletions(-)
diff --git a/tracking/particles.py b/tracking/particles.py
index feea6f4..6cb7c07 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -427,6 +427,8 @@ class Beam:
mpi_switch : bool
Status of MPI parallelisation, should not be changed directly but with
mpi_init() and mpi_close()
+ bunch_index : array of shape (len(self,))
+ Return an array with the positions of the non-empty bunches
Methods
------
@@ -608,6 +610,11 @@ class Beam:
def filling_pattern(self):
"""Return an array with the filling pattern of the beam as bool"""
return self._filling_pattern
+
+ @property
+ def bunch_index(self):
+ """Return an array with the positions of the non-empty bunches."""
+ return np.where(self.filling_pattern == True)[0]
@property
def bunch_current(self):
@@ -647,7 +654,8 @@ class Beam:
"""Return an array with the mean position of alive particles for each
bunches"""
bunch_mean = np.zeros((6,self.ring.h))
- for index, bunch in enumerate(self):
+ for idx, bunch in enumerate(self.not_empty):
+ index = self.bunch_index[idx]
bunch_mean[:,index] = bunch.mean
return bunch_mean
@@ -656,7 +664,8 @@ class Beam:
"""Return an array with the standard deviation of the position of alive
particles for each bunches"""
bunch_std = np.zeros((6,self.ring.h))
- for index, bunch in enumerate(self):
+ for idx, bunch in enumerate(self.not_empty):
+ index = self.bunch_index[idx]
bunch_std[:,index] = bunch.std
return bunch_std
@@ -665,7 +674,8 @@ class Beam:
"""Return an array with the bunch emittance of alive particles for each
bunches and each plane"""
bunch_emit = np.zeros((3,self.ring.h))
- for index, bunch in enumerate(self):
+ for idx, bunch in enumerate(self.not_empty):
+ index = self.bunch_index[idx]
bunch_emit[:,index] = bunch.emit
return bunch_emit
diff --git a/tracking/rf.py b/tracking/rf.py
index 49a743b..bb2f42d 100644
--- a/tracking/rf.py
+++ b/tracking/rf.py
@@ -187,7 +187,7 @@ class CavityResonator():
self.bunch_index = beam.mpi.bunch_num # Number of the tracked bunch in this processor
self.distance = beam.distance_between_bunches
- self.valid_bunch_index = np.where(beam.filling_pattern == True)[0]
+ self.valid_bunch_index = beam.bunch_index
self.tracking = True
self.nturn = 0
--
GitLab