Skip to content

Add constant wake potentials to wakefield pass #976

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Aug 20, 2025
Merged

Add constant wake potentials to wakefield pass #976

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
65c63aa
Add constant wake to wakefieldpass and python functions
Aug 18, 2025
2b011fa
missing kcx,kcy
Aug 18, 2025
a1118a1
missing double definition
Aug 18, 2025
70ac54d
Correct memory allocation
Aug 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 30 additions & 7 deletions atintegrators/WakeFieldPass.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ struct elem
double *waketableQX;
double *waketableQY;
double *waketableZ;
double *waketableCX;
double *waketableCY;
double *turnhistory;
double *z_cuts;
};
Expand All @@ -39,11 +41,13 @@ void WakeFieldPass(double *r_in,int num_particles,double circumference,int nbunc
double *waketableDY = Elem->waketableDY;
double *waketableQX = Elem->waketableQX;
double *waketableQY = Elem->waketableQY;
double *waketableCX = Elem->waketableCX;
double *waketableCY = Elem->waketableCY;
double *waketableZ = Elem->waketableZ;
double *turnhistory = Elem->turnhistory;
double *z_cuts = Elem->z_cuts;

size_t sz = 5*nslice*nbunch*sizeof(double) + num_particles*sizeof(int);
size_t sz = 7*nslice*nbunch*sizeof(double) + num_particles*sizeof(int);
int c;

int *pslice;
Expand All @@ -52,6 +56,8 @@ void WakeFieldPass(double *r_in,int num_particles,double circumference,int nbunc
double *kx2;
double *ky2;
double *kz;
double *kcx;
double *kcy;

void *buffer = atMalloc(sz);
double *dptr = (double *) buffer;
Expand All @@ -62,6 +68,8 @@ void WakeFieldPass(double *r_in,int num_particles,double circumference,int nbunc
kx2 = dptr; dptr += nslice*nbunch;
ky2 = dptr; dptr += nslice*nbunch;
kz = dptr; dptr += nslice*nbunch;
kcx = dptr; dptr += nslice*nbunch;
kcy = dptr; dptr += nslice*nbunch;

iptr = (int *) dptr;
pslice = iptr; iptr += num_particles;
Expand All @@ -71,8 +79,8 @@ void WakeFieldPass(double *r_in,int num_particles,double circumference,int nbunc
slice_bunch(r_in,num_particles,nslice,nturns,nbunch,bunch_spos,bunch_currents,
turnhistory,pslice,z_cuts);
compute_kicks(nslice*nbunch,nturns,nelem,turnhistory,waketableT,waketableDX,
waketableDY,waketableQX,waketableQY,waketableZ,
normfact,kx,ky,kx2,ky2,kz);
waketableDY,waketableQX,waketableQY,waketableZ,waketableCX,waketableCY,
normfact,kx,ky,kx2,ky2,kz,kcx,kcy);

/*apply kicks*/
/* OpenMP not efficient. Too much shared data ?
Expand All @@ -84,8 +92,8 @@ void WakeFieldPass(double *r_in,int num_particles,double circumference,int nbunc
int islice=pslice[c];
if (!atIsNaN(r6[0])) {
r6[4] += kz[islice];
r6[1] += (kx[islice]+r6[0]*kx2[islice])*(1+r6[4]);
r6[3] += (ky[islice]+r6[2]*ky2[islice])*(1+r6[4]);
r6[1] += (kx[islice]+r6[0]*kx2[islice]+kcx[islice])*(1+r6[4]);
r6[3] += (ky[islice]+r6[2]*ky2[islice]+kcy[islice])*(1+r6[4]);
}
}
atFree(buffer);
Expand All @@ -107,6 +115,8 @@ ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem,
double *waketableQX;
double *waketableQY;
double *waketableZ;
double *waketableCX;
double *waketableCY;
double *turnhistory;
double *z_cuts;
int i;
Expand All @@ -123,6 +133,8 @@ ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem,
waketableDY=atGetOptionalDoubleArray(ElemData,"_wakeDY"); check_error();
waketableQX=atGetOptionalDoubleArray(ElemData,"_wakeQX"); check_error();
waketableQY=atGetOptionalDoubleArray(ElemData,"_wakeQY"); check_error();
waketableCX=atGetOptionalDoubleArray(ElemData,"_wakeCX"); check_error();
waketableCY=atGetOptionalDoubleArray(ElemData,"_wakeCY"); check_error();
waketableZ=atGetOptionalDoubleArray(ElemData,"_wakeZ"); check_error();
z_cuts=atGetOptionalDoubleArray(ElemData,"ZCuts"); check_error();

Expand All @@ -143,6 +155,8 @@ ExportMode struct elem *trackFunction(const atElem *ElemData,struct elem *Elem,
Elem->waketableQX=waketableQX;
Elem->waketableQY=waketableQY;
Elem->waketableZ=waketableZ;
Elem->waketableCX=waketableCX;
Elem->waketableCY=waketableCY;
Elem->turnhistory=turnhistory;
Elem->z_cuts=z_cuts;
}
Expand Down Expand Up @@ -182,6 +196,8 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
double *waketableQX;
double *waketableQY;
double *waketableZ;
double *waketableCX;
double *waketableCY;
double *turnhistory;
double *z_cuts;

Expand All @@ -198,6 +214,9 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
waketableQX=atGetOptionalDoubleArray(ElemData,"_wakeQX"); check_error();
waketableQY=atGetOptionalDoubleArray(ElemData,"_wakeQY"); check_error();
waketableZ=atGetOptionalDoubleArray(ElemData,"_wakeZ"); check_error();
waketableCX=atGetOptionalDoubleArray(ElemData,"_wakeCX"); check_error();
waketableCY=atGetOptionalDoubleArray(ElemData,"_wakeCY"); check_error();

z_cuts=atGetOptionalDoubleArray(ElemData,"ZCuts"); check_error();

Elem->nslice=nslice;
Expand All @@ -212,6 +231,8 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
Elem->waketableDY=waketableDY;
Elem->waketableQX=waketableQX;
Elem->waketableQY=waketableQY;
Elem->waketableCX=waketableCX;
Elem->waketableCY=waketableCY;
Elem->waketableZ=waketableZ;
Elem->turnhistory=turnhistory;
Elem->z_cuts=z_cuts;
Expand Down Expand Up @@ -241,13 +262,15 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])

if (nlhs>1) {
/* list of optional fields */
plhs[1] = mxCreateCellMatrix(6,1); /* No optional fields */
plhs[1] = mxCreateCellMatrix(8,1); /* No optional fields */
mxSetCell(plhs[0],0,mxCreateString("_wakeDX"));
mxSetCell(plhs[0],1,mxCreateString("_wakeDY"));
mxSetCell(plhs[0],2,mxCreateString("_wakeQX"));
mxSetCell(plhs[0],3,mxCreateString("_wakeQY"));
mxSetCell(plhs[0],4,mxCreateString("_wakeZ"));
mxSetCell(plhs[0],5,mxCreateString("ZCuts"));
mxSetCell(plhs[0],5,mxCreateString("_wakeCX"));
mxSetCell(plhs[0],6,mxCreateString("_wakeCY"));
mxSetCell(plhs[0],7,mxCreateString("ZCuts"));
}
}
else {
Expand Down
12 changes: 10 additions & 2 deletions atintegrators/atimplib.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,8 +202,9 @@ static void slice_bunch(double *r_in,int num_particles,int nslice,int nturns,
static void compute_kicks(int nslice,int nturns,int nelem,
double *turnhistory,double *waketableT,double *waketableDX,
double *waketableDY,double *waketableQX,double *waketableQY,
double *waketableZ,double *normfact, double *kx,double *ky,
double *kx2,double *ky2,double *kz){
double *waketableZ, double *waketableCX, double *waketableCY,
double *normfact, double *kx,double *ky,
double *kx2,double *ky2,double *kz, double *kcx, double *kcy){
int rank=0;
int size=1;
int i,ii,index;
Expand All @@ -219,6 +220,8 @@ static void compute_kicks(int nslice,int nturns,int nelem,
kx2[i]=0.0;
ky2[i]=0.0;
kz[i]=0.0;
kcx[i]=0.0;
kcy[i]=0.0;
}

#ifdef MPI
Expand All @@ -239,6 +242,9 @@ static void compute_kicks(int nslice,int nturns,int nelem,
if(waketableQX)kx2[i-nslice*(nturns-1)] += normfact[0]*wi*getTableWake(waketableQX,waketableT,ds,index);
if(waketableQY)ky2[i-nslice*(nturns-1)] += normfact[1]*wi*getTableWake(waketableQY,waketableT,ds,index);
if(waketableZ) kz[i-nslice*(nturns-1)] += normfact[2]*wi*getTableWake(waketableZ,waketableT,ds,index);
if(waketableCX)kcx[i-nslice*(nturns-1)] += normfact[0]*wi*getTableWake(waketableCX,waketableT,ds,index);
if(waketableCY)kcy[i-nslice*(nturns-1)] += normfact[1]*wi*getTableWake(waketableCY,waketableT,ds,index);

}
}
}
Expand All @@ -249,6 +255,8 @@ static void compute_kicks(int nslice,int nturns,int nelem,
if(waketableQX)MPI_Allreduce(MPI_IN_PLACE,kx2,nslice,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
if(waketableQY)MPI_Allreduce(MPI_IN_PLACE,ky2,nslice,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
if(waketableZ)MPI_Allreduce(MPI_IN_PLACE,kz,nslice,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
if(waketableCX)MPI_Allreduce(MPI_IN_PLACE,kcx,nslice,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
if(waketableCY)MPI_Allreduce(MPI_IN_PLACE,kcy,nslice,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
MPI_Barrier(MPI_COMM_WORLD);
#endif
};
Expand Down
18 changes: 17 additions & 1 deletion pyat/at/collective/wake_elements.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,9 @@ class WakeElement(Collective, Element):
_wakeQX=lambda v: _array(v),
_wakeQY=lambda v: _array(v),
_wakeZ=lambda v: _array(v),
_wakeT=lambda v: _array(v))
_wakeT=lambda v: _array(v),
_wakeCX=lambda v: _array(v),
_wakeCY=lambda v: _array(v))

def __init__(self, family_name: str, ring: Lattice, wake: Wake, **kwargs):
"""
Expand Down Expand Up @@ -68,6 +70,10 @@ def _build(self, wake):
self._wakeQX = wake.QX
if wake.QY is not None:
self._wakeQY = wake.QY
if wake.CX is not None:
self._wakeCX = wake.CX
if wake.CY is not None:
self._wakeCY = wake.CY

def rebuild_wake(self, wake):
self._build(wake)
Expand Down Expand Up @@ -112,6 +118,16 @@ def WakeQY(self):
"""Quadrupole Y component"""
return getattr(self, '_wakeQY', None)

@property
def WakeCX(self):
"""Constant Dipole X component"""
return getattr(self, '_wakeCX', None)

@property
def WakeCY(self):
"""Constant Dipole Y component"""
return getattr(self, '_wakeCY', None)

@property
def Nslice(self):
"""Number of slices per bunch"""
Expand Down
18 changes: 16 additions & 2 deletions pyat/at/collective/wake_object.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,8 @@ class WakeComponent(Enum):
QX = 3 #: Quadrupole X
QY = 4 #: Quadrupole Y
Z = 5 #: Longitudinal

CX = 6 #: Constant Dipole X
CY = 7 #: Constant Dipole Y

# noinspection PyPep8Naming
class Wake(object):
Expand Down Expand Up @@ -69,7 +70,9 @@ def __init__(self, srange):
WakeComponent.DY: None,
WakeComponent.QX: None,
WakeComponent.QY: None,
WakeComponent.Z: None}
WakeComponent.Z: None,
WakeComponent.CX: None,
WakeComponent.CY: None}

@property
def srange(self):
Expand Down Expand Up @@ -100,6 +103,17 @@ def Z(self):
"""Longitudinal component"""
return self.components[WakeComponent.Z]

@property
def CX(self):
"""Constant Dipole X component"""
return self.components[WakeComponent.CX]

@property
def CY(self):
"""Constant Dipole Y component"""
return self.components[WakeComponent.CY]


def add(self, wtype: WakeType, wcomp: WakeComponent, *args, **kwargs):
"""Add a component to a :py:class:`.Wake`

Expand Down
Loading