Tools : Add Yokoya_elliptic function + misc modif

Add Yokoya_elliptic function to compute Yokoya factors for an elliptic shape.
Add an option to divide by value in read_CST
Add docstrings to read_CST and read_IW2D

collective_effects/data/Yokoya_elliptic_from_Elias_USPAS.csv

+x,dipy,dipx,quady,quadx
+-1.38778e-17,1,1.00307,-0.00270479,0.00429341
+7.94873e-06,0.999988,1.00304,-0.00268229,0.0042735
+0.00131154,0.99802,0.998283,0.00100816,0.00100816
+0.0184543,0.972133,0.935779,0.0495391,-0.0419323
+0.0195168,0.970529,0.931905,0.0521597,-0.0445937
+0.0221637,0.966532,0.922254,0.0586882,-0.0511943
+0.0233454,0.96535,0.917945,0.0616028,-0.0541411
+0.036898,0.9518,0.86594,0.0950297,-0.0879369
+0.0376001,0.951098,0.863246,0.0965459,-0.0896878
+0.0390229,0.949675,0.857042,0.0996185,-0.0932359
+0.0404405,0.948258,0.850861,0.10268,-0.0967708
+0.046659,0.942615,0.823745,0.116109,-0.112278
+0.0579595,0.932362,0.789359,0.140512,-0.140458
+0.0585291,0.931845,0.787626,0.141742,-0.141878
+0.0596446,0.930833,0.784232,0.144151,-0.143994
+0.0639448,0.926931,0.774182,0.153438,-0.152153
+0.0765568,0.919567,0.744707,0.180674,-0.176081
+0.0790209,0.918129,0.737585,0.185995,-0.180757
+0.079368,0.917926,0.736581,0.18656,-0.181415
+0.0796277,0.917774,0.735831,0.186982,-0.181819
+0.098641,0.90472,0.680875,0.217927,-0.211402
+0.101367,0.902849,0.676101,0.222364,-0.215644
+0.101839,0.902525,0.675275,0.222994,-0.216378
+0.102833,0.902214,0.673535,0.224321,-0.217924
+0.121457,0.896392,0.640927,0.24919,-0.2469
+0.126297,0.893069,0.632452,0.255654,-0.254432
+0.126324,0.893051,0.632405,0.255689,-0.254465
+0.135123,0.887009,0.616999,0.267106,-0.265333
+0.143668,0.881143,0.605496,0.278192,-0.275888
+0.147338,0.880067,0.600556,0.282953,-0.28042
+0.148477,0.879733,0.599023,0.28408,-0.281827
+0.164594,0.875006,0.577327,0.300014,-0.295503
+0.169035,0.872782,0.571349,0.304404,-0.299271
+0.174896,0.869847,0.564512,0.310198,-0.304243
+0.177209,0.868689,0.561814,0.311798,-0.306206
+0.182895,0.865841,0.555182,0.31573,-0.31085
+0.201121,0.860497,0.53392,0.328337,-0.325736
+0.203331,0.859849,0.531343,0.329443,-0.32754
+0.203821,0.859705,0.530771,0.329688,-0.327805
+0.205573,0.859461,0.528728,0.330565,-0.328749
+0.226064,0.856601,0.513301,0.340824,-0.339797
+0.231266,0.855831,0.509385,0.343427,-0.342601
+0.237314,0.854935,0.504831,0.346309,-0.345862
+0.246069,0.853639,0.498241,0.350479,-0.349956
+0.246999,0.853501,0.497793,0.350922,-0.35039
+0.266624,0.850405,0.488348,0.36027,-0.359567
+0.269271,0.849578,0.487075,0.361531,-0.360804
+0.276534,0.847307,0.483579,0.364202,-0.3642
+0.286242,0.844273,0.478907,0.367772,-0.36641
+0.290515,0.843676,0.476851,0.369344,-0.367383
+0.302025,0.84207,0.47313,0.373577,-0.370004
+0.308485,0.841169,0.471041,0.375354,-0.371474
+0.317085,0.841142,0.468261,0.377721,-0.373432
+0.328118,0.841109,0.464694,0.380756,-0.375867
+0.328441,0.841064,0.46459,0.380845,-0.375938
+0.334771,0.840181,0.462842,0.382587,-0.377335
+0.350361,0.838005,0.458536,0.384111,-0.380775
+0.358946,0.837979,0.456166,0.384951,-0.382669
+0.364884,0.837961,0.454526,0.385532,-0.383162
+0.372612,0.837937,0.452392,0.386087,-0.383804
+0.372912,0.837844,0.452309,0.386108,-0.383829
+0.39223,0.831805,0.447911,0.387495,-0.385433
+0.396896,0.831791,0.446849,0.38783,-0.385821
+0.405467,0.831765,0.444897,0.388445,-0.387712
+0.413464,0.83174,0.443077,0.389451,-0.389476
+0.432805,0.831682,0.440122,0.391884,-0.393744
+0.438756,0.831664,0.439213,0.392633,-0.395057
+0.452603,0.831622,0.437097,0.394375,-0.395967
+0.454023,0.831617,0.43688,0.39457,-0.396061
+0.455057,0.831614,0.436727,0.394712,-0.396129
+0.487177,0.828947,0.431971,0.399129,-0.39824
+0.493006,0.828463,0.431108,0.399931,-0.398668
+0.495812,0.828454,0.430692,0.400316,-0.398874
+0.495892,0.828454,0.43068,0.400322,-0.39888
+0.530363,0.828349,0.428076,0.402796,-0.401408
+0.533582,0.828339,0.427833,0.403028,-0.401631
+0.536396,0.828331,0.427621,0.40323,-0.401826
+0.537768,0.828327,0.427517,0.403225,-0.401921
+0.573045,0.828219,0.424852,0.403118,-0.404367
+0.576166,0.82821,0.424616,0.40335,-0.404584
+0.579393,0.8282,0.424372,0.40359,-0.404593
+0.579645,0.828181,0.424353,0.403609,-0.404594
+0.612319,0.825713,0.424254,0.406035,-0.404693
+0.619359,0.825181,0.424233,0.406014,-0.404715
+0.62022,0.825116,0.42423,0.406011,-0.404717
+0.62127,0.825037,0.424075,0.406008,-0.40472
+0.652895,0.822647,0.419392,0.405912,-0.404816
+0.662089,0.821953,0.41803,0.405884,-0.404844
+0.663146,0.821873,0.417939,0.405881,-0.404848
+0.670406,0.821851,0.417316,0.405859,-0.40487
+0.683,0.821813,0.416234,0.405821,-0.405677
+0.69873,0.821765,0.414882,0.407431,-0.406685
+0.699795,0.821762,0.414879,0.40754,-0.406753
+0.711803,0.821725,0.414843,0.40877,-0.407523
+0.720136,0.8217,0.414817,0.408745,-0.408057
+0.733826,0.821658,0.414776,0.408703,-0.408099
+0.740615,0.821008,0.414755,0.408682,-0.408119
+0.756305,0.819504,0.413494,0.408635,-0.408167
+0.763329,0.81883,0.412929,0.408613,-0.408188
+0.766541,0.818523,0.412671,0.408604,-0.408389
+0.779873,0.818482,0.4116,0.408563,-0.409223
+0.79819,0.818427,0.411544,0.408508,-0.410368
+0.808425,0.818395,0.411513,0.409101,-0.411008
+0.809978,0.818511,0.411508,0.409191,-0.411105
+0.814368,0.818837,0.410816,0.409446,-0.41138
+0.829603,0.819969,0.408412,0.41033,-0.411426
+0.8477,0.821313,0.410106,0.411379,-0.411481
+0.847936,0.821296,0.410128,0.411393,-0.411482
+0.861024,0.820389,0.411353,0.411353,-0.411521
+0.861488,0.820357,0.411352,0.411352,-0.411523
+0.88982,0.818392,0.411266,0.411266,-0.409555
+0.893503,0.818137,0.411255,0.411255,-0.409299
+0.903381,0.818107,0.411225,0.411225,-0.408614
+0.912071,0.818081,0.411198,0.411198,-0.410191
+0.931705,0.818021,0.411139,0.411139,-0.413756
+0.937396,0.818004,0.411121,0.411121,-0.41479
+0.94455,0.817982,0.4111,0.4111,-0.416089
+0.946102,0.817977,0.411095,0.411095,-0.416371

collective_effects/tools.py

 # -*- coding: utf-8 -*-
 """
-Tools and utilities functions for the Wakefield and impedances classes.
+Tools and utilities functions for the Wakefield and Impedances classes.
 
 @author: Alexis Gamelin
-@date: 14/01/2020
+@date: 24/04/2020
 """
 
 import pandas as pd
 import numpy as np
 from scipy.integrate import trapz
 from collective_effects.wakefield import Wakefield, Impedance, WakeFunction
-#from ..Tracking.synchrotron import Synchrotron
 
-def read_CST(file, impedance_type='long'):
-    """Read CST file format into an Impedance object."""
+def read_CST(file, impedance_type='long', divide_by=None):
+    """
+    Read CST file format into an Impedance object.
+    
+    Parameters
+    ----------
+    file : str
+        Path to the file to read.
+    impedance_type : str, optional
+        Type of the Impedance object.
+    divide_by : float, optional
+        Divide the impedance by a value. Mainly used to normalize transverse 
+        impedance by displacement.
+        
+    Returns
+    -------
+    result : Impedance object
+        Data from file.
+    """
     df = pd.read_csv(file, comment="#", header = None, sep = "\t", 
                     names = ["Frequency","Real","Imaginary"])
     df["Frequency"] = df["Frequency"]*1e9 
+    if divide_by is not None:
+        df["Real"] = df["Real"]/divide_by
+        df["Imaginary"] = df["Imaginary"]/divide_by
     df.set_index("Frequency", inplace = True)
     result = Impedance(variable = df.index,
                        function = df["Real"] + 1j*df["Imaginary"],
@@ -24,7 +43,21 @@ def read_CST(file, impedance_type='long'):
     return result
 
 def read_IW2D(file, impedance_type='long'):
-    """Read IW2D file format into an Impedance object."""
+    """
+    Read IW2D file format into an Impedance object.
+    
+    Parameters
+    ----------
+    file : str
+        Path to the file to read.
+    impedance_type : str, optional
+        Type of the Impedance object.
+        
+    Returns
+    -------
+    result : Impedance object
+        Data from file.
+    """
     df = pd.read_csv(file, sep = ' ', header = None, 
                      names = ["Frequency","Real","Imaginary"], skiprows=1)
     df.set_index("Frequency", inplace = True)
@@ -61,12 +94,14 @@ def loss_factor(wake, sigma):
         if(wake.impedance_type == "long"):
             kloss = trapz(x = omega, 
                           y = 1/np.pi*wake.data["real"][positive_index]*gaussian_spectrum)
-        if(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"):
+        elif(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"):
             kloss = trapz(x = omega, 
                           y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum)
-        if(wake.impedance_type == "xquad" or wake.impedance_type == "yquad"):
+        elif(wake.impedance_type == "xquad" or wake.impedance_type == "yquad"):
             kloss = trapz(x = omega, 
                           y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum)
+        else:
+            raise TypeError("Impedance type not recognized.")
 
     if isinstance(wake, WakeFunction):
         # To be implemented.
@@ -187,3 +222,60 @@ def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
                         " and ydip impedance type.")
         
     return Zeff
+
+
+def Yokoya_elliptic(x_radius , y_radius):
+    """
+    Compute Yokoya factors for an elliptic beam pipe.
+    Function adapted from N. Mounet IW2D.
+
+    Parameters
+    ----------
+    x_radius : float
+        Horizontal semi-axis of the ellipse in [m].
+    y_radius : float
+        Vertical semi-axis of the ellipse in [m].
+
+    Returns
+    -------
+    yoklong : float
+        Yokoya factor for the longitudinal impedance.
+    yokxdip : float
+        Yokoya factor for the dipolar horizontal impedance.
+    yokydip : float
+        Yokoya factor for the dipolar vertical impedance.
+    yokxquad : float
+        Yokoya factor for the quadrupolar horizontal impedance.
+    yokyquad : float
+        Yokoya factor for the quadrupolar vertical impedance.
+    """
+    if y_radius < x_radius:
+        small_semiaxis = y_radius
+        large_semiaxis = x_radius
+    else:
+        small_semiaxis = x_radius
+        large_semiaxis = y_radius
+
+    # read Yokoya factors interpolation file
+    # BEWARE: columns are ratio, dipy, dipx, quady, quadx
+    yokoya_file = pd.read_csv("collective_effects/data/" +
+                              "Yokoya_elliptic_from_Elias_USPAS.csv")
+    ratio_col = yokoya_file["x"]
+    # compute semi-axes ratio (first column of this file)
+    ratio = (large_semiaxis - small_semiaxis)/(large_semiaxis + small_semiaxis)
+
+    # 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"])
+        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
+        yokxquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])
+    else:
+        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"])        
+
+    return (yoklong, yokxdip, yokydip, yokxquad, yokyquad)