<jupyter_start><jupyter_text>Simulation of a FACET-II type plasma wakefield accelerator<NAME> RadiaSoft, LLC <EMAIL>Developed for a project supported by the United States Department of Energy, Office of Science, Office of High Energy Physics under contract number DE-SC0018718. *** IntroductionThis notebook is meant to model a beam-driven plasma wakefield accelerator using nominal FACET-II parameters with an unchanging drive and witness bunch, in an effort to compute the $m=0$ wake fields from the drive bunch using our technique developed and demonstrated with the hollow channel.***<jupyter_code>## Imports

# standard python libraries
import numpy as np
from scipy import constants

import shutil, os

%matplotlib inline
%load_ext autoreload
%autoreload 2

import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib import animation, rc

import h5py as hdf5

# OpenPMD for looking at the file data
from opmd_viewer import OpenPMDTimeSeries

# Imports for the simulations, and setting up the plots
from fbpic.main import Simulation
from fbpic.openpmd_diag import FieldDiagnostic, ParticleDiagnostic, ParticleChargeDensityDiagnostic, \
     set_periodic_checkpoint, restart_from_checkpoint
from fbpic.lpa_utils.bunch import add_elec_bunch_gaussian

# set the colormap and centre the colorbar

import matplotlib.colors as colors

class MidpointNormalize(colors.Normalize):
    """
    Normalise the colorbar so that diverging bars work there way either side from a prescribed midpoint value)

    e.g. im=ax1.imshow(array, norm=MidpointNormalize(midpoint=0.,vmin=-100, vmax=100))
    """
    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
        self.midpoint = midpoint
        colors.Normalize.__init__(self, vmin, vmax, clip)

    def __call__(self, value, clip=None):
        # I'm ignoring masked values and all kinds of edge cases to make a
        # simple example...
        x, y = [self.vmin, self.midpoint, self.vmax], [0, 0.5, 1]
        return np.ma.masked_array(np.interp(value, x, y), np.isnan(value))<jupyter_output><empty_output><jupyter_text>*** Simulation ParametersIn this section, we develop a simulation of the hollow channel. The hollow channel refers to the fact that a drive and witness bunch travel through a hollow ring of plasma -- this set up is akin to having an extremely small beam pipe. For the FBPIC simulation, this requires creating a ring of quiescent electron plasma, and a drive and witness beam initiated with self-consistent fields. We will neglect the ion motion in the plasma channel for this simulation.The simulation uses a moving window, beginning with the drive and witness bunch outside the hollow plasma channel, then passing through, with the simulation running until the fields reach an approximate steady state. We will then use those fields to compute the wake fields and impedances, and compare to the theoretical results in [ref. 1](References).The hollow plasma channel has a density $n_e$ and a radius $b$, and the local plasma frequency is given by $\omega_p = \sqrt{\frac{4 \pi n_e e^2}{m_e}}$ for the electron charge $e$ and mass $m_e$. The plasma wavenumber is $k_p = \omega_p / c$. Because PIC algorithms do not like hard edges in charge density, we will use a linear ramp on the density of the plasma hollow channel, with a slope $k_{ramp}$ specified below.The baseline parameters for our studies will be the nominal design FACET-II parameters described in [ref. 3](References), and a plasma channel radius of thrice the beam $\sigma_r$.<jupyter_code>## Domain physical parameters



## Beam parameters

# Drive bunch is gaussian

# turn the drive beam on and off, as required.
use_drive_bunch = True

drive_sigma_r = 3.65e-6  # meters
drive_sigma_z = 12.77e-6  # meters
drive_Q = 1.e10*(-1.*constants.elementary_charge)   # Coulombs
drive_N_macro = 4000
drive_gamma = 10.e3

# Witness bunch, also gaussian

# turn the witness beam on and off, as required.
use_witness_bunch = True

witness_sigma_r = 3.65e-6 #meters
witness_sigma_z = 6.38e-6  # meters
witness_Q = 4.3e9*(-1.*constants.elementary_charge)   # Coulombs
witness_N_macro = 4000
witness_gamma = 10.e3

trailing_distance = 150.e-6 # meters


## Plasma channel parameters

n_plasma = 4.e16        # cm^-3

# convert to per cubic meter
n_plasma *= 100**3

# derived plasma quantities
omega_p = np.sqrt(n_plasma*constants.elementary_charge**2/(constants.m_e*constants.epsilon_0))
k_p = omega_p/constants.c

lambda_p = 2.*np.pi/k_p

## Domain parameters

# Domain size, include the whole thing and some trailing distance
domain_length = 1.5*lambda_p  # meters
domain_radius = 1.*lambda_p  # meters

# Grid size, resolve the drive bunch
Delta_z = 0.1*min([drive_sigma_z, lambda_p, witness_sigma_z])  # meters
Delta_r = 0.1*min([drive_sigma_r, lambda_p, witness_sigma_r])  # meters

# Derived quantities
Nz = int(np.rint(domain_length/Delta_z))
Nr = int(np.rint(domain_radius/Delta_r))

dt = (np.sqrt(Delta_z**2 + Delta_r**2)/constants.c)        # sec, small fudge factor 
                                                                # to prevent particles from 
                                                                # covering a full cell per time step
        
# Moving window
window_v = constants.c

# start the ramp after the drive bunch has existed a while

ramp_start = domain_length
ramp_length = 5.*drive_sigma_z


# create the density function for the plasma, which is uniform
def dens_func( z, r ) :
    """Returns relative density at position z and r"""
    # Allocate relative density
    n = np.ones_like(z)
    # Make linear ramp
    n = np.where( z < ramp_start + ramp_length, (z-ramp_start)/ramp_length, n )
    # Supress density before the ramp
    n = np.where( z < ramp_start, 0., n )
    return(n)

# We want to run the simulation just long enough for the fields to form behind the drive bunch, 
# so we will simulate for twenty bunch sigmas

sim_length = (ramp_start + ramp_length + 30.*lambda_p)/constants.c

Nsteps = int(sim_length/dt)-int(sim_length/dt)%100 + 1

# Use only the primary azimuthal mode
Nm = 1

# Define the diagnostics
write_fields = True
write_particles = False
dump_period = 50 #Nsteps-1<jupyter_output><empty_output><jupyter_text>*** The SimulationHere we run the actual simulation of the beams traveling through the hollow channel.***<jupyter_code># remove old data
if os.path.exists('./diags/hdf5'):
    shutil.rmtree('./diags/hdf5')

# Create the simulation
sim = Simulation(Nz, domain_length, Nr, domain_radius, Nm, dt, boundaries='open')
# micromanage the particle species
sim.ptcl = []

# add the gaussian drive bunch
if use_drive_bunch:
    add_elec_bunch_gaussian( sim, 
                            sig_r = drive_sigma_r, 
                            sig_z = drive_sigma_z, 
                            n_emit=0., 
                            gamma0=drive_gamma, 
                            sig_gamma=1.,
                            Q=-1.*drive_Q, 
                            N=drive_N_macro, 
                            tf=0.0, 
                            zf=.8*domain_length, boost=None)
    
    if write_particles:
        sim.diags.append( ParticleChargeDensityDiagnostic(dump_period, sim,
                    species={'driver_bunch': sim.ptcl[0]}) )
if use_witness_bunch:
    add_elec_bunch_gaussian( sim, 
                            sig_r = witness_sigma_r, 
                            sig_z = witness_sigma_z, 
                            n_emit=0., 
                            gamma0=drive_gamma, 
                            sig_gamma=1.,
                            Q=-1.*witness_Q, 
                            N=witness_N_macro, 
                            tf=0.0, 
                            zf=.8*domain_length - trailing_distance, boost=None)

# add the hollow channel electrons
plasma_electrons = sim.add_new_species(q = -1.*constants.elementary_charge,
                                 m = constants.electron_mass,
                                 dens_func = dens_func, 
                                 n = n_plasma, p_nz = 1, p_nr = 1, p_nt = 1)

if write_particles:
    sim.diags.append( ParticleDiagnostic(dump_period,
                    {'electrons': sim.ptcl[-1]}, sim.comm ) )

plasma_ions = sim.add_new_species(q = constants.elementary_charge,
                                 m = constants.proton_mass,
                                 dens_func = dens_func, 
                                 n = n_plasma, p_nz = 1, p_nr = 1, p_nt = 1)

# Set the moving window
sim.set_moving_window(v = window_v)

# Add diagnostics
if write_fields:
    sim.diags.append( FieldDiagnostic(dump_period, sim.fld, sim.comm ) )
    
# remove old data
# run the simulation
sim.step(Nsteps)<jupyter_output>FBPIC (0.11.0)

Running on CPU (40 threads per process) 

Calculating initial space charge field...
Done.

Calculating initial space charge field...
Done.

|███████████████████████████████████| 7201/7201, 0:00:00 left, 797 ms/step[K
Total time taken (with compilation): 1:39:51
Average time per iteration (with compilation): 831 ms<jupyter_text>Now we try to make a video of the plasma bubble from the data to see if it's evolving<jupyter_code># import the list of data files
files_list = os.listdir('./diags/hdf5')
dir_name = './diags/hdf5'
step_list = []
for file_name in files_list:
    step_num = int(file_name[-8:-3])
    step_list.append(step_num)

print step_list

ts = OpenPMDTimeSeries(dir_name)
# create a function that computes F_r and F_z from the fields in a given file
def force_vector(dir_name, step_num):
    """ Computes the radial and longitudinal force/charge from a file's field data
    
    Arguments
    file_name: string that gives the file to open
    
    Returns
    F_r: the radial force/charge
    F_z: the longitudinal force/charge
    """

    ts = OpenPMDTimeSeries(dir_name)
    Ez_field, meta = ts.get_field(field='E', coord='z', iteration=step_num)
    Er_field, meta = ts.get_field(field='E', coord='r', iteration=step_num)
    Bt_field, meta = ts.get_field(field='B', coord='t', iteration=step_num)

    rgrid = meta.r
    zgrid = meta.z
    
    F_r = Er_field - constants.speed_of_light*Bt_field
    F_z = Ez_field
    
    return F_r, F_z, rgrid, zgrid
# A function to compute the curl of the force, on a discrete mesh, at the cell midpoints
def curl(F_r, F_z, rs, zs):
    """ Compute the curl of the force, which is zero in the Panofsky-Wenzel limit
    Arguments:
    F_r: vector for the radial force
    F_z: vector for the longitudinal force
    rs : vector of the radial grid points
    zs : vector of the longitudinal grid points
    
    Returns:
    curl_F : curl of the force
    curl_rs: radial coordinates of the curl vector
    curl_zs: longitudinal coordinates of the curl vector
    """
    
    # finite difference in case of variable grid size

    del_r = rs[1:-1] - rs[0:-2]
    del_z = zs[1:-1] - zs[0:-2]
    
    # To co-locate the forces, linearly interpolate to the mid-points on the off-axes
    F_r = (F_r[0:-2,:]+F_r[1:-1,:])/2.
    F_z = (F_z[:,0:-2]+F_z[0,1:-1])/2.
    
    # finite difference the force for the derivatives
    # dFr_dz in this form lives at the midpoints of the z-axis cell edges
    dFr_dz = (F_r[:,1:-1] - F_r[:,0:-2])/del_z[0]
    # dFz_dr in this form lives at the midpoints of the r-axis cell edges
    dFz_dr = (F_z[1:-1,:] - F_z[0:-2,:])/del_r[0]
    
    curl_F =  (dFr_dz - dFz_dr)
    curl_rs = (rs[1:-1] + rs[0:-2])/2.
    curl_zs = (zs[1:-1] + zs[0:-2])/2.
    
    return curl_F, curl_rs, curl_zs
# Pre-process the ranges of data for the color bars

fravg = 0.
frmin = 0.
frmax = 0.

fzavg = 0.
fzmin = 0.
fzmax = 0.

curlavg = 0.
curlmin = 0.
curlmax = 0.
        
for steps in step_list[0:50]:
    Fr, Fz, rgrid, zgrid = force_vector(dir_name, steps)

    fmin = np.amin(Fr*1.e-9)
    if fmin < frmin:
        frmin = fmin
    fmax = np.amax(Fr*1.e-9)
    if fmax > frmax:
        frmax = fmax
    
    fmin = np.amin(Fz*1.e-9)
    if fmin < fzmin:
        fzmin = fmin
    fmax = np.amax(Fz*1.e-9)
    if fmax > fzmax:
        fzmax = fmax
        
    curl_F, rs, zs = curl(Fr, Fz, rgrid, zgrid)

    cmin = np.amin(curl_F*1.e-9)
    if cmin < curlmin:
        curlmin = cmin
    cmax = np.amax(Fz*1.e-9)
    if cmax > curlmax:
        curlmax = cmax

print 'computed the range of values'
#setup figure
fig = plt.figure()

#set up list of images for animation
ims=[]

plt.xlabel(r'$\zeta = z - c t \quad [\mu m]$')
plt.ylabel(r'$r \quad [\mu m]$')

for steps in step_list:
    Fr, Fz, rgrid, zgrid = force_vector(dir_name, steps)
    r_mid = len(rgrid)/2
    
    extent = [(zgrid[0]-constants.speed_of_light*steps*dt)*1.e6,
             (zgrid[-1]-constants.speed_of_light*steps*dt)*1.e6,
             rgrid[r_mid]*1.e6, rgrid[r_mid+r_mid/2]*1.e6]
    
    im = plt.imshow(Fr[r_mid:r_mid+r_mid/2,:]*1.e-9, extent = extent,
           cmap='RdBu', origin='lower', norm=MidpointNormalize(midpoint=fravg,vmin=frmin, vmax=frmax),
                   animated=True)
        
    ims.append([im])
    
cbar = plt.colorbar(orientation='horizontal')
cbar.set_label(r'$E_r - c B_\theta \quad [GV/m]$')

#run animation
ani = animation.ArtistAnimation(fig,ims, interval=100, blit=True)
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
ani.save('pwfa_Er.mp4', writer=writer)
#setup figure
fig = plt.figure()

#set up list of images for animation
ims=[]

plt.xlabel(r'$\zeta = z - c t \quad [\mu m]$')
plt.ylabel(r'$r \quad [\mu m]$')

for steps in step_list:
    Fr, Fz, rgrid, zgrid = force_vector(dir_name, steps)
    r_mid = len(rgrid)/2
    
    extent = [(zgrid[0]-constants.speed_of_light*steps*dt)*1.e6,
             (zgrid[-1]-constants.speed_of_light*steps*dt)*1.e6,
             rgrid[r_mid]*1.e6, rgrid[r_mid+r_mid/2]*1.e6]
    
    im = plt.imshow(Fz[r_mid:r_mid+r_mid/2,:]*1.e-9, extent = extent,
           cmap='RdBu', origin='lower', norm=MidpointNormalize(midpoint=fzavg,vmin=fzmin, vmax=fzmax),
                   animated=True)
        
    ims.append([im])
    
cbar = plt.colorbar(orientation='horizontal')
cbar.set_label(r'$E_z \quad [GV/m]$')

#run animation
ani = animation.ArtistAnimation(fig,ims, interval=100, blit=True)
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
ani.save('pwfa_Ez.mp4', writer=writer)
# Look at the evolving curl(F)

#setup figure
fig = plt.figure()

#set up list of images for animation
ims=[]

plt.xlabel(r'$\zeta = z - c t \quad [\mu m]$')
plt.ylabel(r'$r \quad [\mu m]$')

for steps in step_list:
    Fr, Fz, rgrid, zgrid = force_vector(dir_name, steps)
    
    curl_F, rs, zs = curl(Fr, Fz, rgrid, zgrid)

    r_mid = len(rgrid)/2
    
    extent = [(zgrid[0]-constants.speed_of_light*steps*dt)*1.e6,
             (zgrid[-1]-constants.speed_of_light*steps*dt)*1.e6,
             rgrid[r_mid]*1.e6, rgrid[r_mid+r_mid/2]*1.e6]
    
    im = plt.imshow(np.abs(curl_F[r_mid:r_mid+r_mid/2,:]), extent = extent,
           cmap='bone_r', origin='lower', #norm=MidpointNormalize(midpoint=curlavg,vmin=curlmin, vmax=curlmax),
                   animated=True)
        
    ims.append([im])
    
cbar = plt.colorbar(orientation='horizontal')
cbar.set_label(r'$\nabla \times \vec{F} \quad [GV/m/m]$')

#run animation
ani = animation.ArtistAnimation(fig,ims, interval=100, blit=True)
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
ani.save('curl.mp4', writer=writer)<jupyter_output><empty_output><jupyter_text>Because the $\nabla \times \vec{F}$ function is taking a finite difference on fields with some small noise, the results end up proportional to noise/grid spacing, which gets quite large. To filter this, we take 3x3 averages, which smoothes away the noise and keeps features on a scale larger than a handful of grid cells.<jupyter_code># A function to filter the noise in the curl simulation
from scipy import signal
def curl_filter(curl_array):
    
    # current implemention takes a three-cell block-average, with the edges taking a truncated average
    kernel = np.array([[1.,1.,1.],
                       [1.,2.,1.],
                       [1.,1.,1.]])/10.
    
    filtered_curl = signal.convolve2d(curl_array, kernel, mode='same')
    
    return filtered_curl
# Look at the evolving filtered curl(F)

# Pre-process the ranges of data for the color bars

curlavg = 0.
curlmin = 0.
curlmax = 0.
        
for steps in step_list[100:150]:
    Fr, Fz, rgrid, zgrid = force_vector(dir_name, steps)
        
    curl_F, rs, zs = curl(Fr, Fz, rgrid, zgrid)
    
    filtered_curl_F = curl_filter(curl_F)

    cmin = np.amin(filtered_curl_F*1.e-9)
    if cmin < curlmin:
        curlmin = cmin
    cmax = np.amax(filtered_curl_F*1.e-9)
    if cmax > curlmax:
        curlmax = cmax

print 'computed the range of values'

#setup figure
fig, ax = plt.subplots()

#set up list of images for animation
ims=[]

ax.set_xlabel(r'$\zeta = z - c t \quad [\mu m]$')
ax.set_ylabel(r'$r \quad [\mu m]$')

for steps in step_list:
    Fr, Fz, rgrid, zgrid = force_vector(dir_name, steps)
    
    curl_F, rs, zs = curl(Fr, Fz, rgrid, zgrid)
    
    filtered_curl_F = curl_filter(curl_F)

    r_mid = len(rgrid)/2
    
    extent = [(zgrid[0]-constants.speed_of_light*steps*dt)*1.e6,
             (zgrid[-1]-constants.speed_of_light*steps*dt)*1.e6,
             rgrid[r_mid]*1.e6, rgrid[r_mid+r_mid/2]*1.e6]
    
    im = ax.imshow(filtered_curl_F[r_mid:r_mid+r_mid/2,:]*1.e-9, extent = extent,
           cmap='RdBu', origin='lower', norm=MidpointNormalize(midpoint=curlavg,vmin=curlmin, vmax=curlmax),
                   animated=True)
        
    ims.append([im])
    
cbar = fig.colorbar(im, orientation='horizontal')
cbar.set_label(r'$\nabla \times \vec{F} \quad [GV/m/m]$')

#run animation
ani = animation.ArtistAnimation(fig,ims, interval=100, blit=True)
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
ani.save('filtered_curl.mp4', writer=writer)<jupyter_output>computed the range of values
{'fraction_non_alphanumeric': 0.11074270557029178, 'fraction_numerical': 0.01928603006189213, 'mean_word_length': 2.9791116974494285, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '109892', 'license': 'Apache-2.0', 'n_tokens_mistral': 6448, 'n_tokens_neox': 6149, 'n_words': 1744, 'path': 'jupyter/panofsky_wenzel/pwfa_videos.ipynb', 'repo_name': 'radiasoft/rsfbpic', 'size': '18117'}