# 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)