def test_E_convective_exelectron():
'''
Tests E-field update, convective (JxB) term only by zeroing/unity-ing other terms.
B-field is kept uniform in order to ensure curl(B) = 0
'''
xmin = 0.0 #const.xmin
xmax = 2 * np.pi #const.xmax
k = 1.0
NX = 32
dx = xmax / NX
# Physical location of nodes
E_nodes = (np.arange(NX + 3) - 0.5) * dx
B_nodes = (np.arange(NX + 3) - 1.0) * dx
qn_input = np.ones(
NX + 3) # Analytic input at node points (number density varying)
B_input = np.ones((NX + 3, 3))
J_input = np.zeros((NX + 3, 3))
J_input[:, 0] = np.sin(1.0 * k * E_nodes)
J_input[:, 1] = np.sin(2.0 * k * E_nodes)
J_input[:, 2] = np.sin(3.0 * k * E_nodes)
E_FD = fields.calculate_E(B_input, J_input, qn_input, DX=dx)
E_FD2 = fields.calculate_E_w_exel(B_input, J_input, qn_input, DX=dx)
E_anal = np.zeros((NX + 3, 3))
E_anal[:, 0] = np.sin(3.0 * k * E_nodes) - np.sin(2.0 * k * E_nodes)
E_anal[:, 1] = np.sin(1.0 * k * E_nodes) - np.sin(3.0 * k * E_nodes)
E_anal[:, 2] = np.sin(2.0 * k * E_nodes) - np.sin(1.0 * k * E_nodes)
plot = True
if plot == True:
marker_size = 50
plt.scatter(E_nodes[1:-2], E_anal[1:-2, 0], s=marker_size, c='k')
plt.scatter(E_nodes[1:-2], E_anal[1:-2, 1], s=marker_size, c='k')
plt.scatter(E_nodes[1:-2], E_anal[1:-2, 2], s=marker_size, c='k')
plt.scatter(E_nodes, E_FD[:, 0], s=marker_size, marker='x', c='b')
plt.scatter(E_nodes, E_FD[:, 1], s=marker_size, marker='x', c='b')
plt.scatter(E_nodes, E_FD[:, 2], s=marker_size, marker='x', c='b')
plt.scatter(E_nodes, E_FD2[:, 0], s=marker_size, marker='+', c='r')
plt.scatter(E_nodes, E_FD2[:, 1], s=marker_size, marker='+', c='r')
plt.scatter(E_nodes, E_FD2[:, 2], s=marker_size, marker='+', c='r')
for kk in range(NX + 3):
plt.axvline(E_nodes[kk], linestyle='--', c='r', alpha=0.2)
plt.axvline(B_nodes[kk], linestyle='--', c='b', alpha=0.2)
plt.axvline(xmin, linestyle='-', c='k', alpha=0.2)
plt.axvline(xmax, linestyle='-', c='k', alpha=0.2)
#plt.gcf().text(0.15, 0.93, '$R^2 = %.4f$' % r2)
plt.xlim(xmin - 1.5 * dx, xmax + 2 * dx)
plt.legend()
return
def numerical_loop(real_time, pos, vel, Ie, W_elec, idx, B, E_int, q_dens, Ji,
Ve, Te, DT, max_inc, data_iter, ch_flag):
'''
Does the number crunching for a short snippet. Logs number of time variable changes in ch_flag as powers of 2
(-1 = half, 2 = 4 times slower)
Array values are mutable: Don't have to be returned. Only integer values
'''
qq = 0
while qq < data_iter:
# Check timestep used for this iteration
vel, qq, DT, max_inc, data_iter, ch_flag = aux.check_timestep(
qq, DT, pos, vel, B, E_int, q_dens, Ie, W_elec, max_inc, data_iter,
idx)
# Add timestep to counter
real_time += DT
# Main loop
pos, vel, Ie, W_elec, q_dens_adv, Ji = particles.advance_particles_and_moments(
pos, vel, Ie, W_elec, idx, B, E_int, DT)
q_dens = 0.5 * (q_dens + q_dens_adv)
B = fields.push_B(B, E_int, DT)
E_half, Ve, Te = fields.calculate_E(B, Ji, q_dens)
q_dens = q_dens_adv.copy()
# Predictor-Corrector: Advance fields to start of next timestep
E_int, B = fields.predictor_corrector(B, E_int, E_half, pos, vel,
q_dens_adv, Ie, W_elec, idx, DT)
# Increment loop variable
qq += 1
return DT, ch_flag, max_inc, data_iter, real_time
def test_CAM_CL():
NX = 32
B = np.zeros((NX + 3, 3))
rho_i = np.ones((NX + 3))
J_i = np.zeros((NX + 3, 3))
DT = 0.1
subcycles = 8
E, V, T = fields.calculate_E(B, J_i, rho_i)
B, E_half, Ve_half, Te_half = fields.cyclic_leapfrog(B, rho_i, J_i, DT, subcycles, half_cycle=True)
return
def initialize():
pos, vel, Ie, W_elec, idx = init.initialize_particles()
B, E_int = init.initialize_fields()
DT, max_inc, data_iter = aux.set_timestep(vel)
if data_iter == 0:
data_iter = max_inc
q_dens, Ji = sources.collect_moments(vel, Ie, W_elec, idx)
E_int, Ve, Te = fields.calculate_E(B, Ji, q_dens)
vel = particles.velocity_update(pos, vel, Ie, W_elec, idx, B, E_int,
-0.5 * DT)
return pos, vel, Ie, W_elec, idx, B, E_int, q_dens, Ji, Ve, Te, DT, max_inc, data_iter
B, E_int, E_half, Ve, Te, Te0 = init.initialize_fields()
q_dens, q_dens_adv, Ji, ni, nu = init.initialize_source_arrays()
old_particles, old_fields, temp3De, temp3Db, temp1D,\
v_prime, S, T, mp_flux = init.initialize_tertiary_arrays()
print('Collecting initial moments...')
# Collect initial moments and save initial state
sources.collect_moments(vel, Ie, W_elec, idx, q_dens, Ji, ni, nu)
if te0_equil == 1:
init.set_equilibrium_te0(q_dens, Te0)
DT, max_inc, part_save_iter, field_save_iter, B_damping_array, E_damping_array\
= init.set_timestep(vel, Te0)
fields.calculate_E(B, Ji, q_dens, E_int, Ve, Te, Te0, temp3De, temp3Db,
temp1D, E_damping_array, 0, DT, 0)
print('Saving initial conditions...')
if save_particles == 1:
save.save_particle_data(0, DT, part_save_iter, 0, pos, vel, idx)
if save_fields == 1:
save.save_field_data(0, DT, field_save_iter, 0, Ji, E_int,\
B, Ve, Te, q_dens, B_damping_array, E_damping_array)
# Retard velocity
print('Retarding velocity...')
particles.velocity_update(pos, vel, Ie, W_elec, Ib, W_mag, idx, Ep, Bp, B,
E_int, v_prime, S, T, temp_N, -0.5 * DT)
qq = 1
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
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%25 == 0 and particle_open == 1:
# num_spare = (idx < 0).sum()
# if num_spare < inject_rate.sum() * DT * 5.0:
# # Change this to dynamically expand particle arrays later on (adding more particles)
# print('WARNING :: Less than 5 time increments of spare particles remaining. Widening arrays...')
# pos, vel, idx, Ie, W_elec, Ib, W_mag, Ep, Bp, v_prime, S, T, temp_N, old_particles = \
# increase_particle_array_size(pos, vel, idx, Ie, W_elec, Ib, W_mag, Ep, Bp, v_prime,
# S, T, temp_N, old_particles)
# =============================================================================
# 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
# Push B from N to N + 1/2
fields.push_B(B, E_int, temp3Db, DT, qq, B_damping_array, half_flag=1)
# Calculate E at N + 1/2
fields.calculate_E(B, Ji, q_dens, E_half, Ve, Te, Te0, temp3De, temp3Db,
temp1D, E_damping_array)
###################################
### 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)
# 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 test_E_hall():
'''
Tests E-field update, hall term (B x curl(B)) only by selection of inputs
'''
# =============================================================================
# grids = [32, 64, 128, 256, 512, 1024, 2048]
# err_curl = np.zeros(len(grids))
# err_interp = np.zeros(len(grids))
# err_hall = np.zeros(len(grids))
# =============================================================================
#for NX, ii in zip(grids, range(len(grids))):
NX = 32 #const.NX
xmin = 0.0 #const.xmin
xmax = 2*np.pi #const.xmax
dx = xmax / NX
k = 1.0
marker_size = 70
E_nodes = (np.arange(NX + 3) - 0.5) * dx # Physical location of nodes
B_nodes = (np.arange(NX + 3) - 1.0) * dx
## INPUTS ##
Bx = np.sin(1.0*k*B_nodes)
By = np.sin(2.0*k*B_nodes)
Bz = np.sin(3.0*k*B_nodes)
Bxe = np.sin(1.0*k*E_nodes)
Bye = np.sin(2.0*k*E_nodes)
Bze = np.sin(3.0*k*E_nodes)
dBy = 2.0 * k * np.cos(2.0*k*E_nodes)
dBz = 3.0 * k * np.cos(3.0*k*E_nodes)
B_input = np.zeros((NX + 3, 3))
B_input[:, 0] = Bx
B_input[:, 1] = By
B_input[:, 2] = Bz
Be_input = np.zeros((NX + 3, 3))
Be_input[:, 0] = Bxe
Be_input[:, 1] = Bye
Be_input[:, 2] = Bze
B_center = aux.interpolate_to_center_cspline3D(B_input, DX=dx)
## TEST CURL B (AGAIN JUST TO BE SURE)
curl_B_FD = fields.get_curl_B(B_input, DX=dx)
curl_B_anal = np.zeros((NX + 3, 3))
curl_B_anal[:, 1] = -dBz
curl_B_anal[:, 2] = dBy
## ELECTRIC FIELD CALCULATION ##
E_FD = fields.calculate_E( B_input, np.zeros((NX + 3, 3)), np.ones(NX + 3), DX=dx)
E_FD2 = fields.calculate_E_w_exel(B_input, np.zeros((NX + 3, 3)), np.ones(NX + 3), DX=dx)
E_anal = np.zeros((NX + 3, 3))
E_anal[:, 0] = - (Bye * dBy + Bze * dBz)
E_anal[:, 1] = Bxe * dBy
E_anal[:, 2] = Bxe * dBz
E_anal /= const.mu0
# =============================================================================
# ## Calculate errors ##
# err_curl[ii] = np.abs(curl_B_FD[1: -2, :] - curl_B_anal[1: -2, :]).max()
# err_interp[ii] = np.abs(B_center[1: -2, :] - Be_input[1: -2, :] ).max()
# err_hall[ii] = np.abs(E_FD[1: -2, :] - E_anal[1: -2, :] ).max()
#
# for ii in range(len(grids) - 1):
# order_curl = np.log(err_curl[ii] / err_curl[ii + 1] ) / np.log(2)
# order_interp = np.log(err_interp[ii] / err_interp[ii + 1]) / np.log(2)
# order_hall = np.log(err_hall[ii] / err_hall[ii + 1] ) / np.log(2)
#
# print ''
# print 'Grid reduction: {} -> {}'.format(grids[ii], grids[ii + 1])
# print 'Curl order: {}, \nC-spline Interpolation order: {}'.format(order_curl, order_interp)
# print 'E-field Hall order: {}'.format(order_hall)
# =============================================================================
plot = True
if plot == True:
# =============================================================================
# # Plot curl test
# plt.figure()
# plt.scatter(E_nodes, curl_B_anal[:, 1], s=marker_size, c='k')
# plt.scatter(E_nodes, curl_B_anal[:, 2], s=marker_size, c='k')
# plt.scatter(E_nodes, curl_B_FD[:, 1], s=marker_size, marker='x')
# plt.scatter(E_nodes, curl_B_FD[:, 2], s=marker_size, marker='x')
#
# # Plot center-interpolated B test
# plt.figure()
# plt.scatter(B_nodes, B_input[:, 0], s=marker_size, c='b')
# plt.scatter(B_nodes, B_input[:, 1], s=marker_size, c='b')
# plt.scatter(B_nodes, B_input[:, 2], s=marker_size, c='b')
# plt.scatter(E_nodes, B_center[:, 0], s=marker_size, c='r')
# plt.scatter(E_nodes, B_center[:, 1], s=marker_size, c='r')
# plt.scatter(E_nodes, B_center[:, 2], s=marker_size, c='r')
# =============================================================================
# Plot E-field test solutions
plt.scatter(E_nodes, E_anal[:, 0], s=marker_size, marker='o', c='k')
plt.scatter(E_nodes, E_anal[:, 1], s=marker_size, marker='o', c='k')
plt.scatter(E_nodes, E_anal[:, 2], s=marker_size, marker='o', c='k')
plt.scatter(E_nodes, E_FD[:, 0], s=marker_size, c='b', marker='+')
plt.scatter(E_nodes, E_FD[:, 1], s=marker_size, c='b', marker='+')
plt.scatter(E_nodes, E_FD[:, 2], s=marker_size, c='b', marker='+')
plt.scatter(E_nodes, E_FD2[:, 0], s=marker_size, c='r', marker='x')
plt.scatter(E_nodes, E_FD2[:, 1], s=marker_size, c='r', marker='x')
plt.scatter(E_nodes, E_FD2[:, 2], s=marker_size, c='r', marker='x')
for kk in range(NX + 3):
plt.axvline(E_nodes[kk], linestyle='--', c='r', alpha=0.2)
plt.axvline(B_nodes[kk], linestyle='--', c='b', alpha=0.2)
plt.axvline(xmin, linestyle='-', c='k', alpha=0.2)
plt.axvline(xmax, linestyle='-', c='k', alpha=0.2)
plt.xlim(xmin - 1.5*dx, xmax + 2*dx)
return
qq = 0
while qq < max_inc:
part, qq, DT, max_inc, data_dump_iter, plot_dump_iter, change_flag = aux.check_timestep(
qq, DT, part, B, E, dns_int, max_inc, data_dump_iter,
plot_dump_iter)
if change_flag == 1:
print(
'Timestep halved. Syncing particle velocity/position with DT = {}'
.format(DT))
part, dns_int, dns_half, J_plus, J_minus, G, L = sources.init_collect_moments(
part, 0.5 * DT)
B = fields.cyclic_leapfrog(B, dns_int, J_minus, DT)
E = fields.calculate_E(B, J_minus, dns_half)
J = sources.push_current(J_plus, E, B, L, G, DT)
E = fields.calculate_E(B, J, dns_half)
part = particles.velocity_update(part, B, E, DT)
part, dns_int, J_plus, J_minus, G, L = sources.collect_moments(
part, DT)
dns_int = 0.5 * (dns_int + dns_half)
J = 0.5 * (J_plus + J_minus)
B = fields.cyclic_leapfrog(B, dns_int, J, DT)
if qq % data_dump_iter == 0 and generate_data == 1: # Save data, if flagged
pas.save_data(DT, data_dump_iter, qq, part, J, E, B, dns_int)
def main_loop(pos, vel, idx, Ie, W_elec, Ib, W_mag, \
B, E_int, E_half, q_dens, q_dens_adv, Ji, ni, nu, \
Ve, Te, temp3D, temp3D2, temp1D, old_particles, old_fields,\
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
qq, DT, max_inc, part_save_iter, field_save_iter \
= check_timestep(pos, vel, B, E_int, q_dens, Ie, W_elec, Ib, W_mag, temp3D, \
qq, DT, max_inc, part_save_iter, field_save_iter, idx)
# Move particles, collect moments
particles.advance_particles_and_moments(pos, vel, Ie, W_elec, Ib, W_mag, idx, \
B, E_int, DT, q_dens_adv, Ji, ni, nu, temp1D)
# Average N, N + 1 densities (q_dens at N + 1/2)
q_dens *= 0.5
q_dens += 0.5 * q_dens_adv
# Push B from N to N + 1/2
fields.push_B(B, E_int, temp3D, DT, qq, half_flag=1)
# Calculate E at N + 1/2
fields.calculate_E(B, Ji, q_dens, E_half, Ve, Te, temp3D, temp3D2, temp1D)
###################################
### PREDICTOR CORRECTOR SECTION ###
###################################
# Store old values
old_particles[0, :] = pos
old_particles[1:4, :] = vel
old_particles[4, :] = Ie
old_particles[5:8, :] = W_elec
old_fields[:, 0:3] = B
old_fields[:, 3:6] = Ji
old_fields[:, 6:9] = Ve
old_fields[:, 9] = Te
# Predict fields
E_int *= -1.0
E_int += 2.0 * E_half
fields.push_B(B, E_int, temp3D, DT, qq, 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, \
B, E_int, DT, q_dens, Ji, ni, nu, temp1D, 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, temp3D, DT, qq + 1, half_flag=1)
fields.calculate_E(B, Ji, q_dens, E_int, Ve, Te, temp3D, temp3D2, temp1D)
# 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, :]
B[:] = old_fields[:, 0:3]
Ji[:] = old_fields[:, 3:6]
Ve[:] = old_fields[:, 6:9]
Te[:] = old_fields[:, 9]
fields.push_B(B, E_int, temp3D, DT, qq,
half_flag=0) # Advance the original B
q_dens[:] = q_dens_adv
return qq, DT, max_inc, part_save_iter, field_save_iter
if __name__ == '__main__':
start_time = timer()
# Initialize simulation: Allocate memory and set time parameters
pos, vel, Ie, W_elec, Ib, W_mag, idx = init.initialize_particles()
B, E_int, E_half, Ve, Te, = init.initialize_fields()
q_dens, q_dens_adv, Ji, ni, nu = init.initialize_source_arrays()
old_particles, old_fields, temp3D, temp3D2, temp1D = init.initialize_tertiary_arrays(
)
DT, max_inc, part_save_iter, field_save_iter = init.set_timestep(vel)
# Collect initial moments and save initial state
sources.collect_moments(vel, Ie, W_elec, idx, q_dens, Ji, ni, nu, temp1D)
fields.calculate_E(B, Ji, q_dens, E_int, Ve, Te, temp3D, temp3D2, temp1D)
if save_particles == 1:
save.save_particle_data(DT, part_save_iter, 0, pos, vel)
if save_fields == 1:
save.save_field_data(DT, field_save_iter, 0, Ji, E_int, B, Ve, Te,
q_dens)
particles.assign_weighting_TSC(pos, Ib, W_mag, E_nodes=False)
particles.velocity_update(vel, Ie, W_elec, Ib, W_mag, idx, B, E_int,
-0.5 * DT)
qq = 1
print('Starting main loop...')
while qq < max_inc:
from simulation_parameters_1D import generate_data, NX
if __name__ == '__main__':
start_time = timer()
pos, vel, Ie, W_elec, idx = init.initialize_particles()
B, E_int = init.initialize_fields()
DT, max_inc, data_iter = aux.set_timestep(vel)
print 'Timestep: %.4fs, %d iterations total' % (DT, max_inc)
if generate_data == 1:
save.store_run_parameters(DT, data_iter)
q_dens, Ji = sources.collect_moments(vel, Ie, W_elec, idx)
E_int, Ve, Te = fields.calculate_E(B, Ji, q_dens)
vel = particles.velocity_update(pos, vel, Ie, W_elec, idx, B, E_int,
-0.5 * DT)
qq = 0
while qq < max_inc:
# Check timestep
vel, qq, DT, max_inc, data_iter, ch_flag \
= aux.check_timestep(qq, DT, pos, vel, B, E_int, q_dens, Ie, W_elec, max_inc, data_iter, idx)
if ch_flag == 1:
print 'Timestep halved. Syncing particle velocity with DT = {}'.format(
DT)
elif ch_flag == 2:
print 'Timestep Doubled. Syncing particle velocity with DT = {}'.format(
DT)
rho_int, rho_half, J, J_plus, L,
G, 0.5 * DT)
# Debug: Test if everything is the same if J is replaced with J_minus at each loop.
# Yes it is after loop 0 and loop 1 up until collect_moments()
# Disable this at some point and see if it improves (or even changes) anything.
# =============================================================================
# if qq > 0:
# J[:, :] = J_minus[:, :]
# =============================================================================
#######################
###### MAIN LOOP ######
#######################
fields.cyclic_leapfrog(B, B2, rho_int, J, temp3d, DT, subcycles)
E, Ve, Te = fields.calculate_E(B, J, rho_half)
sources.push_current(J_plus, J, E, B, L, G, DT)
E, Ve, Te = fields.calculate_E(B, J, rho_half)
particles.velocity_update(pos, vel, Ie, W_elec, Ib, W_mag, idx, Ep, Bp,
B, E, v_prime, S, T, temp_N, DT)
# Store pc(1/2) here while pc(3/2) is collected
rho_int[:] = rho_half[:]
sources.collect_moments(pos, vel, Ie, W_elec, idx, ni, nu_plus,
nu_minus, rho_half, J_minus, J_plus, L, G, DT)
rho_int += rho_half
rho_int /= 2.0
J = 0.5 * (J_plus + J_minus)
if adaptive_timestep == 1:
pos, qq, DT, max_inc, part_save_iter, field_save_iter, change_flag, subcycles = aux.check_timestep(qq, DT, pos, vel, B, E, dns_int, max_inc, part_save_iter, field_save_iter, subcycles)
if change_flag == 1:
print('Timestep halved. Syncing particle velocity/position with DT = {}'.format(DT))
pos, Ie, W_elec, dns_int, dns_half, J_plus, J_minus, G, L = sources.init_collect_moments(pos, vel, Ie, W_elec, idx, 0.5*DT)
elif change_flag == 2:
print('Timestep doubled. Syncing particle velocity/position with DT = {}'.format(DT))
pos, Ie, W_elec, dns_int, dns_half, J_plus, J_minus, G, L = sources.init_collect_moments(pos, vel, Ie, W_elec, idx, 0.5*DT)
#######################
###### MAIN LOOP ######
#######################
B = fields.cyclic_leapfrog(B, dns_int, J_minus, DT, subcycles)
E, Ve, Te = fields.calculate_E(B, J_minus, dns_half)
J = sources.push_current(J_plus, E, B, L, G, DT)
E, Ve, Te = fields.calculate_E(B, J, dns_half)
vel = particles.velocity_update(pos, vel, Ie, W_elec, idx, B, E, J, DT)
# Store pc(1/2) here while pc(3/2) is collected
dns_int = dns_half
pos, Ie, W_elec, dns_half, J_plus, J_minus, G, L = sources.collect_moments(pos, vel, Ie, W_elec, idx, DT)
dns_int = 0.5 * (dns_int + dns_half)
J = 0.5 * (J_plus + J_minus)
B = fields.cyclic_leapfrog(B, dns_int, J, DT, subcycles)
E, Ve, Te = fields.calculate_E(B, J, dns_int) # This one's just for output
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, \
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
'''
NC = const.NC
# Check timestep
#print('MAIN Checking timestep')
qq, DT, max_inc, part_save_iter, field_save_iter, B_damping_array, E_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, E_damping_array)
# Move particles, collect moments
#print('MAIN Advancing particles/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)
# Average N, N + 1 densities (q_dens at N + 1/2)
#print('MAIN Averaging density')
q_dens *= 0.5
q_dens += 0.5 * q_dens_adv
if const.disable_waves == 0:
# Push B from N to N + 1/2
fields.push_B(B, E_int, temp3Db, DT, qq, B_damping_array, half_flag=1)
# Calculate E at N + 1/2
fields.calculate_E(B,
Ji,
q_dens,
E_half,
Ve,
Te,
Te0,
temp3De,
temp3Db,
temp1D,
E_damping_array,
qq,
DT,
half_flag=1)
###################################
### PREDICTOR CORRECTOR SECTION ###
###################################
# Store old values
old_particles[0:3, :] = pos
old_particles[3:6, :] = vel
old_particles[6, :] = Ie
old_particles[7:10, :] = W_elec
old_particles[10, :] = 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, pc=1)
# Average N + 1, N + 2 densities (q_dens at N + 3/2)
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 + 1,
DT,
half_flag=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:3, :]
vel[:] = old_particles[3:6, :]
Ie[:] = old_particles[6, :]
W_elec[:] = old_particles[7:10, :]
idx[:] = old_particles[10, :]
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
return qq, DT, max_inc, part_save_iter, field_save_iter