From db49cb81bf895a0c1ae609012f9fc492ba085b52 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Mon, 10 Jan 2022 15:20:26 +0100
Subject: [PATCH] Allow for particle losses in Monitor classes
Add a is_empty property to the Bunch class.
Add a check_empty optional parameter to the Monitor.track_bunch_data method.
ProfileMonitor: if the bunch is empty, no profile is saved.
WakePotentialMonitor: if the bunch is empty, the same from last good turn is saved.
BunchSpectrumMonitor: if the bunch is empty, the incoherent spectrum is nan.
---
tracking/monitors/monitors.py | 26 +++++++++++++++++---------
tracking/particles.py | 7 +++++++
2 files changed, 24 insertions(+), 9 deletions(-)
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index a543fdf..f0f8a4d 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -218,27 +218,32 @@ class Monitor(Element, metaclass=ABCMeta):
except ValueError:
pass
- def track_bunch_data(self, object_to_save):
+ def track_bunch_data(self, object_to_save, check_empty=False):
"""
Track method to use when saving bunch data.
Parameters
----------
object_to_save : Beam or Bunch
+ check_emptiness: bool
+ If True, check if the bunch is empty. If it is, then do nothing.
"""
if self.track_count % self.save_every == 0:
+
if isinstance(object_to_save, Beam):
if (object_to_save.mpi_switch == True):
- if object_to_save.mpi.bunch_num == self.bunch_number:
- self.to_buffer(object_to_save[object_to_save.mpi.bunch_num])
+ bunch = object_to_save[object_to_save.mpi.bunch_num]
else:
- self.to_buffer(object_to_save[self.bunch_number])
+ bunch = object_to_save[self.bunch_number]
elif isinstance(object_to_save, Bunch):
- self.to_buffer(object_to_save)
+ bunch = object_to_save
else:
raise TypeError("object_to_save should be a Beam or Bunch object.")
+
+ if (check_empty == False) or (bunch.is_empty == False):
+ self.to_buffer(bunch)
+
self.track_count += 1
-
class BunchMonitor(Monitor):
"""
@@ -672,7 +677,7 @@ class ProfileMonitor(Monitor):
----------
object_to_save : Bunch or Beam object
"""
- self.track_bunch_data(object_to_save)
+ self.track_bunch_data(object_to_save, check_empty=True)
class WakePotentialMonitor(Monitor):
"""
@@ -1004,8 +1009,11 @@ class BunchSpectrumMonitor(Monitor):
raise TypeError("object_to_save should be a Beam or Bunch object.")
if skip is False:
- for key, value in self.track_dict.items():
- self.positions[value, :, self.save_count] = bunch[key][self.index_sample]
+ try:
+ for key, value in self.track_dict.items():
+ self.positions[value, :, self.save_count] = bunch[key][self.index_sample]
+ except IndexError:
+ self.positions[value, :, self.save_count] = np.nan
self.mean[:, self.save_count] = bunch.mean[self.mean_index]
diff --git a/tracking/particles.py b/tracking/particles.py
index e7939e8..f4351c1 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -76,6 +76,8 @@ class Bunch:
Number of particles in the bunch.
current : float
Bunch current in [A].
+ is_empty : bool
+ Return True if the bunch is empty.
mean : array of shape (6,)
Mean position of alive particles for each coordinates.
std : array of shape (6,)
@@ -192,6 +194,11 @@ class Bunch:
@current.setter
def current(self, value):
self.charge_per_mp = value * self.ring.T0 / self.__len__()
+
+ @property
+ def is_empty(self):
+ """Return True if the bunch is empty."""
+ return ~np.any(self.alive)
@property
def mean(self):
--
GitLab