diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index adc490ba661610646c8ad5831362e407daf7984f..126dbbcafa9e85bb2e76d85b7fa91319420d9c9b 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -10,12 +10,17 @@ from mbtrack2.tracking.monitors.monitors import (Monitor, BunchMonitor,
BeamMonitor,
ProfileMonitor,
WakePotentialMonitor,
- TuneMonitor,
- CavityMonitor)
+ CavityMonitor,
+ BunchSpectrumMonitor,
+ BeamSpectrumMonitor)
from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_phasespacedata,
plot_profiledata,
plot_beamdata,
plot_wakedata,
- plot_tunedata,
- plot_cavitydata)
\ No newline at end of file
+ plot_cavitydata,
+ streak_beamdata,
+ plot_bunchspectrum,
+ streak_bunchspectrum)
+
+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 0d3b9a4cde683c0c301c793ed69d36bbc70a226c..a543fdf91704dda41a9958b3e2cb9ce05735a681 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,52 +834,106 @@ 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.
+ 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
+ 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 tune 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.
If False, open the HDF5 file in standard mode.
+ Attributes
+ ----------
+ fft_resolution : float
+ Return the fft resolution in [Hz].
+ signal_resolution : float
+ Return the signal resolution in [Hz].
+ frequency_samples : array of float
+ Return the fft frequency samples in [Hz].
+
Methods
-------
track(bunch):
- Save tune and/or FFT data.
+ Save spectrum 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):
+ 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:
+ self.n_fft = int(save_every)
+ else:
+ self.n_fft = int(n_fft)
+
+ self.sample_size = int(sample_size)
+ self.store_dict = {"x":0,"y":1,"tau":2}
+
+ if dim == "all":
+ self.track_dict = {"x":0,"y":1,"tau":2}
+ self.mean_index = [True, False, True, False, True, False]
+ elif dim == "tau":
+ self.track_dict = {"tau":0}
+ self.mean_index = [False, False, False, False, True, False]
+ elif dim == "x":
+ self.track_dict = {"x":0}
+ self.mean_index = [True, False, False, False, False, False]
+ 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)
self.ring = ring
self.bunch_number = bunch_number
- group_name = "TuneData_" + str(self.bunch_number)
+ group_name = "BunchSpectrum_" + 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)}
+ dict_buffer = {"incoherent":(3, self.n_fft//2+1, buffer_size),
+ "coherent":(3, self.n_fft//2+1, buffer_size),
+ "mean_incoherent":(3,buffer_size),
+ "std_incoherent":(3,buffer_size)}
+ dict_file = {"incoherent":(3, self.n_fft//2+1, total_size),
+ "coherent":(3, self.n_fft//2+1, total_size),
+ "mean_incoherent":(3,total_size),
+ "std_incoherent":(3,total_size)}
self.monitor_init(group_name, save_every, buffer_size, total_size,
dict_buffer, dict_file, file_name, mpi_mode)
@@ -889,28 +941,48 @@ class TuneMonitor(Monitor):
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))
+ self.save_count = 0
+
+ self.positions = np.zeros((self.size_list, self.sample_size, self.save_every+1))
+ self.mean = np.zeros((self.size_list, self.save_every+1))
index = np.arange(0, int(mp_number))
self.index_sample = sorted(random.sample(list(index), self.sample_size))
+
+ self.incoherent = np.zeros((3, self.n_fft//2+1, self.buffer_size))
+ self.coherent = np.zeros((3, self.n_fft//2+1, self.buffer_size))
- self.save_count = 0
- self.buffer_count = 0
+ self.file[self.group_name].create_dataset(
+ "freq", data=self.frequency_samples)
+
+ @property
+ def fft_resolution(self):
+ """
+ Return the fft resolution in [Hz].
- self.save_tune = save_tune
- self.save_fft = save_fft
- self.save_every = save_every
+ It is defined as the sampling frequency over the number of samples.
+ """
+ return self.ring.f0/self.n_fft
+
+ @property
+ def signal_resolution(self):
+ """
+ Return the signal resolution in [Hz].
- if self.save_fft is True :
- self.n_fft = n_fft
- self.fourier_save = np.zeros((3, self.n_fft//2+1, buffer_size))
+ It is defined as the inverse of the signal length.
+ """
+ return 1/(self.ring.T0*self.save_every)
+
+ @property
+ def frequency_samples(self):
+ """
+ Return the fft frequency samples in [Hz].
+ """
+ return rfftfreq(self.n_fft, self.ring.T0)
def track(self, object_to_save):
"""
- Save tune data.
+ Save spectrum data.
Parameters
----------
@@ -932,9 +1004,10 @@ class TuneMonitor(Monitor):
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]
+ for key, value in self.track_dict.items():
+ self.positions[value, :, self.save_count] = bunch[key][self.index_sample]
+
+ self.mean[:, self.save_count] = bunch.mean[self.mean_index]
self.save_count += 1
@@ -942,7 +1015,7 @@ class TuneMonitor(Monitor):
self.to_buffer(bunch)
self.save_count = 0
- self.track_count += 1
+ self.track_count += 1
def to_buffer(self, bunch):
"""
@@ -952,16 +1025,12 @@ class TuneMonitor(Monitor):
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
+ for key, value in self.track_dict.items():
+ incoherent, mean_incoherent, std_incoherent = self.get_incoherent_spectrum(self.positions[value,:,:])
+ self.incoherent[self.store_dict[key],:,self.buffer_count] = incoherent
+ self.mean_incoherent[self.store_dict[key],self.buffer_count] = mean_incoherent
+ self.std_incoherent[self.store_dict[key],self.buffer_count] = std_incoherent
+ self.coherent[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.mean[value])
self.buffer_count += 1
@@ -977,81 +1046,276 @@ class TuneMonitor(Monitor):
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[self.group_name]["incoherent"][:,:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.incoherent
+ self.file[self.group_name]["mean_incoherent"][:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.mean_incoherent
+ self.file[self.group_name]["std_incoherent"][:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.std_incoherent
+ self.file[self.group_name]["coherent"][:,:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.coherent
self.file.flush()
self.write_count += 1
+
+ def get_incoherent_spectrum(self, positions):
+ """
+ Compute the incoherent spectrum i.e. the average of the absolute value
+ of the FT of the position of every particule of the bunch.
+
+ Returns
+ -------
+ incoherent : array
+ Bunch incoherent spectrum.
+ mean_incoherent : float
+ Mean frequency of the maximum of each individual particle spectrum
+ in [Hz].
+ std_incoherent : float
+ Standard deviation of the frequency of the maximum of each
+ individual particle spectrum in [Hz].
- def get_tune(self, bunch):
"""
- Compute tune by using NAFF algorithm to indentify the fundamental
- harmonic frequency of the particles' motion.
+ fourier = rfft(positions, n=self.n_fft)
+ fourier_abs = np.abs(fourier)
+ max_array = np.argmax(fourier_abs,axis=1)
+ freq_array = self.frequency_samples[max_array]
+ mean_incoherent = np.mean(freq_array)
+ std_incoherent = np.std(freq_array)
+ incoherent = np.mean(fourier_abs, axis=0)
+
+ return incoherent, mean_incoherent, std_incoherent
+
+ def get_coherent_spectrum(self, mean):
+ """
+ Compute the coherent spectrum i.e. the absolute value of the FT of the
+ mean position of the bunch.
+
+ Returns
+ -------
+ coherent : array
+ Bunch coherent spectrum.
+
+ """
+ coherent = np.abs(rfft(mean, n=self.n_fft))
+
+ return coherent
+
+class BeamSpectrumMonitor(Monitor):
+ """
+ Monitor coherent beam spectrum.
+
+ 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.
+ 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.
+
+ Attributes
+ ----------
+ fft_resolution : float
+ Return the fft resolution in [Hz].
+ signal_resolution : float
+ Return the signal resolution in [Hz].
+ frequency_samples : array of float
+ Return the fft frequency samples in [Hz].
+
+ Methods
+ -------
+ track(bunch):
+ Save spectrum data.
+
+ """
+
+ 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:
+ self.n_fft = int(save_every)
+ else:
+ self.n_fft = int(n_fft)
+
+ self.store_dict = {"x":0,"y":1,"tau":2}
+
+ if dim == "all":
+ self.track_dict = {"x":0,"y":1,"tau":2}
+ self.mean_index = [True, False, True, False, True, False]
+ elif dim == "tau":
+ self.track_dict = {"tau":0}
+ self.mean_index = [False, False, False, False, True, False]
+ elif dim == "x":
+ self.track_dict = {"x":0}
+ self.mean_index = [True, False, False, False, False, False]
+ 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)
+
+ self.ring = ring
+ group_name = "BeamSpectrum"
+
+ dict_buffer = {"coherent":(3, self.n_fft//2+1, buffer_size)}
+ dict_file = {"coherent":(3, self.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.save_count = 0
+
+ self.mean = np.zeros((self.size_list, ring.h, self.save_every))
+ self.coherent = np.zeros((3, self.n_fft//2+1, self.buffer_size))
+
+ self.file[self.group_name].create_dataset(
+ "freq", data=self.frequency_samples)
+
+ @property
+ def fft_resolution(self):
+ """
+ Return the fft resolution in [Hz].
+
+ It is defined as the sampling frequency over the number of samples.
+ """
+ return self.ring.f1/self.n_fft
+
+ @property
+ def signal_resolution(self):
+ """
+ Return the signal resolution in [Hz].
+
+ It is defined as the inverse of the signal length.
+ """
+ return 1/(self.ring.T0*self.save_every)
+
+ @property
+ def frequency_samples(self):
+ """
+ Return the fft frequency samples in [Hz].
+ """
+ return rfftfreq(self.n_fft, self.ring.T1)
+
+ def track(self, beam):
+ """
+ Save mean data.
Parameters
----------
- bunch : Bunch object
+ beam : Beam object
"""
+ if (beam.mpi_switch == True):
+ bunch_num = beam.mpi.bunch_num
+ bunch = beam[bunch_num]
+ self.mean[:, bunch_num, self.save_count] = bunch.mean[self.mean_index]
+ else:
+ self.mean[:, :, self.save_count] = beam.bunch_mean[self.mean_index,:]
+
+ self.save_count += 1
+
+ if self.save_count == self.save_every:
+ self.to_buffer(beam)
+ self.save_count = 0
- 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
+ self.track_count += 1
- 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
+ def to_buffer(self, beam):
+ """
+ A method to hold saved data before writing it to the output file.
+
+ """
- tune_single_particle = freq / self.ring.f0
- mean = np.nanmean(tune_single_particle, 0)
- spread = np.nanstd(tune_single_particle, 0)
+ self.time[self.buffer_count] = self.track_count
+
+ for key, value in self.track_dict.items():
+ if (beam.mpi_switch == True):
+ data_core = self.mean[value, beam.mpi.bunch_num, :]
+ full_data = beam.mpi.comm.allgather(data_core)
+ data = np.reshape(full_data, (-1), 'F')
+ else:
+ data = np.reshape(self.mean[value, :, :], (-1), 'F')
+ self.coherent[self.store_dict[key],:,self.buffer_count] = self.get_beam_spectrum(data)
+ self.buffer_count += 1
- return (mean, spread)
+ 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
+
+ self.file[self.group_name]["coherent"][:,:,
+ self.write_count * self.buffer_size:(self.write_count+1) *
+ self.buffer_size] = self.coherent
+
+ self.file.flush()
+ self.write_count += 1
- def get_fft(self):
+ def get_beam_spectrum(self, mean):
"""
- Compute the Fourier transform (using FFT algorithm) of the
- osciallation in horizontal, vertical, and longitudinal plane.
+ Compute the beam coherent spectrum i.e. the absolute value of the FT
+ of the mean position of every bunch.
Returns
-------
- fourier_x_avg, fourier_y_avg, fourier_tau_avg : ndarray
- The average of the transformed input in each plane.
+ coherent : array
+ The beam coherent spectrum.
"""
- 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)
+ coherent = np.abs(rfft(mean, 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)
+ return coherent
class CavityMonitor(Monitor):
"""
@@ -1062,20 +1326,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.
@@ -1087,8 +1351,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
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 780dc4875b6fb2074f1132150677d197f016f7a4..0c648544ec5c7a4f9749e4ef0c3022cf79bc308b 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -3,8 +3,8 @@
Module for plotting the data recorded by the monitor module during the
tracking.
-@author: Watanyu Foosang
-@Date: 10/04/2020
+@author: Watanyu Foosang, Alexis Gamelin
+@Date: 20/07/2021
"""
import numpy as np
@@ -13,17 +13,17 @@ import matplotlib as mpl
import seaborn as sns
import h5py as hp
import random
-from scipy.fft import rfftfreq
+from scipy.stats import gmean
-def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean",
- x_var="time", turn=None, plot_type="normal", cm_lim=None):
+def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
+ x_var="time", turn=None, legend=None):
"""
- Plot data recorded by BeamMonitor.
+ Plot 2D data recorded by BeamMonitor.
Parameters
----------
- filename : str
- Name of the HDF5 file that contains the data.
+ filenames : str or list of str
+ Names of the HDF5 files to be plotted.
dataset : {"current","emit","mean","std"}
HDF5 file's dataset to be plotted. The default is "mean".
dimension : str
@@ -35,18 +35,14 @@ def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean",
stat_var : {"mean", "std"}
Statistical value of the dimension.
x_var : {"time", "index"}
- Choice of the horizontal axis for "normal" plot_type.
- "time" corresponds to turn number.
- "index" corresponds to bunch index.
+ Choice of the horizontal axis:
+ "time" corresponds to turn number.
+ "index" corresponds to bunch index.
turn : int or float, optional
- Choice of the turn to plot for the "normal" plot_type with "index".
+ Choice of the turn to plot when x_var = "index".
If None, the last turn is plotted.
- plot_type : {"normal","streak"}, optional
- Choice of the type of plot.
- "normal" is for the standard line 2D plot.
- "streak" is for the 3D like image.
- cm_lim : list [vmin, vmax], optional
- Colormap scale for the "streak" plot.
+ legend : list of str, optional
+ Legend to add for each file.
Return
------
@@ -54,17 +50,24 @@ def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean",
Figure object with the plot on it.
"""
- file = hp.File(filename, "r")
- data = file["Beam"]
- time = np.array(data["time"])
- if plot_type == "normal":
+ if isinstance(filenames, str):
+ filenames = [filenames]
+ fig, ax = plt.subplots()
+
+ for filename in filenames:
+ file = hp.File(filename, "r")
+ data = file["Beam"]
+ time = np.array(data["time"])
+
if x_var == "time":
x = time
x_label = "Number of turns"
+ bunch_index = data["current"][:,0] != 0
+
if dataset == "current":
- y = np.nansum(data["current"][:],0)*1e3
+ y = np.nansum(data[dataset][bunch_index,:],0)*1e3
y_label = "Total current (mA)"
elif dataset == "emit":
dimension_dict = {"x":0, "y":1, "s":2}
@@ -73,9 +76,9 @@ def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean",
"$\\epsilon_{y}$ (m.rad)",
"$\\epsilon_{s}$ (m.rad)"]
if stat_var == "mean":
- y = np.nanmean(data["emit"][axis,:],0)
+ y = np.nanmean(data[dataset][axis,bunch_index,:],0)
elif stat_var == "std":
- y = np.nanstd(data["emit"][axis,:],0)
+ y = np.nanstd(data[dataset][axis,bunch_index,:],0)
y_label = stat_var + " " + label[axis]
elif dataset == "mean" or dataset == "std":
dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4,
@@ -85,9 +88,9 @@ def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean",
label = ["x (um)", "x' ($\\mu$rad)", "y (um)",
"y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
if stat_var == "mean":
- y = np.nanmean(data[dataset][axis,:],0)*scale[axis]
+ y = np.nanmean(data[dataset][axis,bunch_index,:],0)*scale[axis]
elif stat_var == "std":
- y = np.nanstd(data[dataset][axis,:],0)*scale[axis]
+ y = np.nanstd(data[dataset][axis,bunch_index,:],0)*scale[axis]
label_sup = {"mean":"mean of ", "std":"std of "}
y_label = label_sup[stat_var] + dataset + " " + label[axis]
@@ -122,60 +125,109 @@ def plot_beamdata(filename, dataset="mean", dimension="tau", stat_var="mean",
else:
raise ValueError("x_var should be time or index")
- fig, ax = plt.subplots()
ax.plot(x, y)
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
+ if legend is not None:
+ plt.legend(legend)
+
+ file.close()
- elif plot_type == "streak":
- h = len(data["mean"][0,:,0])
- x = np.arange(h)
- x_label = "Bunch index"
- y = time
- y_label = "Number of turns"
- if dataset == "current":
- z = np.array(data["current"]*1e3).T
- z_label = "Bunch current (mA)"
- elif dataset == "emit":
- dimension_dict = {"x":0, "y":1, "s":2}
- axis = dimension_dict[dimension]
- label = ["$\\epsilon_{x}$ (m.rad)",
- "$\\epsilon_{y}$ (m.rad)",
- "$\\epsilon_{s}$ (m.rad)"]
- z = np.array(data["emit"][axis,:,:]).T
- z_label = label[axis]
- else:
- dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4,
- "delta":5}
- axis = dimension_dict[dimension]
- scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
- label = ["x (um)", "x' ($\\mu$rad)", "y (um)",
- "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
- z = np.array(data[dataset][axis,:,:]).T*scale[axis]
- z_label = label[axis]
+ return fig
+
+def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None):
+ """
+ Plot 3D data recorded by BeamMonitor.
+
+ Parameters
+ ----------
+ filename : str
+ Name of the HDF5 file that contains the data.
+ dataset : {"current","emit","mean","std"}
+ HDF5 file's dataset to be plotted. The default is "mean".
+ dimension : str
+ The dimension of the dataset to plot:
+ for "emit", dimension = {"x","y","s"},
+ for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}.
+ not used if "current".
+ The default is "tau".
+ cm_lim : list [vmin, vmax], optional
+ Colormap scale for the "streak" plot.
+
+ Return
+ ------
+ fig : Figure
+ Figure object with the plot on it.
+
+ """
+
+ file = hp.File(filename, "r")
+ data = file["Beam"]
+ time = np.array(data["time"])
- fig, ax = plt.subplots()
- ax.set_xlabel(x_label)
- ax.set_ylabel(y_label)
- cmap = mpl.cm.cool
- c = ax.imshow(z, cmap=cmap, origin='lower' , aspect='auto',
- extent=[x.min(), x.max(), y.min(), y.max()])
- if cm_lim is not None:
- c.set_clim(vmin=cm_lim[0],vmax=cm_lim[1])
- cbar = fig.colorbar(c, ax=ax)
- cbar.set_label(z_label)
+ h = len(data["mean"][0,:,0])
+ x = np.arange(h)
+ x_label = "Bunch index"
+ y = time
+ y_label = "Number of turns"
+ if dataset == "current":
+ z = (np.array(data["current"])*1e3).T
+ z_label = "Bunch current (mA)"
+ title = z_label
+ elif dataset == "emit":
+ dimension_dict = {"x":0, "y":1, "s":2}
+ axis = dimension_dict[dimension]
+ label = ["$\\epsilon_{x}$ (m.rad)",
+ "$\\epsilon_{y}$ (m.rad)",
+ "$\\epsilon_{s}$ (m.rad)"]
+ z = np.array(data["emit"][axis,:,:]).T
+ z_label = label[axis]
+ title = z_label
+ else:
+ dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4,
+ "delta":5}
+ axis = dimension_dict[dimension]
+ scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
+ label = ["x (um)", "x' ($\\mu$rad)", "y (um)",
+ "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
+ z = np.array(data[dataset][axis,:,:]).T*scale[axis]
+ z_label = label[axis]
+ if dataset == "mean":
+ title = label[axis] + " CM"
+ elif dataset == "std":
+ title = label[axis] + " RMS"
+
+ fig, ax = plt.subplots()
+ ax.set_xlabel(x_label)
+ ax.set_ylabel(y_label)
+ ax.set_title(title)
+
+ if dataset == "mean":
+ cmap = mpl.cm.coolwarm # diverging
+ else:
+ cmap = mpl.cm.inferno # sequential
+
+ c = ax.imshow(z, cmap=cmap, origin='lower' , aspect='auto',
+ extent=[x.min(), x.max(), y.min(), y.max()])
+ if cm_lim is not None:
+ c.set_clim(vmin=cm_lim[0],vmax=cm_lim[1])
+ cbar = fig.colorbar(c, ax=ax)
+ cbar.set_label(z_label)
+
+ file.close()
return fig
-def plot_bunchdata(filename, bunch_number, dataset, dimension="x", x_var="time"):
+def plot_bunchdata(filenames, bunch_number, dataset, dimension="x",
+ legend=None):
"""
Plot data recorded by BunchMonitor.
Parameters
----------
- filename : str
- Name of the HDF5 file that contains the data.
- bunch_number : int
+ filenames : str or list of str
+ Names of the HDF5 files to be plotted.
+ bunch_number : int or list of int
Bunch to plot. This has to be identical to 'bunch_number' parameter in
'BunchMonitor' object.
dataset : {"current", "emit", "mean", "std", "cs_invariant"}
@@ -186,8 +238,8 @@ def plot_bunchdata(filename, bunch_number, dataset, dimension="x", x_var="time")
for "emit", dimension = {"x","y","s"},
for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"},
for "action", dimension = {"x","y"}.
- x_var : {"time", "current"}, optional
- Variable to be plotted on the horizontal axis. The default is "time".
+ legend : list of str, optional
+ Legend to add for each file.
Return
------
@@ -196,60 +248,69 @@ def plot_bunchdata(filename, bunch_number, dataset, dimension="x", x_var="time")
"""
- file = hp.File(filename, "r")
-
- group = "BunchData_{0}".format(bunch_number) # Data group of the HDF5 file
-
- if dataset == "current":
- y_var = file[group][dataset][:]*1e3
- label = "current (mA)"
-
- elif dataset == "emit":
- dimension_dict = {"x":0, "y":1, "s":2}
-
- y_var = file[group][dataset][dimension_dict[dimension]]*1e9
-
- if dimension == "x": label = "hor. emittance (nm.rad)"
- elif dimension == "y": label = "ver. emittance (nm.rad)"
- elif dimension == "s": label = "long. emittance (nm.rad)"
+ if isinstance(filenames, str):
+ filenames = [filenames]
+ if isinstance(bunch_number, int):
+ ll = []
+ for i in range(len(filenames)):
+ ll.append(bunch_number)
+ bunch_number = ll
- elif dataset == "mean" or dataset == "std":
- dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
- scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
- axis_index = dimension_dict[dimension]
+ fig, ax = plt.subplots()
+
+ for i, filename in enumerate(filenames):
+ file = hp.File(filename, "r")
+ group = "BunchData_{0}".format(bunch_number[i]) # Data group of the HDF5 file
- y_var = file[group][dataset][axis_index]*scale[axis_index]
- if dataset == "mean":
- label_list = ["x ($\\mu$m)", "x' ($\\mu$rad)", "y ($\\mu$m)",
- "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
- else:
- label_list = ["$\\sigma_x$ ($\\mu$m)", "$\\sigma_{x'}$ ($\\mu$rad)",
- "$\\sigma_y$ ($\\mu$m)", "$\\sigma_{y'}$ ($\\mu$rad)",
- "$\\sigma_{\\tau}$ (ps)", "$\\sigma_{\\delta}$"]
-
- label = label_list[axis_index]
-
- elif dataset == "cs_invariant":
- dimension_dict = {"x":0, "y":1}
- axis_index = dimension_dict[dimension]
- y_var = file[group][dataset][axis_index]
- label_list = ['$J_x$ (m)', '$J_y$ (m)']
- label = label_list[axis_index]
-
- if x_var == "current":
- x_axis = file[group]["current"][:] * 1e3
- xlabel = "current (mA)"
- elif x_var == "time":
+ if dataset == "current":
+ y_var = file[group][dataset][:]*1e3
+ label = "current (mA)"
+
+ elif dataset == "emit":
+ dimension_dict = {"x":0, "y":1, "s":2}
+
+ y_var = file[group][dataset][dimension_dict[dimension]]*1e9
+
+ if dimension == "x": label = "hor. emittance (nm.rad)"
+ elif dimension == "y": label = "ver. emittance (nm.rad)"
+ elif dimension == "s": label = "long. emittance (nm.rad)"
+
+
+ elif dataset == "mean" or dataset == "std":
+ dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5}
+ scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
+ axis_index = dimension_dict[dimension]
+
+ y_var = file[group][dataset][axis_index]*scale[axis_index]
+ if dataset == "mean":
+ label_list = ["x ($\\mu$m)", "x' ($\\mu$rad)", "y ($\\mu$m)",
+ "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"]
+ else:
+ label_list = ["$\\sigma_x$ ($\\mu$m)", "$\\sigma_{x'}$ ($\\mu$rad)",
+ "$\\sigma_y$ ($\\mu$m)", "$\\sigma_{y'}$ ($\\mu$rad)",
+ "$\\sigma_{\\tau}$ (ps)", "$\\sigma_{\\delta}$"]
+
+ label = label_list[axis_index]
+
+ elif dataset == "cs_invariant":
+ dimension_dict = {"x":0, "y":1}
+ axis_index = dimension_dict[dimension]
+ y_var = file[group][dataset][axis_index]
+ label_list = ['$J_x$ (m)', '$J_y$ (m)']
+ label = label_list[axis_index]
+
x_axis = file[group]["time"][:]
- xlabel = "number of turns"
+ xlabel = "Number of turns"
+
+ ax.plot(x_axis, y_var)
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(label)
+ if legend is not None:
+ plt.legend(legend)
+
+ file.close()
- fig, ax = plt.subplots()
- ax.plot(x_axis,y_var)
- ax.set_xlabel(xlabel)
- ax.set_ylabel(label)
-
- file.close()
return fig
def plot_phasespacedata(filename, bunch_number, x_var, y_var, turn,
@@ -568,120 +629,398 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
elif streak_plot is True:
return fig2
-def plot_tunedata(filename, bunch_number, f0, plot_tune=True, plot_fft=False,
- dimension='x', min_tune=0, max_tune=0.5, min_turn=None,
- max_turn=None, streak_plot=True, profile_plot=False):
+def plot_bunchspectrum(filenames, bunch_number, dataset="incoherent", dim="tau",
+ turns=None, fs=None, log_scale=True, legend=None,
+ norm=False):
"""
- Plot data recorded by TuneMonitor.
-
+ Plot coherent and incoherent spectrum data.
+
Parameters
----------
- filename : str
- Name of the HDF5 file that contains the data.
- bunch_number : int
+ filenames : str or list of str
+ Names of the HDF5 files to be plotted.
+ bunch_number : int or list of int
Bunch to plot. This has to be identical to 'bunch_number' parameter in
- 'BunchMonitor' object.
- f0 : float
- Revolution frequency of the ring used for the tracking in [Hz].
- plot_tune : bool, optional
- If True, tune data usinf NAFF is plotted.
- plot_fft : bool, optional
- If True, FFT data is plotted.
- dimension : {'x', 'y', 's'}
- Option to plot FFT data in horizontal, vertical, or longitudinal plane.
- min_tune, max_tune : int or float, optional
- The minimum and the maximum tune values to plot FFT data.
- min_turn, max_turn : int or float, optional
- The minimum and the maximum number of turns to plot FFT data.
- streak_plot : bool, optional
- If True, the FFT data is plotted as a streak plot.
- profile_plot : bool, optional.
- If True, the FFT data is plotted as line plots.
-
+ 'BunchSpectrumMonitor' object.
+ dataset : {"mean_incoherent", "coherent", "incoherent"}
+ HDF5 file's dataset to be plotted.
+ The default is "incoherent".
+ dim : {"x","y","tau"}, optional
+ The dimension of the dataset to plot.
+ The default is "tau".
+ turns : array or None, optional
+ Numbers of the turns to plot.
+ If None, all turns are shown.
+ The default is None.
+ fs : float or None, optional
+ If not None, the frequency axis is noramlised by fs.
+ The default is None.
+ log_scale : bool, optional
+ If True, the spectrum plots are shown in y-log scale.
+ The default is True.
+ legend : list of str, optional
+ Legend to add for each file.
+ The default is None.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
+
Return
------
fig : Figure
- Figure objects with the plot on it.
+
+ """
+
+ if isinstance(filenames, str):
+ filenames = [filenames]
+
+ if isinstance(bunch_number, int):
+ ll = []
+ for i in range(len(filenames)):
+ ll.append(bunch_number)
+ bunch_number = ll
+
+ fig, ax = plt.subplots()
+
+ for i, filename in enumerate(filenames):
+ file = hp.File(filename, "r")
+ group = file["BunchSpectrum_{0}".format(bunch_number[i])]
+
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
+
+ if dataset == "mean_incoherent":
+ y_var = group["mean_incoherent"][dim_dict[dim],:]
+ y_err = group["std_incoherent"][dim_dict[dim],:]
+ ax.errorbar(time, y_var, y_err)
+ xlabel = "Turn number"
+ ylabel = "Mean incoherent frequency [Hz]"
+ elif dataset == "incoherent" or dataset == "coherent":
+
+ if turns is None:
+ turn_index = np.where(time == time)[0]
+ else:
+ turn_index = []
+ for turn in turns:
+ idx = np.where(time == turn)[0][0]
+ turn_index.append(idx)
+ turn_index = np.array(turn_index)
+
+ if fs is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/fs
+ xlabel = r"$f/f_{s}$"
+
+ for idx in turn_index:
+ y_var = group[dataset][dim_dict[dim],:,idx]
+ if norm is True:
+ y_var = y_var/gmean(y_var)
+ ax.plot(x_var, y_var)
+
+ if log_scale is True:
+ ax.set_yscale('log')
+
+ ylabel = "FFT amplitude [a.u.]"
+ if dataset == "incoherent":
+ ax.set_title("Incoherent spectrum")
+ elif dataset == "coherent":
+ ax.set_title("Coherent spectrum")
+
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+ if legend is not None:
+ plt.legend(legend)
+ file.close()
+
+ return fig
+
+def streak_bunchspectrum(filename, bunch_number, dataset="incoherent",
+ dim="tau", fs=None, log_scale=True, fmin=None,
+ fmax=None, turns=None, norm=False):
+ """
+ Plot 3D data recorded by the BunchSpectrumMonitor.
+
+ Parameters
+ ----------
+ filenames : str
+ Name of the HDF5 file to be plotted.
+ bunch_number : int
+ Bunch to plot. This has to be identical to 'bunch_number' parameter in
+ 'BunchSpectrumMonitor' object.
+ dataset : {"coherent", "incoherent"}
+ HDF5 file's dataset to be plotted.
+ The default is "incoherent".
+ dim : {"x","y","tau"}, optional
+ The dimension of the dataset to plot.
+ The default is "tau".
+ fs : float or None, optional
+ If not None, the frequency axis is noramlised by fs.
+ log_scale : bool, optional
+ If True, the spectrum plots are shown in y-log scale.
+ The default is True.
+ fmin : float, optional
+ If not None, the plot is limitted to values bigger than fmin.
+ fmax : float, optional
+ If not None, the plot is limitted to values smaller than fmax.
+ turns : array, optional
+ If not None, only the turn numbers in the turns array are plotted.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
+
+ Returns
+ -------
+ fig : Figure
"""
file = hp.File(filename, "r")
+ group = file["BunchSpectrum_{0}".format(bunch_number)]
- group = "TuneData_{0}".format(bunch_number) # Data group of the HDF5 file
- time = file[group]["time"]
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
- fig_to_return = []
+ if turns is None:
+ turn_index = np.where(time == time)[0]
+ tmin = time.min()
+ tmax = time.max()
+ else:
+ tmin = turns.min()
+ tmax = turns.max()
+ turn_index = []
+ for turn in turns:
+ idx = np.where(time == turn)[0][0]
+ turn_index.append(idx)
+ turn_index = np.array(turn_index)
- if plot_tune is True:
- tune = file[group]["tune"]
- tune_spread = file[group]["tune_spread"]
-
- fig1, ax1 = plt.subplots()
- ax1.errorbar(x=time, y=tune[0,:], yerr=tune_spread[0,:])
- ax1.errorbar(x=time, y=tune[1,:], yerr=tune_spread[1,:])
- ax1.set_xlabel("Turn number")
- ax1.set_ylabel("Transverse tunes")
- plt.legend(["x","y"])
- fig_to_return.append(fig1)
-
- fig2, ax2 = plt.subplots()
- ax2.errorbar(x=time, y=tune[2,:], yerr=tune_spread[2,:])
- ax2.set_xlabel("Turn number")
- ax2.set_ylabel("Synchrotron tune")
- fig_to_return.append(fig2)
-
- if plot_fft is True:
-
- T0 = 1/f0
- n_freq = file[group]['fft'].shape[1]
- freq = rfftfreq((n_freq-1)*2, T0)
- tune_fft = freq / f0
-
- dimension_dict = {'x':0, 'y':1, 's':2}
- axis = dimension_dict[dimension]
+ data = group[dataset][dim_dict[dim], :, turn_index]
+
+ if log_scale is True:
+ option = mpl.colors.LogNorm()
+ else:
+ option = None
+
+ if fs is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/fs
+ xlabel = r"$f/f_{s}$"
- fourier_save = file[group]['fft'][axis]
+ if fmin is None:
+ fmin = x_var.min()
+ if fmax is None:
+ fmax = x_var.max()
+
+ ind = (x_var > fmin) & (x_var < fmax)
+ x_var=x_var[ind]
+ data = data[ind,:]
+
+ if norm is True:
+ data = data/gmean(data)
- if max_turn is None:
- max_turn = time[-1]
- if min_turn is None:
- min_turn = time[0]
+ ylabel = "Turn number"
+
+ fig, ax = plt.subplots()
+ if dataset == "incoherent":
+ ax.set_title("Incoherent spectrum")
+ elif dataset == "coherent":
+ ax.set_title("Coherent spectrum")
+
+ cmap = mpl.cm.inferno # sequential
+ c = ax.imshow(data.T, cmap=cmap, origin='lower' , aspect='auto',
+ extent=[x_var.min(), x_var.max(), tmin, tmax],
+ norm=option)
+ cbar = fig.colorbar(c, ax=ax)
+ cbar.ax.set_ylabel("FFT amplitude [a.u.]", rotation=270)
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+
+ return fig
+
+def plot_beamspectrum(filenames, dim="tau", turns=None, f0=None,
+ log_scale=True, legend=None, norm=False):
+ """
+ Plot coherent beam spectrum data.
+
+ Parameters
+ ----------
+ filenames : str or list of str
+ Names of the HDF5 files to be plotted.
+ dim : {"x","y","tau"}, optional
+ The dimension of the dataset to plot.
+ The default is "tau".
+ turns : array or None, optional
+ Numbers of the turns to plot.
+ If None, all turns are shown.
+ The default is None.
+ f0 : float or None, optional
+ If not None, the frequency axis is noramlised by f0.
+ The default is None.
+ log_scale : bool, optional
+ If True, the spectrum plots are shown in y-log scale.
+ The default is True.
+ legend : list of str, optional
+ Legend to add for each file.
+ The default is None.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
+
+ Return
+ ------
+ fig : Figure
+
+ """
+
+ if isinstance(filenames, str):
+ filenames = [filenames]
+
+ fig, ax = plt.subplots()
+
+ for i, filename in enumerate(filenames):
+ file = hp.File(filename, "r")
+ group = file["BeamSpectrum"]
+
+ dataset = "coherent"
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
+
+ if turns is None:
+ turn_index = np.where(time == time)[0]
+ else:
+ turn_index = []
+ for turn in turns:
+ idx = np.where(time == turn)[0][0]
+ turn_index.append(idx)
+ turn_index = np.array(turn_index)
- min_tune_iloc = np.where(tune_fft >= min_tune)[0][0]
- max_tune_iloc = np.where(tune_fft <= max_tune)[0][-1]
- save_every = int(time[1] - time[0])
- min_turn_iloc = min_turn // save_every - 1
- max_turn_iloc = max_turn // save_every
-
-
- if streak_plot is True:
- fig3, ax3 = plt.subplots()
- cmap = plt.get_cmap('Blues')
-
- c = ax3.imshow(np.transpose(np.log(
- fourier_save[min_tune_iloc:max_tune_iloc+1,
- min_turn_iloc:max_turn_iloc+1])),
- cmap=cmap, origin='lower' , aspect='auto',
- extent=[min_tune, max_tune, min_turn, max_turn])
- ax3.set_xlabel('$Q_{}$'.format(dimension))
- ax3.set_ylabel("Turns")
- cbar = fig3.colorbar(c, ax=ax3)
- cbar.set_label("log FFT amplitude")
- fig_to_return.append(fig3)
+ if f0 is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/f0
+ xlabel = r"$f/f_{0}$"
- if profile_plot is True:
- fig4, ax4 = plt.subplots()
- ax4.plot(tune_fft[min_tune_iloc:max_tune_iloc+1],
- fourier_save[min_tune_iloc:max_tune_iloc+1,
- min_turn_iloc:max_turn_iloc+1])
- ax4.set_xlabel('$Q_{}$'.format(dimension))
- ax4.set_ylabel("FFT amplitude")
- ax4.legend(time[min_turn_iloc:max_turn_iloc+1])
- fig_to_return.append(fig4)
+ for idx in turn_index:
+ y_var = group[dataset][dim_dict[dim],:,idx]
+ if norm is True:
+ y_var = y_var/gmean(y_var)
+ ax.plot(x_var, y_var)
- file.close()
-
- return tuple(fig_to_return)
+ if log_scale is True:
+ ax.set_yscale('log')
+
+ ylabel = "FFT amplitude [a.u.]"
+ ax.set_title("Beam coherent spectrum")
+
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+ if legend is not None:
+ plt.legend(legend)
+ file.close()
+
+ return fig
+
+def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None,
+ fmax=None, turns=None, norm=False):
+ """
+ Plot 3D data recorded by the BeamSpectrumMonitor.
+
+ Parameters
+ ----------
+ filenames : str
+ Name of the HDF5 file to be plotted.
+ dim : {"x","y","tau"}, optional
+ The dimension of the dataset to plot.
+ The default is "tau".
+ f0 : float or None, optional
+ If not None, the frequency axis is noramlised by f0.
+ log_scale : bool, optional
+ If True, the spectrum plots are shown in y-log scale.
+ The default is True.
+ fmin : float, optional
+ If not None, the plot is limitted to values bigger than fmin.
+ fmax : float, optional
+ If not None, the plot is limitted to values smaller than fmax.
+ turns : array, optional
+ If not None, only the turn numbers in the turns array are plotted.
+ norm : bool, optional
+ If True, normalise the data of each spectrum by its geometric mean.
+ The default is False.
+
+ Returns
+ -------
+ fig : Figure
+
+ """
+
+ file = hp.File(filename, "r")
+ group = file["BeamSpectrum"]
+ dataset="coherent"
+ time = np.array(group["time"])
+ freq = np.array(group["freq"])
+ dim_dict = {"x":0, "y":1, "tau":2}
+
+ if turns is None:
+ turn_index = np.where(time == time)[0]
+ tmin = time.min()
+ tmax = time.max()
+ else:
+ tmin = turns.min()
+ tmax = turns.max()
+ turn_index = []
+ for turn in turns:
+ idx = np.where(time == turn)[0][0]
+ turn_index.append(idx)
+ turn_index = np.array(turn_index)
+
+ data = group[dataset][dim_dict[dim], :, turn_index]
+
+ if log_scale is True:
+ option = mpl.colors.LogNorm()
+ else:
+ option = None
+
+ if f0 is None:
+ x_var = freq
+ xlabel = "Frequency [Hz]"
+ else:
+ x_var = freq/f0
+ xlabel = r"$f/f_{0}$"
+
+ if fmin is None:
+ fmin = x_var.min()
+ if fmax is None:
+ fmax = x_var.max()
+
+ ind = (x_var > fmin) & (x_var < fmax)
+ x_var=x_var[ind]
+ data = data[ind,:]
+
+ if norm is True:
+ data = data/gmean(data)
+
+ ylabel = "Turn number"
+
+ fig, ax = plt.subplots()
+ ax.set_title("Beam coherent spectrum")
+
+ cmap = mpl.cm.inferno # sequential
+ c = ax.imshow(data.T, cmap=cmap, origin='lower' , aspect='auto',
+ extent=[x_var.min(), x_var.max(), tmin, tmax],
+ norm=option)
+ cbar = fig.colorbar(c, ax=ax)
+ cbar.ax.set_ylabel("FFT amplitude [a.u.]", rotation=270)
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+
+ return fig
def plot_cavitydata(filename, cavity_name, phasor="cavity",
plot_type="bunch", bunch_number=0, turn=None, cm_lim=None):
diff --git a/tracking/monitors/tools.py b/tracking/monitors/tools.py
new file mode 100644
index 0000000000000000000000000000000000000000..9735a6e767e14fe22018747866eb2060ff59515c
--- /dev/null
+++ b/tracking/monitors/tools.py
@@ -0,0 +1,75 @@
+# -*- coding: utf-8 -*-
+"""
+This module defines utilities functions, helping to deals with tracking output
+and hdf5 files.
+
+@author: Alexis Gamelin
+"""
+
+import h5py as hp
+
+def merge_files(files_prefix, files_number, file_name=None):
+ """
+ Merge several hdf5 files into one.
+
+ The function assumes that the files to merge have names in the follwing
+ format:
+ - "files_prefix_0.hdf5"
+ - "files_prefix_1.hdf5"
+ ...
+ - "files_prefix_files_number.hdf5"
+
+ Parameters
+ ----------
+ files_prefix : str
+ Name of the files to merge.
+ files_number : int
+ Number of files to merge.
+ file_name : str, optional
+ Name of the file with the merged data. If None, files_prefix without
+ number is used.
+
+ """
+ if file_name == None:
+ file_name = files_prefix
+ f = hp.File(file_name + ".hdf5", "a")
+
+ ## Create file architecture
+ f0 = hp.File(files_prefix + "_" + str(0) + ".hdf5", "r")
+ for group in list(f0):
+ f.require_group(group)
+ for dataset_name in list(f0[group]):
+ if dataset_name == "freq":
+ f0[group].copy(dataset_name, f[group])
+ continue
+ shape = f0[group][dataset_name].shape
+ dtype = f0[group][dataset_name].dtype
+ shape_needed = list(shape)
+ shape_needed[-1] = shape_needed[-1]*files_number
+ shape_needed = tuple(shape_needed)
+ f[group].create_dataset(dataset_name, shape_needed, dtype)
+
+ f0.close()
+
+ ## Copy data
+ for i in range(files_number):
+ fi = hp.File(files_prefix + "_" + str(i) + ".hdf5", "r")
+ for group in list(fi):
+ for dataset_name in list(fi[group]):
+ shape = fi[group][dataset_name].shape
+ n_slice = int(len(shape) - 1)
+ length = shape[-1]
+ slice_list = []
+ for n in range(n_slice):
+ slice_list.append(slice(None))
+ slice_list.append(slice(length*i,length*(i+1)))
+ if (dataset_name == "freq"):
+ continue
+ if (dataset_name == "time") and (i != 0):
+ f[group][dataset_name][tuple(slice_list)] = f[group][dataset_name][(length*i) - 1] + fi[group][dataset_name]
+ else:
+ f[group][dataset_name][tuple(slice_list)] = fi[group][dataset_name]
+ fi.close()
+ f.close()
+
+
\ No newline at end of file