diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 3f1f9d7b02ae3b87e9dd7566b0360d5983f9b211..69a1b8bda4dca3c1dd53210f9e95f26ccd6ddc24 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -10,7 +10,6 @@ from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor,
BeamMonitor,
ProfileMonitor,
WakePotentialMonitor,
- TuneMonitor,
CavityMonitor,
BunchSpectrumMonitor,
BeamSpectrumMonitor)
@@ -21,4 +20,6 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_wakedata,
plot_tunedata,
plot_cavitydata,
- streak_beamdata)
\ No newline at end of file
+ streak_beamdata)
+
+from mbtrack2.tracking.monitors.tools import merge_files
\ No newline at end of file
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 5f4cf94d19a523d93aa76a5069cfaedc10c30d77..cadd0546c3dd773689eb76680fb7e10996125539 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -10,7 +10,6 @@ during tracking.
import numpy as np
import h5py as hp
-import PyNAFF as pnf
import random
from mbtrack2.tracking.element import Element
from mbtrack2.tracking.particles import Bunch, Beam
@@ -77,7 +76,7 @@ class Monitor(Element, metaclass=ABCMeta):
raise ValueError
def monitor_init(self, group_name, save_every, buffer_size, total_size,
- dict_buffer, dict_file, file_name=None, mpi_mode=True,
+ dict_buffer, dict_file, file_name=None, mpi_mode=False,
dict_dtype=None):
"""
Method called to initialize Monitor subclass.
@@ -250,20 +249,20 @@ class BunchMonitor(Monitor):
----------
bunch_number : int
Bunch to monitor
- file_name : string, optional
- Name of the HDF5 where the data will be stored. Must be specified
- the first time a subclass of Monitor is instancied and must be None
- the following times.
- save_every : int or float, optional
+ save_every : int or float
Set the frequency of the save. The data is saved every save_every
call of the montior.
- buffer_size : int or float, optional
+ buffer_size : int or float
Size of the save buffer.
- total_size : int or float, optional
+ total_size : int or float
Total size of the save. The following relationships between the
parameters must exist:
total_size % buffer_size == 0
number of call to track / save_every == total_size
+ file_name : string, optional
+ Name of the HDF5 where the data will be stored. Must be specified
+ the first time a subclass of Monitor is instancied and must be None
+ the following times.
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -275,8 +274,8 @@ class BunchMonitor(Monitor):
Save data
"""
- def __init__(self, bunch_number, file_name=None, save_every=5,
- buffer_size=500, total_size=2e4, mpi_mode=True):
+ def __init__(self, bunch_number, save_every, buffer_size, total_size,
+ file_name=None, mpi_mode=False):
self.bunch_number = bunch_number
group_name = "BunchData_" + str(self.bunch_number)
@@ -313,20 +312,20 @@ class PhaseSpaceMonitor(Monitor):
mp_number : int or float
Number of macroparticle in the phase space to save. If less than the
total number of macroparticles, a random fraction of the bunch is saved.
- file_name : string, optional
- Name of the HDF5 where the data will be stored. Must be specified
- the first time a subclass of Monitor is instancied and must be None
- the following times.
- save_every : int or float, optional
+ save_every : int or float
Set the frequency of the save. The data is saved every save_every
call of the montior.
- buffer_size : int or float, optional
+ buffer_size : int or float
Size of the save buffer.
- total_size : int or float, optional
+ total_size : int or float
Total size of the save. The following relationships between the
parameters must exist:
total_size % buffer_size == 0
number of call to track / save_every == total_size
+ file_name : string, optional
+ Name of the HDF5 where the data will be stored. Must be specified
+ the first time a subclass of Monitor is instancied and must be None
+ the following times.
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -338,8 +337,8 @@ class PhaseSpaceMonitor(Monitor):
Save data
"""
- def __init__(self, bunch_number, mp_number, file_name=None, save_every=1e3,
- buffer_size=10, total_size=100, mpi_mode=True):
+ def __init__(self, bunch_number, mp_number, save_every, buffer_size,
+ total_size, file_name=None, mpi_mode=False):
self.bunch_number = bunch_number
self.mp_number = int(mp_number)
@@ -402,20 +401,20 @@ class BeamMonitor(Monitor):
----------
h : int
Harmonic number of the ring.
- file_name : string, optional
- Name of the HDF5 where the data will be stored. Must be specified
- the first time a subclass of Monitor is instancied and must be None
- the following times.
- save_every : int or float, optional
+ save_every : int or float
Set the frequency of the save. The data is saved every save_every
call of the montior.
- buffer_size : int or float, optional
+ buffer_size : int or float
Size of the save buffer.
- total_size : int or float, optional
+ total_size : int or float
Total size of the save. The following relationships between the
parameters must exist:
total_size % buffer_size == 0
number of call to track / save_every == total_size
+ file_name : string, optional
+ Name of the HDF5 where the data will be stored. Must be specified
+ the first time a subclass of Monitor is instancied and must be None
+ the following times.
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -427,8 +426,8 @@ class BeamMonitor(Monitor):
Save data
"""
- def __init__(self, h, file_name=None, save_every=5, buffer_size=500,
- total_size=2e4, mpi_mode=True):
+ def __init__(self, h, save_every, buffer_size, total_size, file_name=None,
+ mpi_mode=False):
group_name = "Beam"
dict_buffer = {"mean" : (6, h, buffer_size),
@@ -568,6 +567,16 @@ class ProfileMonitor(Monitor):
----------
bunch_number : int
Bunch to monitor.
+ save_every : int or float
+ Set the frequency of the save. The data is saved every save_every
+ call of the montior.
+ buffer_size : int or float
+ Size of the save buffer.
+ total_size : int or float
+ Total size of the save. The following relationships between the
+ parameters must exist:
+ total_size % buffer_size == 0
+ number of call to track / save_every == total_size
dimensions : str or list of str, optional
Dimensions to save.
n_bin : int or list of int, optional
@@ -576,16 +585,6 @@ class ProfileMonitor(Monitor):
Name of the HDF5 where the data will be stored. Must be specified
the first time a subclass of Monitor is instancied and must be None
the following times.
- save_every : int or float, optional
- Set the frequency of the save. The data is saved every save_every
- call of the montior.
- buffer_size : int or float, optional
- Size of the save buffer.
- total_size : int or float, optional
- Total size of the save. The following relationships between the
- parameters must exist:
- total_size % buffer_size == 0
- number of call to track / save_every == total_size
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -597,8 +596,8 @@ class ProfileMonitor(Monitor):
Save data.
"""
- def __init__(self, bunch_number, dimensions="tau", n_bin=75, file_name=None,
- save_every=5, buffer_size=500, total_size=2e4, mpi_mode=True):
+ def __init__(self, bunch_number, save_every, buffer_size, total_size,
+ dimensions="tau", n_bin=75, file_name=None, mpi_mode=False):
self.bunch_number = bunch_number
group_name = "ProfileData_" + str(self.bunch_number)
@@ -689,20 +688,20 @@ class WakePotentialMonitor(Monitor):
n_bin : int
Number of bin to be used to interpolate the wake potential on a fixed
grid.
- file_name : string, optional
- Name of the HDF5 where the data will be stored. Must be specified
- the first time a subclass of Monitor is instancied and must be None
- the following times.
- save_every : int or float, optional
+ save_every : int or float
Set the frequency of the save. The data is saved every save_every
call of the montior.
- buffer_size : int or float, optional
+ buffer_size : int or float
Size of the save buffer.
- total_size : int or float, optional
+ total_size : int or float
Total size of the save. The following relationships between the
parameters must exist:
total_size % buffer_size == 0
number of call to track / save_every == total_size
+ file_name : string, optional
+ Name of the HDF5 where the data will be stored. Must be specified
+ the first time a subclass of Monitor is instancied and must be None
+ the following times.
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -714,8 +713,8 @@ class WakePotentialMonitor(Monitor):
Save data.
"""
- def __init__(self, bunch_number, wake_types, n_bin, file_name=None,
- save_every=5, buffer_size=500, total_size=2e4, mpi_mode=True):
+ def __init__(self, bunch_number, wake_types, n_bin, save_every,
+ buffer_size, total_size, file_name=None, mpi_mode=False):
self.bunch_number = bunch_number
group_name = "WakePotentialData_" + str(self.bunch_number)
@@ -725,7 +724,7 @@ class WakePotentialMonitor(Monitor):
else:
self.wake_types = wake_types
- self.n_bin = n_bin
+ self.n_bin = n_bin*2
dict_buffer = {}
dict_file = {}
@@ -821,11 +820,10 @@ class WakePotentialMonitor(Monitor):
if self.track_count % self.save_every == 0:
self.to_buffer(wake_potential_to_save)
self.track_count += 1
-
-class TuneMonitor(Monitor):
+
+class BunchSpectrumMonitor(Monitor):
"""
- Monitor tunes and the Fourier transform (using FFT algorithm) of the
- osciallation in horizontal, vertical, and longitudinal plane.
+ Monitor the coherent and incoherent bunch spectrums.
Parameters
----------
@@ -836,254 +834,35 @@ class TuneMonitor(Monitor):
Total number of macro-particles in the bunch.
sample_size : int or float
Number of macro-particles to be used for tune and FFT computation.
- This number cannot exceed mp_number.
- save_tune : bool, optional
- If True, tune data is saved.
- save_fft : bool, optional
- If True, FFT data is saved.
- n_fft : int or float, optional
- The number of points used for FFT computation, if n_fft is bigger than
- save_every zero-padding is applied.
- file_name : string, optional
- Name of the HDF5 where the data will be stored. Must be specified
- the first time a subclass of Monitor is instancied and must be None
- the following times.
- save_every : int or float, optional
- Set the frequency of the save. The tune is computed saved every
+ This number cannot exceed mp_number.
+ save_every : int or float
+ Set the frequency of the save. The spectrums are computed every
save_every call of the montior.
- buffer_size : int or float, optional
+ buffer_size : int or float
Size of the save buffer.
- total_size : int or float, optional
+ total_size : int or float
Total size of the save. The following relationships between the
parameters must exist:
total_size % buffer_size == 0
- number of call to track / save_every == total_size - 1
- mpi_mode : bool, optional
- If True, open the HDF5 file in parallel mode, which is needed to
- allow several cores to write in the same file at the same time.
- If False, open the HDF5 file in standard mode.
-
- Methods
- -------
- track(bunch):
- Save tune and/or FFT data.
-
- """
-
- def __init__(self, ring, bunch_number, mp_number, sample_size, save_tune=True,
- save_fft=False, n_fft=10000, file_name=None, save_every=10,
- buffer_size=5, total_size=10, mpi_mode=True):
-
- self.ring = ring
- self.bunch_number = bunch_number
- group_name = "TuneData_" + str(self.bunch_number)
-
- dict_buffer = {"tune":(3, buffer_size), "tune_spread":(3, buffer_size,),
- "fft":(3, n_fft//2+1, buffer_size)}
- dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,),
- "fft":(3, n_fft//2+1, total_size)}
-
- self.monitor_init(group_name, save_every, buffer_size, total_size,
- dict_buffer, dict_file, file_name, mpi_mode)
-
- self.dict_buffer = dict_buffer
- self.dict_file = dict_file
-
- self.sample_size = int(sample_size)
- self.x = np.zeros((self.sample_size, save_every+1))
- self.y = np.zeros((self.sample_size, save_every+1))
- self.tau = np.zeros((self.sample_size, save_every+1))
-
- index = np.arange(0, int(mp_number))
- self.index_sample = sorted(random.sample(list(index), self.sample_size))
-
- self.save_count = 0
- self.buffer_count = 0
-
- self.save_tune = save_tune
- self.save_fft = save_fft
- self.save_every = save_every
-
- if self.save_fft is True :
- self.n_fft = n_fft
- self.fourier_save = np.zeros((3, self.n_fft//2+1, buffer_size))
-
- def track(self, object_to_save):
- """
- Save tune data.
-
- Parameters
- ----------
- object_to_save : Beam or Bunch object
-
- """
- skip = False
- 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
- 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 skip is False:
- self.x[:, self.save_count] = bunch["x"][self.index_sample]
- self.y[:, self.save_count] = bunch["y"][self.index_sample]
- self.tau[:, self.save_count] = bunch["tau"][self.index_sample]
-
- self.save_count += 1
-
- if self.track_count > 0 and self.track_count % self.save_every == 0:
- self.to_buffer(bunch)
- self.save_count = 0
-
- self.track_count += 1
-
- def to_buffer(self, bunch):
- """
- A method to hold saved data before writing it to the output file.
-
- """
-
- self.time[self.buffer_count] = self.track_count
-
- if self.save_tune is True:
- mean, spread = self.get_tune(bunch)
- self.tune[:, self.buffer_count] = mean
- self.tune_spread[:, self.buffer_count] = spread
-
- if self.save_fft is True:
- fx, fy, ftau = self.get_fft()
- self.fourier_save[0,:,self.buffer_count] = fx
- self.fourier_save[1,:,self.buffer_count] = fy
- self.fourier_save[2,:,self.buffer_count] = ftau
-
- self.buffer_count += 1
-
- if self.buffer_count == self.buffer_size:
- self.write()
- self.buffer_count = 0
-
- def write(self):
- """
- Write data from buffer to output file.
-
- """
- self.file[self.group_name]["time"][self.write_count*self.buffer_size:(
- self.write_count+1)*self.buffer_size] = self.time
-
- if self.save_tune is True:
- self.file[self.group_name]["tune"][:,
- self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.tune
- self.file[self.group_name]["tune_spread"][:,
- self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.tune_spread
-
- if self.save_fft is True:
- self.file[self.group_name]["fft"][:,:,
- self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.fourier_save
-
- self.file.flush()
- self.write_count += 1
-
- def get_tune(self, bunch):
- """
- Compute tune by using NAFF algorithm to indentify the fundamental
- harmonic frequency of the particles' motion.
-
- Parameters
- ----------
- bunch : Bunch object
-
- """
-
- turn = self.save_every
- freq = np.zeros((self.sample_size,3))
-
- for i in range(self.sample_size):
- try:
- freq[i,0] = pnf.naff(self.x[i,:], turns=turn-1, nterms=1)[0][1] \
- / self.ring.T0
- except IndexError:
- freq[i,0] = np.nan
-
- try:
- freq[i,1] = pnf.naff(self.y[i,:], turns=turn-1, nterms=1)[0][1] \
- / self.ring.T0
- except IndexError:
- freq[i,1] = np.nan
-
- try:
- freq[i,2] = pnf.naff(self.tau[i,:], turns=turn-1, nterms=1)[0][1] \
- / self.ring.T0
- except IndexError:
- freq[i,2] = np.nan
-
- tune_single_particle = freq / self.ring.f0
- mean = np.nanmean(tune_single_particle, 0)
- spread = np.nanstd(tune_single_particle, 0)
-
- return (mean, spread)
-
- def get_fft(self):
- """
- Compute the Fourier transform (using FFT algorithm) of the
- osciallation in horizontal, vertical, and longitudinal plane.
-
- Returns
- -------
- fourier_x_avg, fourier_y_avg, fourier_tau_avg : ndarray
- The average of the transformed input in each plane.
-
- """
- fourier_x = rfft(self.x, n=self.n_fft)
- fourier_y = rfft(self.y, n=self.n_fft)
- fourier_tau = rfft(self.tau, n=self.n_fft)
-
- fourier_x_avg = np.mean(abs(fourier_x),axis=0)
- fourier_y_avg = np.mean(abs(fourier_y),axis=0)
- fourier_tau_avg = np.mean(abs(fourier_tau),axis=0)
-
- return (fourier_x_avg, fourier_y_avg, fourier_tau_avg)
-
-class BunchSpectrumMonitor(Monitor):
- """
- Monitor the coherent and incoherent bunch spectrums.
-
- Parameters
- ----------
- ring : Synchrotron object
- bunch_number : int
- Bunch to monitor
- mp_number : int or float
- Total number of macro-particles in the bunch.
- sample_size : int or float
- Number of macro-particles to be used for tune and FFT computation.
- This number cannot exceed mp_number.
+ number of call to track / save_every == total_size - 1
+ dim : str, optional
+ Dimensions in which the spectrums have to be computed.
+ Can be:
+ - "all"
+ - "tau"
+ - "x"
+ - "y"
+ - "xy" or "yx"
+ - "xtau" or "taux"
+ - "ytau" or "tauy"
n_fft : int or float, optional
The number of points used for FFT computation, if n_fft is bigger than
save_every zero-padding is applied.
+ If None, save_every is used.
file_name : string, optional
Name of the HDF5 where the data will be stored. Must be specified
the first time a subclass of Monitor is instancied and must be None
the following times.
- save_every : int or float, optional
- Set the frequency of the save. The spectrums are computed every
- save_every call of the montior.
- buffer_size : int or float, optional
- Size of the save buffer.
- total_size : int or float, optional
- Total size of the save. The following relationships between the
- parameters must exist:
- total_size % buffer_size == 0
- number of call to track / save_every == total_size - 1
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -1105,10 +884,13 @@ class BunchSpectrumMonitor(Monitor):
"""
- def __init__(self, ring, bunch_number, mp_number, sample_size, n_fft,
- save_every, buffer_size, total_size, file_name=None,
- mpi_mode=True, dim="all"):
+ def __init__(self, ring, bunch_number, mp_number, sample_size, save_every,
+ buffer_size, total_size, dim="all", n_fft=None,
+ file_name=None, mpi_mode=False):
+ if n_fft is None:
+ n_fft = int(save_every)
+
self.n_fft = int(n_fft)
self.sample_size = int(sample_size)
self.store_dict = {"x":0,"y":1,"tau":2}
@@ -1128,6 +910,12 @@ class BunchSpectrumMonitor(Monitor):
elif dim == "xy" or dim == "yx":
self.track_dict = {"x":0,"y":1}
self.mean_index = [True, False, True, False, False, False]
+ elif dim == "xtau" or dim == "taux":
+ self.track_dict = {"x":0,"tau":1}
+ self.mean_index = [True, False, False, False, True, False]
+ elif dim == "ytau" or dim == "tauy":
+ self.track_dict = {"y":0,"tau":1}
+ self.mean_index = [False, False, True, False, True, False]
else:
raise ValueError("dim is not correct.")
@@ -1322,23 +1110,34 @@ class BeamSpectrumMonitor(Monitor):
Parameters
----------
ring : Synchrotron object
+ save_every : int or float
+ Set the frequency of the save. The spectrums are computed every
+ save_every call of the montior.
+ buffer_size : int or float
+ Size of the save buffer.
+ total_size : int or float
+ Total size of the save. The following relationships between the
+ parameters must exist:
+ total_size % buffer_size == 0
+ number of call to track / save_every == total_size - 1
+ dim : str, optional
+ Dimensions in which the spectrums have to be computed.
+ Can be:
+ - "all"
+ - "tau"
+ - "x"
+ - "y"
+ - "xy" or "yx"
+ - "xtau" or "taux"
+ - "ytau" or "tauy"
n_fft : int or float, optional
The number of points used for FFT computation, if n_fft is bigger than
save_every zero-padding is applied.
+ If None, save_every is used.
file_name : string, optional
Name of the HDF5 where the data will be stored. Must be specified
the first time a subclass of Monitor is instancied and must be None
the following times.
- save_every : int or float, optional
- Set the frequency of the save. The spectrum is computed saved every
- save_every call of the montior.
- buffer_size : int or float, optional
- Size of the save buffer.
- total_size : int or float, optional
- Total size of the save. The following relationships between the
- parameters must exist:
- total_size % buffer_size == 0
- number of call to track / save_every == total_size - 1
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -1351,20 +1150,21 @@ class BeamSpectrumMonitor(Monitor):
signal_resolution : float
Return the signal resolution in [Hz].
frequency_samples : array of float
- Return the fft frequency samples in [Hz].
+ Return the fft frequency samples in [Hz].
Methods
-------
track(bunch):
- Save tune and/or FFT data.
+ Save spectrum data.
"""
- def __init__(self, ring, n_fft=10000,
- file_name=None, save_every=10,
- buffer_size=5, total_size=10, mpi_mode=True,
- dim="all"):
+ def __init__(self, ring, save_every, buffer_size, total_size, dim="all",
+ n_fft=None, file_name=None, mpi_mode=False):
+ if n_fft is None:
+ n_fft = int(save_every)
+
self.n_fft = int(n_fft)
self.store_dict = {"x":0,"y":1,"tau":2}
@@ -1380,6 +1180,17 @@ class BeamSpectrumMonitor(Monitor):
elif dim == "y":
self.track_dict = {"y":0}
self.mean_index = [False, False, True, False, False, False]
+ elif dim == "xy" or dim == "yx":
+ self.track_dict = {"x":0,"y":1}
+ self.mean_index = [True, False, True, False, False, False]
+ elif dim == "xtau" or dim == "taux":
+ self.track_dict = {"x":0,"tau":1}
+ self.mean_index = [True, False, False, False, True, False]
+ elif dim == "ytau" or dim == "tauy":
+ self.track_dict = {"y":0,"tau":1}
+ self.mean_index = [False, False, True, False, True, False]
+ else:
+ raise ValueError("dim is not correct.")
self.size_list = len(self.track_dict)
@@ -1513,20 +1324,20 @@ class CavityMonitor(Monitor):
cavity_name : str
Name of the CavityResonator object to monitor.
ring : Synchrotron object
- file_name : string, optional
- Name of the HDF5 where the data will be stored. Must be specified
- the first time a subclass of Monitor is instancied and must be None
- the following times.
- save_every : int or float, optional
+ save_every : int or float
Set the frequency of the save. The data is saved every save_every
call of the montior.
- buffer_size : int or float, optional
+ buffer_size : int or float
Size of the save buffer.
- total_size : int or float, optional
+ total_size : int or float
Total size of the save. The following relationships between the
parameters must exist:
total_size % buffer_size == 0
number of call to track / save_every == total_size
+ file_name : string, optional
+ Name of the HDF5 where the data will be stored. Must be specified
+ the first time a subclass of Monitor is instancied and must be None
+ the following times.
mpi_mode : bool, optional
If True, open the HDF5 file in parallel mode, which is needed to
allow several cores to write in the same file at the same time.
@@ -1538,8 +1349,8 @@ class CavityMonitor(Monitor):
Save data
"""
- def __init__(self, cavity_name, ring, file_name=None, save_every=5,
- buffer_size=500, total_size=2e4, mpi_mode=True):
+ def __init__(self, cavity_name, ring, save_every, buffer_size, total_size,
+ file_name=None, mpi_mode=False):
self.cavity_name = cavity_name
self.ring = ring