Update methods for ImpedanceModel and BeamLoadingEquilibrium

The energy_loss and power_loss_spectrum methods for ImpedanceModel can now take any bunch spectrum in optional agrument. Add bunch_spectrum and plot_bunch_spectrum to BeamLoadingEquilibrium.

mbtrack2/impedance/impedance_model.py

@@ -305,7 +305,7 @@ class ImpedanceModel():
         Z_type : str, optional
             Type of impedance to plot.
         component : str, optional
-            Component to plot, can be "real" or "imag". 
+            Component to plot, can be "real" or "imag".
         sigma : float, optional
             RMS bunch length in [s] to use for the spectral density. If equal
             to None, the spectral density is not plotted.
@@ -496,29 +496,51 @@ class ImpedanceModel():
 
         return summary
 
-    def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
+    def energy_loss(self,
+                    M,
+                    bunch_spacing,
+                    I,
+                    sigma=None,
+                    bunch_spectrum=None,
+                    freq_spectrum=None,
+                    n_points=10e6):
         """
-        Compute the beam and bunch loss factor and energy losses for each type 
-        of element in the model assuming Gaussian bunches and constant spacing
-        between bunches.
+        Compute the beam and bunch loss factor and energy losses for each type
+        of element in the model.
+
+        Assumtions:
+            - Constant spacing between bunches.
+            - All bunches have the same bunch distribution.
+            - Gaussian bunches if sigma is given.
 
         Parameters
         ----------
-        sigma : float
-            RMS bunch length in [s].
         M : int
             Number of bunches in the beam.
         bunch_spacing : float
             Time between two bunches in [s].
         I : float
             Total beam current in [A].
+        sigma : float, optional
+            RMS bunch length in [s].
+            If None, freq_spectrum and bunch_spectrum must be given.
+            Default is None.
+        bunch_spectrum : array
+            Bunch spectrum to consider (i.e. FT of bunch profile).
+            Not used if sigma is not None.
+            Default is None.
+        freq_spectrum : array
+            Frequency points corresponding to bunch_spectrum in [Hz].
+            Not used if sigma is not None.
+            Default is None.
         n_points : float, optional
             Number of points used in the frequency spectrums.
+            Default is 10e6.
 
         Returns
         -------
         summary : Dataframe
-            Contains the beam and bunch loss factor and energy loss for the 
+            Contains the beam and bunch loss factor and energy loss for the
             full model and for each type of different component.
 
         """
@@ -532,9 +554,24 @@ class ImpedanceModel():
             fmin = -1 * fmax
         f = np.linspace(fmin, fmax, int(n_points))
 
-        beam_spect = beam_spectrum(f, M, bunch_spacing, sigma=sigma)
-
-        bunch_spect = gaussian_bunch_spectrum(f, sigma)
+        if sigma is None:
+            if bunch_spectrum is None:
+                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.")
+                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)
+                bunch_spect = bunch_spectrum_interp(f)
+                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)
 
         attr_list = self.sum_names
 
@@ -566,42 +603,60 @@ class ImpedanceModel():
         return summary
 
     def power_loss_spectrum(self,
-                            sigma,
                             M,
                             bunch_spacing,
                             I,
+                            sigma=None,
+                            bunch_spectrum=None,
+                            freq_spectrum=None,
                             n_points=10e6,
                             max_overlap=False,
                             plot=False):
         """
-        Compute the power loss spectrum of the summed longitudinal impedance 
+        Compute the power loss spectrum of the summed longitudinal impedance
         as in Eq. (4) of [1].
 
-        Assume Gaussian bunches and constant spacing between bunches.
+        Assumtions:
+            - Constant spacing between bunches.
+            - All bunches have the same bunch distribution.
+            - Gaussian bunches if sigma is given.
 
         Parameters
         ----------
-        sigma : float
-            RMS bunch length in [s].
         M : int
             Number of bunches in the beam.
         bunch_spacing : float
             Time between two bunches in [s].
         I : float
             Total beam current in [A].
+        sigma : float, optional
+            RMS bunch length in [s].
+            If None, freq_spectrum and bunch_spectrum must be given.
+            Default is None.
+        bunch_spectrum : array
+            Bunch spectrum to consider (i.e. FT of bunch profile).
+            Not used if sigma is not None.
+            Default is None.
+        freq_spectrum : array
+            Frequency points corresponding to bunch_spectrum in [Hz].
+            Not used if sigma is not None.
+            Default is None.
         n_points : float, optional
             Number of points used in the frequency spectrum.
+            Default is 10e6.
         max_overlap : bool, optional
-            If True, the bunch spectrum (scaled to the number of bunches) is 
-            used instead of the beam spectrum to compute the maximum value of 
-            the power loss spectrum at each frequency. Should only be used to 
-            maximise the power loss at a given frequency (e.g. for HOMs) and 
+            If True, the bunch spectrum (scaled to the number of bunches) is
+            used instead of the beam spectrum to compute the maximum value of
+            the power loss spectrum at each frequency. Should only be used to
+            maximise the power loss at a given frequency (e.g. for HOMs) and
             not for the full spectrum.
+            Default is False.
         plot : bool, optional
             If True, plots:
-                - the overlap between the real part of the longitudinal impedance 
+                - the overlap between the real part of the longitudinal impedance
                 and the beam spectrum.
                 - the power loss spectrum.
+            Default is False.
 
         Returns
         -------
@@ -612,8 +667,8 @@ class ImpedanceModel():
 
         References
         ----------
-        [1] : L. Teofili, et al. "A Multi-Physics Approach to Simulate the RF 
-        Heating 3D Power Map Induced by the Proton Beam in a Beam Intercepting 
+        [1] : L. Teofili, et al. "A Multi-Physics Approach to Simulate the RF
+        Heating 3D Power Map Induced by the Proton Beam in a Beam Intercepting
         Device", in IPAC'18, 2018, doi:10.18429/JACoW-IPAC2018-THPAK093
 
         """
@@ -629,10 +684,31 @@ class ImpedanceModel():
             double_sided_impedance(impedance)
 
         frequency = np.linspace(fmin, fmax, int(n_points))
+
+        if sigma is None:
+            if bunch_spectrum is None:
+                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.")
+                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)
+                bunch_spectrum = bunch_spectrum_interp(frequency)
+        else:
+            bunch_spectrum = None
+
         if max_overlap is False:
-            spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma)
+            spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma, bunch_spectrum)
         else:
-            spectrum = gaussian_bunch_spectrum(frequency, sigma) * M
+            if bunch_spectrum is None:
+                spectrum = gaussian_bunch_spectrum(frequency, sigma) * M
+            else:
+                spectrum = bunch_spectrum * M
 
         pmax = np.floor(fmax / self.ring.f0)
         pmin = np.floor(fmin / self.ring.f0)