# Equations for ell == 0
problem.add_equation(eq_eval("p = 0"), condition="ntheta == 0")
problem.add_equation(eq_eval("u = 0"), condition="ntheta == 0")
problem.add_equation(eq_eval("tau = 0"), condition="ntheta == 0")
logger.info("Problem built")
# Solver
solver = solvers.InitialValueSolver(problem, ts)
solver.stop_sim_time = t_end
# Analysis
t_list = []
E_list = []
weight_theta = b.local_colatitude_weights(1)
weight_r = b.local_radial_weights(1)
reducer = GlobalArrayReducer(d.comm_cart)
vol_test = np.sum(weight_r * weight_theta +
0 * p['g']) * np.pi / (Lmax + 1) / L_dealias
vol_test = reducer.reduce_scalar(vol_test, MPI.SUM)
vol_correction = 4 * np.pi / 3 / vol_test
# Main loop
start_time = time.time()
while solver.ok:
if solver.iteration % 10 == 0:
E0 = np.sum(vol_correction * weight_r * weight_theta * u['g'].real**2)
E0 = 0.5 * E0 * (np.pi) / (Lmax + 1) / L_dealias
E0 = reducer.reduce_scalar(E0, MPI.SUM)
logger.info("t = %f, E = %e" % (solver.sim_time, E0))
t_list.append(solver.sim_time)
E_list.append(E0)
P.append(M_ell.astype(np.complex128))
LU.append([None])
# Output
if rank == 0:
if not os.path.exists('{:s}/'.format(params.snapshots_dir)):
os.mkdir('{:s}/'.format(params.snapshots_dir))
snapshots = FileHandler(params.snapshots_dir, domain, B, max_writes=1)
snapshots.add_task(T, name='T', index=0)
snapshots.add_task(u, name='u0', index=0)
snapshots.add_task(u, name='u1', index=1)
snapshots.add_task(u, name='u2', index=2)
# CFL
reducer = GlobalArrayReducer(domain.dist.comm_cart)
dr = np.gradient(r[0,0])
dangle = 1 / (params.L_max + 1)
def calculate_dt(dt_old):
local_freq = u['g'][0]/dr + u['g'][1]/dangle + u['g'][2]/dangle
global_freq = reducer.global_max(local_freq)
if global_freq == 0.:
dt = np.inf
else:
dt = 1 / global_freq
dt *= params.safety
if dt > params.dt_max:
dt = params.dt_max
if dt < dt_old*(1+params.threshold) and dt > dt_old*(1-params.threshold):
dt = dt_old
N = N_max - B.N_min(ell - R_max)
# multiply by conversion matrices (may be very important)
# note that M matrices are no longer conversion matrices -- need to divide by Ekman or Prandtl
u_len = u_rhs['c'][ell_local].shape[0]
u_rhs['c'][ell_local] = M[ell_local][:u_len, :u_len].dot(
u_rhs['c'][ell_local]) / Ekman
p_len = p_rhs['c'][ell_local].shape[0]
T_rhs['c'][ell_local] = M[ell_local][
u_len + p_len:u_len + 2 * p_len, u_len + p_len:u_len +
2 * p_len].dot(T_rhs['c'][ell_local]) / Prandtl
NL.pack(u_rhs, p_rhs, T_rhs)
reducer = GlobalArrayReducer(domain.dist.comm_cart)
timing_iter = 10
if args['--run_time_iter'] is None:
iter_end = np.inf
else:
iter_end = int(args['--run_time_iter']) + timing_iter
out_cadence = 1000
report_cadence = 1
def initial_iterations(timing_iter):
t = 0.
iter = 0
T_rhs['c'][ell_local] = M[ell_local][u_len+p_len:u_len+p_len+T_len,u_len+p_len:u_len+p_len+T_len].dot(T_rhs['c'][ell_local])
A_rhs['c'][ell_local] = M[ell_local][u_len+p_len+T_len:u_len+p_len+T_len+A_len,u_len+p_len+T_len:u_len+p_len+T_len+A_len].dot(A_rhs['c'][ell_local])
RHS.pack(u_rhs,p_rhs,T_rhs,A_rhs,pi_rhs)
t_list = []
Ek_list = []
Em_list = []
# timestepping loop
start_time = time.time()
t = 0
iter = 0
reducer = GlobalArrayReducer(domain.dist.comm_cart)
while t < t_end:
nonlinear(state_vector,NL,t)
if iter % 5 == 0:
state_vector.unpack(u,p,T,A,pi)
# H = curl(A)
for ell in range(ell_start,ell_end+1):
ell_local = ell - ell_start
B.curl(ell,1,A['c'][ell_local],H['c'][ell_local])
Ek = np.sum(weight_r*weight_theta* 0.5*u['g']**2)*(np.pi)/((L_max+1)*L_dealias)
Em = np.sum(weight_r*weight_theta* 0.5/Rossby*H['g']**2)*np.pi/((L_max+1)*L_dealias)
problem.add_equation(eq_eval("tau_A = 0"), condition="ntheta == 0")
problem.add_equation(eq_eval("T(r=1) = 0"))
# Solver
solver = solvers.InitialValueSolver(problem, ts)
solver.stop_sim_time = t_end
logger.info("built IVP")
# Analysis
t_list = []
KE_list = []
ME_list = []
weight_theta = b.local_colatitude_weights(1)
weight_r = b.local_radial_weights(1)
reducer = GlobalArrayReducer(d.comm_cart)
vol_test = np.sum(weight_r * weight_theta +
0 * p['g']) * np.pi / (Lmax + 1) / L_dealias
vol_test = reducer.reduce_scalar(vol_test, MPI.SUM)
vol_correction = 4 * np.pi / 3 / vol_test
# Main loop
report_cadence = 10
good_solution = True
start_time = time.time()
while solver.ok and good_solution:
if solver.iteration % report_cadence == 0:
KE = np.sum(vol_correction * weight_r * weight_theta * u['g'].real**2)
KE = 0.5 * KE * (np.pi) / (Lmax + 1) / L_dealias
KE = reducer.reduce_scalar(KE, MPI.SUM)
B = (operators.Curl(A)).evaluate()
L = subproblem.left_perm.T @ subproblem.L_min
shape = L.shape
N0, N1 = BC_rows(Nmax, 2)
tau_columns = np.zeros((shape[0], 4))
tau_columns[:N0, 0] = (C(Nmax))[:, -1]
tau_columns[:N0, 1] = (C(Nmax))[:, -2]
tau_columns[N0:N1, 2] = (C(Nmax))[:, -1]
tau_columns[N0:N1, 3] = (C(Nmax))[:, -2]
L[:, -4:] = tau_columns
subproblem.L_min = subproblem.left_perm @ L
subproblem.L_min.eliminate_zeros()
subproblem.expand_matrices(['M', 'L'])
weight_theta = b.local_colatitude_weights(1)
weight_r = b.local_radial_weights(1)
reducer = GlobalArrayReducer(d.comm_cart)
vol_test = np.sum(weight_r * weight_theta +
0 * T['g']) * np.pi / (Lmax + 1) / L_dealias
vol_test = reducer.reduce_scalar(vol_test, MPI.SUM)
vol_correction = 4 * np.pi / 3 / vol_test
report_cadence = 10
dt = 1e-3
plot_cadence = report_cadence
plot = theta_target in theta
var = T['g']
name = 'T_diff'
if plot:
problem.add_equation(eq_eval("T_c(r=1) = 0"))
logger.info("Problem built")
# Solver
solver = solvers.InitialValueSolver(problem, ts)
solver.stop_sim_time = t_end + dt
# Analysis
t_list = []
E_list = []
T_list = []
T_err_list = []
weight_theta = b.local_colatitude_weights(1)
weight_r = b.local_radial_weights(1)
reducer = GlobalArrayReducer(d.comm_cart)
vol_test = np.sum(weight_r * weight_theta +
0 * T['g']) * np.pi / (Lmax + 1) / L_dealias
vol_test = reducer.reduce_scalar(vol_test, MPI.SUM)
vol_correction = 4 * np.pi / 3 / vol_test
report_cadence = 10
# Main loop
start_time = time.time()
while solver.ok:
if args['--verbose'] and solver.iteration % report_cadence == 0 and solver.iteration > 0:
E0 = np.sum(vol_correction * weight_r * weight_theta * u['g'].real**2)
E0 = 0.5 * E0 * (np.pi) / (Lmax + 1) / L_dealias
E0 = reducer.reduce_scalar(E0, MPI.SUM)
T0 = np.sum(vol_correction * weight_r * weight_theta * T['g'].real**2)