diff --git a/mbtrack2/impedance/wakefield.py b/mbtrack2/impedance/wakefield.py
index 5ef7beada349bbcee1a8936354ebdc362632d324..458130fe004e0ee0224c026a3cf51e5c9902cdc2 100644
--- a/mbtrack2/impedance/wakefield.py
+++ b/mbtrack2/impedance/wakefield.py
@@ -782,7 +782,7 @@ class Impedance(ComplexData):
label = r"$Z_{" + self.component_type + r"}$" + unit
ax.set_ylabel(label)
return ax
-
+
def to_pycolleff(self):
"""
Convenience method to export impedance to pycolleff.
@@ -795,20 +795,22 @@ class Impedance(ComplexData):
"""
# pycolleff is a soft dependency
from pycolleff.longitudinal_equilibrium import ImpedanceSource
+
# avoid circular import
from mbtrack2.utilities.misc import double_sided_impedance
imp = ImpedanceSource()
if self.component_type == "long":
double_sided_impedance(self)
# Negative imag sign convention !
- imp.zl_table = self.data["real"].to_numpy() - 1j*self.data["imag"].to_numpy()
+ imp.zl_table = self.data["real"].to_numpy(
+ ) - 1j * self.data["imag"].to_numpy()
imp.ang_freq_table = self.data.index.to_numpy() * 2 * np.pi
else:
raise NotImplementedError()
-
+
imp.calc_method = ImpedanceSource.Methods.ImpedanceDFT
imp.active_passive = ImpedanceSource.ActivePassive.Passive
-
+
return imp
diff --git a/mbtrack2/instability/instabilities.py b/mbtrack2/instability/instabilities.py
index ec4425e2588b23cbae2348cbe633bb2eb91a371a..0a198eb98c90e9f55f16dc5a0805faa948889837 100644
--- a/mbtrack2/instability/instabilities.py
+++ b/mbtrack2/instability/instabilities.py
@@ -196,7 +196,7 @@ def lcbi_growth_rate_mode(ring,
n1 = np.arange(1, n_max)
omega_p = ring.omega0 * (n0*M + mu + nu_s)
omega_m = ring.omega0 * (n1*M - mu - nu_s)
-
+
if bunch_length is None:
FFp = 1
FFm = 1
@@ -204,7 +204,8 @@ def lcbi_growth_rate_mode(ring,
FFp = np.exp(-(omega_p * bunch_length)**2)
FFm = np.exp(-(omega_m * bunch_length)**2)
- sum_val = np.sum(omega_p * Zr(omega_p) * FFp) - np.sum(omega_m * Zr(omega_m) * FFm)
+ sum_val = np.sum(omega_p * Zr(omega_p) * FFp) - np.sum(
+ omega_m * Zr(omega_m) * FFm)
return factor * sum_val
diff --git a/mbtrack2/tracking/__init__.py b/mbtrack2/tracking/__init__.py
index 62e879ffe210caa1a3e56a445ce7b8742d6097db..a877b4689870ca5a9ca0b3cfcd5862b5f4c231ce 100644
--- a/mbtrack2/tracking/__init__.py
+++ b/mbtrack2/tracking/__init__.py
@@ -14,6 +14,7 @@ from mbtrack2.tracking.element import (
TransverseMapSector,
transverse_map_sector_generator,
)
+from mbtrack2.tracking.excite import Sweep
from mbtrack2.tracking.feedback import ExponentialDamper, FIRDamper
from mbtrack2.tracking.monitors import *
from mbtrack2.tracking.parallel import Mpi
@@ -31,4 +32,3 @@ from mbtrack2.tracking.wakepotential import (
LongRangeResistiveWall,
WakePotential,
)
-from mbtrack2.tracking.excite import Sweep
\ No newline at end of file
diff --git a/mbtrack2/tracking/excite.py b/mbtrack2/tracking/excite.py
index 61075d5e65da8aa0d9ce58087c18faea765688c2..818fceb8cc1702956c3b46b2ebe6a4753992d4d7 100644
--- a/mbtrack2/tracking/excite.py
+++ b/mbtrack2/tracking/excite.py
@@ -2,12 +2,14 @@
"""
Module to deal with different kinds of beam excitation.
"""
-import numpy as np
import matplotlib.pyplot as plt
-from mbtrack2.tracking.element import Element
-from mbtrack2.tracking.particles import Bunch, Beam
+import numpy as np
from scipy.signal import chirp
+from mbtrack2.tracking.element import Element
+from mbtrack2.tracking.particles import Beam, Bunch
+
+
class Sweep(Element):
"""
Element which excite the beam in between two frequencies, i.e. apply
@@ -43,6 +45,7 @@ class Sweep(Element):
Plot the sweep voltage applied.
"""
+
def __init__(self, ring, f0, f1, t1, level, plane, bunch_to_sweep=None):
self.ring = ring
self.t = np.arange(0, t1, ring.T0)
@@ -59,7 +62,7 @@ class Sweep(Element):
else:
raise ValueError("plane should be 'x', 'y' or 'tau'.")
self.bunch_to_sweep = bunch_to_sweep
-
+
def track(self, bunch_or_beam):
"""
Tracking method for this element.
@@ -90,7 +93,7 @@ class Sweep(Element):
self.count = 0
else:
raise TypeError("bunch_or_beam should be a Beam or Bunch object.")
-
+
def _track_bunch(self, bunch, count_step=True):
"""
Tracking method for a bunch for this element.
@@ -112,4 +115,4 @@ class Sweep(Element):
ax.plot(self.t, self.sweep)
ax.xlabel("Time [s]")
ax.ylabel("Sweep voltage [V]")
- return fig
\ No newline at end of file
+ return fig
diff --git a/mbtrack2/tracking/monitors/monitors.py b/mbtrack2/tracking/monitors/monitors.py
index 4e87bf2e34f4edb6356929e64e01061f33ec77cd..ee16f66703a994f8065a8d161252dc308bdee78d 100644
--- a/mbtrack2/tracking/monitors/monitors.py
+++ b/mbtrack2/tracking/monitors/monitors.py
@@ -415,7 +415,7 @@ class PhaseSpaceMonitor(Monitor):
self.dict_buffer = dict_buffer
self.dict_file = dict_file
-
+
if self.sample_number != self.mp_number:
index = np.arange(self.mp_number)
samples_meta = random.sample(list(index), self.sample_number)
@@ -445,8 +445,8 @@ class PhaseSpaceMonitor(Monitor):
self.alive[:, self.buffer_count] = bunch.alive[self.samples]
for i, dim in enumerate(bunch):
- self.particles[:, i,
- self.buffer_count] = bunch.particles[dim][self.samples]
+ self.particles[:, i, self.buffer_count] = bunch.particles[dim][
+ self.samples]
self.buffer_count += 1
@@ -1061,7 +1061,7 @@ class BunchSpectrumMonitor(Monitor):
self.n_fft = int(save_every)
else:
self.n_fft = int(n_fft)
-
+
self.higher_orders = higher_orders
self.sample_size = int(sample_size)
@@ -1109,14 +1109,14 @@ class BunchSpectrumMonitor(Monitor):
"mean_incoherent": (3, total_size),
"std_incoherent": (3, total_size)
}
-
+
if self.higher_orders:
- dict_buffer["coherent_q"] = (3, self.n_fft//2+1, buffer_size)
- dict_buffer["coherent_s"] = (3, self.n_fft//2+1, buffer_size)
- dict_buffer["coherent_o"] = (3, self.n_fft//2+1, buffer_size)
- dict_file["coherent_q"] = (3, self.n_fft//2+1, total_size)
- dict_file["coherent_s"] = (3, self.n_fft//2+1, total_size)
- dict_file["coherent_o"] = (3, self.n_fft//2+1, total_size)
+ dict_buffer["coherent_q"] = (3, self.n_fft // 2 + 1, buffer_size)
+ dict_buffer["coherent_s"] = (3, self.n_fft // 2 + 1, buffer_size)
+ dict_buffer["coherent_o"] = (3, self.n_fft // 2 + 1, buffer_size)
+ dict_file["coherent_q"] = (3, self.n_fft // 2 + 1, total_size)
+ dict_file["coherent_s"] = (3, self.n_fft // 2 + 1, total_size)
+ dict_file["coherent_o"] = (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)
@@ -1130,9 +1130,9 @@ class BunchSpectrumMonitor(Monitor):
(self.size_list, self.sample_size, self.save_every + 1))
self.mean = np.zeros((self.size_list, self.save_every + 1))
if self.higher_orders:
- self.std = np.zeros((self.size_list, self.save_every+1))
- self.skew = np.zeros((self.size_list, self.save_every+1))
- self.kurtosis = np.zeros((self.size_list, self.save_every+1))
+ self.std = np.zeros((self.size_list, self.save_every + 1))
+ self.skew = np.zeros((self.size_list, self.save_every + 1))
+ self.kurtosis = np.zeros((self.size_list, self.save_every + 1))
index = np.arange(0, int(mp_number))
self.index_sample = sorted(random.sample(list(index),
@@ -1141,9 +1141,12 @@ class BunchSpectrumMonitor(Monitor):
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))
if self.higher_orders:
- self.coherent_q = np.zeros((3, self.n_fft//2+1, self.buffer_size))
- self.coherent_s = np.zeros((3, self.n_fft//2+1, self.buffer_size))
- self.coherent_o = np.zeros((3, self.n_fft//2+1, self.buffer_size))
+ self.coherent_q = np.zeros(
+ (3, self.n_fft // 2 + 1, self.buffer_size))
+ self.coherent_s = np.zeros(
+ (3, self.n_fft // 2 + 1, self.buffer_size))
+ self.coherent_o = np.zeros(
+ (3, self.n_fft // 2 + 1, self.buffer_size))
self.file[self.group_name].create_dataset("freq",
data=self.frequency_samples)
@@ -1208,7 +1211,8 @@ class BunchSpectrumMonitor(Monitor):
if self.higher_orders:
self.std[:, self.save_count] = bunch.std[self.mean_index]
self.skew[:, self.save_count] = bunch.skew[self.mean_index]
- self.kurtosis[:, self.save_count] = bunch.kurtosis[self.mean_index]
+ self.kurtosis[:, self.save_count] = bunch.kurtosis[
+ self.mean_index]
self.save_count += 1
@@ -1239,9 +1243,18 @@ class BunchSpectrumMonitor(Monitor):
self.buffer_count] = self.get_coherent_spectrum(
self.mean[value])
if self.higher_orders:
- self.coherent_q[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.std[value])
- self.coherent_s[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.skew[value])
- self.coherent_o[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.kurtosis[value])
+ self.coherent_q[
+ self.store_dict[key], :,
+ self.buffer_count] = self.get_coherent_spectrum(
+ self.std[value])
+ self.coherent_s[
+ self.store_dict[key], :,
+ self.buffer_count] = self.get_coherent_spectrum(
+ self.skew[value])
+ self.coherent_o[
+ self.store_dict[key], :,
+ self.buffer_count] = self.get_coherent_spectrum(
+ self.kurtosis[value])
self.buffer_count += 1
@@ -1277,16 +1290,19 @@ class BunchSpectrumMonitor(Monitor):
self.buffer_size:(self.write_count +
1) *
self.buffer_size] = self.coherent
- if self.higher_orders:
- self.file[self.group_name]["coherent_q"][:,:,
- self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.coherent_q
- self.file[self.group_name]["coherent_s"][:,:,
- self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.coherent_s
- self.file[self.group_name]["coherent_o"][:,:,
- self.write_count * self.buffer_size:(self.write_count+1) *
- self.buffer_size] = self.coherent_o
+ if self.higher_orders:
+ self.file[self.group_name][
+ "coherent_q"][:, :, self.write_count *
+ self.buffer_size:(self.write_count + 1) *
+ self.buffer_size] = self.coherent_q
+ self.file[self.group_name][
+ "coherent_s"][:, :, self.write_count *
+ self.buffer_size:(self.write_count + 1) *
+ self.buffer_size] = self.coherent_s
+ self.file[self.group_name][
+ "coherent_o"][:, :, self.write_count *
+ self.buffer_size:(self.write_count + 1) *
+ self.buffer_size] = self.coherent_o
self.file.flush()
self.write_count += 1
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index e9b673b7739b029d5f15819f7997150273a3906a..4c3934b75a81798c4a429d89e3b1ef5b6a8ab330 100644
--- a/mbtrack2/tracking/particles.py
+++ b/mbtrack2/tracking/particles.py
@@ -11,6 +11,7 @@ import seaborn as sns
from scipy.constants import c, e, m_e, m_p
from scipy.stats import moment
+
class Particle:
"""
Define a particle object.
@@ -250,24 +251,24 @@ class Bunch:
"""
std = [[self[name].std()] for name in self]
return np.squeeze(np.array(std))
-
+
@property
def skew(self):
"""
Return the skew (3rd moment) of the position of alive
particles for each coordinates.
"""
- skew = [[moment(self[name],3)] for name in self]
+ skew = [[moment(self[name], 3)] for name in self]
return np.squeeze(np.array(skew))
-
+
@property
def kurtosis(self):
"""
Return the kurtosis (4th moment) of the position of alive
particles for each coordinates.
"""
- kurtosis = [[moment(self[name],4)] for name in self]
- return np.squeeze(np.array(kurtosis))
+ kurtosis = [[moment(self[name], 4)] for name in self]
+ return np.squeeze(np.array(kurtosis))
@property
def emit(self):
diff --git a/mbtrack2/tracking/rf.py b/mbtrack2/tracking/rf.py
index 41284a1aa657f050bcd946d3aa79dc3b3d8a462d..e39271204bc65fd8f015e43a34f4a21064e032d6 100644
--- a/mbtrack2/tracking/rf.py
+++ b/mbtrack2/tracking/rf.py
@@ -3,11 +3,13 @@
This module handles radio-frequency (RF) cavitiy elements.
"""
+from functools import partial
+
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.legend_handler import HandlerPatch
-from functools import partial
+
from mbtrack2.instability import lcbi_growth_rate
from mbtrack2.tracking.element import Element
@@ -1143,7 +1145,7 @@ class CavityResonator():
ref_frame="beam")
return pos, voltage_rec
-
+
def to_pycolleff(self, Impedance=True):
"""
Convenience method to export a CavityResonator to pycolleff.
@@ -1162,11 +1164,11 @@ class CavityResonator():
"""
# pycolleff is a soft dependency
from pycolleff.longitudinal_equilibrium import ImpedanceSource
-
+
cav = ImpedanceSource()
cav.harm_rf = self.m
cav.Q = self.QL
- RoverQ = self.RL/self.QL
+ RoverQ = self.RL / self.QL
cav.shunt_impedance = RoverQ * cav.Q
cav.ang_freq_rf = self.ring.omega1
cav.ang_freq = cav.harm_rf * cav.ang_freq_rf
@@ -1176,16 +1178,14 @@ class CavityResonator():
if Impedance:
raise NotImplementedError()
else:
- cav.calc_method = ImpedanceSource.Methods.Wake
+ cav.calc_method = ImpedanceSource.Methods.Wake
else:
cav.active_passive = ImpedanceSource.ActivePassive.Passive
if Impedance:
cav.calc_method = ImpedanceSource.Methods.ImpedanceDFT
else:
- cav.calc_method = ImpedanceSource.Methods.Wake
+ cav.calc_method = ImpedanceSource.Methods.Wake
return cav
-
-
class ProportionalLoop():
diff --git a/mbtrack2/tracking/synchrotron.py b/mbtrack2/tracking/synchrotron.py
index fa37923a34e880e1c3ed3f8ceac3e28071dc5fc3..c357198ac9043086bab53c92ec19d329fe606ee5 100644
--- a/mbtrack2/tracking/synchrotron.py
+++ b/mbtrack2/tracking/synchrotron.py
@@ -586,7 +586,7 @@ class Synchrotron:
U0=self.U0,
TimeLag=TimeLag)
return at_simple_ring
-
+
def to_pycolleff(self, I0, Vrf, bunch_number):
"""
Return a pycolleff Ring element from the Synchrotron element data.
@@ -608,7 +608,7 @@ class Synchrotron:
"""
# pycolleff is a soft dependency
from pycolleff.colleff import Ring
-
+
ring = Ring()
ring.version = 'from_mbtrack2'
ring.rf_freq = self.f1
@@ -623,12 +623,11 @@ class Synchrotron:
ring.total_current = I0 # total current [A]
ring.sync_tune = self.synchrotron_tune(Vrf) # synchrotron tune
ring.espread = self.sigma_delta
- ring.bunlen = self.sigma_0*c # [m]
+ ring.bunlen = self.sigma_0 * c # [m]
ring.damptx = self.tau[0] # [s]
ring.dampty = self.tau[1] # [s]
ring.dampte = self.tau[2] # [s]
- ring.en_lost_rad = self.U0 # [eV]
+ ring.en_lost_rad = self.U0 # [eV]
ring.gap_voltage = Vrf # [V]
-
- return ring
+ return ring
diff --git a/mbtrack2/utilities/beamloading.py b/mbtrack2/utilities/beamloading.py
index b51554f48d137699de1a86de397eefbe9f5bd3d8..595fafd316940f8cdfa8f7d0cb53c4aa01f8187a 100644
--- a/mbtrack2/utilities/beamloading.py
+++ b/mbtrack2/utilities/beamloading.py
@@ -8,8 +8,10 @@ import numpy as np
from scipy.constants import c
from scipy.integrate import trapz
from scipy.optimize import root
+
from mbtrack2.utilities.spectrum import gaussian_bunch
+
class BeamLoadingEquilibrium():
"""Class used to compute beam equilibrium profile for a given storage ring
and a list of RF cavities of any harmonic. The class assumes an uniform
@@ -40,6 +42,7 @@ class BeamLoadingEquilibrium():
Physical Review Accelerators and Beams, 21(11), 114404.
"""
+
def __init__(self,
ring,
cavity_list,
@@ -68,8 +71,8 @@ class BeamLoadingEquilibrium():
self.mpi = False
self.N = N
self.z0 = np.linspace(self.B1, self.B2, self.N)
- self.tau0 = self.z0/c
- self.rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0)/c
+ self.tau0 = self.z0 / c
+ self.rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0) / c
# Define constants for scaled potential u(z)
self.u0 = self.ring.U0 / (self.ring.ac * self.ring.sigma_delta**2 *
@@ -77,7 +80,7 @@ class BeamLoadingEquilibrium():
self.ug = np.zeros((self.n_cavity, ))
self.ub = np.zeros((self.n_cavity, ))
self.update_potentials()
-
+
def update_rho(self):
uexp_vals = self.uexp(self.z0)
A = trapz(uexp_vals, x=self.z0)
@@ -137,10 +140,10 @@ class BeamLoadingEquilibrium():
def uexp(self, z):
return np.exp(-1 * self.u(z))
-
+
def integrate_func(self, f, g):
"""Return Integral[f*g]/Integral[f] between B1 and B2"""
- A = trapz(f*g, x=self.z0)
+ A = trapz(f * g, x=self.z0)
B = trapz(f, x=self.z0)
return A / B
@@ -394,7 +397,7 @@ class BeamLoadingEquilibrium():
(2 * np.pi * HC.m**2 * self.F[1] * self.ring.h * I0 * f))
return (eta, RQth, f)
-
+
@property
def R_factor(self):
"""
@@ -407,6 +410,6 @@ class BeamLoadingEquilibrium():
Special Topics-Accelerators and Beams 4.3 (2001): 030701.
"""
- rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0)/c
- R = trapz(rho0**2, self.tau0)/trapz(self.rho0**2, self.tau0)
+ rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0) / c
+ R = trapz(rho0**2, self.tau0) / trapz(self.rho0**2, self.tau0)
return R
diff --git a/mbtrack2/utilities/optics.py b/mbtrack2/utilities/optics.py
index 79026f0d67edbbd8ea934f568376499235b23ec7..4e2f58bf3ec2624b7a796c2790a11ce507535154 100644
--- a/mbtrack2/utilities/optics.py
+++ b/mbtrack2/utilities/optics.py
@@ -120,7 +120,7 @@ class Optics:
if self.periodicity > 1:
periods = np.arange(0, self.periodicity)
- shift = periods*twiss.s_pos[-1]
+ shift = periods * twiss.s_pos[-1]
shift = np.repeat(shift, len(twiss.s_pos))
pos = np.tile(twiss.s_pos.T, self.periodicity) + shift
else: