diff --git a/tracking/element.py b/tracking/element.py
index 2231c04f4a964b3bd753a2510e79d9cb50395ca2..9f18581f12ff118058a37f852ee119ba908168c9 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -153,6 +153,12 @@ class TransverseMap(Element):
self.gamma = self.ring.optics.local_gamma
self.dispersion = self.ring.optics.local_dispersion
self.phase_advance = self.ring.tune[0:2]*2*np.pi
+
+ if self.ring.adts is not None:
+ self.adts_poly = [np.poly1d(self.ring.adts[0]),
+ np.poly1d(self.ring.adts[1]),
+ np.poly1d(self.ring.adts[2]),
+ np.poly1d(self.ring.adts[3])]
@Element.parallel
def track(self, bunch):
@@ -166,9 +172,17 @@ class TransverseMap(Element):
bunch : Bunch or Beam object
"""
- # Compute phase adcence which depends on energy via chromaticity
- phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"])
- phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"])
+ # Compute phase advance which depends on energy via chromaticity and ADTS
+ if self.ring.adts is not None:
+ phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"]+
+ self.adts_poly[0](bunch['x'])+self.adts_poly[2](bunch['y']))
+
+ phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"]+
+ self.adts_poly[1](bunch['x'])+self.adts_poly[3](bunch['y']))
+
+ else:
+ phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"])
+ phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"])
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
matrix = np.zeros((6,6,len(bunch)))
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 01a2e217b660a2c8a2c8d0cf4352a0a72e5bb988..d1bb18eb100b0f4ffa9528b09806dc61a66dd7e1 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -17,6 +17,7 @@ from mbtrack2.tracking.particles import Bunch, Beam
from mbtrack2.tracking.rf import CavityResonator
from scipy.interpolate import interp1d
from abc import ABCMeta
+from scipy.fft import rfft, rfftfreq
from mpi4py import MPI
class Monitor(Element, metaclass=ABCMeta):
@@ -823,7 +824,8 @@ class WakePotentialMonitor(Monitor):
class TuneMonitor(Monitor):
"""
- Monitor tunes in horizontal, vertical, and longitudinal plane.
+ Monitor tunes and the Fourier transform (using FFT algorithm) of the
+ osciallation in horizontal, vertical, and longitudinal plane.
Parameters
----------
@@ -833,8 +835,14 @@ class TuneMonitor(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 computation. This number
- needs not to exceed mp_number.
+ 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.
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
@@ -857,19 +865,22 @@ class TuneMonitor(Monitor):
Methods
-------
track(bunch):
- Save tune data.
+ Save tune and/or FFT data.
"""
- def __init__(self, ring, bunch_number, mp_number, sample_size, 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_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,)}
- dict_file = {"tune":(3, total_size), "tune_spread":(3, total_size,)}
+
+ 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)
@@ -885,11 +896,17 @@ class TuneMonitor(Monitor):
index = np.arange(0, int(mp_number))
self.index_sample = sorted(random.sample(list(index), self.sample_size))
- self.save_every = save_every
self.save_count = 0
-
self.buffer_count = 1
+ 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.
@@ -920,10 +937,9 @@ class TuneMonitor(Monitor):
self.save_count += 1
- if self.track_count > 0:
- if self.track_count % self.save_every == 0:
- self.to_buffer(bunch)
- self.save_count = 0
+ 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
@@ -932,10 +948,19 @@ class TuneMonitor(Monitor):
A method to hold saved data before writing it to the output file.
"""
- mean, spread = self.get_tune(bunch)
+
self.time[self.buffer_count] = self.track_count
- self.tune[:, self.buffer_count] = mean
- self.tune_spread[:, self.buffer_count] = spread
+
+ 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
@@ -951,12 +976,18 @@ class TuneMonitor(Monitor):
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]["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_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
@@ -1000,6 +1031,27 @@ class TuneMonitor(Monitor):
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 CavityMonitor(Monitor):
"""
Monitor a CavityResonator object and save attributes (mean, std, emit and current).
@@ -1077,5 +1129,4 @@ class CavityMonitor(Monitor):
pass
else:
raise TypeError("cavity should be a CavityResonator object.")
- self.track_count += 1
-
\ No newline at end of file
+ self.track_count += 1
diff --git a/tracking/monitors/plotting.py b/tracking/monitors/plotting.py
index 9078150ef0c4aafd41cdd26ee0644fb3c61bbf21..2b06b97826ee5fa94b2343d25d42d4b06336d52a 100644
--- a/tracking/monitors/plotting.py
+++ b/tracking/monitors/plotting.py
@@ -13,6 +13,7 @@ import matplotlib as mpl
import seaborn as sns
import h5py as hp
import random
+from scipy.fft import rfftfreq
def plot_beamdata(filename, dataset, dimension=None, stat_var=None, x_var="time"):
"""
@@ -492,7 +493,9 @@ def plot_wakedata(filename, bunch_number, wake_type="Wlong", start=0,
elif streak_plot is True:
return fig2
-def plot_tunedata(filename, bunch_number):
+def plot_tunedata(filename, bunch_number, ring=None, 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):
"""
Plot data recorded by TuneMonitor.
@@ -503,6 +506,24 @@ def plot_tunedata(filename, bunch_number):
bunch_number : int
Bunch to plot. This has to be identical to 'bunch_number' parameter in
'BunchMonitor' object.
+ ring : Synchrotron object, optional
+ The ring configuration that is used in TuneMonitor. If None, the default
+ value of the revolution period and the revolution frequency are used,
+ which are 1.183 us and 0.845 MHz, respectively.
+ plot_tune : bool, optional
+ If True, tune data 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.
+ bunch_profile : bool, optional.
+ If True, the FFT data is plotted as line profiles.
Return
------
@@ -515,23 +536,96 @@ def plot_tunedata(filename, bunch_number):
group = "TuneData_{0}".format(bunch_number) # Data group of the HDF5 file
time = file[group]["time"]
- tune = file[group]["tune"]
- tune_spread = file[group]["tune_spread"]
-
- fig1, ax1 = plt.subplots()
- ax1.errorbar(x=time[1:], y=tune[0,1:], yerr=tune_spread[0,1:])
- ax1.errorbar(x=time[1:], y=tune[1,1:], yerr=tune_spread[1,1:])
- ax1.set_xlabel("Turn number")
- ax1.set_ylabel("Transverse tunes")
- plt.legend(["x","y"])
-
- fig2, ax2 = plt.subplots()
- ax2.errorbar(x=time[1:], y=tune[2,1:], yerr=tune_spread[2,1:])
- ax2.set_xlabel("Turn number")
- ax2.set_ylabel("Synchrotron tune")
+ if plot_tune is True:
+ tune = file[group]["tune"]
+ tune_spread = file[group]["tune_spread"]
+
+ fig1, ax1 = plt.subplots()
+ ax1.errorbar(x=time[1:], y=tune[0,1:], yerr=tune_spread[0,1:])
+ ax1.errorbar(x=time[1:], y=tune[1,1:], yerr=tune_spread[1,1:])
+ ax1.set_xlabel("Turn number")
+ ax1.set_ylabel("Transverse tunes")
+ plt.legend(["x","y"])
+
+ fig2, ax2 = plt.subplots()
+ ax2.errorbar(x=time[1:], y=tune[2,1:], yerr=tune_spread[2,1:])
+ ax2.set_xlabel("Turn number")
+ ax2.set_ylabel("Synchrotron tune")
+
+ if plot_fft is True:
+ if ring is None:
+ T0 = 1.183e-06
+ f0 = 0.845e6
+ else:
+ T0 = ring.T0
+ f0 = ring.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]
+
+ fourier_save = file[group]['fft'][axis]
+
+ if max_turn is None:
+ max_turn = time[-1]
+ if min_turn is None:
+ min_turn = time[1]
+
+ 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
+ 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")
+
+ 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])
+
file.close()
- return (fig1, fig2)
+
+
+ if plot_tune is True and plot_fft is True:
+ if streak_plot is True and profile_plot is True:
+ return fig1, fig2, fig3, fig4
+ elif streak_plot is True:
+ return fig1, fig2, fig3
+ elif profile_plot is True:
+ return fig1, fig2, fig4
+
+ elif plot_tune is True:
+ return fig1, fig2
+
+ elif plot_fft is True:
+ if streak_plot is True and profile_plot is True:
+ return fig3, fig4
+ elif streak_plot is True:
+ return fig3
+ elif profile_plot is True:
+ return fig4
def plot_cavitydata(filename, cavity_name, phasor="cavity",
plot_type="bunch", bunch_number=0, turn=None):
@@ -645,4 +739,5 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
cbar.set_label(ylabel)
file.close()
- return fig
\ No newline at end of file
+ return fig
+ return fig
diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py
index f3842c1e72dc656702f20ca3ad6a346b73259144..d71e78b034ec7b102d41528ad6d4f12b3c1bf0b0 100644
--- a/tracking/synchrotron.py
+++ b/tracking/synchrotron.py
@@ -44,6 +44,19 @@ class Synchrotron:
Horizontal and vertical (non-normalized) chromaticities.
U0 : float, optional
Energy loss per turn in [eV].
+ adts : list of arrays or None, optional
+ List that contains arrays of polynomial's coefficients, in decreasing
+ powers, used to determine Amplitude-Dependent Tune Shifts (ADTS).
+ The order of the elements strictly needs to be
+ [coef_xx, coef_yx, coef_xy, coef_yy], where x and y denote the horizontal
+ and the vertical plane, respectively, and coef_PQ means the polynomial's
+ coefficients of the ADTS in plane P due to the offset in plane Q.
+
+ For example, if the tune shift in y due to the offset x is characterized
+ by the equation dQy(x) = 3*x**2 + 2*x + 1, then coef_yx takes the form
+ np.array([3, 2, 1]).
+
+ Use None, to exclude the ADTS calculation.
Attributes
----------
@@ -99,6 +112,7 @@ class Synchrotron:
self.sigma_delta = kwargs.get('sigma_delta') # Equilibrium energy spread
self.sigma_0 = kwargs.get('sigma_0') # Natural bunch length [s]
self.emit = kwargs.get('emit') # X/Y emittances in [m.rad]
+ self.adts = kwargs.get('adts') # Amplitude-Dependent Tune Shift (ADTS)
@property
def h(self):