Apply formatters and update notebook

mbtrack2/impedance/impedance_model.py

@@ -556,19 +556,28 @@ class ImpedanceModel():
 
         if sigma is None:
             if bunch_spectrum is None:
-                raise ValueError("Either sigma or bunch_spectrum and freq_spectrum"
-                                 + "must be specified.")
+                raise ValueError(
+                    "Either sigma or bunch_spectrum and freq_spectrum" +
+                    "must be specified.")
             else:
                 if freq_spectrum is None:
-                    raise ValueError("freq_spectrum must be given if bunch_spectrum"
-                                     + " is specified.")
+                    raise ValueError(
+                        "freq_spectrum must be given if bunch_spectrum" +
+                        " is specified.")
                 if freq_spectrum[0] >= 0:
-                    freq_spectrum = np.concatenate((-1*freq_spectrum[:0:-1], freq_spectrum))
-                    bunch_spectrum = np.concatenate((bunch_spectrum[:0:-1], bunch_spectrum))
-                bunch_spectrum_interp = interp1d(freq_spectrum, bunch_spectrum,
-                                                 bounds_error=False, fill_value=0)
+                    freq_spectrum = np.concatenate(
+                        (-1 * freq_spectrum[:0:-1], freq_spectrum))
+                    bunch_spectrum = np.concatenate(
+                        (bunch_spectrum[:0:-1], bunch_spectrum))
+                bunch_spectrum_interp = interp1d(freq_spectrum,
+                                                 bunch_spectrum,
+                                                 bounds_error=False,
+                                                 fill_value=0)
                 bunch_spect = bunch_spectrum_interp(f)
-                beam_spect = beam_spectrum(f, M, bunch_spacing, bunch_spectrum=bunch_spect)
+                beam_spect = beam_spectrum(f,
+                                           M,
+                                           bunch_spacing,
+                                           bunch_spectrum=bunch_spect)
         else:
             beam_spect = beam_spectrum(f, M, bunch_spacing, sigma=sigma)
             bunch_spect = gaussian_bunch_spectrum(f, sigma)
@@ -687,23 +696,30 @@ class ImpedanceModel():
 
         if sigma is None:
             if bunch_spectrum is None:
-                raise ValueError("Either sigma or bunch_spectrum and freq_spectrum"
-                                 + "must be specified.")
+                raise ValueError(
+                    "Either sigma or bunch_spectrum and freq_spectrum" +
+                    "must be specified.")
             else:
                 if freq_spectrum is None:
-                    raise ValueError("freq_spectrum must be given if bunch_spectrum"
-                                     + " is specified.")
+                    raise ValueError(
+                        "freq_spectrum must be given if bunch_spectrum" +
+                        " is specified.")
                 if freq_spectrum[0] >= 0:
-                    freq_spectrum = np.concatenate((-1*freq_spectrum[:0:-1], freq_spectrum))
-                    bunch_spectrum = np.concatenate((bunch_spectrum[:0:-1], bunch_spectrum))
-                bunch_spectrum_interp = interp1d(freq_spectrum, bunch_spectrum,
-                                                 bounds_error=False, fill_value=0)
+                    freq_spectrum = np.concatenate(
+                        (-1 * freq_spectrum[:0:-1], freq_spectrum))
+                    bunch_spectrum = np.concatenate(
+                        (bunch_spectrum[:0:-1], bunch_spectrum))
+                bunch_spectrum_interp = interp1d(freq_spectrum,
+                                                 bunch_spectrum,
+                                                 bounds_error=False,
+                                                 fill_value=0)
                 bunch_spectrum = bunch_spectrum_interp(frequency)
         else:
             bunch_spectrum = None
 
         if max_overlap is False:
-            spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma, bunch_spectrum)
+            spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma,
+                                     bunch_spectrum)
         else:
             if bunch_spectrum is None:
                 spectrum = gaussian_bunch_spectrum(frequency, sigma) * M

mbtrack2/utilities/beamloading.py

@@ -6,9 +6,9 @@ Module where the BeamLoadingEquilibrium class is defined.
 import matplotlib.pyplot as plt
 import numpy as np
 from scipy.constants import c
+from scipy.fft import rfft, rfftfreq
 from scipy.integrate import trapz
 from scipy.optimize import root
-from scipy.fft import rfftfreq, rfft
 
 from mbtrack2.utilities.spectrum import gaussian_bunch
 
@@ -421,7 +421,7 @@ class BeamLoadingEquilibrium():
     def bunch_spectrum(self):
         """Return the bunch equilibrium spectrum."""
         dft_rho = rfft(self.rho0) * self.sampling / self.N
-        return np.abs(dft_rho)/np.max(np.abs(dft_rho))
+        return np.abs(dft_rho) / np.max(np.abs(dft_rho))
 
     def plot_bunch_spectrum(self):
         """Plot the bunch equilibrium spectrum"""