From fbc055664bac0fcc68bcdb77f7ea46aa86c5ea67 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Wed, 20 Apr 2022 18:23:23 +0200
Subject: [PATCH] [Fix] Monitor behaviour in multi-bunch
Fix Monitor.track_bunch_data method to check that only the bunch which bunch_num is specified can write on the file
Modify the behaviour of BunchMonitor, PhaseSpaceMonitor, ProfileMonitor
Make the same modification within WakePotentialMonitor.
Style change to BunchSpectrumMonitor
---
mbtrack2/tracking/monitors/monitors.py | 45 +++++++++++++++++++-------
1 file changed, 33 insertions(+), 12 deletions(-)
diff --git a/mbtrack2/tracking/monitors/monitors.py b/mbtrack2/tracking/monitors/monitors.py
index 8091b9f..b1099fc 100644
--- a/mbtrack2/tracking/monitors/monitors.py
+++ b/mbtrack2/tracking/monitors/monitors.py
@@ -226,20 +226,25 @@ class Monitor(Element, metaclass=ABCMeta):
check_emptiness: bool
If True, check if the bunch is empty. If it is, then do nothing.
"""
+ save = True
if self.track_count % self.save_every == 0:
if isinstance(object_to_save, Beam):
if (object_to_save.mpi_switch == True):
- bunch = object_to_save[object_to_save.mpi.bunch_num]
+ if object_to_save.mpi.bunch_num == self.bunch_number:
+ bunch = object_to_save[object_to_save.mpi.bunch_num]
+ else:
+ save = False
else:
bunch = object_to_save[self.bunch_number]
elif isinstance(object_to_save, Bunch):
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)
+
+ if save:
+ if (check_empty == False) or (bunch.is_empty == False):
+ self.to_buffer(bunch)
self.track_count += 1
@@ -712,7 +717,7 @@ class WakePotentialMonitor(Monitor):
Methods
-------
- track(wake_potential_to_save)
+ track(object_to_save, wake_potential_to_save)
Save data.
"""
@@ -812,15 +817,31 @@ class WakePotentialMonitor(Monitor):
self.file.flush()
self.write_count += 1
- def track(self, wake_potential_to_save):
+ def track(self, object_to_save, wake_potential_to_save):
"""
Save data.
Parameters
----------
- object_to_save : WakePotential object
- """
- if self.track_count % self.save_every == 0:
+ object_to_save : Bunch or Beam object
+ wake_potential_to_save : WakePotential object
+ """
+ if isinstance(object_to_save, Beam):
+ if (object_to_save.mpi_switch == True):
+ if object_to_save.mpi.bunch_num == self.bunch_number:
+ save = True
+ else:
+ save = False
+ else:
+ raise NotImplementedError("WakePotentialMonitor for Beam " +
+ "objects is only available " +
+ "with MPI mode.")
+ elif isinstance(object_to_save, Bunch):
+ save = True
+ else:
+ raise TypeError("object_to_save should be a Beam or Bunch object.")
+
+ if save and (self.track_count % self.save_every == 0):
self.to_buffer(wake_potential_to_save)
self.track_count += 1
@@ -992,13 +1013,13 @@ class BunchSpectrumMonitor(Monitor):
object_to_save : Beam or Bunch object
"""
- skip = False
+ save = True
if isinstance(object_to_save, Beam):
if (object_to_save.mpi_switch == True):
if object_to_save.mpi.bunch_num == self.bunch_number:
bunch = object_to_save[object_to_save.mpi.bunch_num]
else:
- skip = True
+ save = False
else:
bunch = object_to_save[self.bunch_number]
elif isinstance(object_to_save, Bunch):
@@ -1006,7 +1027,7 @@ class BunchSpectrumMonitor(Monitor):
else:
raise TypeError("object_to_save should be a Beam or Bunch object.")
- if skip is False:
+ if save:
try:
for key, value in self.track_dict.items():
self.positions[value, :, self.save_count] = bunch[key][self.index_sample]
--
GitLab