<commit_before># WARP example PWFA simulation for Sirepo
# This simulation is designed to produce a quick and clean basic PWFA simulation for use with the Sirepo front end.
#
# <NAME> - 03/15/2016
#
#


from __future__ import absolute_import, division, print_function

import warpoptions
warpoptions.ignoreUnknownArgs = True

"""
-Updated script to handle ionization - initialize with fully ionized hydrogen and un-ionized Helium (at a 2:1 atomic density)

Simulate a basic PWFA with FACET parameters, using Warp's Lifschitz
algorithm: 2D r-z with azimuthal modes (0, 1, ...).

We should use only the 0 mode (standard cylindrical symmetry), but
unfortunately that doesn't seem to work. So we use two: (0,1)

The electrons are pre-ionized. No ionization physics here.

It would be good to use Remi Lehi's stencil with low longitudinal
dispersion, but it's not working for me.

Also, the 2nd-order particle shapes are not working. I think its
required to use 1st-order shape in the radial direction, because
of the symmetry axis. Need to learn more about that.

1/8/2015 - Updated and re-organized to group parameters in a more meaningful way.
Also trims parameters which aren't needed for pure PWFA, and moves them to a
section at the bottom of the script which is commented out.

Naming conventions:
    - ALLCAPS_WITH_UNDERSCORE - constants for the purposes of the simulation (variable overides, etc.)
    - lowercase_with_underscore - private instance variables

Simulation parameters:
    - At top level, specify length of plasma, plasma density, pre-ionized/not
    - beam energy, beam charge, beam size (matched or not)
    - laser parameters (currently optional)

"""

# ---------------------------------------------------------------------------------------
# Imports
# ---------------------------------------------------------------------------------------

# Import warp-specific packages
from warp_init_tools import *

# Basic imports
import math


# ---------------------------------------------------------------------------------------
# User defined functions
# ---------------------------------------------------------------------------------------

def match_beam_transverse(eps_n,l_pe):
    """
    Define the beam radius that is matched to the defined plasma, for a given emittance.

    Parameters
    ----------
    eps_n: transverse, normalized emittance in m-rad
    lp_e: plasma wavelength

    Returns
    -------
    r_rms_match: matched transverse RMS beam radius


    """

    # EPSILON_RMS_N = 2.25e-06    # [m-rad]
    #EPSILON_RMS_N = 0.     # [m-rad]

    # Calculate transverse matching of beam to e- plasma
    k_pe = 2.*math.pi / l_pe #plasma wave#
    twiss_beta_match = sqrt(2.) / k_pe
    r_rms_match  = sqrt(twiss_beta_match * eps_n)

    return r_rms_match

def match_beam_longitudinal(l_pe,fac=0.728583):
    """
    Define the matched bunch length for a plasma. Currently adjusted for Facet Li gas @ 1.5e+23 m^-3.

    Parameters
    ----------
    lp_e: plasma wavelength

    Returns
    -------
    l_bunch_fw: matched bunch length for Facet beam (Li gas at ~ 1.5e+23 m^-3)


    """
    # Matched length (FW) for a longitudinally triangular bunch
    l_fw_matched = 2. * l_pe
    # Maximum RMS length for a longitudinally gaussian bunch
    l_rms_max   = l_pe / math.pi

    # For gaussian beam profiles, bunch length is an arbitrary number of e-foldings
    # Resonance requires RMS_LENGTH <= L_RMS_MAX
    # Ad hoc factor of 0.728583 is reasonable & yields Facet beam
    # length when Li gas density is chosen to be 1.5e+23 m^-3
    rms_length = fac * l_rms_max

    # At FACET, sigma_z ~ 20 microns
    # rms_length = 20.0e-06

    long_norm    = 1.0 / ( sqrt(2.*math.pi) * rms_length )
    x_bunch_hw   = 3.0 * rms_length
    bunch_length = 2.0 * x_bunch_hw
    t_bunch_fw   = bunch_length / clight

    return bunch_length


# ---------------------------------------------------------------------------------------
# Basic user-defined parameters
# ---------------------------------------------------------------------------------------

#Top level Physical Parameters - User editable
#---------------------------------------------

# Reference plasma density (in number of particles per m^3)
N_PLASMA = {{electronPlasma_density}} # plasma density (in m^-3)
# Electron plasma wavelength for full ionization
lambda_pe = 3.34e+07 / sqrt(N_PLASMA)

#Number of standard deviations to fill with particles in each direction (from center of bunch)
#Also used to determine maximum grid dimensions by multiplying plasma wavelength
NUM_SIGMA_R = {{simulationGrid_rScale}}
NUM_SIGMA_L = {{simulationGrid_zScale}}

USE_LASER = {{useLaser}}

USE_BEAM = {{useBeam}}


# ---------------------------------------------------------------------------------------
# General simulation parameters
# ---------------------------------------------------------------------------------------
# Dimension of simulation ("3d", "circ", "2d", "1d")
dim = "circ"
# Number of azimuthal modes beyond m=0, for "circ" (not used for "2d" and "3d")
circ_m = 1

# Moving window (0:off, 1:on)
USE_MOVING_WINDOW = 1
# Speed of the moving window (ignored if use_moving_window = 0)
v_moving_window = clight

# Field boundary conditions (longitudinal and transverse respectively)
f_boundz  = openbc
f_boundxy = absorb
# Particles boundary conditions (longitudinal and transverse respectively)
p_boundz  = absorb
p_boundxy = reflect


# Current smoothing parameters
# ----------------------------
# Turn current smoothing on or off (0:off; 1:on)
use_smooth = 1
# Number of passes of smoother and compensator in each direction (x, y, z)
npass_smooth = array([[ 0 , 0 ], [ 0 , 0 ], [ 1 , 1 ]])
# Smoothing coefficients in each direction (x, y, z)
alpha_smooth = array([[ 0.5, 3.], [ 0.5, 3.], [0.5, 3./2]])
# Stride in each direction (x, y, z)
stride_smooth = array([[ 1 , 1 ], [ 1 , 1 ], [ 1 , 1 ]])



# Numerical parameters
# ---------------------
# Field solver (0:Yee, 1:Karkkainen on EF,B, 3:Lehe)
stencil = 0
# Particle shape (1:linear, 2:quadratic, 3:cubic)
depos_order = 2
# Gathering mode (1:from cell centers, 4:from Yee mesh)
efetch = 1
# Particle pusher (0:Boris, 1:Vay)
particle_pusher = 1



# Diagnostics
# -----------
# Period of diagnostics (in number of timesteps)
diag_period = {{incSteps}}
# Whether to write the fields
write_fields = 1
# Whether to write the particles
write_particles = {{isAnimationView}} or USE_BEAM
# Whether to write the diagnostics in parallel
parallel_output = False


# ---------------------------------------------------------------------------------------
#General Plasma Parameters
# ---------------------------------------------------------------------------------------


# Plasma macroparticles
# ---------------------
# Number of macroparticles per cell in each direction
# In Circ, nppcelly is the number of particles along the
# azimuthal direction. Use a multiple of 4*circ_m
plasma_nx = {{simulationGrid_rParticlesPerCell}}
plasma_ny = 4 * circ_m
plasma_nz = {{simulationGrid_zParticlesPerCell}}


# Plasma content and profile
# --------------------------

# Initialize some preexisting plasmas electrons (0:off, 1:on)
# (Can be used in order to neutralize pre-ionized ions, if any,
# or in order to simulate a plasma without having to initialize ions)
USE_PREEXISTING_ELECTRONS = 1
# Relative density of the preexisting electrons (relative to n_plasma)
rel_dens_preexisting_electrons = 1.
# Initialize plasma ions (0:off, 1:on) - Turn off for speed using a pure electron plasma
USE_IONS = 0
# The different elements used. (Only used if use_ions is different than 0.)
# relative_density is the density relative to n_plasma.
# q_start is the ionization state of the ions at the beginning of the simulation
# q_max is the maximum ionization state
# If q_start is not equal to q_max, ionization between states will be computed.
ion_states = { 'Hydrogen': {'relative_density':1., 'q_start':1, 'q_max':1 },
                 'Helium': {'relative_density':2., 'q_start':0, 'q_max':2 } }

# Specify the grid parameters
# --------------------------

#compute needed quantities for initialization
nx = {{simulationGrid_rCount}}
ny = 300
nz = {{simulationGrid_zCount}}
zmin = {{simulationGrid_zMin}}
zmax = {{simulationGrid_zMax}}
xmin = {{simulationGrid_zMin}}
xmax = {{simulationGrid_rMax}}
ymax = xmax
ymin = -ymax

# Specify the plasma parameters - these can be independent of the initial grid dimensions
# ---------------------------------------------------------------------------------------
# Positions between which the plasma is initialized
# (Transversally, the plasma is initialized between -plasma_xmax and
# plasma_xmax, along x, and -plasma_ymax and plasma_ymax along y)
plasma_zmin = 1.e-6
plasma_zmax = {{electronPlasma_length}}
plasma_xmax = xmax
plasma_ymax = ymax

# Define your own profile using these parameters, and passing them to plasma_dens_func()
#The following values are in units of meters
ramp_start = 0.e-6
ramp_length = {{electronPlasma_entranceRampLength}}
exit_ramp_length = {{electronPlasma_exitRampLength}}
ramp_plateau = plasma_zmax - ramp_length - exit_ramp_length


# ---------------------------------------------------------------------------------------
# User-defined initialization functions. Designed to be called Plasma_Injector()
# ---------------------------------------------------------------------------------------


#def plasma_dens_func(x, y, z, ramp_start, ramp_length, plateau_length):
def plasma_dens_func(x, y, z):
    """
    User-defined function: density profile of the plasma

    It should return the relative density with respect to n_plasma,
    at the position x, y, z (i.e. return a number between 0 and 1)

    Parameters
    ----------
    x, y, z: 1darrays of floats
        Arrays with one element per macroparticle

    Currently Deprecated
    --------------------
    ramp_start: beginning position of density profile
    ramp_length: length of plasma ramp
    plateau_length: length of plateau

    Returns
    -------
    n : 1d array of floats
        Array of relative density, with one element per macroparticles
    """
    # Allocate relative density
    n = ones_like(z)
    # # Make linear ramp
    # n = where( z<ramp_start+ramp_length, (z-ramp_start)/ramp_length, n )
    # # Supress density before the ramp
    # n = where( z<ramp_start, 0., n )
    # # Reduce density by half after the ramp
    # n = where( z> ramp_start+ramp_length+ramp_plateau, 0.5*n, n )
    # # Put the density to 0 later
    # n = where( z> ramp_start+ramp_length+2*ramp_plateau, 0., n )

    n = where(z < ramp_length, 0., n)
    #TODO(pjm): taken from previous script, this could be completely wrong, need to review this
    n = where(z > plasma_zmax - exit_ramp_length, sin(0.5 * pi * (plasma_zmax - z) / exit_ramp_length) * n, n)
    return(n)


#def beam_dens_func(x, y, z, beam_zmin, beam_zmax, beam_rmax):
def beam_dens_func(x, y, z):
    """
    User-defined function: density profile of the beam

    It should return the relative density with respect to n_beam,
    at the position x, y, z (i.e. return a number between 0 and 1)

    Parameters
    ----------
    x, y, z: 1darrays of floats
        Arrays with one element per macroparticle

    Currently Deprecated
    --------------------
    beam_zmin: minimum z position of the beam
    beam_zmax: maximum z position of the beam
    beam_rmax: maximum transverse dimension of the beam (> 0)

    Returns
    -------
    n : 1d array of floats
        Array of relative density, with one element per macroparticles
    """
    # Allocate relative density
    n = ones_like(z)
    # Longitudinal profile: parabolic
    n = n*(z - beam_zmin)*(beam_zmax - z) * 4/(beam_zmax - beam_zmin)**2
    # Transverse profile: parabolic
    r = sqrt( x**2 + y**2)
    n = n*(1 - (r/beam_rmax)**2 )
    # Put the density above rmax to 0
    n[r > beam_rmax] = 0.

    return(n)



# ---------------------------------------------------------------------------------------
# Initialization of the simulation (Normal users should not modify this part.)
# ---------------------------------------------------------------------------------------

# Set some general options for warp
set_diagnostics( 0 )
set_boundary_conditions( f_boundz, f_boundxy, p_boundz, p_boundxy )
set_simulation_box( nz, nx, ny, zmin, zmax, xmax, ymax, dim, xmin=xmin, ymin=ymin )
set_moving_window( USE_MOVING_WINDOW, v_moving_window )

# See smoothing.py
set_smoothing_parameters( use_smooth, dim, npass_smooth,
                         alpha_smooth, stride_smooth )

# Creation of the species
# -----------------------

elec = None
ions = None
elec_from_ions = None

# Create the plasma species
# Reference weight for plasma species
plasma_weight = prepare_weights( N_PLASMA, plasma_nx, plasma_ny,
                            plasma_nz, dim, circ_m )
if USE_PREEXISTING_ELECTRONS:
    elec_weight = rel_dens_preexisting_electrons * plasma_weight
    elec = Species(type=Electron, weight=elec_weight, name='electrons')
if USE_IONS:
    ions, elec_from_ions = initialize_ion_dict( ion_states, plasma_weight,
                                                group_elec_by_element=True )

    # Set the numerical parameters only now: they affect the newly created species
set_numerics( depos_order, efetch, particle_pusher, dim)

# Setup the field solver object
# -----------------------------
em = EM3D(
    stencil=stencil,
    npass_smooth=npass_smooth,
    alpha_smooth=alpha_smooth,
    stride_smooth=stride_smooth,
    l_2dxz= (dim in ["2d", "circ"]),
    l_2drz= (dim in ["circ"]),
    l_1dz = (dim =="1d" ),
    l_getrho=True,
    circ_m = (dim =="circ")*circ_m,
    type_rz_depose=1 )
registersolver(em)

# Introduce the laser (only if required)
# --------------------------------------
if USE_LASER:
    laser_a0 = {{laserPulse_amplitude}}
    lambda_laser = {{laserPulse_wavelength}}
    laser_w0 = {{laserPulse_waist}}
    laser_ctau = {{laserPulse_duration}} * clight
    laser_z0 = -2 * laser_ctau
    laser_zfoc = 50.e-6
    laser_polangle = {{laserPulse_polarizationAngle}}
    laser_source_z = 0

    add_laser( em, dim, laser_a0, laser_w0, laser_ctau, laser_z0,
        lambda0=lambda_laser,
        zf=laser_zfoc, theta_pol=laser_polangle, source_z=laser_source_z )


# Introduce the beam (only if required)
# -------------------------------------
if USE_BEAM:

    # Specify the maximum density of the e- drive bunch
    CHARGE = {{electronBeam_charge}} # total beam charge [C]
    n_total = CHARGE / echarge  # number of electrons

    #Beam energy
    beam_energy = {{electronBeam_energy}} #beam energy in GeV
    GAMMA_BEAM = beam_energy*1.e3/0.510998910 #beam gamma

    #The relative RMS longitudinal energy spread
    DE_OVER_E_RMS = 0.0                                             #USER EDIT

    #beam matching parameters
    EPS_N = {{electronBeam_transverseEmittance}} #normalized transverse emittance
    FIXED_R_RMS = {{electronBeam_rmsRadius}} #Specify this for a fixed radius bunch with 0 emittance. If EPS_N specified, this is ignored.
    L_BUNCH = {{electronBeam_bunchLength}} #This is fixed to match the beam to the Facet beam


    # ---------------------------------------------------------------------------------------
    # Relativistic beam parameters
    # ---------------------------------------------------------------------------------------
    # Initialize beam electrons (0:off, 1:on)
    # (Please be aware that initializing a beam in 2D geometry makes very little
    # physical sense, because of the long range of its space-charge fields)

    # Number of macroparticles per cell in each direction - adjusted to be equal in each direction
    beam_nx = 2 * plasma_nx
    beam_ny = 1 * plasma_ny
    beam_nz = 2 * plasma_nz

    #TRANSVERSE BEAM PARAMETERS
    #----------------------------
    # The RMS normalized transverse emittance
    # EPSILON_RMS_N = 2.25e-06    # [m-rad]
    #EPSILON_RMS_N = 0.     # [m-rad]

    # To avoid betatron oscillations of the drive beam, you
    #   must choose the matched RMS radius (see above).

    if EPS_N > 0:
        #use user specified emittance to match beam
        rms_radius = match_beam_transverse(EPS_N,lambda_pe)
    else:
        #match to FACET parameters (sigma_r ~ 15 microns) or user-specified
        rms_radius = FIXED_R_RMS

    rms_perp    = rms_radius / sqrt(2.) #rms_perp is for a given axis (x or y), so has a separate normalization
    perp_norm   = 1.0 / ( sqrt(2.*math.pi) * rms_perp ) #compute the radius with Gaussian normalization - use rms_perp
    max_radius  = NUM_SIGMA_R * rms_radius #Fill a region of radius NUM_SIGMA_R standard deviations

    #ENERGY AND CHARGE OF BEAM
    #-------------------------

    #GAMMA_BEAM = 45009.78 # at FACET, E~23 GeV
    beta_gamma = sqrt(GAMMA_BEAM**2 - 1.)
    beta_beam  = beta_gamma / GAMMA_BEAM

    #Longitudinal momentum of the beam
    beam_uz = beta_gamma
    v_z_avg = beta_beam  * clight
    u_z_avg = beta_gamma * clight

    # The corresponding RMS momenta
    u_z_rms = DE_OVER_E_RMS * u_z_avg
    u_x_rms = clight * EPS_N / rms_perp #if zero emittance, zero RMS transverse momenta
    u_y_rms = u_x_rms

    #LONGITUDINAL BEAM PARAMETERS
    #----------------------------
    # Calculate longitudinal matching of beam to e- plasma

    # At FACET, sigma_z ~ 20 microns
    # RMS_LENGTH = 20.0e-06

    if L_BUNCH > 0.:
        #use user-defined bunch length
        bunch_length = L_BUNCH
    else:
        #match the bunch to Facet beam
        bunch_length = match_beam_longitudinal(lambda_pe)


    l_norm   = 1.0 / ( sqrt(2.*math.pi)*bunch_length)
    t_bunch_fw   = bunch_length / clight


    #PEAK BEAM DENSITY
    #----------------------------
    # The maximum physical number density of electrons in the beam
    beam_density = n_total * l_norm * perp_norm**2
    # Beam density
    n_beam = beam_density


    # Positions between which the beam is initialized
    # (Transversally, the plasma is initialized between -plasma_xmax and
    # plasma_xmax, along x, and -plasma_ymax and plasma_ymax along y)
    beam_zmax = -1.e-6
    beam_zmin = beam_zmax - bunch_length
    beam_xmax = max_radius
    beam_ymax = beam_xmax

    # Define your own profile using these parameters, and passing them to beam_dens_func()
    beam_rmax = beam_xmax

    # Create the beam
    beam_weight = prepare_weights( n_beam, beam_nx, beam_ny,
                                   beam_nz, dim, circ_m )
    beam = Species(type=Electron, weight=beam_weight, name='beam')
    PlasmaInjector( beam, None, w3d, top, dim, beam_nx, beam_ny, beam_nz,
                beam_zmin, beam_zmax, beam_xmax, beam_ymax,
                dens_func = beam_dens_func, uz_m=beam_uz )
    initialize_beam_fields( em, dim, beam, w3d, top )

# Introduce the plasma
# --------------------
# Create an object to store the information about plasma injection
plasma_injector = PlasmaInjector( elec, ions, w3d, top, dim,
        plasma_nx, plasma_ny, plasma_nz, plasma_zmin,
        plasma_zmax, plasma_xmax, plasma_ymax, plasma_dens_func )

# Continuously inject the plasma, if the moving window is on
if {{isAnimationView}} and USE_MOVING_WINDOW :
    installuserinjection( plasma_injector.continuous_injection )

# Setup the diagnostics
# ---------------------
if write_fields:
    diag1 = FieldDiagnostic( period=diag_period, top=top, w3d=w3d, em=em,
                comm_world=comm_world, lparallel_output=parallel_output,
                write_dir={{outputDir}})
    installafterstep( diag1.write )
if write_particles:
    diag2 = ParticleDiagnostic( period=diag_period, top=top, w3d=w3d,
            species={ species.name : species for species in listofallspecies },
            comm_world=comm_world, lparallel_output=parallel_output,
            write_dir={{outputDir}})
    installafterstep( diag2.write )

# ---------------------------------------------------------------------------------------
# Simulation loop (Normal users should not modify this part either.)
# ---------------------------------------------------------------------------------------

inc_steps = {{incSteps}}

if {{isAnimationView}}:
    doit = True
    while doit:
        step(inc_steps)
        doit = w3d.zmmin + top.zgrid < plasma_zmax
<commit_msg>Fix #472: try warp.init_tools else warp_init_tools<commit_after># WARP example PWFA simulation for Sirepo
# This simulation is designed to produce a quick and clean basic PWFA simulation for use with the Sirepo front end.
#
# <NAME> - 03/15/2016
#
#


from __future__ import absolute_import, division, print_function

import warpoptions
warpoptions.ignoreUnknownArgs = True

"""
-Updated script to handle ionization - initialize with fully ionized hydrogen and un-ionized Helium (at a 2:1 atomic density)

Simulate a basic PWFA with FACET parameters, using Warp's Lifschitz
algorithm: 2D r-z with azimuthal modes (0, 1, ...).

We should use only the 0 mode (standard cylindrical symmetry), but
unfortunately that doesn't seem to work. So we use two: (0,1)

The electrons are pre-ionized. No ionization physics here.

It would be good to use Remi Lehi's stencil with low longitudinal
dispersion, but it's not working for me.

Also, the 2nd-order particle shapes are not working. I think its
required to use 1st-order shape in the radial direction, because
of the symmetry axis. Need to learn more about that.

1/8/2015 - Updated and re-organized to group parameters in a more meaningful way.
Also trims parameters which aren't needed for pure PWFA, and moves them to a
section at the bottom of the script which is commented out.

Naming conventions:
    - ALLCAPS_WITH_UNDERSCORE - constants for the purposes of the simulation (variable overides, etc.)
    - lowercase_with_underscore - private instance variables

Simulation parameters:
    - At top level, specify length of plasma, plasma density, pre-ionized/not
    - beam energy, beam charge, beam size (matched or not)
    - laser parameters (currently optional)

"""

# ---------------------------------------------------------------------------------------
# Imports
# ---------------------------------------------------------------------------------------

# Import warp-specific packages
try:
    from warp.init_tools import *
except ImportError:
    from warp_init_tools import *

# Basic imports
import math


# ---------------------------------------------------------------------------------------
# User defined functions
# ---------------------------------------------------------------------------------------

def match_beam_transverse(eps_n,l_pe):
    """
    Define the beam radius that is matched to the defined plasma, for a given emittance.

    Parameters
    ----------
    eps_n: transverse, normalized emittance in m-rad
    lp_e: plasma wavelength

    Returns
    -------
    r_rms_match: matched transverse RMS beam radius


    """

    # EPSILON_RMS_N = 2.25e-06    # [m-rad]
    #EPSILON_RMS_N = 0.     # [m-rad]

    # Calculate transverse matching of beam to e- plasma
    k_pe = 2.*math.pi / l_pe #plasma wave#
    twiss_beta_match = sqrt(2.) / k_pe
    r_rms_match  = sqrt(twiss_beta_match * eps_n)

    return r_rms_match

def match_beam_longitudinal(l_pe,fac=0.728583):
    """
    Define the matched bunch length for a plasma. Currently adjusted for Facet Li gas @ 1.5e+23 m^-3.

    Parameters
    ----------
    lp_e: plasma wavelength

    Returns
    -------
    l_bunch_fw: matched bunch length for Facet beam (Li gas at ~ 1.5e+23 m^-3)


    """
    # Matched length (FW) for a longitudinally triangular bunch
    l_fw_matched = 2. * l_pe
    # Maximum RMS length for a longitudinally gaussian bunch
    l_rms_max   = l_pe / math.pi

    # For gaussian beam profiles, bunch length is an arbitrary number of e-foldings
    # Resonance requires RMS_LENGTH <= L_RMS_MAX
    # Ad hoc factor of 0.728583 is reasonable & yields Facet beam
    # length when Li gas density is chosen to be 1.5e+23 m^-3
    rms_length = fac * l_rms_max

    # At FACET, sigma_z ~ 20 microns
    # rms_length = 20.0e-06

    long_norm    = 1.0 / ( sqrt(2.*math.pi) * rms_length )
    x_bunch_hw   = 3.0 * rms_length
    bunch_length = 2.0 * x_bunch_hw
    t_bunch_fw   = bunch_length / clight

    return bunch_length


# ---------------------------------------------------------------------------------------
# Basic user-defined parameters
# ---------------------------------------------------------------------------------------

#Top level Physical Parameters - User editable
#---------------------------------------------

# Reference plasma density (in number of particles per m^3)
N_PLASMA = {{electronPlasma_density}} # plasma density (in m^-3)
# Electron plasma wavelength for full ionization
lambda_pe = 3.34e+07 / sqrt(N_PLASMA)

#Number of standard deviations to fill with particles in each direction (from center of bunch)
#Also used to determine maximum grid dimensions by multiplying plasma wavelength
NUM_SIGMA_R = {{simulationGrid_rScale}}
NUM_SIGMA_L = {{simulationGrid_zScale}}

USE_LASER = {{useLaser}}

USE_BEAM = {{useBeam}}


# ---------------------------------------------------------------------------------------
# General simulation parameters
# ---------------------------------------------------------------------------------------
# Dimension of simulation ("3d", "circ", "2d", "1d")
dim = "circ"
# Number of azimuthal modes beyond m=0, for "circ" (not used for "2d" and "3d")
circ_m = 1

# Moving window (0:off, 1:on)
USE_MOVING_WINDOW = 1
# Speed of the moving window (ignored if use_moving_window = 0)
v_moving_window = clight

# Field boundary conditions (longitudinal and transverse respectively)
f_boundz  = openbc
f_boundxy = absorb
# Particles boundary conditions (longitudinal and transverse respectively)
p_boundz  = absorb
p_boundxy = reflect


# Current smoothing parameters
# ----------------------------
# Turn current smoothing on or off (0:off; 1:on)
use_smooth = 1
# Number of passes of smoother and compensator in each direction (x, y, z)
npass_smooth = array([[ 0 , 0 ], [ 0 , 0 ], [ 1 , 1 ]])
# Smoothing coefficients in each direction (x, y, z)
alpha_smooth = array([[ 0.5, 3.], [ 0.5, 3.], [0.5, 3./2]])
# Stride in each direction (x, y, z)
stride_smooth = array([[ 1 , 1 ], [ 1 , 1 ], [ 1 , 1 ]])



# Numerical parameters
# ---------------------
# Field solver (0:Yee, 1:Karkkainen on EF,B, 3:Lehe)
stencil = 0
# Particle shape (1:linear, 2:quadratic, 3:cubic)
depos_order = 2
# Gathering mode (1:from cell centers, 4:from Yee mesh)
efetch = 1
# Particle pusher (0:Boris, 1:Vay)
particle_pusher = 1



# Diagnostics
# -----------
# Period of diagnostics (in number of timesteps)
diag_period = {{incSteps}}
# Whether to write the fields
write_fields = 1
# Whether to write the particles
write_particles = {{isAnimationView}} or USE_BEAM
# Whether to write the diagnostics in parallel
parallel_output = False


# ---------------------------------------------------------------------------------------
#General Plasma Parameters
# ---------------------------------------------------------------------------------------


# Plasma macroparticles
# ---------------------
# Number of macroparticles per cell in each direction
# In Circ, nppcelly is the number of particles along the
# azimuthal direction. Use a multiple of 4*circ_m
plasma_nx = {{simulationGrid_rParticlesPerCell}}
plasma_ny = 4 * circ_m
plasma_nz = {{simulationGrid_zParticlesPerCell}}


# Plasma content and profile
# --------------------------

# Initialize some preexisting plasmas electrons (0:off, 1:on)
# (Can be used in order to neutralize pre-ionized ions, if any,
# or in order to simulate a plasma without having to initialize ions)
USE_PREEXISTING_ELECTRONS = 1
# Relative density of the preexisting electrons (relative to n_plasma)
rel_dens_preexisting_electrons = 1.
# Initialize plasma ions (0:off, 1:on) - Turn off for speed using a pure electron plasma
USE_IONS = 0
# The different elements used. (Only used if use_ions is different than 0.)
# relative_density is the density relative to n_plasma.
# q_start is the ionization state of the ions at the beginning of the simulation
# q_max is the maximum ionization state
# If q_start is not equal to q_max, ionization between states will be computed.
ion_states = { 'Hydrogen': {'relative_density':1., 'q_start':1, 'q_max':1 },
                 'Helium': {'relative_density':2., 'q_start':0, 'q_max':2 } }

# Specify the grid parameters
# --------------------------

#compute needed quantities for initialization
nx = {{simulationGrid_rCount}}
ny = 300
nz = {{simulationGrid_zCount}}
zmin = {{simulationGrid_zMin}}
zmax = {{simulationGrid_zMax}}
xmin = {{simulationGrid_zMin}}
xmax = {{simulationGrid_rMax}}
ymax = xmax
ymin = -ymax

# Specify the plasma parameters - these can be independent of the initial grid dimensions
# ---------------------------------------------------------------------------------------
# Positions between which the plasma is initialized
# (Transversally, the plasma is initialized between -plasma_xmax and
# plasma_xmax, along x, and -plasma_ymax and plasma_ymax along y)
plasma_zmin = 1.e-6
plasma_zmax = {{electronPlasma_length}}
plasma_xmax = xmax
plasma_ymax = ymax

# Define your own profile using these parameters, and passing them to plasma_dens_func()
#The following values are in units of meters
ramp_start = 0.e-6
ramp_length = {{electronPlasma_entranceRampLength}}
exit_ramp_length = {{electronPlasma_exitRampLength}}
ramp_plateau = plasma_zmax - ramp_length - exit_ramp_length


# ---------------------------------------------------------------------------------------
# User-defined initialization functions. Designed to be called Plasma_Injector()
# ---------------------------------------------------------------------------------------


#def plasma_dens_func(x, y, z, ramp_start, ramp_length, plateau_length):
def plasma_dens_func(x, y, z):
    """
    User-defined function: density profile of the plasma

    It should return the relative density with respect to n_plasma,
    at the position x, y, z (i.e. return a number between 0 and 1)

    Parameters
    ----------
    x, y, z: 1darrays of floats
        Arrays with one element per macroparticle

    Currently Deprecated
    --------------------
    ramp_start: beginning position of density profile
    ramp_length: length of plasma ramp
    plateau_length: length of plateau

    Returns
    -------
    n : 1d array of floats
        Array of relative density, with one element per macroparticles
    """
    # Allocate relative density
    n = ones_like(z)
    # # Make linear ramp
    # n = where( z<ramp_start+ramp_length, (z-ramp_start)/ramp_length, n )
    # # Supress density before the ramp
    # n = where( z<ramp_start, 0., n )
    # # Reduce density by half after the ramp
    # n = where( z> ramp_start+ramp_length+ramp_plateau, 0.5*n, n )
    # # Put the density to 0 later
    # n = where( z> ramp_start+ramp_length+2*ramp_plateau, 0., n )

    n = where(z < ramp_length, 0., n)
    #TODO(pjm): taken from previous script, this could be completely wrong, need to review this
    n = where(z > plasma_zmax - exit_ramp_length, sin(0.5 * pi * (plasma_zmax - z) / exit_ramp_length) * n, n)
    return(n)


#def beam_dens_func(x, y, z, beam_zmin, beam_zmax, beam_rmax):
def beam_dens_func(x, y, z):
    """
    User-defined function: density profile of the beam

    It should return the relative density with respect to n_beam,
    at the position x, y, z (i.e. return a number between 0 and 1)

    Parameters
    ----------
    x, y, z: 1darrays of floats
        Arrays with one element per macroparticle

    Currently Deprecated
    --------------------
    beam_zmin: minimum z position of the beam
    beam_zmax: maximum z position of the beam
    beam_rmax: maximum transverse dimension of the beam (> 0)

    Returns
    -------
    n : 1d array of floats
        Array of relative density, with one element per macroparticles
    """
    # Allocate relative density
    n = ones_like(z)
    # Longitudinal profile: parabolic
    n = n*(z - beam_zmin)*(beam_zmax - z) * 4/(beam_zmax - beam_zmin)**2
    # Transverse profile: parabolic
    r = sqrt( x**2 + y**2)
    n = n*(1 - (r/beam_rmax)**2 )
    # Put the density above rmax to 0
    n[r > beam_rmax] = 0.

    return(n)



# ---------------------------------------------------------------------------------------
# Initialization of the simulation (Normal users should not modify this part.)
# ---------------------------------------------------------------------------------------

# Set some general options for warp
set_diagnostics( 0 )
set_boundary_conditions( f_boundz, f_boundxy, p_boundz, p_boundxy )
set_simulation_box( nz, nx, ny, zmin, zmax, xmax, ymax, dim, xmin=xmin, ymin=ymin )
set_moving_window( USE_MOVING_WINDOW, v_moving_window )

# See smoothing.py
set_smoothing_parameters( use_smooth, dim, npass_smooth,
                         alpha_smooth, stride_smooth )

# Creation of the species
# -----------------------

elec = None
ions = None
elec_from_ions = None

# Create the plasma species
# Reference weight for plasma species
plasma_weight = prepare_weights( N_PLASMA, plasma_nx, plasma_ny,
                            plasma_nz, dim, circ_m )
if USE_PREEXISTING_ELECTRONS:
    elec_weight = rel_dens_preexisting_electrons * plasma_weight
    elec = Species(type=Electron, weight=elec_weight, name='electrons')
if USE_IONS:
    ions, elec_from_ions = initialize_ion_dict( ion_states, plasma_weight,
                                                group_elec_by_element=True )

    # Set the numerical parameters only now: they affect the newly created species
set_numerics( depos_order, efetch, particle_pusher, dim)

# Setup the field solver object
# -----------------------------
em = EM3D(
    stencil=stencil,
    npass_smooth=npass_smooth,
    alpha_smooth=alpha_smooth,
    stride_smooth=stride_smooth,
    l_2dxz= (dim in ["2d", "circ"]),
    l_2drz= (dim in ["circ"]),
    l_1dz = (dim =="1d" ),
    l_getrho=True,
    circ_m = (dim =="circ")*circ_m,
    type_rz_depose=1 )
registersolver(em)

# Introduce the laser (only if required)
# --------------------------------------
if USE_LASER:
    laser_a0 = {{laserPulse_amplitude}}
    lambda_laser = {{laserPulse_wavelength}}
    laser_w0 = {{laserPulse_waist}}
    laser_ctau = {{laserPulse_duration}} * clight
    laser_z0 = -2 * laser_ctau
    laser_zfoc = 50.e-6
    laser_polangle = {{laserPulse_polarizationAngle}}
    laser_source_z = 0

    add_laser( em, dim, laser_a0, laser_w0, laser_ctau, laser_z0,
        lambda0=lambda_laser,
        zf=laser_zfoc, theta_pol=laser_polangle, source_z=laser_source_z )


# Introduce the beam (only if required)
# -------------------------------------
if USE_BEAM:

    # Specify the maximum density of the e- drive bunch
    CHARGE = {{electronBeam_charge}} # total beam charge [C]
    n_total = CHARGE / echarge  # number of electrons

    #Beam energy
    beam_energy = {{electronBeam_energy}} #beam energy in GeV
    GAMMA_BEAM = beam_energy*1.e3/0.510998910 #beam gamma

    #The relative RMS longitudinal energy spread
    DE_OVER_E_RMS = 0.0                                             #USER EDIT

    #beam matching parameters
    EPS_N = {{electronBeam_transverseEmittance}} #normalized transverse emittance
    FIXED_R_RMS = {{electronBeam_rmsRadius}} #Specify this for a fixed radius bunch with 0 emittance. If EPS_N specified, this is ignored.
    L_BUNCH = {{electronBeam_bunchLength}} #This is fixed to match the beam to the Facet beam


    # ---------------------------------------------------------------------------------------
    # Relativistic beam parameters
    # ---------------------------------------------------------------------------------------
    # Initialize beam electrons (0:off, 1:on)
    # (Please be aware that initializing a beam in 2D geometry makes very little
    # physical sense, because of the long range of its space-charge fields)

    # Number of macroparticles per cell in each direction - adjusted to be equal in each direction
    beam_nx = 2 * plasma_nx
    beam_ny = 1 * plasma_ny
    beam_nz = 2 * plasma_nz

    #TRANSVERSE BEAM PARAMETERS
    #----------------------------
    # The RMS normalized transverse emittance
    # EPSILON_RMS_N = 2.25e-06    # [m-rad]
    #EPSILON_RMS_N = 0.     # [m-rad]

    # To avoid betatron oscillations of the drive beam, you
    #   must choose the matched RMS radius (see above).

    if EPS_N > 0:
        #use user specified emittance to match beam
        rms_radius = match_beam_transverse(EPS_N,lambda_pe)
    else:
        #match to FACET parameters (sigma_r ~ 15 microns) or user-specified
        rms_radius = FIXED_R_RMS

    rms_perp    = rms_radius / sqrt(2.) #rms_perp is for a given axis (x or y), so has a separate normalization
    perp_norm   = 1.0 / ( sqrt(2.*math.pi) * rms_perp ) #compute the radius with Gaussian normalization - use rms_perp
    max_radius  = NUM_SIGMA_R * rms_radius #Fill a region of radius NUM_SIGMA_R standard deviations

    #ENERGY AND CHARGE OF BEAM
    #-------------------------

    #GAMMA_BEAM = 45009.78 # at FACET, E~23 GeV
    beta_gamma = sqrt(GAMMA_BEAM**2 - 1.)
    beta_beam  = beta_gamma / GAMMA_BEAM

    #Longitudinal momentum of the beam
    beam_uz = beta_gamma
    v_z_avg = beta_beam  * clight
    u_z_avg = beta_gamma * clight

    # The corresponding RMS momenta
    u_z_rms = DE_OVER_E_RMS * u_z_avg
    u_x_rms = clight * EPS_N / rms_perp #if zero emittance, zero RMS transverse momenta
    u_y_rms = u_x_rms

    #LONGITUDINAL BEAM PARAMETERS
    #----------------------------
    # Calculate longitudinal matching of beam to e- plasma

    # At FACET, sigma_z ~ 20 microns
    # RMS_LENGTH = 20.0e-06

    if L_BUNCH > 0.:
        #use user-defined bunch length
        bunch_length = L_BUNCH
    else:
        #match the bunch to Facet beam
        bunch_length = match_beam_longitudinal(lambda_pe)


    l_norm   = 1.0 / ( sqrt(2.*math.pi)*bunch_length)
    t_bunch_fw   = bunch_length / clight


    #PEAK BEAM DENSITY
    #----------------------------
    # The maximum physical number density of electrons in the beam
    beam_density = n_total * l_norm * perp_norm**2
    # Beam density
    n_beam = beam_density


    # Positions between which the beam is initialized
    # (Transversally, the plasma is initialized between -plasma_xmax and
    # plasma_xmax, along x, and -plasma_ymax and plasma_ymax along y)
    beam_zmax = -1.e-6
    beam_zmin = beam_zmax - bunch_length
    beam_xmax = max_radius
    beam_ymax = beam_xmax

    # Define your own profile using these parameters, and passing them to beam_dens_func()
    beam_rmax = beam_xmax

    # Create the beam
    beam_weight = prepare_weights( n_beam, beam_nx, beam_ny,
                                   beam_nz, dim, circ_m )
    beam = Species(type=Electron, weight=beam_weight, name='beam')
    PlasmaInjector( beam, None, w3d, top, dim, beam_nx, beam_ny, beam_nz,
                beam_zmin, beam_zmax, beam_xmax, beam_ymax,
                dens_func = beam_dens_func, uz_m=beam_uz )
    initialize_beam_fields( em, dim, beam, w3d, top )

# Introduce the plasma
# --------------------
# Create an object to store the information about plasma injection
plasma_injector = PlasmaInjector( elec, ions, w3d, top, dim,
        plasma_nx, plasma_ny, plasma_nz, plasma_zmin,
        plasma_zmax, plasma_xmax, plasma_ymax, plasma_dens_func )

# Continuously inject the plasma, if the moving window is on
if {{isAnimationView}} and USE_MOVING_WINDOW :
    installuserinjection( plasma_injector.continuous_injection )

# Setup the diagnostics
# ---------------------
if write_fields:
    diag1 = FieldDiagnostic( period=diag_period, top=top, w3d=w3d, em=em,
                comm_world=comm_world, lparallel_output=parallel_output,
                write_dir={{outputDir}})
    installafterstep( diag1.write )
if write_particles:
    diag2 = ParticleDiagnostic( period=diag_period, top=top, w3d=w3d,
            species={ species.name : species for species in listofallspecies },
            comm_world=comm_world, lparallel_output=parallel_output,
            write_dir={{outputDir}})
    installafterstep( diag2.write )

# ---------------------------------------------------------------------------------------
# Simulation loop (Normal users should not modify this part either.)
# ---------------------------------------------------------------------------------------

inc_steps = {{incSteps}}

if {{isAnimationView}}:
    doit = True
    while doit:
        step(inc_steps)
        doit = w3d.zmmin + top.zgrid < plasma_zmax

{'fraction_non_alphanumeric': 0.18081552444116925, 'fraction_numerical': 0.01353475804470646, 'mean_word_length': 2.215972825657635, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7555223', 'n_tokens_mistral': 12812, 'n_tokens_neox': 11860, 'n_words': 3961}