def init_problem(self):
# Initial conditions
self.u = self.solver.state['u']['g']
self.v = self.solver.state['v']['g']
self.w = self.solver.state['w']['g']
z = self.domain.grid(2)
b = self.solver.state['b']
bz = self.solver.state['bz']
# Random perturbations, initialized globally for same results in parallel
gshape = self.domain.dist.grid_layout.global_shape(scales=1)
slices = self.domain.dist.grid_layout.slices(scales=1)
rand = np.random.RandomState(seed=23)
noise = rand.standard_normal(gshape)[slices]
# Linear background + perturbations damped at walls
zb, zt = self.z_basis.interval
pert = 1e-3 * noise * (zt - z) * (z - zb)
b['g'] = -self.F * (z - pert)
b.differentiate('z', out=bz)
# Flow properties
self.CFL = flow_tools.CFL(self.solver,
initial_dt=1e-4,
cadence=5,
safety=0.5,
max_change=1.5,
min_change=0.5,
max_dt=0.05)
self.CFL.add_velocities(('u', 'v', 'w'))
self.flow = flow_tools.GlobalFlowProperty(self.solver, cadence=10)
self.flow.add_property("sqrt(u*u + v*v + w*w) / R", name='Re')
def init_dedalus_cfl(self, sim):
initial_dt = sim.params.time_stepping.deltat0
if sim.params.ONLY_COARSE_OPER:
self.CFL = None
return
self.CFL = flow_tools.CFL(
sim.dedalus_solver,
initial_dt=initial_dt,
cadence=10,
safety=1,
max_change=1.5,
min_change=0.5,
max_dt=0.125,
threshold=0.05,
)
self.CFL.add_velocities(("vx", "vz"))
def run(self, initial_dt=1e-4, sim_time=100):
# Integration parameters
self.solver.stop_sim_time = self.solver.sim_time + sim_time
self.solver.stop_wall_time = np.inf # 60 * 60.
self.solver.stop_iteration = np.inf
# CFL
CFL = flow_tools.CFL(self.solver,
initial_dt=initial_dt,
cadence=5,
safety=1.5,
max_change=1.5,
min_change=0.5,
max_dt=0.05)
CFL.add_velocities(('u', 'v', 'w'))
flow = flow_tools.GlobalFlowProperty(self.solver, cadence=10)
flow.add_property("0.5*(u*u + v*v + w*w)", name='KE')
self.CFL = CFL
self.flow = flow
try:
logger.info('Starting loop')
start_run_time = time.time()
while self.solver.ok:
dt = CFL.compute_dt()
self.solver.step(dt)
if (self.solver.iteration - 1) % 100 == 0:
logger.info(
'Iteration: %i, Time: %e, dt: %e' %
(self.solver.iteration, self.solver.sim_time, dt))
logger.info('Average KE = %e' % flow.volume_average('KE'))
except:
logger.error('Exception raised, triggering end of main loop.')
raise
finally:
end_run_time = time.time()
logger.info('Iterations: %i' % self.solver.iteration)
logger.info('Sim end time: %f' % self.solver.sim_time)
logger.info('Run time: %.2f sec' % (end_run_time - start_run_time))
logger.info('Run time: %f cpu-hr' %
((end_run_time - start_run_time) / hour *
self.domain.dist.comm_cart.size))
def test_flow_tools_cfl(x_basis_class, Nx, Nz, timestepper, dtype, safety,
z_velocity_mag, dealias):
# Bases
Lx = 2
Lz = 1
c = coords.CartesianCoordinates('x', 'z')
d = distributor.Distributor((c, ))
xb = x_basis_class(c.coords[0], size=Nx, bounds=(0, Lx), dealias=dealias)
x = xb.local_grid(1)
zb = basis.ChebyshevT(c.coords[1],
size=Nz,
bounds=(0, Lz),
dealias=dealias)
z = zb.local_grid(1)
b = (xb, zb)
# Fields
u = field.Field(name='u', dist=d, tensorsig=(c, ), bases=b, dtype=dtype)
# Problem
ddt = operators.TimeDerivative
problem = problems.IVP([u])
problem.add_equation((ddt(u), 0))
# Solver
solver = solvers.InitialValueSolver(problem, timestepper)
# cfl initialization
dt = 1
cfl = flow_tools.CFL(solver, dt, safety=safety, cadence=1)
cfl.add_velocity(u)
# Test Fourier CFL
fourier_velocity = lambda x: np.sin(4 * np.pi * x / Lx)
chebyshev_velocity = lambda z: -z_velocity_mag * z
u['g'][0] = fourier_velocity(x)
u['g'][1] = chebyshev_velocity(z)
solver.step(dt)
solver.step(
dt
) #need two timesteps to get past stored_dt per compute_timestep logic
dt = cfl.compute_timestep()
op = operators.AdvectiveCFL(u, c)
cfl_freq = np.abs(u['g'][0] / op.cfl_spacing(u)[0])
cfl_freq += np.abs(u['g'][1] / op.cfl_spacing(u)[1])
cfl_freq = np.max(cfl_freq)
dt_comparison = safety * (cfl_freq)**(-1)
assert np.allclose(dt, dt_comparison)
u['g'] = 0 * xyone #eta['g'] * (-9 + 6*yy**2)/4. * np.exp(-yy**2/2)
v['g'] = 0 * xyone #2 * deta['g'] * yy * np.exp(-yy**2/2)
s['g'] = xyone #2 * deta['g'] * yy * np.exp(-yy**2/2)
#v['g'] = np.zeros(y.size)#amp*np.sin(2.0*np.pi*x/Lx)*np.exp(-(y*y)/(sigma*sigma))
#s['g'] = np.ones(y.size)
u.differentiate('y', out=uy)
v.differentiate('y', out=vy)
s.differentiate('y', out=sy)
# Time Stepping
solver.stop_sim_time = 50.01
solver.stop_wall_time = np.inf
solver.stop_iteration = np.inf #h500 #np.inf
#dt = 0.025
initial_dt = 0.2 * Lx / nx
cfl = flow_tools.CFL(solver, initial_dt, safety=0.5)
cfl.add_velocities(('u', 'v'))
#logger.info('Solver built')
dt = cfl.compute_dt()
data_dir = 'eq_beta_solodoch'
if domain.distributor.rank == 0:
if not os.path.exists('{:s}/'.format(data_dir)):
os.mkdir('{:s}/'.format(data_dir))
# Analysis
snapshots = solver.evaluator.add_file_handler(os.path.join(
data_dir, 'snapshots'),
sim_dt=1.,
max_writes=np.inf)
snapshots.add_system(solver.state)
def FC_polytrope(Rayleigh=1e4,
Prandtl=1,
aspect_ratio=4,
Taylor=None,
theta=0,
nz=128,
nx=None,
ny=None,
threeD=False,
mesh=None,
n_rho_cz=3,
epsilon=1e-4,
gamma=5 / 3,
run_time=23.5,
run_time_buoyancies=None,
run_time_iter=np.inf,
fixed_T=False,
fixed_flux=False,
mixed_flux_T=False,
const_mu=True,
const_kappa=True,
dynamic_diffusivities=False,
split_diffusivities=False,
restart=None,
start_new_files=False,
rk222=False,
safety_factor=0.2,
max_writes=20,
no_slip=False,
data_dir='./',
out_cadence=0.1,
no_coeffs=False,
no_volumes=False,
no_join=False,
verbose=False):
import dedalus.public as de
from dedalus.tools import post
from dedalus.extras import flow_tools
import time
import os
import sys
from stratified_dynamics import polytropes
from tools.checkpointing import Checkpoint
checkpoint_min = 30
initial_time = time.time()
logger.info("Starting Dedalus script {:s}".format(sys.argv[0]))
if nx is None:
nx = int(np.round(nz * aspect_ratio))
if threeD and ny is None:
ny = nx
if threeD:
atmosphere = polytropes.FC_polytrope_3d(nx=nx, ny=ny, nz=nz, mesh=mesh, constant_kappa=const_kappa, constant_mu=const_mu,\
epsilon=epsilon, gamma=gamma, n_rho_cz=n_rho_cz, aspect_ratio=aspect_ratio,\
fig_dir=data_dir)
else:
if dynamic_diffusivities:
atmosphere = polytropes.FC_polytrope_2d_kappa_mu(nx=nx, nz=nz, constant_kappa=const_kappa, constant_mu=const_mu,\
epsilon=epsilon, gamma=gamma, n_rho_cz=n_rho_cz, aspect_ratio=aspect_ratio,\
fig_dir=data_dir)
else:
atmosphere = polytropes.FC_polytrope_2d(nx=nx, nz=nz, constant_kappa=const_kappa, constant_mu=const_mu,\
epsilon=epsilon, gamma=gamma, n_rho_cz=n_rho_cz, aspect_ratio=aspect_ratio,\
fig_dir=data_dir)
if epsilon < 1e-4:
ncc_cutoff = 1e-14
elif epsilon > 1e-1:
ncc_cutoff = 1e-6
else:
ncc_cutoff = 1e-10
if threeD:
atmosphere.set_IVP_problem(Rayleigh,
Prandtl,
Taylor=Taylor,
theta=theta,
ncc_cutoff=ncc_cutoff,
split_diffusivities=split_diffusivities)
else:
atmosphere.set_IVP_problem(Rayleigh,
Prandtl,
ncc_cutoff=ncc_cutoff,
split_diffusivities=split_diffusivities)
bc_dict = {
'stress_free': False,
'no_slip': False,
'fixed_flux': False,
'mixed_flux_temperature': False,
'fixed_temperature': False
}
if no_slip:
bc_dict['no_slip'] = True
else:
bc_dict['stress_free'] = True
if fixed_flux:
bc_dict['fixed_flux'] = True
elif mixed_flux_T:
bc_dict['mixed_flux_temperature'] = True
else:
bc_dict['fixed_temperature'] = True
atmosphere.set_BC(**bc_dict)
problem = atmosphere.get_problem()
if atmosphere.domain.distributor.rank == 0:
if not os.path.exists('{:s}/'.format(data_dir)):
os.mkdir('{:s}/'.format(data_dir))
if rk222:
logger.info("timestepping using RK222")
ts = de.timesteppers.RK222
cfl_safety_factor = safety_factor * 2
else:
logger.info("timestepping using RK443")
ts = de.timesteppers.RK443
cfl_safety_factor = safety_factor * 4
# Build solver
solver = problem.build_solver(ts)
#Check atmosphere
logger.info("thermal_time = {:g}, top_thermal_time = {:g}".format(atmosphere.thermal_time,\
atmosphere.top_thermal_time))
logger.info("full atm HS check")
atmosphere.check_atmosphere(make_plots=False,
rho=atmosphere.get_full_rho(solver),
T=atmosphere.get_full_T(solver))
if restart is None or start_new_files:
mode = "overwrite"
else:
mode = "append"
logger.info('checkpointing in {}'.format(data_dir))
checkpoint = Checkpoint(data_dir)
if restart is None:
atmosphere.set_IC(solver)
dt = None
else:
logger.info("restarting from {}".format(restart))
dt = checkpoint.restart(restart, solver)
checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min * 60, mode=mode)
if run_time_buoyancies != None:
solver.stop_sim_time = solver.sim_time + run_time_buoyancies * atmosphere.buoyancy_time
else:
solver.stop_sim_time = 100 * atmosphere.thermal_time
solver.stop_iteration = solver.iteration + run_time_iter
solver.stop_wall_time = run_time * 3600
report_cadence = 1
output_time_cadence = out_cadence * atmosphere.buoyancy_time
Hermitian_cadence = 100
logger.info("stopping after {:g} time units".format(solver.stop_sim_time))
logger.info("output cadence = {:g}".format(output_time_cadence))
if threeD:
analysis_tasks = atmosphere.initialize_output(
solver,
data_dir,
sim_dt=output_time_cadence,
coeffs_output=not (no_coeffs),
mode=mode,
max_writes=max_writes,
volumes_output=not (no_volumes))
else:
analysis_tasks = atmosphere.initialize_output(
solver,
data_dir,
sim_dt=output_time_cadence,
coeffs_output=not (no_coeffs),
mode=mode,
max_writes=max_writes)
#Set up timestep defaults
max_dt = output_time_cadence
if dt is None: dt = max_dt
cfl_cadence = 1
cfl_threshold = 0.1
CFL = flow_tools.CFL(solver,
initial_dt=dt,
cadence=cfl_cadence,
safety=cfl_safety_factor,
max_change=1.5,
min_change=0.5,
max_dt=max_dt,
threshold=cfl_threshold)
if threeD:
CFL.add_velocities(('u', 'v', 'w'))
else:
CFL.add_velocities(('u', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("Re_rms", name='Re')
if verbose:
flow.add_property("Pe_rms", name='Pe')
flow.add_property("Nusselt_AB17", name='Nusselt')
start_iter = solver.iteration
start_sim_time = solver.sim_time
try:
start_time = time.time()
start_iter = solver.iteration
logger.info('starting main loop')
good_solution = True
first_step = True
while solver.ok and good_solution:
dt = CFL.compute_dt()
# advance
solver.step(dt)
effective_iter = solver.iteration - start_iter
Re_avg = flow.grid_average('Re')
if threeD and effective_iter % Hermitian_cadence == 0:
for field in solver.state.fields:
field.require_grid_space()
# update lists
if effective_iter % report_cadence == 0:
log_string = 'Iteration: {:5d}, Time: {:8.3e} ({:8.3e}), dt: {:8.3e}, '.format(
solver.iteration - start_iter, solver.sim_time,
(solver.sim_time - start_sim_time) /
atmosphere.buoyancy_time, dt)
if verbose:
log_string += '\n\t\tRe: {:8.5e}/{:8.5e}'.format(
Re_avg, flow.max('Re'))
log_string += '; Pe: {:8.5e}/{:8.5e}'.format(
flow.grid_average('Pe'), flow.max('Pe'))
log_string += '; Nu: {:8.5e}/{:8.5e}'.format(
flow.grid_average('Nusselt'), flow.max('Nusselt'))
else:
log_string += 'Re: {:8.3e}/{:8.3e}'.format(
Re_avg, flow.max('Re'))
logger.info(log_string)
if not np.isfinite(Re_avg):
good_solution = False
logger.info(
"Terminating run. Trapped on Reynolds = {}".format(
Re_avg))
if first_step:
if verbose:
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.spy(solver.pencils[0].L,
markersize=1,
markeredgewidth=0.0)
fig.savefig(data_dir + "sparsity_pattern.png", dpi=1200)
import scipy.sparse.linalg as sla
LU = sla.splu(solver.pencils[0].LHS.tocsc(),
permc_spec='NATURAL')
fig = plt.figure()
ax = fig.add_subplot(1, 2, 1)
ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
ax = fig.add_subplot(1, 2, 2)
ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
fig.savefig(data_dir + "sparsity_pattern_LU.png", dpi=1200)
logger.info("{} nonzero entries in LU".format(LU.nnz))
logger.info("{} nonzero entries in LHS".format(
solver.pencils[0].LHS.tocsc().nnz))
logger.info("{} fill in factor".format(
LU.nnz / solver.pencils[0].LHS.tocsc().nnz))
first_step = False
start_time = time.time()
except:
logger.error('Exception raised, triggering end of main loop.')
finally:
end_time = time.time()
# Print statistics
elapsed_time = end_time - start_time
elapsed_sim_time = solver.sim_time
N_iterations = solver.iteration - 1
logger.info('main loop time: {:e}'.format(elapsed_time))
logger.info('Iterations: {:d}'.format(N_iterations))
logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
if N_iterations > 0:
logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
N_iterations))
if not no_join:
logger.info('beginning join operation')
try:
final_checkpoint = Checkpoint(
data_dir, checkpoint_name='final_checkpoint')
final_checkpoint.set_checkpoint(solver,
wall_dt=1,
mode="append")
solver.step(dt) #clean this up in the future...works for now.
post.merge_process_files(data_dir + '/final_checkpoint/',
cleanup=False)
except:
print('cannot save final checkpoint')
logger.info(data_dir + '/checkpoint/')
post.merge_process_files(data_dir + '/checkpoint/', cleanup=False)
for task in analysis_tasks.keys():
logger.info(analysis_tasks[task].base_path)
post.merge_process_files(analysis_tasks[task].base_path,
cleanup=False)
if (atmosphere.domain.distributor.rank == 0):
logger.info('main loop time: {:e}'.format(elapsed_time))
if start_iter > 1:
logger.info('Iterations (this run): {:d}'.format(N_iterations -
start_iter))
logger.info('Iterations (total): {:d}'.format(N_iterations -
start_iter))
logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
if N_iterations > 0:
logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
N_iterations))
N_TOTAL_CPU = atmosphere.domain.distributor.comm_cart.size
# Print statistics
print('-' * 40)
total_time = end_time - initial_time
main_loop_time = end_time - start_time
startup_time = start_time - initial_time
n_steps = solver.iteration - 1
print(' startup time:', startup_time)
print('main loop time:', main_loop_time)
print(' total time:', total_time)
if n_steps > 0:
print(' iterations:', n_steps)
print(' loop sec/iter:', main_loop_time / n_steps)
print(' average dt:', solver.sim_time / n_steps)
print(
" N_cores, Nx, Nz, startup main loop, main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid"
)
print(
'scaling:', ' {:d} {:d} {:d}'.format(N_TOTAL_CPU, nx, nz),
' {:8.3g} {:8.3g} {:8.3g} {:8.3g} {:8.3g}'.format(
startup_time, main_loop_time, main_loop_time / n_steps,
main_loop_time / n_steps / (nx * nz),
N_TOTAL_CPU * main_loop_time / n_steps / (nx * nz)))
print('-' * 40)
def FC_TriLayer_convection(input_dict):
"""
A driver function for running FC convection in Dedalus.
Parameters
----------
input_dict : dict
A dictionary of strings whose keys are all of the options specified above in the FC_poly.py docstring
"""
import os
from collections import OrderedDict
from dedalus.tools.config import config
from logic.functions import mpi_makedirs
#Get info on data directory, create directories, setup logging
# (Note: this order of imports is bad form, but gets logs properly outputting)
data_dir, case_name = name_case(input_dict)
mpi_makedirs(data_dir)
logdir = os.path.join(data_dir, 'logs')
mpi_makedirs(logdir)
file_level = config['logging']['file_level'] = 'debug'
filename = config['logging']['filename'] = os.path.join(
logdir, 'dedalus_log')
from dedalus import public as de
from dedalus.extras import flow_tools
from logic.atmospheres import TriLayerIH
from logic.domains import DedalusDomain
from logic.experiments import CompressibleConvection
from logic.problems import DedalusIVP
from logic.equations import KappaMuFCE
from logic.ae_tools import FCAcceleratedEvolutionIVP
from logic.outputs import initialize_output, ae_initialize_output
from logic.checkpointing import Checkpoint
from logic.field_averager import AveragerFCAE, AveragerFCStructure
import logging
logger = logging.getLogger(__name__)
logger.info("Running polytrope case: {}".format(case_name))
logger.info("Saving run in: {}".format(data_dir))
# Read in command line args & process them
# Atmosphere params
Ra, Pr, n_rho_rzT, n_rho_cz, n_rho_rzB, eps, aspect = [
float(input_dict[k])
for k in ('--Rayleigh', '--Prandtl', '--n_rho_rzT', '--n_rho_cz',
'--n_rho_rzB', '--epsilon', '--aspect')
]
threeD = input_dict['--3D']
# BCs
thermal_BC = OrderedDict((('flux', False), ('temp', False),
('flux_temp', False), ('temp_flux', False)))
velocity_BC = OrderedDict((('stress_free', False), ('no_slip', False)))
thermal_BC[input_dict['--thermal_BC']] = True
velocity_BC[input_dict['--velocity_BC']] = True
# Coeff resolutions
nx, ny, nz = [int(input_dict[n]) for n in ['--nx', '--ny', '--nz']]
# 3D processor mesh.
mesh = input_dict['--mesh']
if mesh is not None:
mesh = mesh.split(',')
mesh = [int(mesh[0]), int(mesh[1])]
# Stop conditions
run_time_wall, run_time_buoy, run_time_therm = [
input_dict[t]
for t in ['--run_time_wall', '--run_time_buoy', '--run_time_therm']
]
run_time_wall = float(run_time_wall)
if run_time_buoy is not None:
run_time_buoy = float(run_time_buoy)
run_time_therm = None
logger.info(
"Terminating run after {} buoyancy times".format(run_time_buoy))
else:
run_time_therm = float(run_time_therm)
logger.info(
"Terminating run after {} thermal times".format(run_time_therm))
# Initialize atmosphere class
atmo_kwargs = OrderedDict([('epsilon', eps), ('n_rho_rzT', n_rho_rzT),
('n_rho_cz', n_rho_cz),
('n_rho_rzB', n_rho_rzB), ('gamma', 5. / 3),
('R', 1)])
atmosphere = TriLayerIH(**atmo_kwargs)
# Initialize domain class
resolution = (nz, nx, ny)
if not threeD: resolution = resolution[:2]
domain_args = (atmosphere, resolution, aspect)
domain_kwargs = OrderedDict((('mesh', mesh), ))
de_domain = DedalusDomain(*domain_args, **domain_kwargs)
#Set diffusivities
experiment_args = (de_domain, atmosphere, Ra, Pr)
experiment_kwargs = {}
experiment = CompressibleConvection(*experiment_args, **experiment_kwargs)
#Set up problem and equations
if args['--ae']:
problem_type = FCAcceleratedEvolutionIVP
else:
problem_type = DedalusIVP
de_problem = problem_type(de_domain)
equations = KappaMuFCE(thermal_BC, velocity_BC, atmosphere, de_domain,
de_problem)
atmosphere.save_atmo_file('./', de_domain)
# Build solver, set stop times
de_problem.build_solver(de.timesteppers.RK222)
#Setup checkpointing and initial conditions
checkpoint = Checkpoint(data_dir)
dt, mode = experiment.set_IC(
de_problem.solver,
eps,
checkpoint,
restart=args['--restart'],
seed=int(args['--seed']),
checkpoint_dt=float(args['--checkpoint_buoy']) *
atmosphere.atmo_params['t_buoy'])
if run_time_buoy is None:
stop_sim_time = run_time_therm * atmosphere.atmo_params['t_therm']
else:
stop_sim_time = run_time_buoy * atmosphere.atmo_params['t_buoy']
if args['--run_time_restarted']:
stop_sim_time += de_problem.solver.sim_time
de_problem.set_stop_condition(stop_wall_time=run_time_wall * 3600,
stop_sim_time=stop_sim_time)
#Set up outputs
output_dt = float(args['--output_dt']) * atmosphere.atmo_params['t_buoy']
if args['--ae_outs']:
analysis_tasks = ae_initialize_output(
de_domain,
de_problem,
data_dir,
output_dt=output_dt,
output_vol_dt=atmosphere.atmo_params['t_buoy'],
mode=mode) # volumes_output=True)
else:
analysis_tasks = initialize_output(
de_domain,
de_problem,
data_dir,
output_dt=output_dt,
output_vol_dt=atmosphere.atmo_params['t_buoy'],
mode=mode) # volumes_output=True)
# Ensure good initial dt and setup CFL
max_dt = min(
(0.2, float(args['--output_dt']))) * atmosphere.atmo_params['t_buoy']
if dt is None:
dt = max_dt
cfl_safety = 0.1
CFL = flow_tools.CFL(de_problem.solver,
initial_dt=dt,
cadence=1,
safety=cfl_safety * 2,
max_change=1.5,
min_change=0.5,
max_dt=max_dt,
threshold=0.1)
if threeD:
CFL.add_velocities(('u', 'v', 'w'))
else:
CFL.add_velocities(('u', 'w'))
# Solve the IVP.
if args['--ae']:
task_args = (thermal_BC, )
pre_loop_args = ((AveragerFCAE, AveragerFCStructure), (True, False),
data_dir, atmo_kwargs, CompressibleConvection,
experiment_args, experiment_kwargs)
task_kwargs = {}
pre_loop_kwargs = {
'sim_time_start':
int(args['--ae_start_time']) * atmosphere.atmo_params['t_buoy'],
'min_bvp_time':
10 * atmosphere.atmo_params['t_buoy'],
'between_ae_wait_time':
20 * atmosphere.atmo_params['t_buoy'],
'later_bvp_time':
20 * atmosphere.atmo_params['t_buoy'],
'ae_convergence':
1e-2,
'bvp_threshold':
1e-2,
'min_bvp_threshold':
5e-3
}
solve_args = (dt, CFL, data_dir, analysis_tasks)
solve_kwargs = {
'task_args': task_args,
'pre_loop_args': pre_loop_args,
'task_kwargs': task_kwargs,
'pre_loop_kwargs': pre_loop_kwargs,
'time_div': atmosphere.atmo_params['t_buoy'],
'track_fields': ['Pe_rms', 'dissipation'],
'threeD': threeD,
'Hermitian_cadence': 100,
'no_join': args['--no_join'],
'mode': mode
}
de_problem.solve_IVP(*solve_args, **solve_kwargs)
else:
de_problem.solve_IVP(dt,
CFL,
data_dir,
analysis_tasks,
time_div=atmosphere.atmo_params['t_buoy'],
track_fields=['Pe_rms', 'dissipation'],
threeD=threeD,
Hermitian_cadence=100,
no_join=args['--no_join'],
mode=mode)
def FC_convection(Rayleigh=1e6, Prandtl=1, stiffness=1e4, m_rz=3, gamma=5/3,
n_rho_cz=3.5, n_rho_rz=1,
nz_cz=128, nz_rz=128,
nx = None,
width=None,
single_chebyshev=False,
rk222=False,
superstep=False,
dense=False, nz_dense=64,
oz=False,
fixed_flux=False,
run_time=23.5, run_time_buoyancies=np.inf, run_time_iter=np.inf,
dynamic_diffusivities=False,
max_writes=20,out_cadence=0.1, no_coeffs=False, no_join=False,
restart=None, data_dir='./', verbose=False, label=None):
def format_number(number, no_format_min=0.1, no_format_max=10):
if number > no_format_max or number < no_format_min:
try:
mantissa = "{:e}".format(number).split("+")[0].split("e")[0].rstrip("0") or "0"
power = "{:e}".format(number).split("+")[1].lstrip("0") or "0"
except:
mantissa = "{:e}".format(number).split("-")[0].split("e")[0].rstrip("0") or "0"
power = "{:e}".format(number).split("-")[1].lstrip("0") or "0"
power = "-"+power
if mantissa[-1]==".":
mantissa = mantissa[:-1]
mantissa += "e"
else:
mantissa = "{:f}".format(number).rstrip("0") or "0"
if mantissa[-1]==".":
mantissa = mantissa[:-1]
power = ""
number_string = mantissa+power
return number_string
# save data in directory named after script
if data_dir[-1] != '/':
data_dir += '/'
data_dir += sys.argv[0].split('.py')[0]
if fixed_flux:
data_dir += '_flux'
if dynamic_diffusivities:
data_dir += '_dynamic'
if oz:
data_dir += '_oz'
data_dir += "_nrhocz{}_Ra{}_S{}".format(format_number(n_rho_cz), format_number(Rayleigh), format_number(stiffness))
if width:
data_dir += "_erf{}".format(format_number(width))
if label:
data_dir += "_{}".format(label)
data_dir += '/'
from dedalus.tools.config import config
config['logging']['filename'] = os.path.join(data_dir,'logs/dedalus_log')
config['logging']['file_level'] = 'DEBUG'
import mpi4py.MPI
if mpi4py.MPI.COMM_WORLD.rank == 0:
if not os.path.exists('{:s}/'.format(data_dir)):
os.makedirs('{:s}/'.format(data_dir))
logdir = os.path.join(data_dir,'logs')
if not os.path.exists(logdir):
os.mkdir(logdir)
logger = logging.getLogger(__name__)
logger.info("saving run in: {}".format(data_dir))
import dedalus.public as de
from dedalus.tools import post
from dedalus.extras import flow_tools
from dedalus.core.future import FutureField
from stratified_dynamics import multitropes
from tools.checkpointing import Checkpoint
checkpoint_min = 30
initial_time = time.time()
logger.info("Starting Dedalus script {:s}".format(sys.argv[0]))
constant_Prandtl=True
mixed_temperature_flux=None
if oz:
stable_top=True
if not fixed_flux:
mixed_temperature_flux=True
else:
stable_top=False
# Set domain
if nx is None:
nx = nz_cz*4
if single_chebyshev:
nz = nz_cz
nz_list = [nz_cz]
else:
if dense:
nz = nz_rz+nz_dense+nz_cz
#nz_list = [nz_rz, int(nz_dense/2), int(nz_dense/2), nz_cz]
nz_list = [nz_rz, nz_dense, nz_cz]
else:
nz = nz_rz+nz_cz
nz_list = [nz_rz, nz_cz]
if dynamic_diffusivities:
atmosphere = multitropes.FC_multitrope_2d_kappa_mu(nx=nx, nz=nz_list, stiffness=stiffness, m_rz=m_rz, gamma=gamma,
n_rho_cz=n_rho_cz, n_rho_rz=n_rho_rz,
verbose=verbose, width=width,
constant_Prandtl=constant_Prandtl,
stable_top=stable_top)
else:
atmosphere = multitropes.FC_multitrope(nx=nx, nz=nz_list, stiffness=stiffness, m_rz=m_rz, gamma=gamma,
n_rho_cz=n_rho_cz, n_rho_rz=n_rho_rz,
verbose=verbose, width=width,
constant_Prandtl=constant_Prandtl,
stable_top=stable_top)
atmosphere.set_IVP_problem(Rayleigh, Prandtl)
atmosphere.set_BC(mixed_temperature_flux=mixed_temperature_flux, fixed_flux=fixed_flux)
problem = atmosphere.get_problem()
if atmosphere.domain.distributor.rank == 0:
if not os.path.exists('{:s}/'.format(data_dir)):
os.makedirs('{:s}/'.format(data_dir))
if rk222:
logger.info("timestepping using RK222")
ts = de.timesteppers.RK222
cfl_safety_factor = 0.2*2
else:
logger.info("timestepping using RK443")
ts = de.timesteppers.RK443
cfl_safety_factor = 0.2*4
# Build solver
solver = problem.build_solver(ts)
# initial conditions
if restart is None:
mode = "overwrite"
else:
mode = "append"
logger.info("checkpointing in {}".format(data_dir))
checkpoint = Checkpoint(data_dir)
if restart is None:
atmosphere.set_IC(solver)
dt = None
else:
logger.info("restarting from {}".format(restart))
dt = checkpoint.restart(restart, solver)
checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min*60, mode=mode)
logger.info("thermal_time = {:g}, top_thermal_time = {:g}".format(atmosphere.thermal_time, atmosphere.top_thermal_time))
max_dt = atmosphere.min_BV_time
max_dt = atmosphere.buoyancy_time*out_cadence
if dt is None: dt = max_dt/5
report_cadence = 1
output_time_cadence = out_cadence*atmosphere.buoyancy_time
solver.stop_sim_time = solver.sim_time + run_time_buoyancies*atmosphere.buoyancy_time
solver.stop_iteration = solver.iteration + run_time_iter
solver.stop_wall_time = run_time*3600
logger.info("output cadence = {:g}".format(output_time_cadence))
analysis_tasks = atmosphere.initialize_output(solver, data_dir, coeffs_output=not(no_coeffs), sim_dt=output_time_cadence, max_writes=max_writes, mode=mode)
cfl_cadence = 1
CFL = flow_tools.CFL(solver, initial_dt=dt, cadence=cfl_cadence, safety=cfl_safety_factor,
max_change=1.5, min_change=0.5, max_dt=max_dt, threshold=0.1)
if superstep:
CFL_traditional = flow_tools.CFL(solver, initial_dt=max_dt, cadence=cfl_cadence, safety=cfl_safety_factor,
max_change=1.5, min_change=0.5, max_dt=max_dt, threshold=0.1)
CFL_traditional.add_velocities(('u', 'w'))
vel_u = FutureField.parse('u', CFL.solver.evaluator.vars, CFL.solver.domain)
delta_x = atmosphere.Lx/nx
CFL.add_frequency(vel_u/delta_x)
vel_w = FutureField.parse('w', CFL.solver.evaluator.vars, CFL.solver.domain)
mean_delta_z_cz = atmosphere.Lz_cz/nz_cz
CFL.add_frequency(vel_w/mean_delta_z_cz)
else:
CFL.add_velocities(('u', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("Re_rms", name='Re')
try:
logger.info("starting main loop")
start_time = time.time()
start_iter = solver.iteration
good_solution = True
first_step = True
while solver.ok and good_solution:
dt = CFL.compute_dt()
# advance
solver.step(dt)
effective_iter = solver.iteration - start_iter
# update lists
if effective_iter % report_cadence == 0:
Re_avg = flow.grid_average('Re')
log_string = 'Iteration: {:5d}, Time: {:8.3e} ({:8.3e}), '.format(solver.iteration, solver.sim_time, solver.sim_time/atmosphere.buoyancy_time)
log_string += 'dt: {:8.3e}'.format(dt)
if superstep:
dt_traditional = CFL_traditional.compute_dt()
log_string += ' (vs {:8.3e})'.format(dt_traditional)
log_string += ', '
log_string += 'Re: {:8.3e}/{:8.3e}'.format(Re_avg, flow.max('Re'))
logger.info(log_string)
if not np.isfinite(Re_avg):
good_solution = False
logger.info("Terminating run. Trapped on Reynolds = {}".format(Re_avg))
if first_step:
if verbose:
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.spy(solver.pencils[0].L, markersize=0.5, markeredgewidth=0.0)
fig.savefig(data_dir+"sparsity_pattern.png", dpi=2400)
#fig.savefig(data_dir+"sparsity_pattern.svg", format="svg")
import scipy.sparse.linalg as sla
LU = sla.splu(solver.pencils[0].LHS.tocsc(), permc_spec='NATURAL')
fig = plt.figure()
ax = fig.add_subplot(1,2,1)
ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
ax = fig.add_subplot(1,2,2)
ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
fig.savefig(data_dir+"sparsity_pattern_LU.png", dpi=1200)
#fig.savefig(data_dir+"sparsity_pattern_LU.svg", format="svg")
logger.info("{} nonzero entries in LU".format(LU.nnz))
logger.info("{} nonzero entries in LHS".format(solver.pencils[0].LHS.tocsc().nnz))
logger.info("{} fill in factor".format(LU.nnz/solver.pencils[0].LHS.tocsc().nnz))
first_step = False
start_time = time.time()
except:
logger.error('Exception raised, triggering end of main loop.')
raise
finally:
end_time = time.time()
# Print statistics
elapsed_time = end_time - start_time
elapsed_sim_time = solver.sim_time
N_iterations = solver.iteration - 1
logger.info('main loop time: {:e}'.format(elapsed_time))
logger.info('Iterations: {:d}'.format(N_iterations))
logger.info('iter/sec: {:g}'.format(N_iterations/(elapsed_time)))
if N_iterations > 0:
logger.info('Average timestep: {:e}'.format(elapsed_sim_time / N_iterations))
logger.info('beginning join operation')
try:
final_checkpoint = Checkpoint(data_dir, checkpoint_name='final_checkpoint')
final_checkpoint.set_checkpoint(solver, wall_dt=1, mode="append")
solver.step(dt) #clean this up in the future...works for now.
post.merge_process_files(data_dir+'/final_checkpoint/')
except:
print('cannot save final checkpoint')
if not(no_join):
logger.info(data_dir+'/checkpoint/')
post.merge_process_files(data_dir+'/checkpoint/')
for task in analysis_tasks:
logger.info(analysis_tasks[task].base_path)
post.merge_process_files(analysis_tasks[task].base_path)
if (atmosphere.domain.distributor.rank==0):
N_TOTAL_CPU = atmosphere.domain.distributor.comm_cart.size
# Print statistics
print('-' * 40)
total_time = end_time-initial_time
main_loop_time = end_time - start_time
startup_time = start_time-initial_time
n_steps = solver.iteration-1
print(' startup time:', startup_time)
print('main loop time:', main_loop_time)
print(' total time:', total_time)
if n_steps > 0:
print(' iterations:', n_steps)
print(' loop sec/iter:', main_loop_time/n_steps)
print(' average dt:', solver.sim_time/n_steps)
print(" N_cores, Nx, Nz, startup main loop, main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid")
print('scaling:',
' {:d} {:d} {:d}'.format(N_TOTAL_CPU,nx,nz),
' {:8.3g} {:8.3g} {:8.3g} {:8.3g} {:8.3g}'.format(startup_time,
main_loop_time,
main_loop_time/n_steps,
main_loop_time/n_steps/(nx*nz),
N_TOTAL_CPU*main_loop_time/n_steps/(nx*nz)))
print('-' * 40)
return data_dir
def Rayleigh_Benard(Rayleigh=1e6,
Prandtl=1,
nz=64,
nx=None,
ny=None,
aspect=4,
fixed_flux=False,
fixed_T=False,
mixed_flux_T=True,
stress_free=False,
no_slip=True,
restart=None,
run_time=23.5,
run_time_buoyancy=None,
run_time_iter=np.inf,
run_time_therm=1,
max_writes=20,
max_slice_writes=20,
output_dt=0.2,
data_dir='./',
coeff_output=True,
verbose=False,
no_join=False,
do_bvp=False,
num_bvps=10,
bvp_convergence_factor=1e-3,
bvp_equil_time=10,
bvp_resolution_factor=1,
bvp_transient_time=30,
bvp_final_equil_time=None,
min_bvp_time=50,
first_bvp_time=20,
first_bvp_convergence_factor=1e-2,
threeD=False,
seed=42,
mesh=None,
overwrite=False):
import os
from dedalus.tools.config import config
config['logging']['filename'] = os.path.join(data_dir, 'logs/dedalus_log')
config['logging']['file_level'] = 'DEBUG'
import mpi4py.MPI
if mpi4py.MPI.COMM_WORLD.rank == 0:
if not os.path.exists('{:s}/'.format(data_dir)):
os.makedirs('{:s}/'.format(data_dir))
logdir = os.path.join(data_dir, 'logs')
if not os.path.exists(logdir):
os.mkdir(logdir)
logger = logging.getLogger(__name__)
logger.info("saving run in: {}".format(data_dir))
import time
from dedalus import public as de
from dedalus.extras import flow_tools
from dedalus.tools import post
# input parameters
logger.info("Ra = {}, Pr = {}".format(Rayleigh, Prandtl))
# Parameters
Lz = 1.
Lx = aspect * Lz
Ly = aspect * Lz
if nx is None:
nx = int(nz * aspect)
if ny is None:
ny = int(nz * aspect)
if threeD:
logger.info("resolution: [{}x{}x{}]".format(nx, ny, nz))
equations = BoussinesqEquations3D(nx=nx,
ny=ny,
nz=nz,
Lx=Lx,
Ly=Ly,
Lz=Lz,
mesh=mesh)
else:
logger.info("resolution: [{}x{}]".format(nx, nz))
equations = BoussinesqEquations2D(nx=nx, nz=nz, Lx=Lx, Lz=Lz)
equations.set_IVP(Rayleigh, Prandtl)
bc_dict = {
'fixed_flux': None,
'fixed_temperature': None,
'mixed_flux_temperature': None,
'mixed_temperature_flux': None,
'stress_free': None,
'no_slip': None
}
if mixed_flux_T:
bc_dict['mixed_flux_temperature'] = True
elif fixed_T:
bc_dict['fixed_temperature'] = True
elif fixed_flux:
bc_dict['fixed_flux'] = True
if stress_free:
bc_dict['stress_free'] = True
elif no_slip:
bc_dict['no_slip'] = True
supercrit = Rayleigh / RA_CRIT
equations.set_BC(**bc_dict)
# Build solver
ts = de.timesteppers.RK443
cfl_safety = 0.8
solver = equations.problem.build_solver(ts)
logger.info('Solver built')
checkpoint = Checkpoint(data_dir)
if isinstance(restart, type(None)):
equations.set_IC(solver, seed=seed)
dt = None
mode = 'overwrite'
else:
logger.info("restarting from {}".format(restart))
checkpoint.restart(restart, solver)
if overwrite:
mode = 'overwrite'
else:
mode = 'append'
checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min * 60, mode=mode)
# Integration parameters
# if not isinstance(run_time_therm, type(None)):
# solver.stop_sim_time = run_time_therm*equations.thermal_time + solver.sim_time
# elif not isinstance(run_time_buoyancy, type(None)):
# solver.stop_sim_time = run_time_buoyancy + solver.sim_time
# else:
solver.stop_sim_time = np.inf
solver.stop_wall_time = run_time * 3600.
solver.stop_iteration = run_time_iter
Hermitian_cadence = 100
# Analysis
max_dt = output_dt
analysis_tasks = equations.initialize_output(solver,
data_dir,
coeff_output=coeff_output,
output_dt=output_dt,
mode=mode)
# CFL
CFL = flow_tools.CFL(solver,
initial_dt=0.1,
cadence=1,
safety=cfl_safety,
max_change=1.5,
min_change=0.5,
max_dt=max_dt,
threshold=0.1)
if threeD:
CFL.add_velocities(('u', 'v', 'w'))
else:
CFL.add_velocities(('u', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("Re", name='Re')
# flow.add_property("interp(w, z=0.95)", name='w near top')
# u, v, w = solver.state['u'], solver.state['v'], solver.state['w']
if do_bvp:
if not threeD:
ny = 0
atmo_class = BoussinesqEquations2D
else:
atmo_class = BoussinesqEquations3D
bvp_solver = BoussinesqBVPSolver(atmo_class, nx, ny, nz, \
flow, equations.domain.dist.comm_cart, \
solver, num_bvps, bvp_equil_time, \
threeD=threeD,
bvp_transient_time=bvp_transient_time, \
bvp_run_threshold=bvp_convergence_factor, \
bvp_l2_check_time=1, mesh=mesh,\
first_bvp_time=first_bvp_time,
first_run_threshold=first_bvp_convergence_factor,\
plot_dir='{}/bvp_plots/'.format(data_dir),\
min_avg_dt=1e-10, final_equil_time=bvp_final_equil_time,
min_bvp_time=min_bvp_time)
bc_dict.pop('stress_free')
bc_dict.pop('no_slip')
# print(equations.domain.grid(0), equations.domain.grid(1), equations.domain.grid(2))
first_step = True
# Main loop
try:
logger.info('Starting loop')
Re_avg = 0
continue_bvps = True
not_corrected_times = True
init_time = solver.sim_time
start_iter = solver.iteration
while (solver.ok and np.isfinite(Re_avg)) and continue_bvps:
dt = CFL.compute_dt()
solver.step(dt) #, trim=True)
# Solve for blow-up over long timescales in 3D due to hermitian-ness
effective_iter = solver.iteration - start_iter
if threeD and effective_iter % Hermitian_cadence == 0:
for field in solver.state.fields:
field.require_grid_space()
Re_avg = flow.grid_average('Re')
log_string = 'Iteration: {:5d}, '.format(solver.iteration)
log_string += 'Time: {:8.3e} ({:8.3e} therm), dt: {:8.3e}, '.format(
solver.sim_time, solver.sim_time / equations.thermal_time, dt)
log_string += 'Re: {:8.3e}/{:8.3e}'.format(Re_avg, flow.max('Re'))
logger.info(log_string)
if not_corrected_times and Re_avg > 1:
if not isinstance(run_time_therm, type(None)):
solver.stop_sim_time = run_time_therm * equations.thermal_time + solver.sim_time
elif not isinstance(run_time_buoyancy, type(None)):
solver.stop_sim_time = run_time_buoyancy + solver.sim_time
not_corrected_times = False
if do_bvp:
bvp_solver.update_avgs(dt, Re_avg, min_Re=np.sqrt(supercrit))
if bvp_solver.check_if_solve():
atmo_kwargs = {'nz': nz * bvp_resolution_factor, 'Lz': Lz}
diff_args = [Rayleigh, Prandtl]
bvp_solver.solve_BVP(atmo_kwargs, diff_args, bc_dict)
if bvp_solver.terminate_IVP():
continue_bvps = False
if first_step:
if verbose:
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.spy(solver.pencils[0].L,
markersize=1,
markeredgewidth=0.0)
fig.savefig(data_dir + "sparsity_pattern.png", dpi=1200)
import scipy.sparse.linalg as sla
LU = sla.splu(solver.pencils[0].LHS.tocsc(),
permc_spec='NATURAL')
fig = plt.figure()
ax = fig.add_subplot(1, 2, 1)
ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
ax = fig.add_subplot(1, 2, 2)
ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
fig.savefig(data_dir + "sparsity_pattern_LU.png", dpi=1200)
logger.info("{} nonzero entries in LU".format(LU.nnz))
logger.info("{} nonzero entries in LHS".format(
solver.pencils[0].LHS.tocsc().nnz))
logger.info("{} fill in factor".format(
LU.nnz / solver.pencils[0].LHS.tocsc().nnz))
first_step = False
start_time = time.time()
except:
raise
logger.error('Exception raised, triggering end of main loop.')
finally:
end_time = time.time()
main_loop_time = end_time - start_time
n_iter_loop = solver.iteration - 1
logger.info('Iterations: {:d}'.format(n_iter_loop))
logger.info('Sim end time: {:f}'.format(solver.sim_time))
logger.info('Run time: {:f} sec'.format(main_loop_time))
logger.info('Run time: {:f} cpu-hr'.format(
main_loop_time / 60 / 60 * equations.domain.dist.comm_cart.size))
logger.info('iter/sec: {:f} (main loop only)'.format(n_iter_loop /
main_loop_time))
try:
final_checkpoint = Checkpoint(data_dir,
checkpoint_name='final_checkpoint')
final_checkpoint.set_checkpoint(solver, wall_dt=1, mode=mode)
solver.step(dt) #clean this up in the future...works for now.
post.merge_process_files(data_dir + '/final_checkpoint/',
cleanup=False)
except:
raise
print('cannot save final checkpoint')
finally:
if not no_join:
logger.info('beginning join operation')
post.merge_analysis(data_dir + 'checkpoints')
for task in analysis_tasks:
logger.info(task.base_path)
post.merge_analysis(task.base_path)
logger.info(40 * "=")
logger.info('Iterations: {:d}'.format(n_iter_loop))
logger.info('Sim end time: {:f}'.format(solver.sim_time))
logger.info('Run time: {:f} sec'.format(main_loop_time))
logger.info('Run time: {:f} cpu-hr'.format(
main_loop_time / 60 / 60 *
equations.domain.dist.comm_cart.size))
logger.info('iter/sec: {:f} (main loop only)'.format(
n_iter_loop / main_loop_time))
snapshots.add_task('rho(T)')
snapshots.add_task('u')
snapshots.add_task('w')
checkpoints = solver.evaluator.add_file_handler('checkpoints%s' % ending,
sim_dt=10,
max_writes=1)
checkpoints.add_system(solver.state, layout='c')
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=10)
flow.add_property("sqrt(u*u + w*w) / nu", name='Re')
CFL = flow_tools.CFL(solver,
initial_dt=dt,
cadence=3,
safety=0.35,
max_dt=max_dt,
threshold=0.05)
CFL.add_velocities(('u*heaviside', 'w*heaviside'))
# Main loop
try:
logger.info('Starting loop')
start_time = time.time()
while solver.ok:
dt = CFL.compute_dt()
if solver.sim_time > 1:
dt = min(dt, max_dt)
# Initial condition
def smoothstep(z, d):
return 0.5 * (1 + np.tanh(z / d))
b0 = bz_deep * (model.z + h0) * smoothstep(-model.z - h0, d)
# Add noise ... ?
noise = dedaLES.random_noise(model.domain)
b0 = b0 + b_noise * noise
model.set_fields(b=b0, u=0, v=0, w=0)
# CFL conditions
CFL = flow_tools.CFL(model.solver,
initial_dt=1e-2,
cadence=1,
max_change=1.5,
max_dt=10,
safety=0.5)
CFL.add_velocities(('u', 'v+Vg', 'w'))
#Analysis tasks
snap = model.solver.evaluator.add_file_handler('snapshots', sim_dt=30)
#snap.add_task("interp(b, z=0)", scales=1, name='b surface')
#snap.add_task("interp(u, z=0)", scales=1, name='u surface')
#snap.add_task("interp(v, z=0)", scales=1, name='v surface')
snap.add_task("interp(b, z=-25)", scales=1, name='b 25')
snap.add_task("interp(u, z=-25)", scales=1, name='u 25')
snap.add_task("interp(v, z=-25)", scales=1, name='v 25')
snap.add_task("interp(w, z=-25)", scales=1, name='w 25')
def Rayleigh_Benard(Rayleigh=1e6,
Prandtl=1,
nz=64,
nx=None,
aspect=4,
fixed_flux=False,
fixed_T=True,
viscous_heating=False,
restart=None,
run_time=23.5,
run_time_buoyancy=50,
run_time_iter=np.inf,
max_writes=10,
max_slice_writes=10,
data_dir='./',
coeff_output=True,
verbose=False,
no_join=False):
import os
from dedalus.tools.config import config
config['logging']['filename'] = os.path.join(data_dir, 'logs/dedalus_log')
config['logging']['file_level'] = 'DEBUG'
import mpi4py.MPI
if mpi4py.MPI.COMM_WORLD.rank == 0:
if not os.path.exists('{:s}/'.format(data_dir)):
os.makedirs('{:s}/'.format(data_dir))
logdir = os.path.join(data_dir, 'logs')
if not os.path.exists(logdir):
os.mkdir(logdir)
logger = logging.getLogger(__name__)
logger.info("saving run in: {}".format(data_dir))
import time
from dedalus import public as de
from dedalus.extras import flow_tools
from dedalus.tools import post
# input parameters
logger.info("Ra = {}, Pr = {}".format(Rayleigh, Prandtl))
# Parameters
Lz = 1.
Lx = aspect * Lz
if nx is None:
nx = int(nz * aspect)
logger.info("resolution: [{}x{}]".format(nx, nz))
# Create bases and domain
x_basis = de.Fourier('x', nx, interval=(0, Lx), dealias=3 / 2)
z_basis = de.Chebyshev('z', nz, interval=(0, Lz), dealias=3 / 2)
domain = de.Domain([x_basis, z_basis], grid_dtype=np.float64)
if fixed_flux:
T_bc_var = 'Tz'
elif fixed_T:
T_bc_var = 'T'
# 2D Boussinesq hydrodynamics
problem = de.IVP(domain, variables=['Tz', 'T', 'p', 'u', 'w', 'Oy'])
problem.meta['p', T_bc_var, 'Oy', 'w']['z']['dirichlet'] = True
T0_z = domain.new_field()
T0_z.meta['x']['constant'] = True
T0_z['g'] = -1
T0 = domain.new_field()
T0.meta['x']['constant'] = True
T0['g'] = Lz / 2 - domain.grid(-1)
problem.parameters['T0'] = T0
problem.parameters['T0_z'] = T0_z
problem.parameters['P'] = (Rayleigh * Prandtl)**(-1 / 2)
problem.parameters['R'] = (Rayleigh / Prandtl)**(-1 / 2)
problem.parameters['F'] = F = 1
problem.parameters['Lx'] = Lx
problem.parameters['Lz'] = Lz
problem.substitutions['plane_avg(A)'] = 'integ(A, "x")/Lx'
problem.substitutions['vol_avg(A)'] = 'integ(A)/Lx/Lz'
problem.substitutions['v'] = '0'
problem.substitutions['Ox'] = '0'
problem.substitutions['Oz'] = '(dx(v) )'
problem.substitutions['Kx'] = '( -dz(Oy))'
problem.substitutions['Ky'] = '(dz(Ox) - dx(Oz))'
problem.substitutions['Kz'] = '(dx(Oy) )'
problem.substitutions['vorticity'] = 'Oy'
problem.substitutions['enstrophy'] = 'Oy**2'
problem.substitutions['u_fluc'] = '(u - plane_avg(u))'
problem.substitutions['w_fluc'] = '(w - plane_avg(w))'
problem.substitutions['KE'] = '(0.5*(u*u+w*w))'
problem.substitutions['sigma_xz'] = '(dx(w) + Oy + dx(w))'
problem.substitutions['sigma_xx'] = '(2*dx(u))'
problem.substitutions['sigma_zz'] = '(2*dz(w))'
if viscous_heating:
problem.substitutions[
'visc_heat'] = 'R*(sigma_xz**2 + sigma_xx*dx(u) + sigma_zz*dz(w))'
problem.substitutions['visc_flux_z'] = 'R*(u*sigma_xz + w*sigma_zz)'
else:
problem.substitutions['visc_heat'] = '0'
problem.substitutions['visc_flux_z'] = '0'
problem.substitutions['conv_flux_z'] = '(w*T + visc_flux_z)/P'
problem.substitutions['kappa_flux_z'] = '(-Tz)'
problem.add_equation("Tz - dz(T) = 0")
problem.add_equation(
"dt(T) - P*(dx(dx(T)) + dz(Tz)) + w*T0_z = -(u*dx(T) + w*Tz) - visc_heat"
)
# O == omega = curl(u); K = curl(O)
problem.add_equation("dt(u) + R*Kx + dx(p) = v*Oz - w*Oy ")
problem.add_equation("dt(w) + R*Kz + dz(p) -T = u*Oy - v*Ox ")
problem.add_equation("dx(u) + dz(w) = 0")
problem.add_equation("Oy - dz(u) + dx(w) = 0")
problem.add_bc("right(p) = 0", condition="(nx == 0)")
if fixed_flux:
problem.add_bc("left(Tz) = 0")
problem.add_bc("right(Tz) = 0")
elif fixed_T:
problem.add_bc("left(T) = 0")
problem.add_bc("right(T) = 0")
problem.add_bc("left(Oy) = 0")
problem.add_bc("right(Oy) = 0")
problem.add_bc("left(w) = 0")
problem.add_bc("right(w) = 0", condition="(nx != 0)")
# Build solver
ts = de.timesteppers.RK443
cfl_safety = 1
solver = problem.build_solver(ts)
logger.info('Solver built')
checkpoint = solver.evaluator.add_file_handler(data_dir + 'checkpoints',
wall_dt=8 * 3600,
max_writes=1)
checkpoint.add_system(solver.state, layout='c')
# Initial conditions
x = domain.grid(0)
z = domain.grid(1)
T = solver.state['T']
Tz = solver.state['Tz']
# Random perturbations, initialized globally for same results in parallel
noise = global_noise(domain, scale=1, frac=0.25)
if restart is None:
# Linear background + perturbations damped at walls
zb, zt = z_basis.interval
pert = 1e-3 * noise * (zt - z) * (z - zb)
T['g'] = pert
T.differentiate('z', out=Tz)
else:
logger.info("restarting from {}".format(restart))
checkpoint.restart(restart, solver)
# Integration parameters
solver.stop_sim_time = run_time_buoyancy
solver.stop_wall_time = run_time * 3600.
solver.stop_iteration = run_time_iter
# Analysis
analysis_tasks = []
snapshots = solver.evaluator.add_file_handler(data_dir + 'slices',
sim_dt=0.1,
max_writes=max_slice_writes)
snapshots.add_task("T + T0", name='T')
snapshots.add_task("enstrophy")
snapshots.add_task("vorticity")
analysis_tasks.append(snapshots)
snapshots_small = solver.evaluator.add_file_handler(
data_dir + 'slices_small', sim_dt=0.1, max_writes=max_slice_writes)
snapshots_small.add_task("T + T0", name='T', scales=0.25)
snapshots_small.add_task("enstrophy", scales=0.25)
snapshots_small.add_task("vorticity", scales=0.25)
analysis_tasks.append(snapshots_small)
if coeff_output:
coeffs = solver.evaluator.add_file_handler(data_dir + 'coeffs',
sim_dt=0.1,
max_writes=max_slice_writes)
coeffs.add_task("T+T0", name="T", layout='c')
coeffs.add_task("T - plane_avg(T)", name="T'", layout='c')
coeffs.add_task("w", layout='c')
coeffs.add_task("u", layout='c')
coeffs.add_task("enstrophy", layout='c')
coeffs.add_task("vorticity", layout='c')
analysis_tasks.append(coeffs)
profiles = solver.evaluator.add_file_handler(data_dir + 'profiles',
sim_dt=0.1,
max_writes=max_writes)
profiles.add_task("plane_avg(T+T0)", name="T")
profiles.add_task("plane_avg(T)", name="T'")
profiles.add_task("plane_avg(u)", name="u")
profiles.add_task("plane_avg(w)", name="w")
profiles.add_task("plane_avg(enstrophy)", name="enstrophy")
# This may have an error:
profiles.add_task("plane_avg(conv_flux_z)/plane_avg(kappa_flux_z) + 1",
name="Nu")
profiles.add_task("plane_avg(conv_flux_z) + plane_avg(kappa_flux_z)",
name="Nu_2")
analysis_tasks.append(profiles)
scalar = solver.evaluator.add_file_handler(data_dir + 'scalar',
sim_dt=0.1,
max_writes=max_writes)
scalar.add_task("vol_avg(T)", name="IE")
scalar.add_task("vol_avg(KE)", name="KE")
scalar.add_task("vol_avg(-T*z)", name="PE_fluc")
scalar.add_task("vol_avg(T) + vol_avg(KE) + vol_avg(-T*z)", name="TE")
scalar.add_task("0.5*vol_avg(u_fluc*u_fluc+w_fluc*w_fluc)", name="KE_fluc")
scalar.add_task("0.5*vol_avg(u*u)", name="KE_x")
scalar.add_task("0.5*vol_avg(w*w)", name="KE_z")
scalar.add_task("0.5*vol_avg(u_fluc*u_fluc)", name="KE_x_fluc")
scalar.add_task("0.5*vol_avg(w_fluc*w_fluc)", name="KE_z_fluc")
scalar.add_task("vol_avg(plane_avg(u)**2)", name="u_avg")
scalar.add_task("vol_avg((u - plane_avg(u))**2)", name="u1")
scalar.add_task("vol_avg(conv_flux_z) + 1.", name="Nu")
analysis_tasks.append(scalar)
# workaround for issue #29
problem.namespace['enstrophy'].store_last = True
# CFL
CFL = flow_tools.CFL(solver,
initial_dt=0.1,
cadence=1,
safety=cfl_safety,
max_change=1.5,
min_change=0.5,
max_dt=0.1,
threshold=0.1)
CFL.add_velocities(('u', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("sqrt(u*u + w*w) / R", name='Re')
first_step = True
# Main loop
try:
logger.info('Starting loop')
Re_avg = 0
while solver.ok and np.isfinite(Re_avg):
dt = CFL.compute_dt()
solver.step(dt) #, trim=True)
Re_avg = flow.grid_average('Re')
log_string = 'Iteration: {:5d}, '.format(solver.iteration)
log_string += 'Time: {:8.3e}, dt: {:8.3e}, '.format(
solver.sim_time, dt)
log_string += 'Re: {:8.3e}/{:8.3e}'.format(Re_avg, flow.max('Re'))
logger.info(log_string)
if first_step:
if verbose:
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.spy(solver.pencils[0].L,
markersize=1,
markeredgewidth=0.0)
fig.savefig(data_dir + "sparsity_pattern.png", dpi=1200)
import scipy.sparse.linalg as sla
LU = sla.splu(solver.pencils[0].LHS.tocsc(),
permc_spec='NATURAL')
fig = plt.figure()
ax = fig.add_subplot(1, 2, 1)
ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
ax = fig.add_subplot(1, 2, 2)
ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
fig.savefig(data_dir + "sparsity_pattern_LU.png", dpi=1200)
logger.info("{} nonzero entries in LU".format(LU.nnz))
logger.info("{} nonzero entries in LHS".format(
solver.pencils[0].LHS.tocsc().nnz))
logger.info("{} fill in factor".format(
LU.nnz / solver.pencils[0].LHS.tocsc().nnz))
first_step = False
start_time = time.time()
except:
logger.error('Exception raised, triggering end of main loop.')
raise
finally:
end_time = time.time()
main_loop_time = end_time - start_time
n_iter_loop = solver.iteration - 1
logger.info('Iterations: {:d}'.format(n_iter_loop))
logger.info('Sim end time: {:f}'.format(solver.sim_time))
logger.info('Run time: {:f} sec'.format(main_loop_time))
logger.info('Run time: {:f} cpu-hr'.format(main_loop_time / 60 / 60 *
domain.dist.comm_cart.size))
logger.info('iter/sec: {:f} (main loop only)'.format(n_iter_loop /
main_loop_time))
final_checkpoint = solver.evaluator.add_file_handler(
data_dir + 'final_checkpoint', iter=1)
final_checkpoint.add_system(solver.state, layout='c')
solver.step(dt) #clean this up in the future...works for now.
post.merge_analysis(data_dir + 'final_checkpoint')
if not no_join:
logger.info('beginning join operation')
post.merge_analysis(data_dir + 'checkpoints')
for task in analysis_tasks:
logger.info(task.base_path)
post.merge_analysis(task.base_path)
logger.info(40 * "=")
logger.info('Iterations: {:d}'.format(n_iter_loop))
logger.info('Sim end time: {:f}'.format(solver.sim_time))
logger.info('Run time: {:f} sec'.format(main_loop_time))
logger.info('Run time: {:f} cpu-hr'.format(main_loop_time / 60 / 60 *
domain.dist.comm_cart.size))
logger.info('iter/sec: {:f} (main loop only)'.format(n_iter_loop /
main_loop_time))
load_writes.add_task('T')
load_writes.add_task('phi')
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence = 1)
flow.add_property("sqrt(vx*vx + vy*vy + vz*vz) / nu", name = 'Re_k')
flow.add_property("sqrt(vx*vx + vy*vy + vz*vz) / eta", name = 'Re_m')
flow.add_property("sqrt(vx*vx + vy*vy + vz*vz)", name = 'flow_speed')
flow.add_property("sqrt(vx*vx + vy*vy + vz*vz) / sqrt(T)", name = 'Ma')
flow.add_property("sqrt(Bx*Bx + By*By + Bz*Bz) / sqrt(rho)", name = 'Al_v')
flow.add_property("T", name = 'temp')
char_time = 50. # this should be set to a characteristic time in the problem (the alfven crossing time of the tube, for example)
CFL_safety = 0.3
CFL = flow_tools.CFL(solver, initial_dt = dt, cadence = 1, safety = CFL_safety,
max_change = 1.5, min_change = 0.005, max_dt = output_cadence, threshold = 0.05)
CFL.add_velocities(('vx', 'vy', 'vz'))
CFL.add_velocities(('Va_x', 'Va_y', 'Va_z'))
CFL.add_velocities(( 'Cs', 'Cs', 'Cs'))
good_solution = True
# Main loop
try:
logger.info('Starting loop')
logger_string = 'kappa: {:.3g}, mu: {:.3g}, eta: {:.3g}, dt: {:.3g}'.format(kappa, mu, eta, dt)
logger.info(logger_string)
start_time = time.time()
while solver.ok and good_solution:
dt = CFL.compute_dt()
solver.step(dt)
)
bp.add_equation("omega + Lap(psi) = 0", condition='nx != 0 or ny != 0')
bp.add_equation("psi = 0", condition='nx == 0 and ny == 0')
# Build solver
dt = 0.01
ts = de.timesteppers.RK443
IVP = bp.build_solver(ts)
IVP.stop_sim_time = 200.
IVP.stop_wall_time = np.inf
IVP.stop_iteration = 10000000
logger.info('Solver built')
CFL = flow_tools.CFL(IVP,
initial_dt=dt,
cadence=10,
safety=0.8,
max_change=2.0,
max_dt=dt)
CFL.add_velocities(('u', 'v'))
if mhd:
CFL.add_velocities(('Bx', 'By'))
if domain.distributor.rank == 0:
if not os.path.exists('{:s}/'.format(data_dir)):
os.mkdir('{:s}/'.format(data_dir))
# Analysis
snapshots = IVP.evaluator.add_file_handler(os.path.join(data_dir, 'snapshots'),
sim_dt=1.,
max_writes=50)
snapshots.add_system(IVP.state)
def main():
args = get_args()
# Parameters
Lx, Lz = (args.lx, args.lz)
Prandtl = args.prandtl
Rayleigh = args.rayleigh
seed = args.seed
# Create bases and domain
x_basis = de.Fourier('x', args.res_x, interval=(0, Lx), dealias=3/2) # 256
z_basis = de.Chebyshev('z', args.res_z, interval=(-Lz/2, Lz/2), dealias=3/2) # 64
domain = de.Domain([x_basis, z_basis], grid_dtype=np.float64)
# 2D Boussinesq hydrodynamics
problem = de.IVP(domain, variables=['p','b','u','w','bz','uz','wz'])
problem.meta['p','b','u','w']['z']['dirichlet'] = True
problem.parameters['P'] = (Rayleigh * Prandtl)**(-1/2)
problem.parameters['R'] = (Rayleigh / Prandtl)**(-1/2)
problem.parameters['F'] = F = 1
problem.add_equation("dx(u) + wz = 0")
# problem.add_equation("dt(b) - P*(dx(dx(b)) + dz(bz)) - F*w = -(u*dx(b) + w*bz)")
problem.add_equation("dt(b) - P*(dx(dx(b)) + dz(bz)) = -(u*dx(b) + w*bz)")
problem.add_equation("dt(u) - R*(dx(dx(u)) + dz(uz)) + dx(p) = -(u*dx(u) + w*uz)")
problem.add_equation("dt(w) - R*(dx(dx(w)) + dz(wz)) + dz(p) - b = -(u*dx(w) + w*wz)")
problem.add_equation("bz - dz(b) = 0")
problem.add_equation("uz - dz(u) = 0")
problem.add_equation("wz - dz(w) = 0")
problem.add_bc("left(b) = 0.5")
problem.add_bc("left(u) = 0")
problem.add_bc("left(w) = 0")
problem.add_bc("right(b) = -0.5")
problem.add_bc("right(u) = 0")
problem.add_bc("right(w) = 0", condition="(nx != 0)")
problem.add_bc("right(p) = 0", condition="(nx == 0)")
# Build solver
solver = problem.build_solver(de.timesteppers.RK222)
logger.info('Solver built')
# Initial conditions or restart
if not pathlib.Path('restart.h5').exists():
# Initial conditions
x = domain.grid(0)
z = domain.grid(1)
b = solver.state['b']
bz = solver.state['bz']
# Random perturbations, initialized globally for same results in parallel
gshape = domain.dist.grid_layout.global_shape(scales=1)
slices = domain.dist.grid_layout.slices(scales=1)
rand = np.random.RandomState(seed=seed)
noise = rand.standard_normal(gshape)[slices]
# Linear background + perturbations damped at walls
zb, zt = z_basis.interval
pert = 1e-3 * noise * (zt - z) * (z - zb)
b['g'] += F * pert
b.differentiate('z', out=bz)
# Timestepping and output
dt = args.dt
stop_sim_time = args.stop_sim_time
fh_mode = 'overwrite'
else:
# Restart
write, last_dt = solver.load_state('restart.h5', -1)
# Timestepping and output
dt = last_dt
stop_sim_time = args.stop_sim_time
fh_mode = 'append'
# Integration parameters
solver.stop_sim_time = stop_sim_time
# Analysis
snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.05, max_writes=250, mode=fh_mode)
snapshots.add_system(solver.state)
# CFL
CFL = flow_tools.CFL(solver, initial_dt=dt, cadence=1, safety=1,
max_change=1.5, min_change=0., max_dt=0.125, threshold=0.05)
CFL.add_velocities(('u', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=10)
flow.add_property("sqrt(u*u + w*w) / R", name='Re')
# Main loop
try:
logger.info('Starting loop')
start_time = time.time()
while solver.proceed:
dt = CFL.compute_dt()
dt = solver.step(dt)
if (solver.iteration-1) % 10 == 0:
logger.info('Iteration: %i, Time: %e, dt: %e' %(solver.iteration, solver.sim_time, dt))
logger.info('Max Re = %f' %flow.max('Re'))
except:
logger.error('Exception raised, triggering end of main loop.')
raise
finally:
end_time = time.time()
logger.info('Iterations: %i' %solver.iteration)
logger.info('Sim end time: %f' %solver.sim_time)
logger.info('Run time: %.2f sec' %(end_time-start_time))
logger.info('Run time: %f cpu-hr' %((end_time-start_time)/60/60*domain.dist.comm_cart.size))
def simulate(self,
initial_conditions=None,
scheme=de.timesteppers.RK443,
sim_time=2,
wall_time=np.inf,
tight=False):
"""Load initial conditions, run the simulation, and merge results.
Parameters:
initial_conditions (None, str): determines from what source to
draw the initial conditions. Valid options are as follows:
- None: use trivial conditions (T_ = 1 - z, T = 1 - z + eps).
- 'resume': use the most recent state file in the
records directory (load both model and DA system).
- An .h5 filename: load state variables for the model and
reset the data assimilation state variables to zero.
scheme (de.timesteppers): The kind of solver to use. Options are
RK443 (de.timesteppers.RK443), RK111, RK222, RKSMR, etc.
sim_time (float): The maximum amount of simulation time allowed
(in seconds) before ending the simulation.
wall_time (float): The maximum amound of computing time allowed
(in seconds) before ending the simulation.
tight (bool): If True, set a low cadence and min_dt for refined
simulation. If False, set a higher cadence and min_dt for a
more coarse (but faster) simulation.
"""
if not self.problem:
raise TypeError("problem not initialized (run setup())")
self.logger.debug("\n")
self.logger.debug("NEW SIMULATION")
solver = self.problem.build_solver(scheme)
self.logger.info("Solver built")
N = int(self.problem.parameters['N'])
# Initial conditions --------------------------------------------------
if initial_conditions is None: # "Trivial" conditions.
dt = 1e-4
eps = 1e-4
k = 3.117
x, z = self.problem.domain.grids(scales=1)
T, T_ = solver.state['T'], solver.state['T_']
# Start T from rest plus a small perturbation.
T['g'] = 1 - z + eps * np.sin(k * x) * np.sin(2 * np.pi * z)
T.differentiate('z', out=solver.state['Tz'])
# Start T_ from rest.
T_['g'] = 1 - z
T_.differentiate('z', out=solver.state['Tz_'])
self.logger.info("Using trivial initial conditions")
elif isinstance(initial_conditions, str): # Load data from a file.
# Resume: load the state of the last (merged) state file.
resume = initial_conditions == "resume"
if resume:
initial_conditions = self._get_merged_file("states")
if not initial_conditions.endswith(".h5"):
raise ValueError(
"'{}' is not an h5 file".format(initial_conditions))
# Load the data from the specified h5 file into the system.
self.logger.info("Loading initial conditions from {}".format(
initial_conditions))
with h5py.File(initial_conditions, 'r') as infile:
dt = infile["scales/timestep"][-1] * .01 # initial dt
errs = []
tasks = ["T", "Tz", "psi", "psiz", "zeta", "zetaz"]
if resume: # Only load assimilating variables to resume.
tasks += ["T_", "Tz_", "psi_", "psiz_", "zeta_", "zetaz_"]
solver.sim_time = infile["scales/sim_time"][-1]
niters = infile["scales/iteration"][-1]
solver.initial_iteration = niters
solver.iteration = niters
for name in tasks:
# Get task data from the h5 file (recording failures).
try:
data = infile["tasks/" + name][-1, :, :]
except KeyError as e:
errs.append("tasks/" + name)
continue
# Determine the chunk belonging to this process.
chunk = data.shape[1] // SIZE
subset = data[:, RANK * chunk:(RANK + 1) * chunk]
# Change the corresponding state variable.
scale = solver.state[name]['g'].shape[0] / \
self.problem.parameters["xsize"]
solver.state[name].set_scales(1)
solver.state[name]['g'] = subset
solver.state[name].set_scales(scale)
if errs:
raise KeyError("Missing keys in '{}': '{}'".format(
initial_conditions, "', '".join(errs)))
# Initial conditions for assimilating system: T_0 = P_4(T0).
if not resume:
G = self.problem.domain.new_field()
G['c'] = solver.state['T']['c'].copy()
solver.state['T_']['g'] = BoussinesqDataAssimilation2D.P_N(
G, 4, True)
solver.state['T_'].differentiate('z', out=solver.state['Tz_'])
# Driving / projection function arguments -----------------------------
dT = solver.state['T_'] - solver.state['T']
self.problem.parameters["driving"].args = [dT, N]
self.problem.parameters["driving"].original_args = [dT, N]
# Stopping Parameters -------------------------------------------------
solver.stop_sim_time = sim_time # Length of simulation.
solver.stop_wall_time = wall_time # Real time allowed to compute.
solver.stop_iteration = np.inf # Maximum iterations allowed.
# State snapshots -----------------------------------------------------
# Save the entire state in states/ files. USE iter, NOT sim_dt.
# NOTE: This is where BoussinesqDataAssimilation2Dmovie differs.
snaps = solver.evaluator.add_file_handler(os.path.join(
self.records_dir, "states"),
iter=1,
max_writes=5000,
mode="append")
snaps.add_task("T")
snaps.add_task("T_")
snaps.add_task("driving", name="P_N")
# Control Flow --------------------------------------------------------
if tight: # Tighter control flow (slower but safer).
cfl = flow_tools.CFL(solver,
initial_dt=dt,
cadence=1,
safety=1,
max_change=1.5,
min_change=0.01,
max_dt=0.01,
min_dt=1e-10)
else: # Looser control flow (faster but risky).
cfl = flow_tools.CFL(solver,
initial_dt=dt,
cadence=10,
safety=1,
max_change=1.5,
min_change=0.5,
max_dt=0.01,
min_dt=1e-6)
cfl.add_velocities(('v', 'w'))
cfl.add_velocities(('v_', 'w_'))
# Flow properties (print during run; not recorded in the records files)
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("sqrt(v **2 + w **2) / Ra", name='Re')
flow.add_property("sqrt(v_**2 + w_**2) / Ra", name='Re_')
# MAIN COMPUTATION LOOP -----------------------------------------------
try:
self.logger.info("Starting main loop")
start_time = time.time()
while solver.ok:
# Recompute time step and iterate.
dt = cfl.compute_dt()
dt = solver.step(dt)
# Print info to the screen every 10 iterations.
if solver.iteration % 10 == 0:
info = "Iteration {:>5d}, Time: {:.7f}, dt: {:.2e}".format(
solver.iteration, solver.sim_time, dt)
Re = flow.max("Re")
Re_ = flow.max("Re_")
info += ", Max Re = {:f}".format(Re)
info += ", Max Re_ = {:f}".format(Re_)
self.logger.info(info)
# Make sure the simulation hasn't blown up.
if np.isnan(Re) or np.isnan(Re_):
raise ValueError("Reynolds number went to infinity!!"
"\nRe = {}, Re_ = {}".format(Re, Re_))
except BaseException as e:
self.logger.error("Exception raised, triggering end of main loop.")
raise
finally:
total_time = time.time() - start_time
cpu_hr = total_time / 60 / 60 * SIZE
self.logger.info("Iterations: {:d}".format(solver.iteration))
self.logger.info("Sim end time: {:.3e}".format(solver.sim_time))
self.logger.info("Run time: {:.3e} sec".format(total_time))
self.logger.info("Run time: {:.3e} cpu-hr".format(cpu_hr))
self.logger.debug("END OF SIMULATION\n")
model.set_tracer_gradient_bc("b", "bottom", gradient="initial_N2")
model.build_solver(timestepper='SBDF3')
# Initial condition
noise = noise_amp * dedaLES.random_noise(model.domain) * model.z * (Lz - model.z) / Lz**2
initial_b = initial_N2*model.z
model.set_fields(
u = noise,
v = noise,
b = initial_b + np.sqrt(initial_N2) * noise
#w = noise,
)
model.stop_at(sim_time=run_time)
CFL = flow_tools.CFL(model.solver, initial_dt=initial_dt, cadence=CFL_cadence, max_change=1.5, safety=0.5)
CFL.add_velocities(('u', 'v', 'w'))
stats = flow_tools.GlobalFlowProperty(model.solver, cadence=stats_cadence)
stats.add_property("w*b", name="buoyancyflux")
stats.add_property("ε + ε_sgs", name="epsilon")
stats.add_property("χ + χ_sgs", name="chi")
stats.add_property("w*w", name="wsquared")
stats.add_property("sqrt(u*u + v*v + w*w) / ν", name='Re')
analysis = model.solver.evaluator.add_file_handler(identifier(model, closure=closure), iter=analysis_cadence, max_writes=max_writes)
analysis.add_system(model.solver.state, layout='g')
#visualization
analysis.add_task("interp(b, y=0)", scales=1, name='b midplane')
analysis.add_task("interp(u, y=0)", scales=1, name='u midplane')
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("vol_avg(rho_fluc*phi)", name='PE_fluc')
flow.add_property("Re_rms", name='Re')
flow.add_property("Ma_rms", name='Ma')
flow.add_property("integ(rho_full*s1*Cp)", name='integ_s')
flow.add_property(
"((dz(T_full)/T_full)**2 - (dz(T0)/T0)**2)",
name='dissipation') # good for seeing if the solution isn't resolved.
################################################
# CFL and sim stop conditions
CFL = flow_tools.CFL(solver,
initial_dt=dt,
cadence=1,
safety=safety * float(args['--safety_factor']),
max_change=1.5,
min_change=0.5,
max_dt=max_dt,
threshold=0.05)
CFL.add_velocities(('u', 'v', 'w'))
if N_buoyancy is None:
solver.stop_sim_time = buoyancy_time * 50 + solver.initial_sim_time
else:
solver.stop_sim_time = N_buoyancy * buoyancy_time + solver.initial_sim_time
solver.stop_wall_time = float(args['--run_time']) * 3600
solver.stop_iteration = np.inf + solver.initial_iteration
good_solution = True
Hermitian_cadence = 100
u.differentiate('x', out=ux)
v.differentiate('y', out=vy)
v.differentiate('x', out=vx)
u.differentiate('y', out=uy)
h.differentiate('x', out=hx)
h.differentiate('y', out=hy)
solver.stop_sim_time = 100 * 86400
solver.stop_wall_time = np.inf
solver.stop_iteration = np.inf
initial_dt = dt_max
#print(initial_dt)
cfl = flow_tools.CFL(solver,
initial_dt=initial_dt,
cadence=10,
safety=CFLfac,
max_change=1.5,
min_change=0.5,
max_dt=dt_max,
threshold=0.5)
#initial_dt = 0.2*Lx/nx
#cfl = flow_tools.CFL(solver,initial_dt=initial_dt,cadence=10,safety=0.5,threshold=0.5)
cfl.add_velocities(('u', 'v'))
#analysis = solver.evaluator.add_file_handler(exp_name, sim_dt=900, max_writes=300)
analysis = solver.evaluator.add_file_handler(exp_name, iter=1, max_writes=300)
analysis.add_task('h', layout='g')
analysis.add_task('u', layout='g')
analysis.add_task('v', layout='g')
analysis.add_task('y', layout='g')
analysis.add_task('x', layout='g')
#analysis.add_task('Q',layout='g')
def set_up_cfl_fine(self, **kwargs):
"""Helper function to set up dedalus CFL, include all
arguments except solver, which will be the coarse solver"""
self.fine_CFL = flow_tools.CFL(self.fine_solver, **kwargs)
self.dt_fine = kwargs['initial_dt']
name='vs') # Save in horizontal wavenumber, vertical physical space
snap.add_task('w', layout=domain.dist.layouts[1], scales=1, name='ws')
snap.add_task(
'u', layout=domain.dist.layouts[1], scales=1,
name='us') # Save in horizontal wavenumber, vertical physical space
snap.add_task('b', layout=domain.dist.layouts[1], scales=1, name='bs')
snap.add_task('NL(u, uz)', layout=domain.dist.layouts[1], scales=1, name='NLu')
snap.add_task('NL(v, vz)', layout=domain.dist.layouts[1], scales=1, name='NLv')
snap.add_task('NL(w, wz)', layout=domain.dist.layouts[1], scales=1, name='NLw')
# CFL
CFL = flow_tools.CFL(solver,
initial_dt=1e-4,
cadence=10,
safety=0.8,
max_change=1.5,
min_change=0,
max_dt=30,
threshold=0.05) #Adjust these for your problem
CFL.add_velocities(('u', 'w'))
# Flow properties
# Main loop
end_init_time = time.time()
logger.info('Initialization time: %f' % (end_init_time - start_init_time))
try:
logger.info('Starting loop')
start_run_time = time.time()
while solver.ok:
dt = CFL.compute_dt()
ts = de.SBDF2
cfl_safety_factor = 0.3
else:
ts = de.RK443
cfl_safety_factor = 0.4
if args['--safety']:
cfl_safety_factor = float(args['--safety'])
solver = problem.build_solver(ts)
solver.stop_iteration = run_time_iter
Δt = max_Δt = float(args['--max_dt'])
cfl = flow_tools.CFL(solver,
Δt,
safety=cfl_safety_factor,
cadence=1,
threshold=0.1,
max_change=1.5,
min_change=0.5,
max_dt=max_Δt)
cfl.add_velocity(u)
ρ = np.exp(Υ0 + Υ).evaluate()
h = cP * (T + T0)
KE = 0.5 * ρ * dot(u, u)
IE = cP * Ma2 * h * (s + s0)
Re = (ρ / μ) * np.sqrt(dot(u, u))
ω = -div(skew(u))
KE.store_last = True
IE.store_last = True
Re.store_last = True
ω.store_last = True
max_dt = np.min((0.1 * true_t_ff, 1))
if dt is None: dt = max_dt
analysis_tasks = initialize_rotating_output(solver,
data_dir,
aspect,
output_dt=0.1 * true_t_ff,
slice_output_dt=1 * true_t_ff,
vol_output_dt=10 * true_t_ff,
mode=mode,
volumes=True)
# CFL
CFL = flow_tools.CFL(solver,
initial_dt=dt,
cadence=1,
safety=cfl_safety,
max_change=1.5,
min_change=0.5,
max_dt=max_dt,
threshold=0.1)
CFL.add_velocities(('u', 'v', 'w'))
### 8. Setup flow tracking for terminal output, including rolling averages
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("Re", name='Re')
flow.add_property("Ro", name='Ro')
flow.add_property("true_Ro", name='true_Ro')
flow.add_property("vel_rms**2/2", name='KE')
flow.add_property("Nu", name='Nu')
flow.add_property("-1 + (left(T1_z) + right(T1_z) ) / 2", name='Tz_excess')
flow.add_property("T0+T1", name='T')
# Integration parameters
solver.stop_sim_time = 100
solver.stop_wall_time = 60 * 60.
solver.stop_iteration = np.inf
# Analysis
snap = solver.evaluator.add_file_handler('snapshots', sim_dt=0.2, max_writes=10)
snap.add_task("interp(p, z=0)", scales=1, name='p midplane')
snap.add_task("interp(b, z=0)", scales=1, name='b midplane')
snap.add_task("interp(u, z=0)", scales=1, name='u midplane')
snap.add_task("interp(v, z=0)", scales=1, name='v midplane')
snap.add_task("interp(w, z=0)", scales=1, name='w midplane')
snap.add_task("integ(b, 'z')", name='b integral x4', scales=4)
# CFL
CFL = flow_tools.CFL(solver, initial_dt=1e-4, cadence=5, safety=0.5,
max_change=1.5, min_change=0.5, max_dt=0.05)
CFL.add_velocities(('u', 'v', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=10)
flow.add_property("sqrt(u*u + v*v + w*w) / R", name='Re')
# Main loop
end_init_time = time.time()
logger.info('Initialization time: %f' %(end_init_time-start_init_time))
try:
logger.info('Starting loop')
start_run_time = time.time()
while solver.ok:
dt = CFL.compute_dt()
solver.step(dt)
def run(subs):
print("\n\nBegin Linearized Pure Hydro, Const/Static Background\n")
print("Using substitutions: ", subs)
print()
# set lowest level of all loggers to INFO, i.e. dont follow DEBUG stuff
root = logging.root
for h in root.handlers:
h.setLevel("INFO")
# name logger using the module name
logger = logging.getLogger(__name__)
# Input parameters
gamma = 1.4
kappa = 1.e-4
c_v = 1.
k_z = 3.13 # vertical wavenumber
R_in = 1.
R_out = 3.
# background is const
rho0 = 5.
T0 = 10.
p0 = (gamma - 1) * c_v * rho0 * T0
height = 2. * numpy.pi / k_z
# Problem Domain
r_basis = de.Chebyshev('r', 32, interval=(R_in, R_out), dealias=3 / 2)
z_basis = de.Fourier('z', 32, interval=(0., height), dealias=3 / 2)
th_basis = de.Fourier('th',
32,
interval=(0., 2. * numpy.pi),
dealias=3 / 2)
domain = de.Domain([r_basis, th_basis, z_basis], grid_dtype=numpy.float64)
# alias the grids
z = domain.grid(2, scales=domain.dealias)
th = domain.grid(1, scales=domain.dealias)
r = domain.grid(0, scales=domain.dealias)
# Equations
# incompressible, axisymmetric, Navier-Stokes in Cylindrical
# expand the non-constant-coefficients (NCC) to precision of 1.e-8
TC = de.EVP(domain,
variables=['rho', 'p', 'T', 'u', 'v', 'w', 'Tr'],
ncc_cutoff=1.e-8,
eigenvalue='omega')
TC.parameters['gamma'] = gamma
TC.parameters['kappa'] = kappa
TC.parameters['p0'] = p0
TC.parameters['T0'] = T0
TC.parameters['rho0'] = rho0
# multiply equations through by r**2 or r to avoid 1/r**2 & 1/r terms
if (subs):
# substitute for dt()
TC.substitutions['dt(A)'] = '-1.j*omega*A'
# r*div(U)
TC.substitutions['r_div(A,B,C)'] = 'A + r*dr(A) + dth(B) + r*dz(C)'
# r-component of gradient
TC.substitutions['grad_1(A)'] = 'dr(A)'
# th-component of gradient
TC.substitutions['r_grad_2(A)'] = 'r*dth(A)'
# z-component of gradient
TC.substitutions['grad_3(A)'] = 'dz(A)'
# r*r*Laplacian(scalar)
TC.substitutions['r2_Lap(f, fr)'] = \
'r*r*dr(fr) + r*dr(f) + dth(f) + r*r*dz(dz(f))'
# equations using substituions
TC.add_equation("p/p0 - rho/rho0 - T/T0 = 0")
TC.add_equation("r*r*dt(p) + gamma*p0*r*r_div(u,v,w) - " + \
"(gamma-1)*kappa*r2_Lap(T, Tr) = 0")
TC.add_equation("r*dt(rho) + rho0*r_div(u,v,w) = 0")
TC.add_equation("rho0*dt(u) + grad_1(p) = 0")
TC.add_equation("r*rho0*dt(v) + r_grad_2(p) = 0")
TC.add_equation("rho0*dt(w) + grad_3(p) = 0")
TC.add_equation("Tr - dr(T) = 0")
else:
print("\nNon-Substitutions Version is not coded up\n")
sys.exit(2)
# Boundary Conditions
TC.add_bc("left(u) = 0")
TC.add_bc("left(v) = 0")
TC.add_bc("left(w) = 0")
TC.add_bc("right(u) = 0", condition="nz != 0")
TC.add_bc("right(v) = 0")
TC.add_bc("right(w) = 0")
TC.add_bc("integ(p,'r') = 0", condition="nz == 0")
###############################################################
# Force break since the following has not been edited yet...
###############################################################
print("\nThe rest of the script needs to be changed still\n")
sys.exit(2)
# Timestepper
dt = max_dt = 1.
Omega_1 = TC.parameters['V_l'] / R_in
period = 2. * numpy.pi / Omega_1
logger.info('Period: %f' % (period))
ts = de.timesteppers.RK443
IVP = TC.build_solver(ts)
IVP.stop_sim_time = 15. * period
IVP.stop_wall_time = numpy.inf
IVP.stop_iteration = 10000000
# alias the state variables
p = IVP.state['p']
u = IVP.state['u']
v = IVP.state['v']
w = IVP.state['w']
ur = IVP.state['ur']
vr = IVP.state['vr']
wr = IVP.state['wr']
# new field
phi = Field(domain, name='phi')
for f in [phi, p, u, v, w, ur, vr, wr]:
f.set_scales(domain.dealias, keep_data=False)
v['g'] = v_analytic
#p['g'] = p_analytic
v.differentiate(1, vr)
# incompressible perturbation, arbitrary vorticity
phi['g'] = 1.e-3 * numpy.random.randn(*v['g'].shape)
phi.differentiate(1, u)
u['g'] *= -1. * numpy.sin(numpy.pi * (r - R_in))
phi.differentiate(0, w)
w['g'] *= numpy.sin(numpy.pi * (r - R_in))
u.differentiate(1, ur)
w.differentiate(1, wr)
# Time step size
CFL = flow_tools.CFL(IVP,
initial_dt=1.e-3,
cadence=5,
safety=0.3,
max_change=1.5,
min_change=0.5)
CFL.add_velocities(('u', 'w'))
# Analysis
# integrated energy every 10 steps
analysis1 = IVP.evaluator.add_file_handler("scalar_data", iter=10)
analysis1.add_task("integ(0.5 * (u*u + v*v + w*w))", name="total KE")
analysis1.add_task("integ(0.5 * (u*u + w*w))", name="meridional KE")
analysis1.add_task("integ((u*u)**0.5)", name="u_rms")
analysis1.add_task("integ((w*w)**0.5)", name="w_rms")
# Snapshots every half inner rotation period
analysis2 = IVP.evaluator.add_file_handler('snapshots',
sim_dt=0.5 * period,
max_size=2**30)
analysis2.add_system(IVP.state, layout='g')
# Radial profiles every 100 steps
analysis3 = IVP.evaluator.add_file_handler("radial_profiles", iter=100)
analysis3.add_task("integ(r*v, 'z')", name="Angular Momentum")
# MAIN LOOP
dt = CFL.compute_dt()
start_time = time.time()
while IVP.ok:
IVP.step(dt)
if (IVP.iteration % 10 == 0):
logger.info('Iteration: %i, Time: %e, dt: %e' % \
(IVP.iteration, IVP.sim_time, dt))
dt = CFL.compute_dt()
end_time = time.time()
logger.info('Total time: %f' % (end_time - start_time))
logger.info('Iterations: %i' % (IVP.iteration))
logger.info('Average timestep: %f' % (IVP.sim_time / IVP.iteration))
logger.info('Period: %f' % (period))
logger.info('\n\tSimulation Complete\n')
return
u.differentiate('r', out=ur)
v.differentiate('r', out=vr)
w.differentiate('r', out=wr)
#Setting Simulation Runtime
omega1 = 1 / eta - 1.
period = 2 * np.pi / omega1
solver.stop_sim_time = 6 * period
solver.stop_wall_time = 24 * 3600. #np.inf
solver.stop_iteration = 2000
#CFL stuff
CFL = flow_tools.CFL(solver,
initial_dt=1e-2,
cadence=5,
safety=0.3,
max_change=1.5,
min_change=0.5,
max_dt=1,
threshold=0.1)
CFL.add_velocities(('u', 'v', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=10)
flow.add_property("abs(DivU)", name='divu')
flow.add_property("integ(r*KE)", name='KE')
flow.add_property("integ(r*enstrophy)", name='enstrophy')
dt = CFL.compute_dt()
# Main loop
geo_factor = 1
# initial conditions
x, z = domain.grids(scales=1)
u = solver.state['u']
w = solver.state['w']
p = solver.state['p']
B = solver.state['B'] # zero for everything
solver.stop_sim_time = stop_time
solver.stop_wall_time = np.inf
solver.stop_iteration = np.inf
# CFL conditions
initial_dt = 0.8 * Lz / nz
cfl = flow_tools.CFL(solver,
initial_dt,
safety=0.8,
max_change=1.5,
min_change=0.5,
max_dt=400)
# too large of a timestep makes things rather diffusive
cfl.add_velocities(('u', 'w'))
# analysis
# fields to record
analysis = solver.evaluator.add_file_handler(sim_name,
sim_dt=50,
max_writes=300)
analysis.add_task('B', name='buoyancy')
analysis.add_task('w', name='vertical velocity')
analysis.add_task('u', name='horizontal velocity')
analysis.add_task('p', name='pressure')
analysis.add_task('Nsq')
sim_dt=param.scalar_sim_dt,
max_writes=100)
an2.add_task('integ(((u1+u2)**2 + (w1+w2)**2) /(a0+a1+a2)/2)',
name='KE',
layout='g')
an2.add_task('integ((u2*u2+w2*w2)/(a0+a1+a2)/2)', name='KE_pert', layout='g')
an2.add_task('-integ(u2*dx(t2xx) + u2*dz(t2xz) + w2*dx(t2xz) + w2*dz(t2zz))',
name='D',
layout='g')
# Monitoring
flow = flow_tools.GlobalFlowProperty(solver, cadence=param.CFL['cadence'])
flow.add_property('integ((u2*u2+w2*w2)/(a0+a1+a2)/2)', name='KE_pert')
# CFL calculator
CFL = flow_tools.CFL(solver, **param.CFL)
CFL.add_velocities(('u1+u2', 'w1+w2'))
# Main loop
dt_floor = 2**(np.log2(param.CFL['min_dt']) // 1)
try:
logger.info('Starting loop')
start_time = time.time()
while solver.ok:
dt = CFL.compute_dt()
dt = dt_floor * (dt // dt_floor)
dt = solver.step(dt, trim=param.trim)
if (solver.iteration - 1) % param.CFL['cadence'] == 0:
logger.info('Iteration: %i, Time: %e, dt: %e' %
(solver.iteration, solver.sim_time, dt))
logger.info('Ave KE = %e' % flow.max('KE_pert'))
# snapshots.add_task('w')
# snapshots.add_task('u')
profiles = solver.evaluator.add_file_handler('profiles_nom', sim_dt=0.1, max_writes=100, mode=fh_mode)
profiles.add_task("P*integ(dz(b) - 1, 'x') / Lx", name='diffusive_flux')
profiles.add_task("integ(w*b, 'x') / Lx", name='adv_flux')
profiles.add_task("integ(integ(w**2 + u**2, 'x'), 'z') / Lx", name='ke')
profiles.add_task("integ(integ((dx(w) - uz)**2, 'x'), 'z') / Lx", name='enst')
profiles.add_task("integ(((P*integ(dz(b) - 1, 'x') / Lx) - (integ(w*b, 'x') / Lx) ) / (P*integ(dz(b) - 1, 'x') / Lx), 'z')", name='nu')
checkpoints = solver.evaluator.add_file_handler('checkpoints_nom', sim_dt=20, max_writes=50, mode=fh_mode)
checkpoints.add_system(solver.state)
# snapshots.add_task('b')
# CFL
CFL = flow_tools.CFL(solver, initial_dt=dt, cadence=10, safety=0.5,
max_change=1.5, min_change=0.5, max_dt=0.125, threshold=0.05)
CFL.add_velocities(('u', 'w'))
# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=10)
flow.add_property("sqrt(u*u + w*w) / R", name='Re')
# Main loop
try:
logger.info('Starting loop')
start_time = time.time()
while solver.proceed:
dt = CFL.compute_dt()
dt = solver.step(dt)
if (solver.iteration-1) % 10 == 0:
logger.info('Iteration: %i, Time: %e, dt: %e' %(solver.iteration, solver.sim_time, dt))
psi = solver.state['psi']
psi['g'] = init_solver.state['init_psi']['g']
# plt.imshow(init_vorticity['g'])
# plt.show()
# plt.imshow(psi['g'])
# plt.show()
# Now you are ready to go.
# Anytime you ask for zeta, u, or v they will be non-zero because psy is non-zero.
cfl = flow_tools.CFL(solver,
initial_dt=dt,
cadence=10,
safety=.5,
max_change=1.5,
min_change=0.5)
cfl.add_velocities(('u', 'v'))
dout = solver.evaluator.add_dictionary_handler(iter=10)
fout = solver.evaluator.add_file_handler('analysis_tasks', sim_dt=1e-3)
outputs = [fout]
if PLOTTING:
outputs.append(dout)
for output in outputs:
output.add_system(solver.state)
output.add_task('zeta', scales=1, name='zeta')
output.add_task('u', scales=1, name='u')
