import numpy as np
import dedalus.public as d3
import logging
import matplotlib.pyplot as plt
import h5py
from dedalus.extras.plot_tools import plot_bot_2d
logger = logging.getLogger(__name__)
# Parameters
Lx, Ly = (10, 20*np.pi)
Nx, Ny = (128, 121)
dealias = 3/2
stop_sim_time = 30
# timestepper = d3.RK222
timestepper = d3.RKSMR
max_timestep = 1e-2
dtype = np.float64
# dtype = np.complex128
# Bases
coords = d3.CartesianCoordinates('x', 'y')
dist = d3.Distributor(coords, dtype=dtype)
xbasis = d3.RealFourier(coords['x'], size=Nx, bounds=(-Lx/2, Lx/2), dealias=dealias)
# ybasis = d3.RealFourier(coords['y'], size=Ny, bounds=(-Ly/2, Ly/2), dealias=dealias)
ybasis = d3.Chebyshev(coords['y'], size=Ny, bounds=(-Ly/2, Ly/2), dealias=dealias)
# Fields
x, y = dist.local_grids(xbasis, ybasis)
ex, ey = coords.unit_vector_fields(dist)
# u_vec = dist.VectorField(coords, name='u_vec', bases=(xbasis,ybasis))
u = dist.Field(name='u', bases=(xbasis,ybasis))
v = dist.Field(name='v', bases=(xbasis,ybasis))
eta = dist.Field(name='eta', bases=(xbasis,ybasis))
# f = dist.Field(name='f', bases=(xbasis,ybasis))
f = dist.Field(name='f', bases=(ybasis))
omega = dist.Field(name='omega')
# f['g'] = np.ones(Ny)*np.sin(2*np.pi*y/Ly)
f['g'] = np.sin(2*np.pi*y/Ly)
tau1 = dist.Field(name='tau1',bases=xbasis)
tau2 = dist.Field(name='tau2',bases=xbasis)
tau3 = dist.Field(name='tau3',bases=xbasis)
tau4 = dist.Field(name='tau4',bases=xbasis)
tau5 = dist.Field(name='tau5',bases=xbasis)
lift_basis = ybasis.derivative_basis(1) # Chebyshev U basis
lift = lambda A: d3.Lift(A, lift_basis, -1) # Shortcut for multiplying by U_{N-1}(y)
# Substitutions
dx = lambda A: d3.Differentiate(A, coords['x'])
dy = lambda A: d3.Differentiate(A, coords['y'])
# Problem
problem = d3.IVP([u, v, eta, tau1, tau2,tau3,tau4,tau5], namespace=locals())
problem.add_equation("dt(u) + dx(eta) - f*v + 0.00*lap(u) + lift(tau4) + lift(tau5) = 0")
problem.add_equation("dt(v) + dy(eta) + f*u - 0.000*lap(v) + lift(tau1) + lift(tau3) + lift(tau2)= 0")
problem.add_equation("dt(eta) + dx(u) + dy(v) + lift(tau2) = 0")
problem.add_equation("v(y=-Ly/2) = 0")
problem.add_equation("v(y=Ly/2) = 0")
problem.add_equation("dy(v)(y=-Ly/2) = 0")
problem.add_equation("dy(v)(y=Ly/2) = 0")
problem.add_equation("u(y=-Ly/2) = u(y=Ly/2)")
# problem = d3.IVP([u, v, eta], namespace=locals())
# problem.add_equation("dt(u) + dx(eta) - f*v = 0")
# problem.add_equation("dt(v) + dy(eta) + f*u = 0")
# problem.add_equation("dt(eta) + dx(u) + dy(v) = 0")
# Initial Conditions
# umode = np.load('umode.npy')
# vmode = np.load('vmode.npy')
# hmode = np.load('hmode.npy')
hmode = np.random.rand(1, 121)
umode = np.random.rand(1, 121)
vmode = np.random.rand(1, 121)
kx = 4*np.pi/Lx
u['g'] = umode[0].real*np.cos(kx*x)
v['g'] = vmode[0].real*np.cos(kx*x)
eta['g'] = hmode[0].real*np.cos(kx*x)
# print(np.shape(hmode[0]))
# print(np.shape(eta['g']))
# print(np.shape(u['g']))
# print(u['g'])
# Solver
solver = problem.build_solver(timestepper)
solver.stop_sim_time = stop_sim_time
# Analysis
snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.1, max_writes=10)
snapshots.add_task(eta, name='eta')
snapshots.add_task(u, name='u')
snapshots.add_task(v, name='v')
# CFL
CFL = d3.CFL(solver, initial_dt=max_timestep, cadence=10, safety=0.5, threshold=0.05,
max_change=1.5, min_change=0.5, max_dt=max_timestep)
CFL.add_velocity(u*ex + v*ey)
# Main loop
try:
logger.info('Starting main loop')
while solver.proceed:
timestep = CFL.compute_timestep()
solver.step(timestep)
if (solver.iteration-1) % 10 == 0:
# max_w = np.sqrt(flow.max('w2'))
logger.info('Iteration=%i, Time=%e, dt=%e' %(solver.iteration, solver.sim_time, timestep))
except:
logger.error('Exception raised, triggering end of main loop.')
raise
finally:
solver.log_stats()