Add power loss spectrum function to evaluate heating at a given frequency

collective_effects/impedance_model.py

@@ -8,10 +8,11 @@ import pandas as pd
 import numpy as np
 import matplotlib.pyplot as plt
 from scipy.integrate import trapz
+from scipy.interpolate import interp1d
 from mpl_toolkits.axes_grid1.inset_locator import inset_axes
 from mbtrack2.collective_effects.utilities import (beam_spectrum, gaussian_bunch_spectrum, 
                                                    beam_loss_factor, spectral_density,
-                                                   effective_impedance)
+                                                   effective_impedance, double_sided_impedance)
 from mbtrack2.collective_effects.wakefield import WakeField
 from mbtrack2.tracking.element import Element
 
@@ -422,3 +423,102 @@ class ImpedanceModel(Element):
         summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"]*1e12*Q**2/(self.ring.T0)*M
                 
         return summary
+    
+    def power_loss_spectrum(self, sigma, M, bunch_spacing, I, n_points=10e6, 
+                            max_overlap=False,plot=False):
+        """
+        Compute the power loss spectrum of the summed longitudinal impedance 
+        as in Eq. (4) of [1].
+        
+        Assume Gaussian bunches and constant spacing between bunches.
+
+        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].
+        n_points : float, optional
+            Number of points used in the frequency spectrum.
+        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 
+            not for the full spectrum.
+        plot : bool, optional
+            If True, plots:
+                - the overlap between the real part of the longitudinal impedance 
+                and the beam spectrum.
+                - the power loss spectrum.
+
+        Returns
+        -------
+        pf0 : array
+            Frequency points.
+        power_loss : array
+            Power loss spectrum in [W].
+            
+        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 
+        Device", in IPAC'18, 2018, doi:10.18429/JACoW-IPAC2018-THPAK093
+
+        """
+        
+        impedance = self.sum.Zlong
+        fmax = impedance.data.index.max()
+        fmin = impedance.data.index.min()
+        
+        Q = I*self.ring.T0/M
+            
+        if fmin >= 0:
+            fmin = -1*fmax
+            
+        frequency = np.linspace(fmin, fmax, int(n_points))
+        if max_overlap is False:
+            spectrum =  beam_spectrum(frequency, M, bunch_spacing, sigma)
+        else:
+            spectrum = gaussian_bunch_spectrum(frequency, sigma)*M
+        
+        pmax = np.floor(fmax/self.ring.f0)
+        pmin = np.floor(fmin/self.ring.f0)
+        
+        if pmin >= 0:
+            double_sided_impedance(impedance)
+            pmin = -1*pmax
+        
+        p = np.arange(pmin+1,pmax)    
+        pf0 = p*self.ring.f0
+        ReZ = np.real(impedance(pf0))
+        spectral_density = np.abs(spectrum)**2
+        # interpolation of the spectrum is needed to avoid problems liked to 
+        # division by 0
+        # computing the spectrum directly to the frequency points gives
+        # wrong results
+        spect = interp1d(frequency, spectral_density)
+        power_loss = (self.ring.f0 * Q)**2 * ReZ * spect(pf0)
+        
+        if plot is True:
+            fig, ax = plt.subplots()
+            twin = ax.twinx()
+            p1, = ax.plot(pf0, ReZ, color="r",label="Re[Z] [Ohm]")
+            p2, = twin.plot(pf0, spect(pf0)*(self.ring.f0 * Q)**2, color="b", 
+                            label="Beam spectrum [a.u.]")
+            ax.set_xlabel("Frequency [Hz]")
+            ax.set_ylabel("Re[Z] [Ohm]")
+            twin.set_ylabel("Beam spectrum [a.u.]")
+            plots = [p1, p2]
+            ax.legend(handles=plots, loc="best")
+            
+            fig, ax = plt.subplots()
+            ax.plot(pf0, power_loss)
+            ax.set_xlabel("Frequency [Hz]")
+            ax.set_ylabel("Power loss [W]")
+        
+        return pf0, power_loss