# Parameters
Lx, Lz = (2.*math.sqrt(2), 1.)
Prandtl = 6.8
Rayleigh = 5e3
# Create bases and domain
x_basis = de.Fourier('x', 128, interval=(0, Lx))
z_basis = de.Chebyshev('z', 128, interval=(0, Lz))
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','vort'])
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)) - w = -(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_equation("uz - dx(w) + vort = 0")
problem.add_bc("left(b) = 0")
problem.add_bc("left(u) = 0")
problem.add_bc("left(w) = 0")
problem.add_bc("right(b) = 0")
problem.add_bc("right(u) = 0")
problem.add_bc("right(w) = 0", condition="(nx != 0)")
problem.add_bc("integ(p, 'z') = 0", condition="(nx == 0)")
# Build solver
solver = problem.build_solver(de.timesteppers.RK443)
logger.info('Solver built')
# Initial conditions
x = domain.grid(0)
z = domain.grid(1)
b = solver.state['b']
bz = solver.state['bz']
u = solver.state['u']
w = solver.state['w']
vort = solver.state['vort']
# 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=42)
noise = rand.standard_normal(gshape)[slices]
# Linear background + perturbations damped at walls
#zb, zt = z_basis.interval
bpert = noise
#bpert = rd.array_gen_from_file('tmpr_init.dat')
b['g'] = bpert
u['g'] = rd.array_gen_from_file('uvel_init.dat')
w['g'] = rd.array_gen_from_file('wvel_init.dat')
# b.differentiate('z', out=bz)
# Integration parameters
solver.stop_sim_time = 2000
solver.stop_wall_time = 20000 * 60.
solver.stop_iteration = np.inf
# CFL
CFL = flow_tools.CFL(solver, initial_dt=2.5e-3, cadence=5, safety=0.1,
max_change=1.5, min_change=0.5, max_dt=0.0025)
CFL.add_velocities(('u', 'w'))