# Store data for final plot
u.set_scales(1, keep_data=True)
u_list = [np.copy(u['g'])]
t_list = [solver.sim_time]
# Main loop
dt = 2e-3
while solver.ok:
solver.step(dt)
if solver.iteration % 20 == 0:
u.set_scales(1, keep_data=True)
u_list.append(np.copy(u['g']))
t_list.append(solver.sim_time)
if solver.iteration % 100 == 0:
logger.info('Iteration: %i, Time: %e, dt: %e' %(solver.iteration, solver.sim_time, dt))
# Create space-time plot
u_array = np.array(u_list)
t_array = np.array(t_list)
xmesh, ymesh = quad_mesh(x=x, y=t_array)
plt.figure()
plt.pcolormesh(xmesh, ymesh, u_array, cmap='RdBu_r')
plt.axis(pad_limits(xmesh, ymesh))
plt.colorbar()
plt.xlabel('x')
plt.ylabel('t')
plt.title('KdV-Burgers, (a,b)=(%g,%g)' %(problem.parameters['a'], problem.parameters['b']))
plt.savefig('kdv_burgers.png')
n_rows = 3
xmesh, zmesh = quad_mesh(x=x[:, 0], y=z[0])
for t_idx, sim_time in list(enumerate(sim_times[:-1])):
fig = plt.figure(dpi=200)
idx = 1
for var in mesh_vars:
axes = fig.add_subplot(n_rows, n_cols, idx, title=var)
var_dat = state_vars[var]
p = axes.pcolormesh(xmesh,
zmesh,
var_dat[t_idx].T,
vmin=var_dat.min(),
vmax=var_dat.max())
axes.axis(pad_limits(xmesh, zmesh))
cb = fig.colorbar(p, ax=axes)
cb.ax.set_yticklabels(cb.ax.get_yticklabels(), rotation=30)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
idx += 1
for var in z_vars + slice_vars:
axes = fig.add_subplot(n_rows, n_cols, idx, title=var)
z_pts = (zmesh[1:, 0] + zmesh[:-1, 0]) / 2
if var in slice_vars:
var_dat = state_vars[var][:, 0, :]
p = axes.plot(get_analytical_sponge(var, z_pts, sim_time),
z_pts)
else:
var_dat = state_vars[var]
def plot(name, params):
slice_suffix = '(x=0)'
SAVE_FMT_STR = 't_%d.png'
snapshots_dir = SNAPSHOTS_DIR % name
path = '{s}/{s}_s1'.format(s=snapshots_dir)
matplotlib.rcParams.update({'font.size': 6})
plot_vars = ['uz']
c_vars = ['uz_c']
f_vars = ['uz_f']
# z_vars = ['F_z', 'E'] # sum these over x
z_vars = []
slice_vars = ['%s%s' % (i, slice_suffix) for i in ['uz']]
n_cols = 4
n_rows = 1
plot_stride = 1
if os.path.exists('%s.mp4' % name):
print('%s.mp4 already exists, not regenerating' % name)
return
sim_times, domain, state_vars = load(name, params)
x = domain.grid(0, scales=params['INTERP_X'])
z = domain.grid(1, scales=params['INTERP_Z'])
xmesh, zmesh = quad_mesh(x=x[:, 0], y=z[0])
for var in z_vars:
state_vars[var] = np.sum(state_vars[var], axis=1)
for var in slice_vars:
state_vars[var] = state_vars[var.replace(slice_suffix, '')][:, 0, :]
uz_est = params['F'] * get_uz_f_ratio(params)
for t_idx, sim_time in list(enumerate(sim_times))[::plot_stride]:
fig = plt.figure(dpi=200)
idx = 1
for var in plot_vars:
axes = fig.add_subplot(n_rows, n_cols, idx, title=var)
var_dat = state_vars[var]
p = axes.pcolormesh(xmesh, zmesh, var_dat[t_idx].T)
# vmin=var_dat.min(), vmax=var_dat.max())
axes.axis(pad_limits(xmesh, zmesh))
cb = fig.colorbar(p, ax=axes)
cb.ax.set_yticklabels(cb.ax.get_yticklabels(), rotation=30)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
idx += 1
for var in z_vars + slice_vars:
axes = fig.add_subplot(n_rows, n_cols, idx, title=var)
var_dat = state_vars[var]
z_pts = (zmesh[1:, 0] + zmesh[:-1, 0]) / 2
p = axes.plot(var_dat[t_idx], z_pts, linewidth=0.5)
if var == 'uz%s' % slice_suffix:
p = axes.plot(uz_est * np.exp(
(z_pts - params['Z0']) / (2 * params['H'])),
z_pts,
'orange',
linewidth=0.5)
p = axes.plot(-uz_est * np.exp(
(z_pts - params['Z0']) / (2 * params['H'])),
z_pts,
'orange',
linewidth=0.5)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
xlims = [var_dat.min(), var_dat.max()]
axes.set_xlim(*xlims)
p = axes.plot(xlims, [params['SPONGE_LOW']] * len(xlims),
'r:',
linewidth=0.5)
p = axes.plot(xlims, [params['SPONGE_HIGH']] * len(xlims),
'r:',
linewidth=0.5)
p = axes.plot(xlims, [params['Z0'] + 3 * params['S']] * len(xlims),
'b--',
linewidth=0.5)
p = axes.plot(xlims, [params['Z0'] - 3 * params['S']] * len(xlims),
'b--',
linewidth=0.5)
idx += 1
for var in c_vars:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title='%s (kx=kx_d)' % var)
var_dat = state_vars[var]
kx_idx = round(params['KX'] / (2 * np.pi / params['XMAX']))
p = axes.semilogx(var_dat[t_idx][kx_idx],
range(len(var_dat[t_idx][kx_idx])),
linewidth=0.5)
idx += 1
for var in f_vars:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title='%s (Cheb. summed)' % var)
var_dat = state_vars[var.replace('_f', '_c')]
summed_dat = np.sum(np.abs(var_dat[t_idx]), 1)
p = axes.semilogx(summed_dat,
range(len(summed_dat)),
linewidth=0.5)
idx += 1
fig.suptitle(
'Config: %s (t=%.2f, kx=%.2f, kz=%.2f, omega=%.2f)' %
(name, sim_time, params['KX'], params['KZ'], params['OMEGA']))
fig.subplots_adjust(hspace=0.7, wspace=0.6)
savefig = SAVE_FMT_STR % (t_idx // plot_stride)
plt.savefig('%s/%s' % (snapshots_dir, savefig))
print('Saved %s/%s' % (snapshots_dir, savefig))
plt.close()
os.system('ffmpeg -y -framerate 12 -i %s/%s %s.mp4' %
(snapshots_dir, SAVE_FMT_STR, name))
def plot_for_sigma(SIGMA):
SCALE = 4
XMAX = 30
N_X = 128
T_F = 100
NUMSTEPS = 5e2
DT = T_F / NUMSTEPS
# Bases and domain
x_basis = de.Chebyshev('x', N_X, interval=(0, XMAX), dealias=3 / 2)
domain = de.Domain([x_basis], np.float64)
# Problem
x = domain.grid(0)
problem = de.IVP(domain, variables=['y', 'y_x', 'y_t'])
problem.parameters['SIGMA'] = SIGMA
problem.add_equation(
'dt(y_t) - dx(y_x)' +
'= cos(x - t) * exp(-(x - 10)**2 / (2 * SIGMA**2)) / SIGMA')
problem.add_equation('dx(y) - y_x = 0')
problem.add_equation('dt(y) - y_t = 0')
problem.add_bc('right(y_x + y_t) = 0')
problem.add_bc('left(y_x - y_t) = 0')
# Build solver
solver = problem.build_solver(de.timesteppers.RK443)
solver.stop_sim_time = T_F
solver.stop_wall_time = 100000 # should never get hit
solver.stop_iteration = 100000 # should never get hit
# Initial conditions
y = solver.state['y']
y_x = solver.state['y_x']
y_t = solver.state['y_t']
gshape = domain.dist.grid_layout.global_shape(scales=1)
y['g'] = np.zeros(gshape)
y_x['g'] = np.zeros(gshape)
y_t['g'] = np.zeros(gshape)
# Store data for final plot
t_list = [solver.sim_time]
y.set_scales(SCALE, keep_data=True)
y_x.set_scales(SCALE, keep_data=True)
y_t.set_scales(SCALE, keep_data=True)
y_list = [np.copy(y['g'])]
yx_list = [np.copy(y_x['g'])]
yt_list = [np.copy(y_t['g'])]
# Main loop
while solver.ok:
solver.step(DT)
if solver.iteration % (NUMSTEPS // 500) == 0:
y.set_scales(SCALE, keep_data=True)
y_x.set_scales(SCALE, keep_data=True)
y_t.set_scales(SCALE, keep_data=True)
y_list.append(np.copy(y['g']))
yx_list.append(np.copy(y_x['g']))
yt_list.append(np.copy(y_t['g']))
t_list.append(solver.sim_time)
if solver.iteration % (NUMSTEPS // 2) == 0:
logger.info('Iteration: %i, Time: %.3f/%3f, dt: %.3e',
solver.iteration, solver.sim_time, T_F, DT)
# Create space-time plot
_x = np.array(domain.grid(0, scales=SCALE))
t_array = np.array(t_list)
y_array = np.array(y_list)
yx_array = np.array(yx_list)
yt_array = np.array(yt_list)
print('Max:', y_array.max())
xmesh, tmesh = quad_mesh(x=_x, y=t_array)
fig = plt.figure()
PLOT_CFG = True
if PLOT_CFG:
axes1 = fig.add_subplot(1, 1, 1, title='y')
p1 = axes1.pcolormesh(xmesh, tmesh, y_array, cmap='YlGnBu')
axes1.axis(pad_limits(xmesh, tmesh))
fig.colorbar(p1, ax=axes1)
else:
plt.plot(t_array, y_array[:, 0], 'g-')
fig.suptitle('Wave Equation, rad BC')
fig.subplots_adjust(wspace=0.4, hspace=0.4)
plt.savefig('1d_rad_%d.png' % SIGMA, dpi=600)
plt.clf()