Skip to content
Snippets Groups Projects
Commit a7c73e5a authored by Alexis GAMELIN's avatar Alexis GAMELIN
Browse files

Misc monitor changes 

Remove TuneMonitor (remplaced by BunchSpectrumMonitor).
Remove default values for save_very, buffer_size, total_size.
Set default value for mpi_mode to False.
Add "xtau"/"ytau" choices for Bunch/BeamSpectrumMonitor classes.
Set default value for n_fft to None to avoid zero-padding by default.
Fix WakePotentialMontior n_bin (x2) problem.
parent 7d5bd52c
No related branches found
No related tags found
No related merge requests found
Show whitespace changes
Inline Side-by-side
tracking/monitors/__init__.py
+ 3
− 2
View file @ a7c73e5a
......@@ -10,7 +10,6 @@ from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor,
BeamMonitor,
ProfileMonitor,
WakePotentialMonitor,
TuneMonitor,
CavityMonitor,
BunchSpectrumMonitor,
BeamSpectrumMonitor)
......@@ -22,3 +21,5 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_tunedata,
plot_cavitydata,
streak_beamdata)
from mbtrack2.tracking.monitors.tools import merge_files
\ No newline at end of file
tracking/monitors/monitors.py
+ 130
− 319
View file @ a7c73e5a
......@@ -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 = {}
......@@ -822,10 +821,9 @@ class WakePotentialMonitor(Monitor):
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
----------
......@@ -837,253 +835,34 @@ class TuneMonitor(Monitor):
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
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.
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,9 +884,12 @@ 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)
......@@ -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.
......@@ -1356,14 +1155,15 @@ class BeamSpectrumMonitor(Monitor):
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
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment

The place to be for your source code