Repository: atcollab/at (language: python)

Issue:
Cavity with beam loading
Hello,
I wrote this pass method in 2016 and I updated it a few times. I put it in the repository because other people (@lcarver) are interested on it.
I use the matlab function [mxGetComplexDoubles](https://www.mathworks.com/help/matlab/apiref/mxgetcomplexdoubles.html), calling it with the function atGetComplexArray, to get a complex field (BeamVoltage). It works in matlab, but I have no idea of the similar function in python, so the pass method does not work in pyAT at the moment.
@willrogers or @lfarv maybe can help me to have a version working also in pyAT.


Discussion / hints:
@carmignani : I implemented the access to complex attributes in python.
But I have a few questions on this integrator:
- Can you give a simple example to test the integrator ?
- To allow `atradoff` to turnoff the cavity, it would be better to rename the integrator to `RFBeamLoadingCavityPass` since `atradoff` looks for *CavityPass at the end of the passmethod name to turn the element into `IdentityPass` or `DriftPass`,
- It looks like the passmethod modifies the `VoltageBeam` attribute. This is forbidden in Matlab (a function cannot modify its arguments), and it will probably fail in python. Is this desired, and is there another way to achieve the result ?
Hello Laurent,
I will give a simple example for the integrator.
Ok for the change of name.
It is not really needed to change the VoltageBeam field of the element in the matlab/python memory, but it was convenient to have its value at the end of the simulation. It's a trick to be able to split the simulation and restart it from the exact VoltageBeam needed at that turn.
I did not do many tests with the mxGetComplexDoubles function, because when I did the real simulations I was using the old matlab functions mxGetPr and mxGetPi, but I think this trick works with matlab.
If you think it is better to remove this feature, we can decide to pass the value of VoltageBeam and not the pointer as it is done now.
After many trials, I discovered that though the `complex.h` header is available in Microsoft Visual C,  [the standard arithmetic operators are not defined on complex types](https://docs.microsoft.com/en-us/cpp/c-runtime-library/complex-math-support?view=msvc-160). Unless we rewrite the code to avoid complex numbers or switch to C++, we have to skip this integrator on Windows.

Fix patch:
diff --git a/atintegrators/BeamLoadingCavityPass.c b/atintegrators/BeamLoadingCavityPass.c
new file mode 100644
index 000000000..7d696d611
--- /dev/null
+++ b/atintegrators/BeamLoadingCavityPass.c
@@ -0,0 +1,312 @@
+#include "atelem.c"
+#include "atimplib.c"
+#include "attrackfunc.c"
+
+/*
+ * BeamLoadingCavity pass method by Simon White.  
+ *
+ */
+ 
+#define C0     2.99792458e8 
+
+struct elem
+{
+  int nslice;
+  int nturnsw;
+  int mode;
+  double normfact;
+  double phasegain;
+  double voltgain;
+  double *turnhistory;
+  double *z_cuts;
+  double Length;
+  double Voltage;
+  double Energy;
+  double Frequency;
+  double HarmNumber;
+  double TimeLag;
+  double Qfactor;
+  double Rshunt;
+  double Beta;
+  double phis;
+  double *vbunch;
+  double *vbeam_phasor;
+  double *vbeam;
+  double *vcav;
+  double *vgen;
+}; 
+   
+
+void BeamLoadingCavityPass(double *r_in,int num_particles,int nbunch,
+                           double *bunch_spos,double *bunch_currents,
+                           double circumference,int nturn,struct elem *Elem) {
+    /*
+     * r_in - 6-by-N matrix of initial conditions reshaped into
+     * 1-d array of 6*N elements
+     */   
+    long nslice = Elem->nslice;
+    long nturnsw = Elem->nturnsw;
+    long mode = Elem->mode;
+    double normfact = Elem->normfact;  
+    double le = Elem->Length;
+    double energy = Elem->Energy;
+    double rffreq = Elem->Frequency;
+    double harmn = Elem->HarmNumber;
+    double tlag = Elem->TimeLag;
+    double qfactor = Elem->Qfactor;
+    double rshunt = Elem->Rshunt;
+    double beta = Elem->Beta;
+    double *turnhistory = Elem->turnhistory;
+    double *z_cuts = Elem->z_cuts;
+    double *vbunch = Elem->vbunch;
+    double phasegain = Elem->phasegain;
+    double voltgain = Elem->voltgain;
+    double *vbeam_phasor = Elem->vbeam_phasor;
+    double *vbeamk = Elem->vbeam;
+    double *vcavk = Elem->vcav;
+    double *vgenk = Elem->vgen;
+
+    size_t sz = nslice*nbunch*sizeof(double) + num_particles*sizeof(int);
+    int c;
+    int *pslice;
+    double *kz;
+    double freqres = rffreq/(1-tan(vgenk[1])/(2*qfactor));
+    double vgen = vgenk[0];
+    double psi = vgenk[1];
+    
+    void *buffer = atMalloc(sz);
+    double *dptr = (double *) buffer;
+    int *iptr;
+    kz = dptr;
+    dptr += nslice*nbunch;
+    iptr = (int *) dptr;
+    pslice = iptr; iptr += num_particles;
+
+    trackRFCavity(r_in,le,vgen/energy,rffreq,harmn,tlag,-psi,nturn,circumference/C0,num_particles);
+    rotate_table_history(nturnsw,nslice*nbunch,turnhistory,circumference);
+    slice_bunch(r_in,num_particles,nslice,nturnsw,nbunch,bunch_spos,bunch_currents,
+                turnhistory,pslice,z_cuts);
+    if(mode==2){
+        compute_kicks_phasor(nslice,nbunch,nturnsw,turnhistory,normfact,kz,freqres,
+                             qfactor,rshunt,vbeam_phasor,circumference,energy,beta,
+                             vbeamk,vbunch);                        
+    }else if(mode==1){
+        compute_kicks_longres(nslice,nbunch,nturnsw,turnhistory,normfact,kz,freqres,
+                              qfactor,rshunt,beta,vbeamk,energy,vbunch);
+    }
+    /*apply kicks and RF*/
+    /* OpenMP not efficient. Too much shared data ?
+    #pragma omp parallel for if (num_particles > OMP_PARTICLE_THRESHOLD) default(none) \
+    shared(r_in,num_particles,pslice,kz) private(c)
+    */   
+    for (c=0; c<num_particles; c++) {
+        double *r6 = r_in+c*6;
+        int islice=pslice[c];
+        if (!atIsNaN(r6[0])) {         
+            r6[4] += kz[islice]; 
+        }
+    }
+    update_vgen(vbeamk,vcavk,vgenk,phasegain,voltgain); 
+    atFree(buffer);
+}
+
+
+#if defined(MATLAB_MEX_FILE) || defined(PYAT)
+ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem,
+        double *r_in, int num_particles, struct parameters *Param)
+{
+    double rl = Param->RingLength;
+    int nturn=Param->nturn;
+    if (!Elem) {
+        long nslice,nturns,mode;
+        double wakefact;
+        double normfact, phasegain, voltgain;
+        double *turnhistory;
+        double *z_cuts;
+        double Energy, Frequency, TimeLag, Length;
+        double qfactor,rshunt,beta;
+        double *vbunch;
+        double *vbeam_phasor;
+        double *vbeam;
+        double *vgen;
+        double *vcav;
+
+        /*attributes for RF cavity*/
+        Length=atGetDouble(ElemData,"Length"); check_error();
+        Energy=atGetDouble(ElemData,"Energy"); check_error();
+        Frequency=atGetDouble(ElemData,"Frequency"); check_error();
+        TimeLag=atGetOptionalDouble(ElemData,"TimeLag",0); check_error();
+        /*attributes for resonator*/
+        nslice=atGetLong(ElemData,"_nslice"); check_error();
+        nturns=atGetLong(ElemData,"_nturns"); check_error();
+        mode=atGetLong(ElemData,"_mode"); check_error();
+        wakefact=atGetDouble(ElemData,"_wakefact"); check_error();
+        qfactor=atGetDouble(ElemData,"Qfactor"); check_error();
+        rshunt=atGetDouble(ElemData,"Rshunt"); check_error();
+        beta=atGetDouble(ElemData,"_beta"); check_error();
+        normfact=atGetDouble(ElemData,"NormFact"); check_error();
+        phasegain=atGetDouble(ElemData,"PhaseGain"); check_error();
+        voltgain=atGetDouble(ElemData,"VoltGain"); check_error();
+        turnhistory=atGetDoubleArray(ElemData,"_turnhistory"); check_error();
+        vbunch=atGetDoubleArray(ElemData,"_vbunch"); check_error();
+        vbeam=atGetDoubleArray(ElemData,"_vbeam"); check_error();
+        vcav=atGetDoubleArray(ElemData,"_vcav"); check_error();
+        vgen=atGetDoubleArray(ElemData,"_vgen"); check_error();
+        vbeam_phasor=atGetDoubleArray(ElemData,"_vbeam_phasor"); check_error(); 
+        /*optional attributes*/
+        z_cuts=atGetOptionalDoubleArray(ElemData,"ZCuts"); check_error();
+       
+        Elem = (struct elem*)atMalloc(sizeof(struct elem));
+        
+        Elem->Length=Length;
+        Elem->mode=mode;
+        Elem->Frequency=Frequency;
+        Elem->HarmNumber=round(Frequency*rl/C0);
+        Elem->Energy = Energy;
+        Elem->TimeLag=TimeLag;   
+        Elem->nslice=nslice;
+        Elem->nturnsw=nturns;
+        Elem->normfact=normfact*wakefact;
+        Elem->turnhistory=turnhistory;
+        Elem->Qfactor = qfactor;
+        Elem->Rshunt = rshunt;
+        Elem->Beta = beta;
+        Elem->z_cuts=z_cuts;
+        Elem->vbunch = vbunch;
+        Elem->vbeam = vbeam;
+        Elem->vgen = vgen;
+        Elem->vcav = vcav;
+        Elem->phasegain = phasegain;
+        Elem->voltgain = voltgain;
+        Elem->vbeam_phasor = vbeam_phasor;
+    }
+    if(num_particles<Param->nbunch){
+        atError("Number of particles has to be greater or equal to the number of bunches.");
+    }else if (num_particles%Param->nbunch!=0){
+        atWarning("Number of particles not a multiple of the number of bunches: uneven bunch load.");
+    }
+    #ifdef _MSC_VER
+    if(Elem->mode==1){
+        atError("Beam loading Phasor mode not implemented in Windows.");
+    }
+    #endif
+    BeamLoadingCavityPass(r_in,num_particles,Param->nbunch,Param->bunch_spos,
+                          Param->bunch_currents,rl,nturn,Elem);
+    return Elem;
+}
+
+MODULE_DEF(BeamLoadingCavityPass)       /* Dummy module initialisation */
+#endif /*defined(MATLAB_MEX_FILE) || defined(PYAT)*/
+
+#if defined(MATLAB_MEX_FILE)
+
+void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
+{ 	
+  if(nrhs == 2)
+  {
+      long nslice,nturns,mode;
+      double wakefact;
+      double normfact, phasegain, voltgain;
+      double *turnhistory;
+      double *z_cuts;
+      double Energy, Frequency, TimeLag, Length;
+      double qfactor,rshunt,beta;
+      double *vbunch;
+      double *vbeam_phasor;
+      double *vbeam;
+      double *vgen;
+      double *vcav;
+      /*attributes for RF cavity*/
+      Length=atGetDouble(ElemData,"Length"); check_error();
+      Energy=atGetDouble(ElemData,"Energy"); check_error();
+      Frequency=atGetDouble(ElemData,"Frequency"); check_error();
+      TimeLag=atGetOptionalDouble(ElemData,"TimeLag",0); check_error();
+      /*attributes for resonator*/
+      nslice=atGetLong(ElemData,"_nslice"); check_error();
+      nturns=atGetLong(ElemData,"_nturns"); check_error();
+      mode=atGetLong(ElemData,"_mode"); check_error();
+      wakefact=atGetDouble(ElemData,"_wakefact"); check_error();
+      qfactor=atGetDouble(ElemData,"Qfactor"); check_error();
+      rshunt=atGetDouble(ElemData,"Rshunt"); check_error();
+      beta=atGetDouble(ElemData,"_beta"); check_error();
+      normfact=atGetDouble(ElemData,"NormFact"); check_error();
+      phasegain=atGetDouble(ElemData,"PhaseGain"); check_error();
+      voltgain=atGetDouble(ElemData,"VoltGain"); check_error();
+      turnhistory=atGetDoubleArray(ElemData,"_turnhistory"); check_error();
+      vbunch=atGetDoubleArray(ElemData,"_vbunch"); check_error();
+      vbeam=atGetDoubleArray(ElemData,"_vbeam"); check_error();
+      vcav=atGetDoubleArray(ElemData,"_vcav"); check_error();
+      vgen=atGetDoubleArray(ElemData,"_vgen"); check_error();
+      vbeam_phasor=atGetDoubleArray(ElemData,"_vbeam_phasor"); check_error(); 
+      /*optional attributes*/
+      z_cuts=atGetOptionalDoubleArray(ElemData,"ZCuts"); check_error();
+      Elem = (struct elem*)atMalloc(sizeof(struct elem));
+      Elem->Length=Length;
+      Elem->mode=mode;
+      Elem->Frequency=Frequency;
+      Elem->HarmNumber=1;
+      Elem->Energy = Energy;
+      Elem->TimeLag=TimeLag;   
+      Elem->nslice=nslice;
+      Elem->nturnsw=nturns;
+      Elem->normfact=normfact*wakefact;
+      Elem->turnhistory=turnhistory;
+      Elem->Qfactor = qfactor;
+      Elem->Rshunt = rshunt;
+      Elem->Beta = beta;
+      Elem->z_cuts=z_cuts;
+      Elem->vbunch = vbunch;
+      Elem->vbeam = vbeam;
+      Elem->vgen = vgen;
+      Elem->vcav = vcav;
+      Elem->phasegain = phasegain;
+      Elem->voltgain = voltgain;
+      Elem->vbeam_phasor = vbeam_phasor;
+      if (mxGetM(prhs[1]) != 6) mexErrMsgIdAndTxt("AT:WrongArg","Second argument must be a 6 x N matrix");
+      /* ALLOCATE memory for the output array of the same size as the input  */
+      plhs[0] = mxDuplicateArray(prhs[1]);
+      r_in = mxGetDoubles(plhs[0]);
+      double *bspos = malloc(sizeof(double));
+      double *bcurr = malloc(sizeof(double));
+      bspos[0] = 0.0;
+      bcurr[0] = 0.0;
+      BeamLoadingCavityPass(r_in,num_particles,1,bspos,bcur,1,0,Elem);
+      free(bspos);
+      free(bcur);
+  }
+  else if (nrhs == 0)
+  {   /* return list of required fields */
+      plhs[0] = mxCreateCellMatrix(19,1);
+      mxSetCell(plhs[0],0,mxCreateString("Length"));
+      mxSetCell(plhs[0],1,mxCreateString("Energy"));
+      mxSetCell(plhs[0],2,mxCreateString("Frequency"));
+      mxSetCell(plhs[0],3,mxCreateString("_nslice"));
+      mxSetCell(plhs[0],4,mxCreateString("_nturns"));
+      mxSetCell(plhs[0],5,mxCreateString("_mode"));
+      mxSetCell(plhs[0],6,mxCreateString("_wakefact"));
+      mxSetCell(plhs[0],7,mxCreateString("Qfactor"));
+      mxSetCell(plhs[0],8,mxCreateString("Rshunt"));
+      mxSetCell(plhs[0],9,mxCreateString("_beta"));
+      mxSetCell(plhs[0],10,mxCreateString("NormFact"));
+      mxSetCell(plhs[0],11,mxCreateString("PhaseGain"));
+      mxSetCell(plhs[0],12,mxCreateString("VoltGain"));
+      mxSetCell(plhs[0],13,mxCreateString("_turnhistory"));
+      mxSetCell(plhs[0],14,mxCreateString("_vbunch"));
+      mxSetCell(plhs[0],15,mxCreateString("_vbeam"));
+      mxSetCell(plhs[0],16,mxCreateString("_vcav"));
+      mxSetCell(plhs[0],17,mxCreateString("_vgen"));
+      mxSetCell(plhs[0],18,mxCreateString("_vbeam_phasor"));
+      if(nlhs>1) /* optional fields */
+      {
+          plhs[1] = mxCreateCellMatrix(2,1);
+          mxSetCell(plhs[1],0,mxCreateString("TimeLag"));
+          mxSetCell(plhs[0],1,mxCreateString("ZCuts"));
+      }
+  }
+  else
+  {
+      mexErrMsgIdAndTxt("AT:WrongArg","Needs 0 or 2 arguments");
+  }
+  
+}
+#endif /* MATLAB_MEX_FILE */
diff --git a/atintegrators/CavityPass.c b/atintegrators/CavityPass.c
index 14762679f..81cf917a5 100644
--- a/atintegrators/CavityPass.c
+++ b/atintegrators/CavityPass.c
@@ -2,7 +2,7 @@
 
  * Accelerator Toolbox
  * Revision 3/10/04
- * A.Terebilo terebilo@ssrl.slac.stanford.edu
+ * A.Terebilo
  */
 
 #include "atelem.c"
diff --git a/atintegrators/CorrectorPass.c b/atintegrators/CorrectorPass.c
index 9b8b6b86e..a1d4d964c 100644
--- a/atintegrators/CorrectorPass.c
+++ b/atintegrators/CorrectorPass.c
@@ -1,7 +1,7 @@
 /* CorrectorPass.c 
    Accelerator Toolbox 
    Revision 10/11/01
-   A.Terebilo terebilo@ssrl.slac.stanford.edu
+   A.Terebilo
    CorrectorPass expects an element to have a fields:
    Length, KickAngle
 */
diff --git a/atintegrators/GWigSymplecticPass.c b/atintegrators/GWigSymplecticPass.c
index 2dbf3c185..a943be136 100644
--- a/atintegrators/GWigSymplecticPass.c
+++ b/atintegrators/GWigSymplecticPass.c
@@ -6,10 +6,10 @@
  *---------------------------------------------------------------------------
  * Modification Log:
  * -----------------
- * .02  2003-06-18     J. Li, jing@fel.duke.edu
+ * .02  2003-06-18     J. Li
  *				Cleanup the code
  *
- * .01  2003-04-20     YK Wu, wu@fel.duke.edu
+ * .01  2003-04-20     YK Wu
  *				GWiggler interface
  *
  *---------------------------------------------------------------------------
diff --git a/atintegrators/GWigSymplecticRadPass.c b/atintegrators/GWigSymplecticRadPass.c
index 396f47ad2..bc4dfdb9d 100644
--- a/atintegrators/GWigSymplecticRadPass.c
+++ b/atintegrators/GWigSymplecticRadPass.c
@@ -6,10 +6,10 @@
  *---------------------------------------------------------------------------
  * Modification Log:
  * -----------------
- * .02  2003-06-18     J. Li, jing@fel.duke.edu
+ * .02  2003-06-18     J. Li
  *				Cleanup the code
  *
- * .01  2003-04-20     YK Wu, wu@fel.duke.edu
+ * .01  2003-04-20     YK Wu
  *				GWiggler interface
  *
  *---------------------------------------------------------------------------
diff --git a/atintegrators/IdTablePass.c b/atintegrators/IdTablePass.c
index 6d2d0c395..1739e85c9 100644
--- a/atintegrators/IdTablePass.c
+++ b/atintegrators/IdTablePass.c
@@ -1,7 +1,7 @@
 /* IdTablePass.c
  * Accelerator Toolbox
  * Created: 13/11/08
- * Z.Mart?? zeus@cells.es
+ * Zeus Marti
  *
  * Based in the matlab routine:
  * WigTablePass.m - The tracking table is described in
diff --git a/atintegrators/IdentityPass.c b/atintegrators/IdentityPass.c
index b0a0b005b..1ba8085d4 100644
--- a/atintegrators/IdentityPass.c
+++ b/atintegrators/IdentityPass.c
@@ -1,7 +1,7 @@
 /* IdentityPass.c 
    Accelerator Toolbox
    Revision 7/16/03
-   A.Terebilo terebilo@slac.stanford.edu
+   A.Terebilo
 */
 
 #include "atelem.c"
diff --git a/atintegrators/ImpedanceTablePass.c b/atintegrators/ImpedanceTablePass.c
index 3548c913e..d2bc29b2d 100755
--- a/atintegrators/ImpedanceTablePass.c
+++ b/atintegrators/ImpedanceTablePass.c
@@ -4,7 +4,7 @@
 
 /*
  * Impedance pass method by Simon White.  
- * User may contact simon.white@esrf.fr for questions and comments.
+ * 
  */
 
 struct elem
diff --git a/atintegrators/QuadLinearPass.c b/atintegrators/QuadLinearPass.c
index 4f02006a8..48732f2d3 100644
--- a/atintegrators/QuadLinearPass.c
+++ b/atintegrators/QuadLinearPass.c
@@ -1,7 +1,7 @@
 /* QuadLinearPass.c 
    Accelerator Toolbox 
    Revision 6/27/00
-   A.Terebilo terebilo@ssrl.slac.stanford.edu
+   A.Terebilo
 */
 
 
diff --git a/atintegrators/RFCavityPass.c b/atintegrators/RFCavityPass.c
index 2c6fd32be..ad96133b4 100755
--- a/atintegrators/RFCavityPass.c
+++ b/atintegrators/RFCavityPass.c
@@ -5,13 +5,8 @@
  *  Nicola Carmignani
  */
 
-#include <math.h>
-#include "atelem.c"
-#include "atlalib.c"
-#include "driftkickrad.c"
+#include "attrackfunc.c"
 
-#define TWOPI  6.28318530717959
-#define C0  	2.99792458e8 
 
 struct elem 
 {
@@ -32,29 +27,7 @@ void RFCavityPass(double *r_in, double le, double nv, double freq, double h, dou
    1-d array of 6*N elements
 */
 {
-    int c;
-
-    if (le == 0) {
-        for (c = 0; c<num_particles; c++) {
-            double *r6 = r_in+c*6;
-            if(!atIsNaN(r6[0]))
-                r6[4] += -nv*sin(TWOPI*freq*((r6[5]-lag)/C0 - (h/freq-T0)*nturn) - philag);
-        }
-    }
-    else {
-        double halflength = le/2;
-        for (c = 0;c<num_particles;c++) {
-            double *r6 = r_in+c*6;
-            if(!atIsNaN(r6[0]))  {
-                /* Propagate through a drift equal to half cavity length */
-                drift6(r6, halflength);
-                /* Longitudinal momentum kick */
-                r6[4] += -nv*sin(TWOPI*freq*((r6[5]-lag)/C0 - (h/freq-T0)*nturn) - philag);
-                /* Propagate through a drift equal to half cavity length */
-                drift6(r6, halflength);
-            }
-        }
-    }
+    trackRFCavity(r_in, le, nv, freq, h, lag, philag, nturn, T0, num_particles);
 }
 
 #if defined(MATLAB_MEX_FILE) || defined(PYAT)
diff --git a/atintegrators/SolenoidLinearPass.c b/atintegrators/SolenoidLinearPass.c
index c321a7048..5cd1698cb 100644
--- a/atintegrators/SolenoidLinearPass.c
+++ b/atintegrators/SolenoidLinearPass.c
@@ -1,7 +1,7 @@
 /* SolenoidLinearPass.c 
    Accelerator Toolbox 
    Revision 7/17/03
-   A.Terebilo terebilo@slac.stanford.edu
+   A.Terebilo
 */
 
 #include "atelem.c"
diff --git a/atintegrators/ThinMPolePass.c b/atintegrators/ThinMPolePass.c
index 06ba0c952..4fa095565 100644
--- a/atintegrators/ThinMPolePass.c
+++ b/atintegrators/ThinMPolePass.c
@@ -1,6 +1,6 @@
 /* ThinMPolePass
    Accelerator Toolbox 
-   A.Terebilo terebilo@ssrl.slac.stanford.edu
+   A.Terebilo
 */
 
 #include "atelem.c"
@@ -94,9 +94,9 @@ MODULE_DEF(ThinMPolePass)        /* Dummy module initialisation */
 void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
 {
     if (nrhs == 2 ) {
-        double *r_in;
-        const mxArray *ElemData = prhs[0];
-        int num_particles = mxGetN(prhs[1]);
+        double *r_in;
+        const mxArray *ElemData = prhs[0];
+        int num_particles = mxGetN(prhs[1]);
         double bax = 0.0, bay = 0.0;
         int MaxOrder;
         double *PolynomA, *PolynomB, *BendingAngle, *R1, *R2, *T1, *T2;
diff --git a/atintegrators/VariableThinMPolePass.c b/atintegrators/VariableThinMPolePass.c
index f87de0928..a52cfc33d 100644
--- a/atintegrators/VariableThinMPolePass.c
+++ b/atintegrators/VariableThinMPolePass.c
@@ -1,6 +1,6 @@
 /* VariableThinMPolePass
    Accelerator Toolbox 
-   S.White simon.white@esrf.fr
+   S.White
 */
 
 #include "atelem.c"
diff --git a/atintegrators/atimplib.c b/atintegrators/atimplib.c
index b79e8442d..6a745246d 100644
--- a/atintegrators/atimplib.c
+++ b/atintegrators/atimplib.c
@@ -1,12 +1,17 @@
 #include "atelem.c"
 #include <math.h>
 #include <float.h>
+#include <complex.h>
 #ifdef MPI
 #include <mpi.h>
 #include <mpi4py/mpi4py.h>
 #endif
 
 
+#define TWOPI  6.28318530717959
+#define C0     2.99792458e8 
+
+
 int binarySearch(double *array,double value,int upper,int lower,int nStep){
     int pivot = (int)(lower+upper)/2;
     if ((upper-lower)<=1){
@@ -122,7 +127,7 @@ void slice_bunch(double *r_in,int num_particles,int nslice,int nturns,
                  int nbunch,double *bunch_spos,double *bunch_currents,
                  double *turnhistory,int *pslice,double *z_cuts){
     
-    int i,ii,ib,is;
+    int i,ii,ib;
     double *rtmp;
     
     double *smin = malloc(nbunch*sizeof(double));
@@ -255,4 +260,183 @@ void compute_kicks(int nslice,int nturns,int nelem,
 };
 
 
+double *wakefunc_long_resonator(double ds, double freqres, double qfactor, double rshunt, double beta) {
+
+    double omega, alpha, omegabar;
+    double *wake = malloc(2*sizeof(double));
+    wake[0] = 0.0;
+    wake[1] = 0.0;
+    double dt;
+    
+    omega = TWOPI * freqres;
+    alpha = omega / (2 * qfactor);
+    omegabar = sqrt(fabs(omega*omega - alpha*alpha));
+    
+    dt = -ds/(beta * C0);      
+             
+    if (dt==0) {
+        wake[0] = rshunt * alpha;
+        wake[1] = 0.0;
+    } else if (dt<0) {
+        if (qfactor > 0.5) {
+            wake[0] = 2 * rshunt * alpha * exp(alpha * dt) * (cos(omegabar * dt) + \
+                      alpha / omegabar * sin(omegabar*dt));
+            wake[1] = 2 * rshunt * alpha * exp(alpha * dt) * (sin(omegabar * dt) - \
+                      alpha / omegabar * cos(omegabar*dt));
+        } else if (qfactor == 0.5) {
+            wake[0] = 2 * rshunt * alpha * exp(alpha * dt) * (1. + alpha * dt);
+            wake[1] = 0.0;
+        } else if (qfactor < 0.5) {
+            wake[0] = 2 * rshunt * alpha * exp(alpha * dt) * (cosh(omegabar * dt) + \
+                      alpha / omegabar * sinh(omegabar * dt)); 
+            wake[1] = 2 * rshunt * alpha * exp(alpha * dt) * (sinh(omegabar * dt) - \
+                      alpha / omegabar * cosh(omegabar*dt));
+        }       
+    } else {
+        wake[0] = 0.0;
+        wake[1] = 0.0;
+    }            
+    return wake;
+}
+
+
+void compute_kicks_longres(int nslice,int nbunch,int nturns, double *turnhistory,double normfact,
+                           double *kz,double freq, double qfactor, double rshunt,
+                           double beta, double *vbeamk, double energy, double *vbunch) {
+
+    int rank=0;
+    int size=1;
+    int i,ii,ib;
+    double ds,wi;
+    double *turnhistoryZ = turnhistory+nslice*nbunch*nturns*2;
+    double *turnhistoryW = turnhistory+nslice*nbunch*nturns*3;
+    double *wake = malloc(2*sizeof(double));
+    double vba, vbp;
+    double *vbr = vbunch;
+    double *vbi = vbunch+nbunch;
+
+    for (i=0;i<nslice*nbunch;i++) {
+        ib = (int)(i/nslice);
+        kz[i]=0.0;
+        vbr[ib] = 0.0;
+        vbi[ib] = 0.0;
+    }
+
+    vbeamk[0] = 0.0;
+    vbeamk[1] = 0.0;
+    wake[0] = 0.0;
+    wake[1] = 0.0;
+
+
+    #ifdef MPI
+    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+    MPI_Comm_size(MPI_COMM_WORLD, &size);
+    #endif
+    for(i=nslice*nbunch*(nturns-1);i<nslice*nbunch*nturns;i++){  
+        ib = (int)((i-nslice*nbunch*(nturns-1))/nslice);
+        if(turnhistoryW[i]>0.0 && rank==(i+size)%size){
+            for (ii=0;ii<nslice*nbunch*nturns;ii++){
+                ds = turnhistoryZ[i]-turnhistoryZ[ii];
+                if(turnhistoryW[ii]>0.0 && ds>=0){
+                    wi = turnhistoryW[ii];
+                    wake = wakefunc_long_resonator(ds,freq,qfactor,rshunt,beta);       
+                    kz[i-nslice*nbunch*(nturns-1)] += normfact*wi*wake[0];
+                    vbeamk[0] += normfact*wi*wake[0]*energy/nbunch;
+                    vbeamk[1] -= normfact*wi*wake[1]*energy/nbunch;
+                    vbr[ib] += normfact*wi*wake[0]*energy;
+                    vbi[ib] -= normfact*wi*wake[1]*energy;
+                }            
+            }
+        }
+    }
+    #ifdef MPI
+    MPI_Allreduce(MPI_IN_PLACE,kz,nslice*nbunch,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
+    MPI_Allreduce(MPI_IN_PLACE,vbeamk,2,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
+    MPI_Allreduce(MPI_IN_PLACE,vbr,nbunch,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
+    MPI_Allreduce(MPI_IN_PLACE,vbi,nbunch,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
+    MPI_Barrier(MPI_COMM_WORLD);
+    #endif
+    vba = sqrt(vbeamk[0]*vbeamk[0]+vbeamk[1]*vbeamk[1]);   
+    vbp = atan2(vbeamk[1],vbeamk[0]);
+    vbeamk[0] = vba;
+    vbeamk[1] = vbp;
+    for(i=0;i<nbunch;i++){
+        double vr = vbr[i];
+        double vi = vbi[i];
+        vbr[i] = sqrt(vr*vr+vi*vi); 
+        vbi[i] = atan2(vi,vr);
+    }
+    free(wake);
+};
+
+
+void compute_kicks_phasor(int nslice, int nbunch, int nturns, double *turnhistory,
+                          double normfact, double *kz,double freq, double qfactor,
+                          double rshunt, double *vbeam, double circumference,
+                          double energy, double beta, double *vbeamk, double *vbunch){  
+    #ifndef _MSC_VER  
+    int i,ib;
+    double wi;
+    double dt =0.0;
+    double *turnhistoryZ = turnhistory+nslice*nbunch*nturns*2;
+    double *turnhistoryW = turnhistory+nslice*nbunch*nturns*3;
+    double omr = TWOPI*freq;
+    double complex vbeamc = vbeam[0]*cexp(I*vbeam[1]);
+    double complex vbeamkc = 0.0;
+    double kick = rshunt*omr/(2*qfactor);
+    double bc = beta*C0;
+    double *vbr = vbunch;
+    double *vbi = vbunch+nbunch;
+    
+    for (i=0;i<nslice*nbunch;i++) {
+        ib = (int)(i/nslice);
+        kz[i]=0.0;
+        vbr[ib] = 0.0;
+        vbi[ib] = 0.0;
+    }
+    
+    for(i=nslice*nbunch*(nturns-1);i<nslice*nbunch*nturns;i++){
+        ib = (int)((i-nslice*nbunch*(nturns-1))/nslice);
+        wi = turnhistoryW[i];
+        if(i==nslice*nbunch*(nturns-1)){
+            dt = (circumference+turnhistoryZ[i])/bc;
+        }else{
+            dt = (turnhistoryZ[i]-turnhistoryZ[i-1])/bc;
+        }      
+        vbeamc *= cexp((I*omr-omr/(2*qfactor))*dt);
+        vbeamkc += vbeamc+normfact*wi*kick*energy;
+        vbr[ib] += creal(vbeamc+normfact*wi*kick*energy);
+        vbi[ib] += cimag(vbeamc+normfact*wi*kick*energy);
+        kz[i-nslice*nbunch*(nturns-1)] = creal(vbeamc/energy+normfact*wi*kick);
+        vbeamc += 2*normfact*wi*kick*energy;    
+    }
+    dt = -turnhistoryZ[nslice*nbunch*nturns-1]/bc;
+    vbeamc = vbeamc*cexp((I*omr-omr/(2*qfactor))*dt);
+    vbeam[0] = cabs(vbeamc);
+    vbeam[1] = carg(vbeamc); 
+    vbeamkc = vbeamkc/nslice/nbunch;
+    vbeamk[0] = cabs(vbeamkc);
+    vbeamk[1] = carg(vbeamkc);   
+    
+    for(i=0;i<nbunch;i++){
+        double vr = vbr[i]/nslice;
+        double vi = vbi[i]/nslice;
+        vbr[i] = sqrt(vr*vr+vi*vi); 
+        vbi[i] = atan2(vi,vr);
+    }
+    #endif    
+};
+
 
+void update_vgen(double *vbeam,double *vcav,double *vgen,double voltgain,double phasegain){
+    double vbeamr = vbeam[0]*cos(vbeam[1]);
+    double vbeami = vbeam[0]*sin(vbeam[1]);
+    double vcavr = vcav[0]*cos(vcav[1]);
+    double vcavi = vcav[0]*sin(vcav[1]);   
+    double vgenr = vcavr - vbeamr;
+    double vgeni = vcavi - vbeami; 
+    double vga = sqrt(vgenr*vgenr+vgeni*vgeni);   
+    double vgp = atan2(vgeni,vgenr)-vcav[1];
+    vgen[0] += (vga-vgen[0])*voltgain;
+    vgen[1] += (vgp-vgen[1])*phasegain;
+}
diff --git a/atintegrators/attrackfunc.c b/atintegrators/attrackfunc.c
new file mode 100644
index 000000000..941c8911d
--- /dev/null
+++ b/atintegrators/attrackfunc.c
@@ -0,0 +1,39 @@
+#include <math.h>
+#include "atelem.c"
+#include "driftkickrad.c"
+
+#define TWOPI  6.28318530717959
+#define C0  	2.99792458e8 
+
+void trackRFCavity(double *r_in, double le, double nv, double freq, double h, double lag, double philag,
+                  int nturn, double T0, int num_particles)
+/* le - physical length
+   nv - peak voltage (V) normalized to the design enegy (eV)
+   r is a 6-by-N matrix of initial conditions reshaped into
+   1-d array of 6*N elements
+*/
+{
+    int c;
+
+    if (le == 0) {
+        for (c = 0; c<num_particles; c++) {
+            double *r6 = r_in+c*6;
+            if(!atIsNaN(r6[0]))
+                r6[4] += -nv*sin(TWOPI*freq*((r6[5]-lag)/C0 - (h/freq-T0)*nturn) - philag);
+        }
+    }
+    else {
+        double halflength = le/2;
+        for (c = 0;c<num_particles;c++) {
+            double *r6 = r_in+c*6;
+            if(!atIsNaN(r6[0]))  {
+                /* Propagate through a drift equal to half cavity length */
+                drift6(r6, halflength);
+                /* Longitudinal momentum kick */
+                r6[4] += -nv*sin(TWOPI*freq*((r6[5]-lag)/C0 - (h/freq-T0)*nturn) - philag);
+                /* Propagate through a drift equal to half cavity length */
+                drift6(r6, halflength);
+            }
+        }
+    }
+}
diff --git a/atintegrators/gwig.c b/atintegrators/gwig.c
index be0ff1b59..162aec500 100644
--- a/atintegrators/gwig.c
+++ b/atintegrators/gwig.c
@@ -2,15 +2,15 @@
  *----------------------------------------------------------------------------
  * Modification Log:
  * -----------------
- * .04  2003-04-29      YK Wu, Duke University, wu@fel.duke.edu 
+ * .04  2003-04-29      YK Wu, Duke University
  *			using scientific notation for constants. 
  *                      Checked with TRACY pascal code.
  *                      Computing differential pathlength only.
  *
- * .03  2003-04-28      YK Wu, Duke University, wu@fel.duke.edu 
+ * .03  2003-04-28      YK Wu, Duke University 
  *			Convert to C code and cross-checked with the pascal version;
  *
- * .02  2001-12-xx      Y. K. Wu, Duke University, wu@fel.duke.edu
+ * .02  2001-12-xx      Y. K. Wu, Duke University
  *                      Implementing DA version of the wiggler integrator for Pascal.
  *                      Gauge is disabled !!! (Dec. 4, 2001)
  *  
diff --git a/atintegrators/gwigR.c b/atintegrators/gwigR.c
index 98d93969e..fd1995222 100644
--- a/atintegrators/gwigR.c
+++ b/atintegrators/gwigR.c
@@ -2,15 +2,15 @@
  *----------------------------------------------------------------------------
  * Modification Log:
  * -----------------
- * .04  2003-04-29      YK Wu, Duke University, wu@fel.duke.edu 
+ * .04  2003-04-29      YK Wu, Duke University 
  *			using scientific notation for constants. 
  *                      Checked with TRACY pascal code.
  *                      Computing differential pathlength only.
  *
- * .03  2003-04-28      YK Wu, Duke University, wu@fel.duke.edu 
+ * .03  2003-04-28      YK Wu, Duke University 
  *			Convert to C code and cross-checked with the pascal version;
  *
- * .02  2001-12-xx      Y. K. Wu, Duke University, wu@fel.duke.edu
+ * .02  2001-12-xx      Y. K. Wu, Duke University
  *                      Implementing DA version of the wiggler integrator for Pascal.
  *                      Gauge is disabled !!! (Dec. 4, 2001)
  *  
diff --git a/atintegrators/interpolate.c b/atintegrators/interpolate.c
index d4766ed69..09dc5f9da 100644
--- a/atintegrators/interpolate.c
+++ b/atintegrators/interpolate.c
@@ -1,7 +1,7 @@
 /* Interpolate.c 
    Accelerator Toolbox 
    Created: 14/11/08
-   Z.Martí zeus@cells.es
+   Z.Martí
  
   
   
diff --git a/pyat/at/collective/__init__.py b/pyat/at/collective/__init__.py
index cc1986d37..264492fa5 100644
--- a/pyat/at/collective/__init__.py
+++ b/pyat/at/collective/__init__.py
@@ -4,3 +4,4 @@
 from .wake_elements import *
 from .wake_functions import *
 from .wake_object import *
+from .beam_loading import *
diff --git a/pyat/at/collective/beam_loading.py b/pyat/at/collective/beam_loading.py
new file mode 100644
index 000000000..751955702
--- /dev/null
+++ b/pyat/at/collective/beam_loading.py
@@ -0,0 +1,281 @@
+import numpy
+from enum import IntEnum
+from ..lattice import Lattice
+from at.lattice import RFCavity, Collective
+from at.lattice.elements import _array
+from at.lattice.utils import Refpts, uint32_refpts, make_copy
+from at.physics import get_timelag_fromU0
+from at.constants import clight
+from typing import Sequence, Optional, Union
+
+
+class BLMode(IntEnum):
+    WAKE = 1
+    PHASOR = 2
+
+
+def add_beamloading(ring: Lattice, qfactor: Union[float, Sequence[float]],
+                    rshunt: Union[float, Sequence[float]],
+                    cavpts: Refpts = None, copy: Optional[bool] = False,
+                    **kwargs):
+    r"""Function to add beam loading to a cavity element, the cavity
+    element is changed to a beam loading element that combines the energy
+    kick from both the cavity and the resonator
+
+    Parameters:
+        ring:           Lattice object
+        qfactor:        Q factor. Scalar or array of float are accepted
+        rshunt:         Shunt impedance, [:math:`\Omega`] for longitudinal
+                        Scalar or array of float are accepted
+
+    Keyword Arguments:
+        cavpts (Refpts):    refpts of the cavity. If None (default) apply
+                            to all cavity in the ring
+        Nslice (int):       Number of slices per bunch. Default: 101
+        Nturns (int):       Number of turn for the wake field. Default: 1
+        ZCuts:              Limits for fixed slicing, default is adaptive
+        NormFact (float):   Normalization factor
+        mode (BLMode):  method for beam loading calculation BLMode.PHASOR
+            (default) uses the phasor method, BLMode.WAKE uses the wake
+            function. For high Q resonator, the phasor method should be
+            used
+        copy:       If True, returns a shallow copy of ring with new
+                    beam loading elements. Otherwise, modify ring in-place
+    """
+    @make_copy(copy)
+    def apply(ring, cavpts, newelems):
+        for ref, elem in zip(cavpts, newelems):
+            ring[ref] = elem
+
+    if cavpts is None:
+        cavpts = ring.get_refpts(RFCavity)
+    else:
+        cavpts = uint32_refpts(cavpts, len(ring))
+    qfactor = numpy.broadcast_to(qfactor, (len(cavpts), ))
+    rshunt = numpy.broadcast_to(rshunt, (len(cavpts), ))
+    new_elems = []
+    for ref, qf, rs in zip(cavpts, qfactor, rshunt):
+        cav = ring[ref]
+        assert isinstance(cav, RFCavity), \
+            'Beam loading can only be assigned to a cavity element'
+        new_elems.append(BeamLoadingElement.build_from_cav(cav, ring, qf,
+                                                           rs, **kwargs))
+    return apply(ring, cavpts, new_elems)
+
+
+def remove_beamloading(ring, cavpts: Refpts = None,
+                       copy: Optional[bool] = False):
+    """Function to remove beam loading from a cavity element, the beam
+    loading element is changed to a beam loading element that
+    combines the energy kick from both the cavity and the resonator
+
+    Parameters:
+        ring:           Lattice object
+
+    Keyword Arguments:
+        cavpts (Refpts):    refpts of the beam loading Elements.
+                            If None (default) apply to all elements
+        copy:       If True, returns a shallow copy of ring with new
+                      cavity elements. Otherwise, modify ring in-place
+    """
+    @make_copy(copy)
+    def apply(ring, cavpts, newelems):
+        for ref, elem in zip(cavpts, newelems):
+            ring[ref] = elem
+
+    if cavpts is None:
+        cavpts = ring.get_refpts(BeamLoadingElement)
+    else:
+        cavpts = uint32_refpts(cavpts, len(ring))
+    new_elems = []
+    for ref in cavpts:
+        bl = ring[ref]
+        assert isinstance(bl, BeamLoadingElement), \
+            'Cannot remove beam loading: not a beam loading element'
+        family_name = bl.FamName.replace('_BL', '')
+        harm = numpy.round(bl.Frequency*ring.circumference/clight)
+        new_elems.append(RFCavity(family_name, bl.Length, bl.Voltage,
+                                  bl.Frequency, harm, bl.Energy,
+                                  TimeLag=getattr(bl, 'TimeLag', 0.0),
+                                  PhaseLag=getattr(bl, 'PhaseLag', 0.0)))
+    return apply(ring, cavpts, new_elems)
+
+
+class BeamLoadingElement(RFCavity, Collective):
+    """Class to generate a beamloading element, inherits from Element
+       additional argument are ring, cavity, qfactor, rshunt
+    """
+    default_pass = {False: 'DriftPass', True: 'BeamLoadingCavityPass'}
+    _conversions = dict(RFCavity._conversions,
+                        Rshunt=float, Qfactor=float, NormFact=float,
+                        PhaseGain=float, VoltGain=float, _mode=int,
+                        _beta=float, _wakefact=float, _nslice=int,
+                        ZCuts=lambda v: _array(v),
+                        _nturns=int, _phis=float,
+                        _turnhistory=lambda v: _array(v),
+                        _vbunch=lambda v: _array(v),
+                        _vbeam_phasor=lambda v: _array(v, shape=(2,)),
+                        _vbeam=lambda v: _array(v, shape=(2,)),
+                        _vcav=lambda v: _array(v, shape=(2,)),
+                        _vgen=lambda v: _array(v, shape=(2,)))
+
+    def __init__(self, family_name: str, length: float, voltage: float,
+                 frequency: float, ring: Lattice, qfactor: float,
+                 rshunt: float, mode: Optional[BLMode] = BLMode.PHASOR,
+                 **kwargs):
+        r"""
+        Parameters:
+            ring:            Lattice object
+            length:          Length of the cavity
+            voltage:         Cavity voltage [V]
+            frequency:       Cavity frequency [Hz]
+            qfactor:         Q factor
+            rshunt:          Shunt impedance, [:math:`\Omega`]
+
+        Keyword Arguments:
+            Nslice (int):       Number of slices per bunch. Default: 101
+            Nturns (int):       Number of turn for the wake field. Default: 1
+            ZCuts:              Limits for fixed slicing, default is adaptive
+            NormFact (float):   Normalization factor
+            mode (BLMode):  method for beam loading calculation BLMode.PHASOR
+                (default) uses the phasor method, BLMode.WAKE uses the wake
+                function. For high Q resonator, the phasor method should be
+                used
+        Returns:
+            bl_elem (Element): beam loading element
+        """
+        kwargs.setdefault('PassMethod', self.default_pass[True])
+        assert isinstance(mode, BLMode), \
+            'Beam loading mode has to be an instance of BLMode'
+        zcuts = kwargs.pop('ZCuts', None)
+        phil = kwargs.pop('phil', 0)
+        energy = ring.energy
+        harmonic_number = numpy.round(frequency*ring.circumference/clight)
+        self.Rshunt = rshunt
+        self.Qfactor = qfactor
+        self.NormFact = kwargs.pop('NormFact', 1.0)
+        self.PhaseGain = kwargs.pop('PhaseGain', 1.0)
+        self.VoltGain = kwargs.pop('VoltGain', 1.0)
+        self._mode = int(mode)
+        self._beta = ring.beta
+        self._wakefact = - ring.circumference/(clight *
+                                               ring.energy*ring.beta**3)
+        self._nslice = kwargs.pop('Nslice', 101)
+        self._nturns = kwargs.pop('Nturns', 1)
+        self._turnhistory = None    # Defined here to avoid warning
+        self._vbunch = None
+        if zcuts is not None:
+            self.ZCuts = zcuts
+        super(BeamLoadingElement, self).__init__(family_name, length,
+                                                 voltage, frequency,
+                                                 harmonic_number,
+                                                 energy, **kwargs)
+        _, ts = get_timelag_fromU0(ring)
+        self._phis = 2*numpy.pi*self.Frequency*(ts+self.TimeLag)/clight
+        self._vbeam_phasor = numpy.zeros(2)
+        self._vbeam = numpy.zeros(2)
+        self._vgen = numpy.zeros(2)
+        self._vcav = numpy.array([self.Voltage,
+                                  numpy.pi/2-self._phis-phil])
+        self.clear_history(ring=ring)
+
+    def is_compatible(self, other):
+        return False
+
+    def clear_history(self, ring=None):
+        if ring is None:
+            current = 0.0
+            nbunch = 1
+        else:
+            current = ring.beam_current
+            nbunch = ring.nbunch
+        tl = self._nturns*self._nslice*nbunch
+        self._turnhistory = numpy.zeros((tl, 4), order='F')
+        self._vbunch = numpy.zeros((nbunch, 2), order='F')
+        self._init_bl_params(current)
+
+    def _init_bl_params(self, current):
+        theta = -self._vcav[1]+numpy.pi/2
+        vb = 2*current*self.Rshunt
+        a = self.Voltage*numpy.cos(theta-self._phis)
+        b = self.Voltage*numpy.sin(theta-self._phis)-vb*numpy.cos(theta)
+        psi = numpy.arcsin(b/numpy.sqrt(a**2+b**2))
+        vgen = self.Voltage*numpy.cos(psi) + \
+            vb*numpy.cos(psi)*numpy.sin(self._phis)
+        self._vgen = numpy.array([vgen, psi])
+        self._vbeam = numpy.array([2*current*self.Rshunt*numpy.cos(psi),
+                                   numpy.pi-psi])
+
+    @property
+    def ResFrequency(self):
+        """Resonator frequency"""
+        return self.Frequency/(1-numpy.tan(self.Vgen[1])/(2*self.Qfactor))
+
+    @property
+    def Vbeam(self):
+        """Beam phasor (amplitude, phase)"""
+        return self._vbeam
+
+    @property
+    def Vbunch(self):
+        """Bunch phasor (amplitude, phase)"""
+        return self._vbunch
+
+    @property
+    def Vcav(self):
+        """Cavity phasor (amplitude, phase)"""
+        return self._vcav
+
+    @property
+    def Vgen(self):
+        """Generator phasor (amplitude, phase)"""
+        return self._vgen
+
+    @Vgen.setter
+    def Vgen(self, value):
+        self._vgen = value
+
+    @staticmethod
+    def build_from_cav(cav: RFCavity, ring: Sequence,
+                       qfactor: float, rshunt: float,
+                       mode: Optional[BLMode] = BLMode.PHASOR,
+                       **kwargs):
+        r"""Function to build the BeamLoadingElement from a cavity
+        the FamName, Length, Voltage, Frequency and HarmNumber are
+        taken from the cavity element
+
+        Parameters:
+            ring:            Lattice object
+            qfactor:         Q factor
+            rshunt:          Shunt impedance, [:math:`\Omega`]
+
+        Keyword Arguments:
+            Nslice (int):       Number of slices per bunch. Default: 101
+            Nturns (int):       Number of turn for the wake field. Default: 1
+            ZCuts:              Limits for fixed slicing, default is adaptive
+            NormFact (float):   Normalization factor
+            mode (BLMode):  method for beam loading calculation BLMode.PHASOR
+                (default) uses the phasor method, BLMode.WAKE uses the wake
+                function. For high Q resonator, the phasor method should be
+                used
+        Returns:
+            bl_elem (Element): beam loading element
+        """
+        _CAV_ATTRIBUTES = ['Length', 'Voltage', 'Frequency']
+        _EXCL_ATTRIBUTES = ['PassMethod', 'Energy', 'HarmNumber']
+        cav_attrs = dict(cav.items())
+        family_name = cav_attrs.pop('FamName') + '_BL'
+        [cav_attrs.pop(k) for k in _EXCL_ATTRIBUTES]
+        cav_args = [cav_attrs.pop(k, getattr(cav, k)) for k in
+                    _CAV_ATTRIBUTES]
+        return BeamLoadingElement(family_name, *cav_args, ring,
+                                  qfactor, rshunt, mode=mode,
+                                  **cav_attrs, **kwargs)
+
+    def __repr__(self):
+        """Simplified __repr__ to avoid errors due to arguments
+        not defined as attributes
+        """
+        attrs = dict((k, v) for (k, v) in self.items()
+                     if not k.startswith('_'))
+        return '{0}({1})'.format(self.__class__.__name__, attrs)
diff --git a/pyat/at/collective/haissinski.py b/pyat/at/collective/haissinski.py
index ec72f505f..b0cc92ce5 100644
--- a/pyat/at/collective/haissinski.py
+++ b/pyat/at/collective/haissinski.py
@@ -63,7 +63,8 @@ def __init__(self, wake_object, ring, m=12, kmax=1,
         s = wake_object.srange/self.sigma_l
         self.ds = numpy.diff(s)[0]
         self.s = -s
-        self.wtot_fun = interp1d(self.s, -wake_object.Z)
+        self.wtot_fun = interp1d(self.s, -wake_object.Z,
+                                 bounds_error=False, fill_value = 0)
 
         if m % 2 != 0:
             raise AttributeError('m must be even and int')
@@ -102,6 +103,8 @@ def precompute_S(self):
         Equation 16
         '''
         print('Computing integrated wake potential')
+        print(self.q_array)
+        print(self.ds)
         sr = numpy.arange(2*numpy.amin(self.q_array),
                           numpy.abs(2*numpy.amin(self.q_array))
                           + self.ds, self.ds)
diff --git a/pyat/at/collective/wake_elements.py b/pyat/at/collective/wake_elements.py
index 39907e284..ab4aa5baa 100644
--- a/pyat/at/collective/wake_elements.py
+++ b/pyat/at/collective/wake_elements.py
@@ -13,9 +13,9 @@ class WakeElement(Collective, Element):
     _BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES
     default_pass = {False: 'IdentityPass', True: 'WakeFieldPass'}
     _conversions = dict(Element._conversions, _nslice=int, _nturns=int,
-                        _nelem=int, NumParticles=float, Circumference=float,
+                        _nelem=int, _wakeFact=float,
                         NormFact=lambda v: _array(v, (3,)),
-                        WakeFact=float,
+                        ZCuts=lambda v: _array(v),
                         _wakeDX=lambda v: _array(v),
                         _wakeDY=lambda v: _array(v),
                         _wakeQX=lambda v: _array(v),
@@ -123,11 +123,11 @@ def Nslice(self, nslice):
     @property
     def Nturns(self):
         """Number of turn for the wake field"""
-        return self._nslice
+        return self._nturns
 
     @Nturns.setter
-    def Nturns(self, nslice):
-        self._nslice = nslice
+    def Nturns(self, nturns):
+        self._nturns = nturns
         self.clear_history()
 
     @property
@@ -261,11 +261,6 @@ def __init__(self, family_name: str, ring: Lattice, srange,
                                                    WakeComponent.Z, frequency,
                                                    qfactor, rshunt, **kwargs)
 
-    def rebuild_wake(self):
-        wake = Wake.long_resonator(self.WakeT, self._resfrequency,
-                                   self._qfactor, self._rshunt, self._beta)
-        self._build(wake)
-
 
 class ResWallElement(WakeElement):
     """Class to generate a resistive wall element, inherits from WakeElement
@@ -349,15 +344,3 @@ def Yokoya(self):
     def Yokoya(self, yokoya):
         self._yokoya = yokoya
         self.rebuild_wake()
-
-
-def set_wake_turnhistory(ring):
-    """Function to set the shape of the turn history
-    based on the number of slices, turns and bunches
-    """
-    welems = ring.get_elements(WakeElement)
-    for w in welems:
-        w.clear_history(ring=ring)
-
-
-Lattice.set_wake_turnhistory = set_wake_turnhistory
diff --git a/pyat/at/collective/wake_object.py b/pyat/at/collective/wake_object.py
index 56c72429e..58d4edda7 100644
--- a/pyat/at/collective/wake_object.py
+++ b/pyat/at/collective/wake_object.py
@@ -141,6 +141,10 @@ def _readwakefile(self, filename, scol=0, wcol=1, sfact=1, wfact=1,
 
     def _resonator(self, wcomp, frequency, qfactor, rshunt, beta,
                    yokoya_factor=1):
+        if numpy.any(self._srange < 0):
+            warnings.warn(AtWarning('You are adding a wake function with '
+                                    'negative srange. This may cause issues '
+                                    'interpolating around zero.\n'))
         if wcomp is WakeComponent.Z:
             return long_resonator_wf(self._srange, frequency,
                                      qfactor, rshunt, beta)
@@ -152,6 +156,10 @@ def _resonator(self, wcomp, frequency, qfactor, rshunt, beta,
             raise AtError('Invalid WakeComponent: {}'.format(wcomp))
 
     def _reswall(self, wcomp, length, rvac, conduct, beta, yokoya_factor=1):
+        if numpy.any(self._srange < 0):
+            warnings.warn(AtWarning('You are adding a wake function with '
+                                    'negative srange. This may cause issues '
+                                    'interpolating around zero.\n'))
         if wcomp is WakeComponent.Z:
             raise AtError('Resitive wall not available '
                           'for WakeComponent: {}'.format(wcomp))
@@ -181,7 +189,7 @@ def resonator(srange, wakecomp,
         Several resonators can be built in one step by setting ``nelems``> 1
         and giving array parameters
         """
-        wake = Wake(srange)
+        wake = Wake(numpy.unique(srange.clip(0)))
         try:
             wakecomp = numpy.broadcast_to(wakecomp, (nelems, ))
             frequency = numpy.broadcast_to(frequency, (nelems, ))
@@ -243,7 +251,7 @@ def resistive_wall(srange, wakecomp: WakeComponent,
             yokoya_factor:  Yokoya factor
             nelems:
         """
-        wake = Wake(srange)
+        wake = Wake(numpy.unique(srange.clip(0)))
         try:
             wakecomp = numpy.broadcast_to(wakecomp, (nelems, ))
             length = numpy.broadcast_to(length, (nelems, ))
diff --git a/pyat/at/lattice/elements.py b/pyat/at/lattice/elements.py
index d0a870950..01a507443 100644
--- a/pyat/at/lattice/elements.py
+++ b/pyat/at/lattice/elements.py
@@ -228,6 +228,10 @@ class Collective(_DictLongtMotion):
     def _get_collective(self):
         # noinspection PyUnresolvedReferences
         return self.PassMethod != self.default_pass[False]
+        
+    @abc.abstractmethod
+    def clear_history(self):
+        pass
 
 
 class Element(object):
diff --git a/pyat/at/lattice/lattice_object.py b/pyat/at/lattice/lattice_object.py
index 4e8685516..10ab6b136 100644
--- a/pyat/at/lattice/lattice_object.py
+++ b/pyat/at/lattice/lattice_object.py
@@ -193,12 +193,13 @@ def __init__(self, *args,
         kwargs.setdefault('_particle', Particle())
         # dummy initialization in case the harmonic number is not there
         kwargs.setdefault('_fillpattern', numpy.ones(1))
-        kwargs.setdefault('beam_current', 0.0)
         # Remove temporary keywords
         frequency = kwargs.pop('_frequency', None)
         cell_length = kwargs.pop('_length', None)
         cell_h = kwargs.pop('_harmnumber', math.nan)
         ring_h = kwargs.pop('harmonic_number', periodicity*cell_h)
+        bcurrent = kwargs.pop('beam_current', 0.0)
+        kwargs.setdefault('_beam_current', bcurrent)
 
         if 'energy' in kwargs:
             kwargs.pop('_energy', None)
@@ -540,6 +541,15 @@ def particle(self, particle: Union[str, Particle]):
         else:
             self._particle = Particle(particle)
 
+    def set_wake_turnhistory(self):
+        """Function to reset the shape of the turn history
+        in collective elements based on the number of slices,
+        turns and bunches
+        """
+        for e in self:
+            if e.is_collective:
+                e.clear_history(ring=self)
+
     def set_fillpattern(self, bunches: Union[int, numpy.ndarray] = 1):
         """Function to generate the filling pattern lof the ring.
         The filling pattern is computed as:
@@ -602,6 +612,16 @@ def bunch_list(self) -> numpy.ndarray:
         """Indices of filled bunches"""
         return numpy.flatnonzero(self._fillpattern)
 
+    @property
+    def beam_current(self):
+        return self._beam_current
+
+    @beam_current.setter
+    def beam_current(self, value, clear_history=True):
+        self._beam_current = value
+        if clear_history:
+            self.set_wake_turnhistory()
+
     @property
     def bunch_currents(self) -> numpy.ndarray:
         """Bunch currents [A]"""
@@ -611,7 +631,12 @@ def bunch_currents(self) -> numpy.ndarray:
     @property
     def bunch_spos(self) -> numpy.ndarray:
         """Bunch position around the ring [m]"""
-        bs = self.circumference/len(self._fillpattern)
+        try:
+            circ = self.beta*clight* \
+                self.harmonic_number/self.rf_frequency
+        except AtError:
+            circ = self.circumference       
+        bs = circ/len(self._fillpattern)
         allpos = bs*numpy.arange(len(self._fillpattern))
         return allpos[self._fillpattern > 0]
 
diff --git a/pyat/at/physics/energy_loss.py b/pyat/at/physics/energy_loss.py
index 58b4f61cf..eb3616557 100644
--- a/pyat/at/physics/energy_loss.py
+++ b/pyat/at/physics/energy_loss.py
@@ -6,6 +6,7 @@
 from scipy.optimize import least_squares
 from at.lattice import Lattice, Dipole, Wiggler, RFCavity, Refpts
 from at.lattice import check_radiation, AtError, AtWarning
+from at.lattice import QuantumDiffusion, Collective
 from at.lattice import checktype, set_value_refpts, get_cells, refpts_len
 from at.constants import clight, Cgamma
 from at.tracking import lattice_pass
@@ -68,13 +69,7 @@ def dipole_i2(dipole: Dipole):
     def tracking(ring):
         """Losses from tracking
         """
-        ringtmp = ring.radiation_off(dipole_pass=None,
-                                     quadrupole_pass=None,
-                                     wiggler_pass=None,
-                                     sextupole_pass=None,
-                                     octupole_pass=None,
-                                     copy=True)
-
+        ringtmp = ring.disable_6d(RFCavity, QuantumDiffusion, Collective, copy=True)
         o6 = numpy.squeeze(lattice_pass(ringtmp, numpy.zeros(6),
                            refpts=len(ringtmp)))
         if numpy.isnan(o6[0]):
@@ -140,7 +135,6 @@ def eq(x, freq, rfv, tl0, u0):
     freq = numpy.array([cav.Frequency for cav in ring.select(cavpts)])
     rfv = numpy.array([cav.Voltage for cav in ring.select(cavpts)])
     tl0 = numpy.array([cav.TimeLag for cav in ring.select(cavpts)])
-
     try:
         frf = singlev(freq)
         tml = singlev(tl0)
@@ -148,15 +142,17 @@ def eq(x, freq, rfv, tl0, u0):
         ctmax = clight/numpy.amin(freq)/2
         tt0 = tl0[numpy.argmin(freq)]
         bounds = (-ctmax, ctmax)
-        args = (freq,rfv,tl0,u0)
+        args = (freq, rfv, tl0, u0)
         r = []
         for i in range(divider):
             fact = (i+1)/divider
-            r.append(least_squares(eq,bounds[0]*fact+tt0, args=args, bounds=bounds+tt0))
-            r.append(least_squares(eq,bounds[1]*fact+tt0, args=args, bounds=bounds+tt0))
+            r.append(least_squares(eq, bounds[0]*fact+tt0,
+                                   args=args, bounds=bounds+tt0))
+            r.append(least_squares(eq, bounds[1]*fact+tt0,
+                                   args=args, bounds=bounds+tt0))
         res = numpy.array([abs(ri.fun[0]) for ri in r])
         ok = res < 1.0e-6
-        vals = numpy.array([abs(ri.x[0]).round(decimals=6) for ri in r])    
+        vals = numpy.array([abs(ri.x[0]).round(decimals=6) for ri in r])
         if not numpy.any(ok):
             raise AtError('No solution found for Phis, please check '
                           'RF settings')
diff --git a/pyat/examples/CollectiveEffects/BeamLoading.py b/pyat/examples/CollectiveEffects/BeamLoading.py
new file mode 100644
index 000000000..26524b71b
--- /dev/null
+++ b/pyat/examples/CollectiveEffects/BeamLoading.py
@@ -0,0 +1,168 @@
+import numpy
+import matplotlib.pyplot as plt
+import matplotlib.cm as cmap
+import at
+from at.collective import BeamLoadingElement, add_beamloading, BLMode
+from mpi4py import MPI
+from at.tracking.utils import get_bunches_std_mean
+from at.constants import qe
+from at.physics import harmonic_analysis
+
+
+
+
+def analytical_qs(ring, I):
+
+    E0 = ring.energy * qe
+    alpha_c = at.get_mcf(ring.radiation_off(copy=True))
+    omega_rf = 2 * numpy.pi * ring.get_rf_frequency()
+    Vc = ring.get_rf_voltage()
+
+    envel = at.envelope_parameters(ring.radiation_on(copy=True))
+    U0 = envel.U0 * qe
+
+    Q0 = bl_elem.Qfactor * (1 + beta)
+    Rsh = 2 * bl_elem.Rshunt * (1 + beta)  # Difference in definition
+
+    phi_s = numpy.arccos(U0 / (qe * Vc))  # sync. phase in radian
+    psi = - numpy.arctan(Rsh * I * numpy.sin(phi_s) /
+                         (Vc * (1 + beta)))  # tuning angle
+
+    omega_res = (omega_rf /
+                 (1 + Rsh * I * numpy.sin(phi_s) /
+                  (2 * Q0 * Vc)))  # resonant freq.
+
+    h = ring.harmonic_number  # harmonic number
+    omega0 = omega_rf / h  # rev. freq.
+    T0 = 2 * numpy.pi / omega0  # rev. time
+
+    omega_s0 = numpy.sqrt(qe * Vc * omega_rf * alpha_c * numpy.sin(phi_s) /
+                          (E0 * T0))  # Synch. freq. for a single particle
+
+    K = alpha_c * qe * I / (E0 * T0) * (Rsh / 2) / (1 + beta)
+
+    # Synch. freq. with beam loading
+    tilde_omega_s = numpy.sqrt(omega_s0**2 +
+                               K * omega_rf * numpy.sin(2 * psi) + 1j * 0)
+    return tilde_omega_s / omega0
+
+
+comm = MPI.COMM_WORLD
+size = comm.Get_size()
+rank = comm.Get_rank()
+
+print(size, rank)
+
+ring = at.load_lattice('../../../machine_data/esrf.m')
+ring.radiation_on(cavity_pass='RFCavityPass')
+ring.set_rf_frequency()
+
+
+# In fact, the following resonator parameters are for the EBS machine
+# Not the old ESRF SR, however for benchmarking this doesn't matter
+beta = 2.8
+q0 = 37500  # unloaded
+qfactor = q0 / (1 + beta)  # loaded
+rshunt = 145 * qfactor * 13  # loaded
+o6, _ = ring.find_orbit6()
+
+ring.set_rf_voltage(8.e6)
+
+Nbunches = 16
+ring.set_fillpattern(Nbunches)
+
+_, fring = at.fast_ring(ring)
+fring.pop(-1)  # drop diffusion element
+
+# Here we specify whether we want to use PHASOR or WAKE
+# beam loading models.
+mode = 'PHASOR'
+if mode == 'WAKE':
+    blm = BLMode.WAKE
+else:
+    blm = BLMode.PHASOR
+
+# Npart must be at least Nbunches per core
+Npart = Nbunches
+
+# Now we give the fring and convert it
+# into a beam loaded cavity.
+add_beamloading(fring, qfactor, rshunt, Nslice=1,
+                          Nturns=50, mode=blm,
+                          VoltGain=0.1, PhaseGain=0.1)
+                          
+bl_elem = fring[at.get_refpts(fring, BeamLoadingElement)[0]]
+
+# Specify some simulation parameters
+kickTurn = 500
+Nturns = 2**14 + kickTurn
+current = 150e-3
+fring.beam_current = current
+
+part = numpy.zeros((6, Npart))
+part = (part.T + o6).T
+
+if rank == 0:
+
+    print('\n')
+    print('Current:', current)
+
+    z_all = numpy.zeros((Nturns, Nbunches))
+    dp_all = numpy.zeros((Nturns, Nbunches))
+
+for i in numpy.arange(Nturns):
+    # Apply a kick to ensure the coherent tune is large
+    if i == kickTurn:
+        part[4, :] += 3e-3
+
+    _ = at.lattice_pass(fring, part, nturns=1)
+
+    # Gather particles over all cores (compatible with MPI on or off)
+    allPartsg = comm.gather(part)
+    if rank == 0:
+        allPartsg = numpy.concatenate(allPartsg, axis=1)
+        stds, means = get_bunches_std_mean(allPartsg, Nbunches)
+
+        dp_all[i] = numpy.array([x[4] for x in means])
+        z_all[i] = numpy.array([x[5] for x in means])
+        if i % 1000 == 0:
+            # Print the turn number and number of lost particles
+            print(i, numpy.sum(numpy.isnan(allPartsg)))
+
+
+# Now we compute the tune of each bunch, and compare with the analytical
+if rank == 0:
+    qscoh = numpy.zeros(Nbunches)
+    for ib in numpy.arange(Nbunches):
+        qs = harmonic_analysis.get_tunes_harmonic(dp_all[kickTurn:, ib],
+                                                      num_harmonics=20,
+                                                      fmin=1e-5, fmax=0.1)
+        qscoh[ib] = qs
+
+    qs_mn, qs_std = numpy.array([numpy.mean(qscoh), numpy.std(qscoh)])
+
+    qs_theory = analytical_qs(ring, current)
+
+    print('Analytical:', numpy.real(qs_theory))
+    print('Simulated:', qs_mn, 'pm', qs_std)
+
+    freq = numpy.fft.rfftfreq(Nturns - kickTurn)
+    for i in numpy.arange(Nbunches):
+        fftdat = numpy.abs(numpy.fft.rfft(dp_all[kickTurn:, i]))
+        if i == 0:
+            plt.semilogy(freq, fftdat,
+                         color=cmap.jet(float(i + 1) / Nbunches),
+                         label='Bunch by Bunch FFT')
+        else:
+            plt.semilogy(freq, fftdat,
+                         color=cmap.jet(float(i + 1) / Nbunches))
+
+    plt.xlim([0, 8e-3])
+    plt.xlabel('Tune')
+    plt.ylabel('FFT Amp [A.U.]')
+    plt.axvline(numpy.mean(qscoh), label='Coherent Tune',
+                linestyle='dashed', color='k')
+    plt.axvline(numpy.real(qs_theory), label='Analytical Beam Loaded Tune',
+                linestyle='dashed', color='r')
+    plt.legend()
+    plt.show()
diff --git a/pyat/examples/CollectiveEffects/LongDistribution.py b/pyat/examples/CollectiveEffects/LongDistribution.py
index 9ca0ef18b..2e8760ef6 100644
--- a/pyat/examples/CollectiveEffects/LongDistribution.py
+++ b/pyat/examples/CollectiveEffects/LongDistribution.py
@@ -15,7 +15,7 @@
 m = 50 #30 is quite coarse, 70 or 80 is very fine. 50 is middle
 kmax = 8
 
-srange = Wake.build_srange(-0.36, 0.36, 1.0e-5, 1.0e-2, ring.circumference, ring.circumference)
+srange = Wake.build_srange(0, 0.36, 1.0e-5, 1.0e-2, ring.circumference, ring.circumference)
 
 # Now we initialise the Haissinski class, and solve, then we normalise the distribution and shift the charge center to be at 0
 wobj = Wake.long_resonator(srange, freq, qfactor, Rs, ring.beta)
@@ -40,9 +40,9 @@
 ax1.set_xlabel('z [mm]')
 ax1.set_ylabel(r'$\rho(z)$')
 plt.show()
-
 '''
 
+
 # Now we set up and run the tracking. The final distribution is an average of the last numAve turns 
 welem = LongResonatorElement('wake', ring, srange, freq, qfactor, Rs, Nslice=300)
 welem.Current = current
diff --git a/pyat/examples/CollectiveEffects/RobinsonInstability.py b/pyat/examples/CollectiveEffects/RobinsonInstability.py
index 2b15eda2d..77f0df84e 100644
--- a/pyat/examples/CollectiveEffects/RobinsonInstability.py
+++ b/pyat/examples/CollectiveEffects/RobinsonInstability.py
@@ -28,7 +28,7 @@
 bucket_size = clight/freqres
 current = 0.1   # A
 
-welem = LongResonatorElement('wake', ring, srange, fr, qfactor, rshunt, Nturns=wturns, Nslice=10)
+welem = LongResonatorElement('wake', ring, srange, fr, qfactor, rshunt, Nturns=wturns, Nslice=1)
 
 welem.Current = current
 fring.append(welem)


Test patch:
diff --git a/atmat/attests/pytests.m b/atmat/attests/pytests.m
index 2866f6776..52460d388 100644
--- a/atmat/attests/pytests.m
+++ b/atmat/attests/pytests.m
@@ -90,7 +90,7 @@ function orbit6(testCase,lat2)
             porbit6=double(porbit6)';
             % Matlab
             [~,morbit6]=findorbit6(lattice.m);
-            testCase.verifyEqual(morbit6,porbit6,AbsTol=1.E-15 );
+            testCase.verifyEqual(morbit6,porbit6,AbsTol=5.E-15 );
         end
 
         function m44(testCase,lat2,dp)