def increase_particle_array_size(pos, vel, idx, Ie, W_elec, Ib, W_mag, Ep, Bp,
v_prime, S, T, temp_N, old_particles):
'''
Allocates more memory in case spare particles run out. Super inefficient way of doing it, so
when this function is called, it'll probably double the size of the program in memory. Maybe
numba/numpy has a native function that will let us do this?
'''
old_N = pos.shape[0]
inc_size = 5 * nsp_ppc.sum()
# Initialize 'increased' arrays
ipos = np.zeros((old_N + inc_size), dtype=nb.float64)
ivel = np.zeros((3, old_N + inc_size), dtype=nb.float64)
iidx = np.zeros((old_N + inc_size), dtype=nb.uint16)
iIe = np.zeros((old_N + inc_size), dtype=nb.uint16)
iIb = np.zeros((old_N + inc_size), dtype=nb.uint16)
iW_elec = np.zeros((3, old_N + inc_size), dtype=nb.float64)
iW_mag = np.zeros((3, old_N + inc_size), dtype=nb.float64)
iEp = np.zeros((3, old_N + inc_size), dtype=nb.float64)
iBp = np.zeros((3, old_N + inc_size), dtype=nb.float64)
iv_prime = np.zeros((3, old_N + inc_size), dtype=nb.float64)
iS = np.zeros((3, old_N + inc_size), dtype=nb.float64)
iT = np.zeros((3, old_N + inc_size), dtype=nb.float64)
itemp_N = np.zeros((old_N + inc_size), dtype=nb.float64)
iold_particles = np.zeros((old_particles.shape[0], old_N + inc_size),
dtype=nb.float64)
# Fill new arrays with existing values except W_mag, Ib, Ep, Bp, S, T, v_prime, temp_N and old_particles
# Since these will either be zeroed or overwritten
ipos[:old_N] = pos[:]
ivel[:, :old_N] = vel[:, :]
iidx[:old_N] = idx[:]
iIe[:old_N] = Ie[:]
iW_elec[:, :old_N] = W_elec[:, :]
return ipos, ivel, iidx, iIe, iIb, iW_elec, iW_mag, iEp, iBp, iv_prime, iS, iT, itemp_N, iold_particles
def main_loop(pos, vel, idx, Ie, W_elec, Ib, W_mag, Ep, Bp, v_prime, S, T,temp_N, \
B, E_int, E_half, q_dens, q_dens_adv, Ji, ni, nu, mp_flux, \
Ve, Te, Te0, temp3De, temp3Db, temp1D, old_particles, old_fields,\
B_damping_array, E_damping_array, qq, DT, max_inc, part_save_iter, field_save_iter):
'''
Main loop separated from __main__ function, since this is the actual computation bit.
Could probably be optimized further, but I wanted to njit() it.
The only reason everything's not njit() is because of the output functions.
Future: Come up with some way to loop until next save point
Thoughts: declare a variable steps_to_go. Output all time variables at return
to resync everything, and calculate steps to next stop.
If no saves, steps_to_go = max_inc
'''
# Check timestep
if qq % 20 == 0:
#print('Checking timestep')
qq, DT, max_inc, part_save_iter, field_save_iter, damping_array \
= check_timestep(pos, vel, B, E_int, q_dens, Ie, W_elec, Ib, W_mag, temp3De, Ep, Bp, v_prime, S, T,temp_N,\
qq, DT, max_inc, part_save_iter, field_save_iter, idx, B_damping_array)
# Check number of spare particles every 25 steps
if qq % 50 == 0 and particle_open == 1:
#print('Checking number of spare particles')
num_spare = (idx < 0).sum()
if num_spare < nsp_ppc.sum():
print(
'WARNING :: Less than one cell of spare particles remaining.')
if num_spare < inject_rate.sum() * DT * 5.0:
# Change this to dynamically expand particle arrays later on (adding more particles)
# Can do it by cell lots (i.e. add a cell's worth each time)
print(
'WARNING :: No space particles remaining. Exiting simulation.'
)
raise IndexError
# Move particles, collect moments
particles.advance_particles_and_moments(pos, vel, Ie, W_elec, Ib, W_mag, idx, Ep, Bp, v_prime, S, T,temp_N,\
B, E_int, DT, q_dens_adv, Ji, ni, nu, mp_flux)
# Average N, N + 1 densities (q_dens at N + 1/2)
q_dens *= 0.5
q_dens += 0.5 * q_dens_adv
# Compute fields at N + 1/2
fields.push_B(B, E_int, temp3Db, DT, qq, B_damping_array, half_flag=1)
fields.calculate_E(B, Ji, q_dens, E_half, Ve, Te, Te0, temp3De, temp3Db,
temp1D, E_damping_array, qq, DT, 0)
###################################
### PREDICTOR CORRECTOR SECTION ###
###################################
# Store old values
mp_flux_old = mp_flux.copy()
old_particles[0, :] = pos
old_particles[1:4, :] = vel
old_particles[4, :] = Ie
old_particles[5:8, :] = W_elec
old_particles[8, :] = idx
old_fields[:, 0:3] = B
old_fields[:NC, 3:6] = Ji
old_fields[:NC, 6:9] = Ve
old_fields[:NC, 9] = Te
# Predict fields
E_int *= -1.0
E_int += 2.0 * E_half
fields.push_B(B, E_int, temp3Db, DT, qq, B_damping_array, half_flag=0)
# Advance particles to obtain source terms at N + 3/2
particles.advance_particles_and_moments(pos, vel, Ie, W_elec, Ib, W_mag, idx, Ep, Bp, v_prime, S, T,temp_N,\
B, E_int, DT, q_dens, Ji, ni, nu, mp_flux, pc=1)
q_dens *= 0.5
q_dens += 0.5 * q_dens_adv
# Compute predicted fields at N + 3/2
fields.push_B(B, E_int, temp3Db, DT, qq + 1, B_damping_array, half_flag=1)
fields.calculate_E(B, Ji, q_dens, E_int, Ve, Te, Te0, temp3De, temp3Db,
temp1D, E_damping_array, qq, DT, 1)
# Determine corrected fields at N + 1
E_int *= 0.5
E_int += 0.5 * E_half
# Restore old values: [:] allows reference to same memory (instead of creating new, local instance)
pos[:] = old_particles[0, :]
vel[:] = old_particles[1:4, :]
Ie[:] = old_particles[4, :]
W_elec[:] = old_particles[5:8, :]
idx[:] = old_particles[8, :]
B[:] = old_fields[:, 0:3]
Ji[:] = old_fields[:NC, 3:6]
Ve[:] = old_fields[:NC, 6:9]
Te[:] = old_fields[:NC, 9]
fields.push_B(B, E_int, temp3Db, DT, qq, B_damping_array,
half_flag=0) # Advance the original B
q_dens[:] = q_dens_adv
mp_flux = mp_flux_old.copy()
return qq, DT, max_inc, part_save_iter, field_save_iter
def inject_particles(pos, vel, idx, flux, dt, pc):
'''
Basic injection routine. Flux is set by either:
-- Measuring outgoing
-- Set as initial condition by Maxwellian
This code injects particles to equal that outgoing flux. Not sure
how this is going to go with conserving moments, but we'll see.
If this works, might be able to use Daughton conditions later. But should
at least keep constant under static conditions.
NOTE: Finding a random particle like that, do I have to change the timestep?
What if the particle is just a little too fast? Check later. Or put in cap
(3*vth?)
'''
N_retries = 10 # Set number of times to try and reinitialize a particle
# This is so when the flux is low (but non-zero) the code
# doesn't start getting exponentially longer at the tail
# looking for a particle with almost zero velocity.
# Check number of spare particles
num_spare = (idx < 0).sum()
if num_spare < 2 * nsp_ppc.sum():
print(
'WARNING :: Less than two cells worth of spare particles remaining.'
)
if num_spare == 0:
print(
'WARNING :: No space particles remaining. Exiting simulation.')
raise IndexError
# Create particles one at a time until flux drops below a certain limit
# or number of retries is reached
acc = 0
n_created = np.zeros((2, Nj), dtype=np.float64)
for ii in range(2):
for jj in range(Nj):
while flux[ii, jj] > 0:
# Loop until not-too-energetic particle found
vx = 3e8
new_particle = 0
# Find a vx that'll fit in remaining flux
for n_tried in range(N_retries):
vx = np.random.normal(0.0, vth_par[jj])
if vx <= flux[ii, jj]:
new_particle = 1
break
# If successful, load particle
if new_particle == 1:
# Successful particle found, set parameters and subtract flux
# Find first negative idx to initialize as new particle
for kk in nb.prange(acc, pos.shape[1]):
if idx[kk] < 0:
acc = kk + 1
break
# Decide direction of vx, and placement of particle
# This could probably be way optimized later.
if ii == 0:
vel[0, kk] = np.abs(vx)
max_pos = vel[0, kk] * dt
if abs(max_pos) < 0.5 * dx:
pos[0,
kk] = xmin + np.random.uniform(0, 1) * max_pos
else:
pos[0, kk] = xmin + max_pos
else:
vel[0, kk] = -np.abs(vx)
max_pos = vel[0, kk] * dt
if abs(max_pos) < 0.5 * dx:
pos[0,
kk] = xmax + np.random.uniform(0, 1) * max_pos
else:
pos[0, kk] = xmax + max_pos
vel[1, kk] = np.random.normal(0.0, vth_perp[jj])
vel[2, kk] = np.random.normal(0.0, vth_perp[jj])
gyangle = init.get_gyroangle_single(vel[:, kk])
rL = np.sqrt(vel[1, kk]**2 +
vel[2, kk]**2) / (qm_ratios[idx[kk]] * B_xmax)
pos[1, kk] = rL * np.cos(gyangle)
pos[2, kk] = rL * np.sin(gyangle)
idx[kk] = jj
flux[ii, jj] -= abs(vx)
n_created[ii, jj] += 1
else:
#print('Number of retries reached, stopping injection...')
break
#print('Particles created:\n', n_created)
return
# Reflect/reinject script
# =============================================================================
# for ii in nb.prange(pos.shape[1]):
# sp = idx[ii]
# reflag = 0
# # Reflect cold
# if temp_type[sp] == 0:
# if pos[0, ii] > xmax:
# pos[0, ii] = 2*xmax - pos[0, ii]; reflag=1
# elif pos[0, ii] < xmin:
# pos[0, ii] = 2*xmin - pos[0, ii]; reflag=1
#
# if reflag==1:
# vel[0, ii] *= -1.0
#
# # Reinitialize hot
# else:
# if pos[0, ii] > xmax or pos[0, ii] < xmin:
# vel[0, ii] = np.random.normal(0, vth_par[sp]) * (-np.sign(pos[0, ii]))
# vel[1, ii] = np.random.normal(0, vth_perp[sp])
# vel[2, ii] = np.random.normal(0, vth_perp[sp])
#
# max_pos = abs(vel[0, ii])*DT
# pos[0, ii] = np.sign(pos[0, ii]) * (xmax - max_pos)
# =============================================================================
def position_update(pos, vel, idx, pos_old, DT, Ie, W_elec):
'''
Updates the position of the particles using x = x0 + vt.
Also updates particle nearest node and weighting.
INPUT:
pos -- Particle position array (Also output)
vel -- Particle velocity array (Also output for reflection)
idx -- Particle index array (Also output for reflection)
DT -- Simulation time step
Ie -- Particle leftmost to nearest node array (Also output)
W_elec -- Particle weighting array (Also output)
Note: This function also controls what happens when a particle leaves the
simulation boundary.
acc is an array accumulator that should prevent starting the search from the
beginning of the array each time a particle needs to be reinitialized
Is it super inefficient to do two loops for the disable/enable bit? Or worse
to have arrays all jumbled and array access all horrible? Is there a more
efficient way to code it?
To Do: Code with one loop, same acc, etc. Assume that there is enough
disabled particles to fulfill the needs for a single timestep (especially
since all the 'spare' particles are at the end. Spare particles should only
then be accessed when there was not sufficient numbers of 'normal' particles
in domain, or they were accessed out of order. Might save a bit of time. But
also need to check it against the 2-loop version.
IDEA: INSTEAD OF STORING POS_OLD, USE Ie INSTEAD, SINCE IT STORES THE CLOSEST
CELL CENTER - IF IT WAS IN LHS BOUNDARY CELL Ie[ii] == ND (+1?)
To do: Increase the size of particle arrays rather than Exception if particles
run out. Maybe have this check outside in main() so all particle-related arrays
can be copied and extended dynamically if required. Just need to check that
nothing relies purely on N.
'''
# L/R boundary, species
n_flux = np.zeros(2)
pos_old[:, :] = pos
pos[0, :] += vel[0, :] * DT
pos[1, :] += vel[1, :] * DT
pos[2, :] += vel[2, :] * DT
if particle_periodic == 1:
for ii in nb.prange(pos.shape[1]):
if pos[0, ii] > xmax:
pos[0, ii] += xmin - xmax
n_flux[1] += 1
elif pos[0, ii] < xmin:
pos[0, ii] += xmax - xmin
n_flux[0] += 1
else:
# Disable loop: Remove particles that leave the simulation space
n_deleted = 0
for ii in nb.prange(pos.shape[1]):
if (pos[0, ii] < xmin or pos[0, ii] > xmax):
pos[:, ii] *= 0.0
vel[:, ii] *= 0.0
idx[ii] = -1
n_deleted += 1
# Check number of spare particles
num_spare = (idx < 0).sum()
if num_spare < 2 * nsp_ppc.sum():
print('WARNING :: Less than two cells worth of spare particles remaining.')
if num_spare == 0:
print('WARNING :: No space particles remaining. Exiting simulation.')
raise IndexError
acc = 0; n_created = 0
for ii in nb.prange(pos.shape[1]):
# If particle goes from cell 1 to cell 2
if (pos_old[0, ii] < xmin + dx):
if (pos[0, ii] >= xmin + dx) and (pos[0, ii] <= xmin + 2*dx):
# Find first negative idx to initialize as new particle
for kk in nb.prange(acc, pos.shape[1]):
if idx[kk] < 0:
acc = kk + 1
break
# Create new particle with pos_old, vel of this particle
pos[0, kk] = pos[0, ii] - dx
pos[1, kk] = pos[1, ii]
pos[2, kk] = pos[2, ii]
vel[:, kk] = vel[:, ii]
idx[kk] = idx[ii]
n_created += 1
# If particle goes from cell NX to cell NX - 1
elif (pos_old[0, ii] > xmax - dx):
if (pos[0, ii] <= xmax - dx) and (pos[0, ii] > xmax - 2*dx):
# Find first negative idx to initialize as new particle
for kk in nb.prange(acc, pos.shape[1]):
if idx[kk] < 0:
acc = kk + 1
break
# Create new particle with pos_old, vel of this particle
pos[0, kk] = pos[0, ii] + dx
pos[1, kk] = pos[1, ii]
pos[2, kk] = pos[2, ii]
vel[:, kk] = vel[:, ii]
idx[kk] = idx[ii]
n_created += 1
print(n_flux)
#print(n_created, 'created ; ', n_deleted, 'deleted')
assign_weighting_TSC(pos, idx, Ie, W_elec)
return
def position_update(pos, vel, idx, pos_old, DT, Ie, W_elec):
'''
Updates the position of the particles using x = x0 + vt.
Also updates particle nearest node and weighting.
INPUT:
pos -- Particle position array (Also output)
vel -- Particle velocity array (Also output for reflection)
idx -- Particle index array (Also output for reflection)
DT -- Simulation time step
Ie -- Particle leftmost to nearest node array (Also output)
W_elec -- Particle weighting array (Also output)
Note: This function also controls what happens when a particle leaves the
simulation boundary.
NOTE :: Check the > < and >= <= equal positions when more awake, just make sure
they're right (not under or over triggering the condition)
acc is an array accumulator that should prevent starting the search from the
beginning of the array each time a particle needs to be reinitialized
Is it super inefficient to do two loops for the disable/enable bit? Or worse
to have arrays all jumbled and array access all horrible? Is there a more
efficient way to code it?
To Do: Code with one loop, same acc, etc. Assume that there is enough
disabled particles to fulfill the needs for a single timestep (especially
since all the 'spare' particles are at the end. Spare particles should only
then be accessed when there was not sufficient numbers of 'normal' particles
in domain, or they were accessed out of order. Might save a bit of time. But
also need to check it against the 2-loop version.
IDEA: Instead of searching for negative indices, assume all negatives are at
end of array. Sort arrays every X timesteps so 'disabled' particles are always
at end.
OR: Generate array of disabled indices at each timestep and call that instead of
searching each time.
IDEA: INSTEAD OF STORING POS_OLD, USE Ie INSTEAD, SINCE IT STORES THE CLOSEST
CELL CENTER - IF IT WAS IN LHS BOUNDARY CELL Ie[ii] == ND (+1?)
I THINK I FIXED THE ERROR :: POOR SPECIFICATION OF WHETHER OR NOT THE PARTICLE WAS
IN CELL 2 OR NOT.
'''
pos_old[:, :] = pos
pos[0, :] += vel[0, :] * DT
pos[1, :] += vel[1, :] * DT
pos[2, :] += vel[2, :] * DT
if particle_periodic == 1:
for ii in nb.prange(pos.shape[1]):
if pos[0, ii] > xmax:
pos[0, ii] += xmin - xmax
elif pos[0, ii] < xmin:
pos[0, ii] += xmax - xmin
else:
# Disable loop: Remove particles that leave the simulation space
n_deleted = 0
for ii in nb.prange(pos.shape[1]):
if (pos[0, ii] < xmin or pos[0, ii] > xmax):
pos[:, ii] *= 0.0
vel[:, ii] *= 0.0
idx[ii] = -1
n_deleted += 1
# Check number of spare particles
num_spare = (idx < 0).sum()
if num_spare < 2 * nsp_ppc.sum():
print(
'WARNING :: Less than two cells worth of spare particles remaining.'
)
if num_spare == 0:
print(
'WARNING :: No space particles remaining. Exiting simulation.'
)
raise IndexError
acc = 0
n_created = 0
for ii in nb.prange(pos.shape[1]):
# If particle goes from cell 1 to cell 2
if (pos_old[0, ii] < xmin + dx):
if (pos[0, ii] >= xmin + dx) and (pos[0, ii] <= xmin + 2 * dx):
# Find first negative idx to initialize as new particle
for kk in nb.prange(acc, pos.shape[1]):
if idx[kk] < 0:
acc = kk + 1
break
# Create new particle with pos_old, vel of this particle
pos[0, kk] = pos[0, ii] - dx
pos[1, kk] = pos[1, ii]
pos[2, kk] = pos[2, ii]
vel[:, kk] = vel[:, ii]
idx[kk] = idx[ii]
n_created += 1
# If particle goes from cell NX to cell NX - 1
elif (pos_old[0, ii] > xmax - dx):
if (pos[0, ii] <= xmax - dx) and (pos[0, ii] > xmax - 2 * dx):
# Find first negative idx to initialize as new particle
for kk in nb.prange(acc, pos.shape[1]):
if idx[kk] < 0:
acc = kk + 1
break
# Create new particle with pos_old, vel of this particle
pos[0, kk] = pos[0, ii] + dx
pos[1, kk] = pos[1, ii]
pos[2, kk] = pos[2, ii]
vel[:, kk] = vel[:, ii]
idx[kk] = idx[ii]
n_created += 1
#print(n_created, 'created ; ', n_deleted, 'deleted')
assign_weighting_TSC(pos, idx, Ie, W_elec)
return