def check_charge_conservation(sim, rho_ions):
"""
Check that the relation div(E) - rho/epsilon_0 is satisfied, with a
relative precision close to the machine precision (directly in spectral space)
Parameters
----------
sim: Simulation object
rho_ions: list of 2d complex arrays (one per mode)
The density of the ions (which are not explicitly present in the `sim`
object, since they are motionless)
"""
# Create a global field object across all subdomains, and copy the fields
global_Nz, _ = sim.comm.get_Nz_and_iz(local=False,
with_damp=False,
with_guard=False)
global_zmin, global_zmax = sim.comm.get_zmin_zmax(local=False,
with_damp=False,
with_guard=False)
global_fld = Fields(global_Nz,
global_zmax,
sim.fld.Nr,
sim.fld.rmax,
sim.fld.Nm,
sim.fld.dt,
zmin=global_zmin,
n_order=sim.fld.n_order,
use_cuda=False)
# Gather the fields of the interpolation grid
for m in range(sim.fld.Nm):
# Gather E
for field in ['Er', 'Et', 'Ez']:
local_array = getattr(sim.fld.interp[m], field)
gathered_array = sim.comm.gather_grid_array(local_array)
setattr(global_fld.interp[m], field, gathered_array)
# Gather rho
global_fld.interp[m].rho = \
sim.comm.gather_grid_array( sim.fld.interp[m].rho + rho_ions[m] )
# Loop over modes and check charge conservation in spectral space
if sim.comm.rank == 0:
global_fld.interp2spect('E')
global_fld.interp2spect('rho_prev')
for m in range(global_fld.Nm):
spect = global_fld.spect[m]
# Calculate div(E) in spectral space
divE = spect.kr * (spect.Ep - spect.Em) + 1.j * spect.kz * spect.Ez
# Calculate rho/epsilon_0 in spectral space
rho_eps0 = spect.rho_prev / epsilon_0
# Calculate relative RMS error
rel_err = np.sqrt( np.sum(abs(divE - rho_eps0)**2) \
/ np.sum(abs(rho_eps0)**2) )
print('Relative error on divE in mode %d: %e' % (m, rel_err))
assert rel_err < 1.e-11
def get_space_charge_fields(sim, ptcl, direction='forward'):
"""
Add the space charge field from `ptcl` the interpolation grid
This assumes that all the particles being passed have the same gamma.
Parameters
----------
sim : a Simulation object
Contains the values of the fields, and the MPI communicator
ptcl : a Particles object
The list of the species which are relativistic and
will produce a space charge field. (Do not pass the
particles which are at rest.)
direction : string, optional
Can be either "forward" or "backward".
Propagation direction of the beam.
"""
if sim.comm.rank == 0:
print("Calculating initial space charge field...")
# Calculate the mean gamma by computing weighted sum on each subdomain
w_sum_local = ptcl.w.sum()
w_gamma_sum_local = (ptcl.w * 1. / ptcl.inv_gamma).sum()
if sim.comm.mpi_comm is None:
w_sum = w_sum_local
w_gamma_sum = w_gamma_sum_local
else:
w_sum = sim.comm.mpi_comm.allreduce(w_sum_local)
w_gamma_sum = sim.comm.mpi_comm.allreduce(w_gamma_sum_local)
# Check that the number of particles is not 0
if w_sum == 0:
warnings.warn(
"Tried to calculate space charge, but found 0 macroparticles in \n"
"the corresponding species. Skipping space charge calculation...\n"
)
return
else:
gamma = w_gamma_sum / w_sum
# Project the charge and currents onto the local subdomain
sim.deposit('rho',
exchange=True,
species_list=[ptcl],
update_spectral=False)
sim.deposit('J', exchange=True, species_list=[ptcl], update_spectral=False)
# Create a global field object across all subdomains, and copy the sources
# (Space-charge calculation is a global operation)
# Note: in the single-proc case, this is also useful in order not to
# erase the pre-existing E and B field in sim.fld
global_Nz, _ = sim.comm.get_Nz_and_iz(local=False,
with_damp=True,
with_guard=False)
global_zmin, global_zmax = sim.comm.get_zmin_zmax(local=False,
with_damp=True,
with_guard=False)
global_fld = Fields(global_Nz,
global_zmax,
sim.fld.Nr,
sim.fld.rmax,
sim.fld.Nm,
sim.fld.dt,
n_order=sim.fld.n_order,
smoother=sim.fld.smoother,
zmin=global_zmin,
use_cuda=False)
# Gather the sources on the interpolation grid of global_fld
for m in range(sim.fld.Nm):
for field in ['Jr', 'Jt', 'Jz', 'rho']:
local_array = getattr(sim.fld.interp[m], field)
gathered_array = sim.comm.gather_grid_array(local_array,
with_damp=True)
setattr(global_fld.interp[m], field, gathered_array)
# Calculate the space-charge fields on the global grid
# (For a multi-proc simulation: only performed by the first proc)
if sim.comm.rank == 0:
# - Convert the sources to spectral space
global_fld.interp2spect('rho_prev')
global_fld.interp2spect('J')
if sim.filter_currents:
global_fld.filter_spect('rho_prev')
global_fld.filter_spect('J')
# - Get the space charge fields in spectral space
for m in range(global_fld.Nm):
get_space_charge_spect(global_fld.spect[m], gamma, direction)
# - Convert the fields back to real space
global_fld.spect2interp('E')
global_fld.spect2interp('B')
# Communicate the results from proc 0 to the other procs
# and add it to the interpolation grid of sim.fld.
# - First find the indices at which the fields should be added
Nz_local, iz_start_local_domain = sim.comm.get_Nz_and_iz(
local=True, with_damp=True, with_guard=False, rank=sim.comm.rank)
_, iz_start_local_array = sim.comm.get_Nz_and_iz(local=True,
with_damp=True,
with_guard=True,
rank=sim.comm.rank)
iz_in_array = iz_start_local_domain - iz_start_local_array
# - Then loop over modes and fields
for m in range(sim.fld.Nm):
for field in ['Er', 'Et', 'Ez', 'Br', 'Bt', 'Bz']:
# Get the local result from proc 0
global_array = getattr(global_fld.interp[m], field)
local_array = sim.comm.scatter_grid_array(global_array,
with_damp=True)
# Add it to the fields of sim.fld
local_field = getattr(sim.fld.interp[m], field)
local_field[iz_in_array:iz_in_array + Nz_local, :] += local_array
if sim.comm.rank == 0:
print("Done.\n")
def add_laser_direct(sim, laser_profile, boost):
"""
Add a laser pulse in the simulation, by directly adding it to the mesh
Note:
-----
Arbitrary laser profiles can be passed through `laser_profile`
(which must provide the *transverse electric field*)
For any profile:
- The field is automatically decomposed into azimuthal modes
- The Ez field is automatically calculated so as to ensure that div(E)=0
- The B field is automatically calculated so as to ensure propagation
in the right direction.
Parameters:
-----------
sim: a Simulation object
The structure that contains the simulation.
laser_profile: a valid laser profile object
Laser profiles can be imported from fbpic.lpa_utils.laser
boost: a BoostConverter object or None
Contains the information about the boost to be applied
"""
if sim.comm.rank == 0:
print("Initializing laser pulse on the mesh...")
# Get the local azimuthally-decomposed laser fields Er and Et on each proc
laser_Er, laser_Et = get_laser_Er_Et(sim, laser_profile, boost)
# Save previous values on the grid, and replace them with the laser fields
# (This is done in preparation for gathering among procs)
saved_Er = []
saved_Et = []
for m in range(sim.fld.Nm):
saved_Er.append(sim.fld.interp[m].Er.copy())
sim.fld.interp[m].Er[:, :] = laser_Er[:, :, m]
saved_Et.append(sim.fld.interp[m].Et.copy())
sim.fld.interp[m].Et[:, :] = laser_Et[:, :, m]
# Create a global field object across all subdomains, and copy the fields
# (Calculating the self-consistent Ez and B is a global operation)
global_Nz, _ = sim.comm.get_Nz_and_iz(local=False,
with_damp=True,
with_guard=False)
global_zmin, global_zmax = sim.comm.get_zmin_zmax(local=False,
with_damp=True,
with_guard=False)
global_fld = Fields(global_Nz,
global_zmax,
sim.fld.Nr,
sim.fld.rmax,
sim.fld.Nm,
sim.fld.dt,
zmin=global_zmin,
n_order=sim.fld.n_order,
use_cuda=False)
# Gather the fields of the interpolation grid
for m in range(sim.fld.Nm):
for field in ['Er', 'Et']:
local_array = getattr(sim.fld.interp[m], field)
gathered_array = sim.comm.gather_grid_array(local_array,
with_damp=True)
setattr(global_fld.interp[m], field, gathered_array)
# Now that the (gathered) laser fields are stored in global_fld,
# copy the saved field back into the local grid
for m in range(sim.fld.Nm):
sim.fld.interp[m].Er[:, :] = saved_Er[m]
sim.fld.interp[m].Et[:, :] = saved_Et[m]
# Calculate the Ez and B fields on the global grid
# (For a multi-proc simulation: only performed by the first proc)
if sim.comm.rank == 0:
calculate_laser_fields(global_fld, laser_profile.propag_direction)
# Communicate the results from proc 0 to the other procs
# and add it to the interpolation grid of sim.fld.
# - First find the indices at which the fields should be added
Nz_local, iz_start_local_domain = sim.comm.get_Nz_and_iz(
local=True, with_damp=True, with_guard=False, rank=sim.comm.rank)
_, iz_start_local_array = sim.comm.get_Nz_and_iz(local=True,
with_damp=True,
with_guard=True,
rank=sim.comm.rank)
iz_in_array = iz_start_local_domain - iz_start_local_array
# - Then loop over modes and fields
for m in range(sim.fld.Nm):
for field in ['Er', 'Et', 'Ez', 'Br', 'Bt', 'Bz']:
# Get the local result from proc 0
global_array = getattr(global_fld.interp[m], field)
local_array = sim.comm.scatter_grid_array(global_array,
with_damp=True)
# Add it to the fields of sim.fld
local_field = getattr(sim.fld.interp[m], field)
local_field[iz_in_array:iz_in_array + Nz_local, :] += local_array
if sim.comm.rank == 0:
print("Done.\n")