mbtrack2/impedance/wakefield.py

@@ -795,13 +795,15 @@ 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()

mbtrack2/instability/instabilities.py

@@ -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
 

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

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)

mbtrack2/tracking/monitors/monitors.py

@@ -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
 
@@ -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
 
@@ -1278,14 +1291,17 @@ class BunchSpectrumMonitor(Monitor):
                                                            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.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.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.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()

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.

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
 
@@ -1186,8 +1188,6 @@ class CavityResonator():
         return cav
 
 
-
-
 class ProportionalLoop():
     """
     Proportional feedback loop to control a CavityResonator amplitude and phase.

mbtrack2/tracking/synchrotron.py

@@ -631,4 +631,3 @@ class Synchrotron:
         ring.gap_voltage = Vrf  # [V]
 
         return ring
-

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,

mbtrack2/utilities/optics.py

@@ -119,9 +119,10 @@ class Optics:
             self.periodicity = periodicity
 
         if self.periodicity > 1:
-            for i in range(self.periodicity - 1):
-                pos = np.append(twiss.s_pos,
-                                twiss.s_pos + twiss.s_pos[-1] * (i+1))
+            periods = np.arange(0, self.periodicity)
+            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:
             pos = twiss.s_pos