From 65c63aacad8ea383a1d4fcc60c9687d0dba6882b Mon Sep 17 00:00:00 2001 From: Lee Carver Date: Mon, 18 Aug 2025 14:29:44 +0200 Subject: [PATCH 1/4] Add constant wake to wakefieldpass and python functions --- atintegrators/WakeFieldPass.c | 33 +++++++++++++++++++++++------ atintegrators/atimplib.c | 10 ++++++++- pyat/at/collective/wake_elements.py | 18 +++++++++++++++- pyat/at/collective/wake_object.py | 18 ++++++++++++++-- 4 files changed, 69 insertions(+), 10 deletions(-) diff --git a/atintegrators/WakeFieldPass.c b/atintegrators/WakeFieldPass.c index 703f2fac34..ee23a9d159 100755 --- a/atintegrators/WakeFieldPass.c +++ b/atintegrators/WakeFieldPass.c @@ -19,6 +19,8 @@ struct elem double *waketableQX; double *waketableQY; double *waketableZ; + double *waketableCX; + double *waketableCY; double *turnhistory; double *z_cuts; }; @@ -39,6 +41,8 @@ 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; @@ -62,6 +66,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; @@ -71,8 +77,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 ? @@ -84,8 +90,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); @@ -107,6 +113,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; @@ -123,6 +131,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(); @@ -143,6 +153,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; } @@ -182,6 +194,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; @@ -198,6 +212,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; @@ -212,6 +229,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; @@ -241,13 +260,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 { diff --git a/atintegrators/atimplib.c b/atintegrators/atimplib.c index 204e8cda41..35579dd1d7 100644 --- a/atintegrators/atimplib.c +++ b/atintegrators/atimplib.c @@ -202,7 +202,8 @@ 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 *waketableZ, double *waketableCX, double *waketableCY, + double *normfact, double *kx,double *ky, double *kx2,double *ky2,double *kz){ int rank=0; int size=1; @@ -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 @@ -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); + } } } @@ -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 }; diff --git a/pyat/at/collective/wake_elements.py b/pyat/at/collective/wake_elements.py index c56c5f7340..aadc95e8d6 100644 --- a/pyat/at/collective/wake_elements.py +++ b/pyat/at/collective/wake_elements.py @@ -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): """ @@ -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) @@ -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""" diff --git a/pyat/at/collective/wake_object.py b/pyat/at/collective/wake_object.py index 9b9025d0e7..4032304c77 100644 --- a/pyat/at/collective/wake_object.py +++ b/pyat/at/collective/wake_object.py @@ -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): @@ -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): @@ -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` From 2b011fa9b755c11c4e906c525e2e68ecb02474a1 Mon Sep 17 00:00:00 2001 From: Lee Carver Date: Mon, 18 Aug 2025 14:31:46 +0200 Subject: [PATCH 2/4] missing kcx,kcy --- atintegrators/atimplib.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atintegrators/atimplib.c b/atintegrators/atimplib.c index 35579dd1d7..f0d724cfa6 100644 --- a/atintegrators/atimplib.c +++ b/atintegrators/atimplib.c @@ -204,7 +204,7 @@ static void compute_kicks(int nslice,int nturns,int nelem, double *waketableDY,double *waketableQX,double *waketableQY, double *waketableZ, double *waketableCX, double *waketableCY, double *normfact, double *kx,double *ky, - double *kx2,double *ky2,double *kz){ + double *kx2,double *ky2,double *kz, double *kcx, double *kcy){ int rank=0; int size=1; int i,ii,index; From a1118a1775bd7b74f5c132be85cbc4e1651ab710 Mon Sep 17 00:00:00 2001 From: Lee Carver Date: Mon, 18 Aug 2025 14:33:53 +0200 Subject: [PATCH 3/4] missing double definition --- atintegrators/WakeFieldPass.c | 2 ++ 1 file changed, 2 insertions(+) diff --git a/atintegrators/WakeFieldPass.c b/atintegrators/WakeFieldPass.c index ee23a9d159..348c8f6a41 100755 --- a/atintegrators/WakeFieldPass.c +++ b/atintegrators/WakeFieldPass.c @@ -56,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; From 70ac54db82f91eea66db753ff4ea487df9ea4806 Mon Sep 17 00:00:00 2001 From: Lee Carver Date: Mon, 18 Aug 2025 14:50:27 +0200 Subject: [PATCH 4/4] Correct memory allocation --- atintegrators/WakeFieldPass.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atintegrators/WakeFieldPass.c b/atintegrators/WakeFieldPass.c index 348c8f6a41..a179efc946 100755 --- a/atintegrators/WakeFieldPass.c +++ b/atintegrators/WakeFieldPass.c @@ -47,7 +47,7 @@ void WakeFieldPass(double *r_in,int num_particles,double circumference,int nbunc 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;