- added Chebyshev, Legendre and Sacherer (Sinusoidal)

 modes to effective impedance computation
 -  for transverse plane modes, the frequency sampling should be
 with fractional part of the transverse tune.
 Otherwise code does not work for single bunch case (M=1)
 - added head_tail_form_factor(), a function that computes only
  the denominator of the effective impedance.
  To be used to compute tune shifts in low intensity approximation
  - added tune_shift_from_effective_impedance() to implement it
  in the future
  - some autorefactoring

mbtrack2/utilities/misc.py

@@ -9,13 +9,14 @@ from pathlib import Path
 from scipy.interpolate import interp1d
 from mbtrack2.impedance.wakefield import Impedance
 from mbtrack2.utilities.spectrum import spectral_density
-    
-def effective_impedance(ring, imp, m, mu, sigma, M, tuneS, xi=None, 
+
+
+def effective_impedance(ring, imp, m, mu, sigma, M, tuneS, xi=None,
                         mode="Hermite"):
     """
     Compute the effective (longitudinal or transverse) impedance. 
     Formulas from Eq. (1) and (2) p238 of [1].
-    
+
     Parameters
     ----------
     ring : Synchrotron object
@@ -42,61 +43,112 @@ def effective_impedance(ring, imp, m, mu, sigma, M, tuneS, xi=None,
     Zeff : float 
         effective impedance in [ohm] or in [ohm/m] depanding on the impedance
         type.
-        
+
     References
     ----------
     [1] : Handbook of accelerator physics and engineering, 3rd printing.
     """
-    
+
     if not isinstance(imp, Impedance):
         raise TypeError("{} should be an Impedance object.".format(imp))
-        
+
     fmin = imp.data.index.min()
     fmax = imp.data.index.max()
     if fmin > 0:
         double_sided_impedance(imp)
-        
-    if mode == "Hermite":
+
+    if mode in ["Hermite", "Legendre", "Sinusoidal", "Sacherer", "Chebyshev"]:
         def h(f):
             return spectral_density(frequency=f, sigma=sigma, m=m,
-                                    mode="Hermite")
+                                    mode=mode)
     else:
         raise NotImplementedError("Not implemanted yet.")
-    
-    pmax = fmax/(ring.f0 * M) - 1
-    pmin = fmin/(ring.f0 * M) + 1
-    
-    p = np.arange(pmin,pmax+1)
-    
+
     if imp.component_type == "long":
+        pmax = fmax/(ring.f0 * M) - 1
+        pmin = fmin/(ring.f0 * M) + 1
+        p = np.arange(pmin, pmax+1)
+
         fp = ring.f0*(p*M + mu + m*tuneS)
-        fp = fp[np.nonzero(fp)] # Avoid division by 0
-        num = np.sum( imp(fp) * h(fp) / (fp*2*np.pi) )
-        den = np.sum( h(fp) )
+        fp = fp[np.nonzero(fp)]  # Avoid division by 0
+        num = np.sum(imp(fp) * h(fp) / (fp*2*np.pi))
+        den = np.sum(h(fp))
         Zeff = num/den
-        
+
     elif imp.component_type == "xdip" or imp.component_type == "ydip":
         if imp.component_type == "xdip":
-            tuneXY = ring.tune[0]
-            if xi is None :
+            tuneXY = ring.tune[0]-np.floor(ring.tune[0])
+            if xi is None:
                 xi = ring.chro[0]
         elif imp.component_type == "ydip":
-            tuneXY = ring.tune[1]
+            tuneXY = ring.tune[1]-np.floor(ring.tune[1])
             if xi is None:
                 xi = ring.chro[1]
+        pmax = fmax/(ring.f0 * M) - 1
+        pmin = fmin/(ring.f0 * M) + 1
+        p = np.arange(pmin, pmax+1)
         fp = ring.f0*(p*M + mu + tuneXY + m*tuneS)
         f_xi = xi/ring.eta*ring.f0
-        num = np.sum( imp(fp) * h(fp - f_xi) )
-        den = np.sum( h(fp - f_xi) )
+        num = np.sum(imp(fp) * h(fp - f_xi))
+        den = np.sum(h(fp - f_xi))
         Zeff = num/den
     else:
         raise TypeError("Effective impedance is only defined for long, xdip"
                         " and ydip impedance type.")
-        
+
     return Zeff
 
 
-def yokoya_elliptic(x_radius , y_radius):
+def head_tail_form_factor(ring, imp, m, sigma, tuneS, xi=None, mode="Hermite", mu=0):
+    M = 1
+    if not isinstance(imp, Impedance):
+        raise TypeError("{} should be an Impedance object.".format(imp))
+
+    fmin = imp.data.index.min()
+    fmax = imp.data.index.max()
+    if fmin > 0:
+        double_sided_impedance(imp)
+
+    if mode in ["Hermite", "Legendre", "Sinusoidal", "Sacherer", "Chebyshev"]:
+        def h(f):
+            return spectral_density(frequency=f, sigma=sigma, m=m,
+                                    mode=mode)
+    else:
+        raise NotImplementedError("Not implemanted yet.")
+
+    pmax = np.floor(fmax/(ring.f0 * M))
+    pmin = np.ceil(fmin/(ring.f0 * M))
+
+    p = np.arange(pmin, pmax+1)
+
+    if imp.component_type == "long":
+        fp = ring.f0*(p*M + mu + m*tuneS)
+        fp = fp[np.nonzero(fp)]  # Avoid division by 0
+        den = np.sum(h(fp))
+
+    elif imp.component_type == "xdip" or imp.component_type == "ydip":
+        if imp.component_type == "xdip":
+            tuneXY = ring.tune[0]-np.floor(ring.tune[0])
+            if xi is None:
+                xi = ring.chro[0]
+        elif imp.component_type == "ydip":
+            tuneXY = ring.tune[1]-np.floor(ring.tune[0])
+            if xi is None:
+                xi = ring.chro[1]
+        fp = ring.f0*(p*M + mu + tuneXY + m*tuneS)
+        f_xi = xi/ring.eta*ring.f0
+        den = np.sum(h(fp - f_xi))
+    else:
+        raise TypeError("Effective impedance is only defined for long, xdip"
+                        " and ydip impedance type.")
+
+    return den
+
+
+def tune_shift_from_effective_impedance(Zeff):
+    pass
+
+def yokoya_elliptic(x_radius, y_radius):
     """
     Compute Yokoya factors for an elliptic beam pipe.
     Function adapted from N. Mounet IW2D.
@@ -127,7 +179,7 @@ def yokoya_elliptic(x_radius , y_radius):
     else:
         small_semiaxis = x_radius
         large_semiaxis = y_radius
-        
+
     path_to_file = Path(__file__).parent
     file = path_to_file / "data" / "Yokoya_elliptic_from_Elias_USPAS.csv"
 
@@ -140,7 +192,7 @@ def yokoya_elliptic(x_radius , y_radius):
 
     # interpolate Yokoya file at the correct ratio
     yoklong = 1
-    
+
     if y_radius < x_radius:
         yokydip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
         yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
@@ -150,10 +202,11 @@ def yokoya_elliptic(x_radius , y_radius):
         yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
         yokydip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
         yokxquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
-        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])        
+        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])
 
     return (yoklong, yokxdip, yokydip, yokxquad, yokyquad)
 
+
 def beam_loss_factor(impedance, frequency, spectrum, ring):
     """
     Compute "beam" loss factor using the beam spectrum, uses a sum instead of 
@@ -172,7 +225,7 @@ def beam_loss_factor(impedance, frequency, spectrum, ring):
     -------
     kloss_beam : float
         Beam loss factor in [V/C].
-        
+
     References
     ----------
     [1] : Handbook of accelerator physics and engineering, 3rd printing. 
@@ -180,24 +233,25 @@ def beam_loss_factor(impedance, frequency, spectrum, ring):
     """
     pmax = np.floor(impedance.data.index.max()/ring.f0)
     pmin = np.floor(impedance.data.index.min()/ring.f0)
-    
+
     if pmin >= 0:
         double_sided_impedance(impedance)
         pmin = -1*pmax
-    
-    p = np.arange(pmin+1,pmax)    
+
+    p = np.arange(pmin+1, pmax)
     pf0 = p*ring.f0
     ReZ = np.real(impedance(pf0))
     spectral_density = np.abs(spectrum)**2
-    # interpolation of the spectrum is needed to avoid problems liked to 
+    # 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)
     kloss_beam = ring.f0 * np.sum(ReZ*spect(pf0))
-    
+
     return kloss_beam
 
+
 def double_sided_impedance(impedance):
     """
     Add negative frequency points to single sided impedance spectrum following
@@ -209,31 +263,30 @@ def double_sided_impedance(impedance):
         Single sided impedance.
     """
     fmin = impedance.data.index.min()
-    
+
     if fmin >= 0:
         negative_index = impedance.data.index*-1
         negative_data = impedance.data.set_index(negative_index)
-        
+
         imp_type = impedance.component_type
-        
+
         if imp_type == "long":
             negative_data["imag"] = -1*negative_data["imag"]
-            
+
         elif (imp_type == "xdip") or (imp_type == "ydip"):
             negative_data["real"] = -1*negative_data["real"]
-        
+
         elif (imp_type == "xquad") or (imp_type == "yquad"):
             negative_data["real"] = -1*negative_data["real"]
-            
+
         else:
             raise ValueError("Wrong impedance type")
-            
+
         try:
             negative_data = negative_data.drop(0)
         except KeyError:
             pass
-            
+
         all_data = impedance.data.append(negative_data)
         all_data = all_data.sort_index()
         impedance.data = all_data
-