Add feedback passmethod

atintegrators/FeedbackPass.c

@@ -0,0 +1,153 @@
+/* IdentityPass.c 
+   Accelerator Toolbox
+   Revision 7/16/03
+   A.Terebilo
+*/
+
+#include "atelem.c"
+#include "atlalib.c"
+#ifdef MPI
+#include <mpi.h>
+#include <mpi4py/mpi4py.h>
+#endif
+
+struct elem 
+{
+  double GX;
+  double GY;
+  double GZ;
+  double *closed_orbit;
+};
+
+void FeedbackPass(double *r_in, int num_particles, struct elem *Elem)
+{	
+    int c;
+    double *mx = atMalloc(sizeof(double));
+    double *my = atMalloc(sizeof(double));
+    double *mz = atMalloc(sizeof(double));
+    long *npart = atMalloc(sizeof(long));
+
+    *mx = 0.0; 
+    *my = 0.0;
+    *mz = 0.0; 
+    *npart = 0;
+
+    double *closed_orbit; 
+    closed_orbit = Elem->closed_orbit;
+
+    double gx = Elem->GX;
+    double gy = Elem->GY;
+    double gz = Elem->GZ;
+
+    for (c = 0; c<num_particles; c++) {	/*Loop over particles  */
+        double *r6 = r_in+c*6;
+
+        if (!atIsNaN(r6[0])) {
+            *npart += 1; 
+            *mx += r6[0];
+            *my += r6[2];
+            *mz += r6[5]; 
+        }
+    }
+
+    #ifdef MPI
+    MPI_Allreduce(MPI_IN_PLACE, npart, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
+    MPI_Allreduce(MPI_IN_PLACE, mx, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); 
+    MPI_Allreduce(MPI_IN_PLACE, my, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); 
+    MPI_Allreduce(MPI_IN_PLACE, mz, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); 
+    MPI_Barrier(MPI_COMM_WORLD);
+    #endif
+
+
+    if (*npart>0.0){
+        *mx /= *npart;
+        *my /= *npart;
+        *mz /= *npart;    
+        for (c = 0; c<num_particles; c++) {	/*Loop over particles  */
+            double *r6 = r_in+c*6;
+            if (!atIsNaN(r6[0])) {
+                r6[0] -= (*mx-closed_orbit[0])*gx;
+                r6[2] -= (*my-closed_orbit[2])*gy;
+                r6[5] -= (*mz-closed_orbit[5])*gz;    
+            }
+        }
+    }
+
+    atFree(mx);
+    atFree(my);
+    atFree(mz);
+    atFree(npart);
+}
+
+#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)
+{
+    if (!Elem) {
+        double GX, GY, GZ, *closed_orbit;
+        GX=atGetDouble(ElemData,"GX"); check_error();
+        GY=atGetDouble(ElemData,"GY"); check_error();
+        GZ=atGetDouble(ElemData,"GZ"); check_error();
+        closed_orbit = atGetDoubleArray(ElemData,"closed_orbit"); check_error();
+
+        Elem = (struct elem*)atMalloc(sizeof(struct elem));
+        Elem->GX=GX;
+        Elem->GY=GY;
+        Elem->GZ=GZ;
+        Elem->closed_orbit = closed_orbit;
+    }
+    FeedbackPass(r_in,num_particles,Elem);
+    return Elem;
+}
+
+MODULE_DEF(FeedbackPass)        /* Dummy module initialisation */
+
+#endif /*defined(MATLAB_MEX_FILE) || defined(PYAT)*/
+
+
+#ifdef MATLAB_MEX_FILE
+
+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]);
+        struct elem El, *Elem=&El;
+
+        double GX,GY,GZ,*closed_orbit;
+
+        GX=atGetDouble(ElemData,"GX"); check_error();
+        GY=atGetDouble(ElemData,"GY"); check_error();
+        GZ=atGetDouble(ElemData,"GZ"); check_error();
+        closed_orbit = atGetDoubleArray(ElemData,"closed_orbit");check_error();
+
+
+        Elem = (struct elem*)atMalloc(sizeof(struct elem));
+        Elem->GX=GX;
+        Elem->GY=GY;
+        Elem->GZ=GZ;
+        Elem->closed_orbit = closed_orbit;
+
+    if (mxGetM(prhs[1]) != 6) mexErrMsgIdAndTxt("AT:WrongArg","Second argument must be a 6 x N matrix: particle array");
+
+        /* ALLOCATE memory for the output array of the same size as the input  */
+        plhs[0] = mxDuplicateArray(prhs[1]);
+        r_in = mxGetDoubles(plhs[0]);
+        FeedbackPass(r_in,num_particles,Elem);
+    }
+    else if (nrhs == 0) {
+        /* list of required fields */
+        plhs[0] = mxCreateCellMatrix(4,1);
+        mxSetCell(plhs[0],0,mxCreateString("GX"));
+        mxSetCell(plhs[0],1,mxCreateString("GY"));
+        mxSetCell(plhs[0],2,mxCreateString("GZ"));
+        mxSetCell(plhs[0],3,mxCreateString("closed_orbit"));
+    }
+    else {
+        mexErrMsgIdAndTxt("AT:WrongArg","Needs 2 or 0 arguments");
+    }
+}
+#endif
+

pyat/at/lattice/elements/basic_elements.py

@@ -18,6 +18,7 @@
     "SimpleRadiation",
     "QuantumDiffusion",
     "EnergyLoss",
+    "Feedback"
 ]
 
 import warnings
@@ -626,5 +627,35 @@ def __init__(self, family_name: str, energy_loss: float, **kwargs):
         kwargs.setdefault("PassMethod", self.default_pass[False])
         super().__init__(family_name, EnergyLoss=energy_loss, **kwargs)
 
+class Feedback(Element):
+    """Transverse and Longitudinal Feedback element"""
+
+    def __init__(self, family_name: str,
+                       GX: float = 0.0, 
+                       GY: float = 0.0,
+                       GZ: float = 0.0, 
+                       closed_orbit: float = np.zeros(6), 
+                       **kwargs):
+        """
+        Args:
+            family_name:    Name of the element
+            GX:             Feedback Gain in Horizontal
+            GY:             Feedback Gain in Vertical
+            GZ:             Feedback Gain in Longitudinal
+            closed_orbit:   6D closed orbit to feedback around
+
+        Default PassMethod: ``FeedbackPass``
+
+        Note that this element does not SET the closed orbit. That 
+        must be handled in the linear maps for x and y, or the timelag
+        for the longitudinal plane. But an accurate closed orbit must be
+        provided in order to have a well behaving feedback. 
+        """
+        kwargs.setdefault("PassMethod", "FeedbackPass")
+        super().__init__(family_name, GX=GX, GY=GY, GZ=GZ, closed_orbit=closed_orbit, **kwargs)
+
+
+
+
 
 Radiative.register(EnergyLoss)

pyat/at/physics/fastring.py

@@ -141,7 +141,9 @@ def simple_ring(
     U0: float = 0.0,
     name: str = "",
     particle: str | Particle = "relativistic",
-    TimeLag: float | Sequence[float] = 0.0
+    TimeLag: float | Sequence[float] = 0.0,
+    orb6: Sequence[float] = np.zeros(6), 
+
 ) -> Lattice:
     """Generates a "simple ring" based on a given dictionary
        of global parameters
@@ -201,7 +203,8 @@ def simple_ring(
           or a Particle object
         TimeLag: Set the timelag of the cavities, Default=0. Can be scalar
           or sequence of scalars (as with harmonic_number and Vrf).      
-
+        orb6: 6d closed orbit. Needed for inclusion of feedbacks.
+
     If the given emitx, emity or espread is 0, then no equlibrium emittance
     is applied in this plane.
     If the given tau is 0, then no radiation damping is applied for this plane.
@@ -265,7 +268,8 @@ def simple_ring(
     # generate the linear tracking element, we set a length
     # which is needed to give the lattice object the correct length
     # (although it is not used for anything else)
-    lin_elem = M66("Linear", m66=Mat66, Length=circumference)
+
+    lin_elem = M66("Linear", m66=Mat66, Length=circumference, T1=-orb6, T2=orb6)
 
     # Generate the simple radiation element
     simplerad = SimpleRadiation(