Add constant wake potentials to wakefield pass

[lcarver]

Testing still ongoing, making PR now for collaborative development

[lcarver]

Here is a test script.

I define a constant wake CX which is a heaviside step function with an amplitude of 1. I make a simplering, but I pop out every single element, so I track only the wake element. I set the wakefact to be equal to 1. I make a gaussian beam with 3mm sigma.

I then slice the bunch and compute the mean xp of each longitudinal slice, which is then appropriately weighted. This is plotted in red below. The shape of the curve is exactly what I would expect, which is the convolution of a gaussian with a step function.

I also compare the amplitude with an analytical gaussian convolved with the CX wake. The final value agrees nicely.

import at
import numpy as np
import matplotlib.pyplot as plt
from at.collective import Wake, WakeType, WakeComponent, WakeElement
from at.constants import qe

ring = at.simple_ring(6e9, 843.977, 992, 0.1, 0.2, 6.5e6, 8.5e-5, betax=1, betay=1, alphax=0.0, alphay=0.0)
current = 1/10
ring.set_beam_current(current)
t0 = 1/ring.revolution_frequency
sigm = at.sigma_matrix(betax=1, alphax=0.0, emitx=100e-12, betay=1, alphay=0.0, emity=10e-12, blength=3e-3, espread=1e-3)
Nparts = 10000
I_per_parts = current/Nparts
parts = at.beam(Nparts,sigm)

srange = np.arange(-0.1,0.1,1e-4)

         
cx_val = 1
wakeCX = np.heaviside(srange,0)*cx_val
wa = Wake(srange)
wa.add(WakeType.TABLE, WakeComponent.CX, srange, wakeCX)
welem = WakeElement('const_wake', ring, wa, Nslice=200)
#welem.set_normfactxy(ring)
welem._wakefact = 1

ring.pop(0)
ring.pop(0)
ring.pop(0)
ring.pop(0)
ring.pop(0)

ring.append(welem)


p_before = parts.copy()
_,out,_ = ring.track(p_before, in_place=False, nturns=1, refpts=None)
p_after = out['rout']

delta_xp = p_after[1,:] - parts[1,:]

# quick slicing

maxZ = np.amax(p_before[5,:])+1e-6
minZ = np.amin(p_before[5,:])-1e-6
Nslices = 200
binedges = np.linspace(minZ, maxZ, Nslices+1)
binmeans = np.array([(binedges[i] + binedges[i+1])/2 for i in range(len(binedges)-1)]) 

slice_counter = np.zeros(Nslices)
slice_xp = np.zeros(Nslices)
for ip in np.arange(Nparts):
    par = p_before[:,ip]
    zpos = par[5]
    dxp = delta_xp[ip]
    #dxp /= (1+par[4])
    
    slice_ind = np.atleast_1d(np.squeeze(np.argwhere(zpos>binedges)))[-1]
    slice_xp[slice_ind] += dxp
    slice_counter[slice_ind] += 1

mean_xp_per_slice = welem._wakefact * slice_xp / slice_counter

sigz = 3e-3
x0 = 0
amp = current / Nparts
gauss = (amp)/(np.sqrt(2*np.pi)*sigz) * np.exp(-(srange-x0)**2/2/sigz**2)
cc = np.convolve(gauss, wakeCX)

print(np.amax(cc), np.amin(cc))
print(np.amax(mean_xp_per_slice[~np.isnan(mean_xp_per_slice)]), np.amin(mean_xp_per_slice[~np.isnan(mean_xp_per_slice)]))

plt.plot(binmeans, mean_xp_per_slice, color='r', label='simulation')
plt.axhline(np.amax(cc), color='k', linestyle='dashed', label='Max of Conv of Gaussian with CX Wake')
plt.xlabel('s [m]')
plt.legend()
plt.show()