def equator_plot(phi, r, data, **kw):
fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(111, polar=True)
r, phi = plot_tools.quad_mesh(r.ravel(), phi.ravel())
r[:, 0] = 0
pcm = ax.pcolormesh(phi, r, data, **kw)
ax.set_aspect(1)
plt.colorbar(pcm)
def top_plot(phi, theta, data, pcm=None, cmap=None):
r = np.sin(theta)
# Plot
fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(111, polar=True)
r, phi = plot_tools.quad_mesh(r.ravel(), phi.ravel())
r[:, 0] = 0
pcm = ax.pcolormesh(phi, r, data, cmap=cmap)
ax.set_aspect(1)
plt.colorbar(pcm)
def plot(name, params):
rank = CW.rank
size = CW.size
populate_globals(params)
slice_suffix = '(x=0)'
mean_suffix = '(mean)'
sub_suffix = ' (- mean)'
res_suffix = ' (res)'
snapshots_dir = SNAPSHOTS_DIR % name
matplotlib.rcParams.update({'font.size': 6})
V_GZ = get_vgz(g, H, KX, KZ)
# z_b = N_Z // 4
z_b = 0
# available cfgs:
# plot_vars: 2D plot
# f_vars: vertically-summed Fourier components
# f2_vars: 2D plot w/ horizontal Fourier transform
# slice_vars: sliced at x=0
# mean_vars: horizontally averaged
# sub_vars: 2D plot, mean-subtracted
# res_vars: subtract analytical solution (only uz/ux)
def get_plot_vars(cfg):
''' unpacks above variables from cfg shorthand '''
ret_vars = [
cfg.get('plot_vars', []),
[i + '_f' for i in cfg.get('f_vars', [])],
cfg.get('f2_vars', []),
[i + mean_suffix for i in cfg.get('mean_vars', [])],
[i + slice_suffix for i in cfg.get('slice_vars', [])],
[i + sub_suffix for i in cfg.get('sub_vars', [])],
[i + res_suffix for i in cfg.get('res_vars', [])]
]
n_cols = cfg.get('n_cols', sum([len(arr) for arr in ret_vars]))
n_rows = cfg.get('n_rows', 1)
ret = [n_cols, n_rows, cfg['save_fmt_str']]
ret.extend(ret_vars)
return ret
plot_cfgs = [
{
'save_fmt_str': 'p_%03i.png',
'slice_vars': ['uz'],
'mean_vars': ['ux', 'S_{px}'],
'res_vars_vars': ['uz'],
},
{
'save_fmt_str': 'm_%03i.png',
'plot_vars': ['ux', 'uz', 'rho1', 'P1'],
},
]
dyn_vars = ['uz', 'ux', 'rho', 'P']
sim_times, domain, state_vars = load(name, params, dyn_vars, plot_stride)
x = domain.grid(0, scales=1)
z = domain.grid(1, scales=1)[:, z_b:]
z0 = z[0]
z1s = np.ones(np.shape(z0))
xmesh, zmesh = quad_mesh(x=x[:, 0], y=z0)
x2mesh, z2mesh = quad_mesh(x=np.arange(N_X // 2), y=z0)
# preprocess
for var in dyn_vars + ['ux_z']:
state_vars[var + mean_suffix] = horiz_mean(state_vars[var], N_X)
for var in dyn_vars:
state_vars[var + slice_suffix] = np.copy(state_vars[var][:, 0, :])
for var in dyn_vars:
# can't figure out how to numpy this together
means = state_vars[var + mean_suffix]
state_vars[var + sub_suffix] = np.copy(state_vars[var])
for idx, _ in enumerate(state_vars[var + sub_suffix]):
mean = means[idx]
state_vars[var + sub_suffix][idx] -= np.tile(mean, (N_X, 1))
_, _, _, _, _, dux2s, duz2s = \
subtract_lins(params, state_vars, sim_times, domain)
for cfg in plot_cfgs:
n_cols, n_rows, save_fmt_str, plot_vars, f_vars,\
f2_vars, mean_vars, slice_vars, sub_vars, res_vars\
= get_plot_vars(cfg)
uz_est = F * get_uz_f_ratio(params)
ux_est = uz_est * KZ / KX
for t_idx, time in list(enumerate(sim_times))[rank::size]:
fig = plt.figure(dpi=400)
uz_anal = get_anal_uz(params, time, x, z)
ux_anal = get_anal_ux(params, time, x, z)
uz_mean = np.outer(np.ones(N_X),
state_vars['uz%s' % mean_suffix][t_idx])
ux_mean = np.outer(np.ones(N_X),
state_vars['ux%s' % mean_suffix][t_idx])
S_px_mean = state_vars['S_{px}%s' % mean_suffix]
z_top = get_front_idx(S_px_mean[t_idx], get_flux_th(params) * 0.2)
z_bot = get_idx(Z0 + 3 * S, z0)
idx = 1
for var in plot_vars + sub_vars + res_vars:
if res_suffix in var:
# divide by analytical profile and normalize
dux2 = dux2s[t_idx]
duz2 = duz2s[t_idx]
if var == 'uz%s' % res_suffix:
var_dat = duz2 / uz_est
title = 'u_z'
elif var == 'ux%s' % res_suffix:
var_dat = dux2 / ux_est
title = 'u_x'
else:
raise ValueError('lol wtf is %s' % var)
# truncate uz_anal at sponge and critical layers
var_dat[:, 0:z_bot] = 0
var_dat[:, z_top:] = 0
err = np.sqrt(np.mean(var_dat**2))
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title=r'$\delta %s$ (RMS = %.4e)' %
(title, err))
vmin = var_dat.min()
vmax = var_dat.max()
else:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title=r'$%s$' % var)
var_dat = state_vars[var][t_idx, :, z_b:]
vmin = state_vars[var].min()
vmax = state_vars[var].max()
p = axes.pcolormesh(xmesh,
zmesh,
var_dat.T,
vmin=vmin,
vmax=vmax)
axes.axis(pad_limits(xmesh, zmesh))
cb = fig.colorbar(p, ax=axes)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
idx += 1
for var in f2_vars:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title=r'$\log %s$ (x-FT)' % var)
var_dat = state_vars[var][:, :, z_b:]
var_dat_t = np.fft.fft(var_dat[t_idx], axis=0)
var_dat_shaped = np.log(
np.abs(2 * var_dat_t.real[:N_X // 2, :]))
p = axes.pcolormesh(x2mesh,
z2mesh,
var_dat_shaped.T,
vmin=var_dat.min(),
vmax=var_dat.max())
axes.axis(pad_limits(x2mesh, z2mesh))
cb = fig.colorbar(p, ax=axes)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
idx += 1
for var in mean_vars + slice_vars:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title=r'$%s$' % var)
var_dat = state_vars[var][:, z_b:]
p = axes.plot(var_dat[t_idx],
z0,
'r-',
linewidth=0.7,
label='Data')
if var == 'uz%s' % slice_suffix:
p = axes.plot(uz_anal[0, :], z0, 'orange', linewidth=0.5)
p = axes.plot(-uz_est * z1s,
z0,
'g',
uz_est * z1s,
z0,
'g',
linewidth=0.5)
if var == 'ux%s' % slice_suffix:
p = axes.plot(ux_anal[0, :], z0, 'orange', linewidth=0.5)
p = axes.plot(-ux_est,
z0,
'g',
ux_est,
z0,
'g',
linewidth=0.5)
if var == 'S_{px}%s' % mean_suffix:
k_damp = get_k_damp(params)
p = axes.plot(uz_est**2 / 2 * abs(KZ / KX) * RHO0 * z1s,
z0,
'orange',
label=r'$x_0z_0$ (Anal.)',
linewidth=0.5)
# compute all of the S_px cross terms (00 is model)
Spx01 = horiz_mean(RHO0 * state_vars['ux'][t_idx] *
(state_vars['uz'][t_idx] - uz_anal),
N_X,
axis=0)
Spx10 = horiz_mean(RHO0 * state_vars['uz'][t_idx] *
(state_vars['ux'][t_idx] - ux_anal),
N_X,
axis=0)
Spx11 = horiz_mean(RHO0 *
(state_vars['ux'][t_idx] - ux_anal) *
(state_vars['uz'][t_idx] - uz_anal),
N_X,
axis=0)
p = axes.plot(Spx01[z_bot:z_top],
z[0, z_bot:z_top],
'g:',
linewidth=0.4,
label=r'$x_0z_1$')
p = axes.plot(Spx10[z_bot:z_top],
z[0, z_bot:z_top],
'b:',
linewidth=0.4,
label=r'$x_1z_0$')
p = axes.plot(Spx11[z_bot:z_top],
z[0, z_bot:z_top],
'k-',
linewidth=0.7,
label=r'$x_1z_1$')
axes.legend()
if var == 'ux%s' % mean_suffix:
# mean flow = E[ux * uz] / V_GZ
p = axes.plot(uz_est**2 * abs(KZ) / KX / (2 * abs(V_GZ)) *
z1s,
z0,
'orange',
linewidth=0.5)
# critical = omega / kx
p = axes.plot(OMEGA / KX * z1s, z0, 'green', linewidth=0.5)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
xlims = [var_dat.min(), var_dat.max()]
axes.set_xlim(*xlims)
axes.set_ylim(z0.min(), z0.max())
p = axes.plot(xlims,
[SPONGE_LOW + SPONGE_WIDTH * SPONGE_LOW]\
* len(xlims),
'r:',
linewidth=0.5)
p = axes.plot(
xlims, [SPONGE_HIGH - SPONGE_WIDTH *
(ZMAX - SPONGE_HIGH)] * len(xlims),
'r:',
linewidth=0.5)
p = axes.plot(xlims, [Z0 + 3 * S] * len(xlims),
'b--',
linewidth=0.5)
p = axes.plot(xlims, [Z0 - 3 * S] * len(xlims),
'b--',
linewidth=0.5)
idx += 1
for var in f_vars:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title=r'$%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(
r'Config: $%s (t=%.2f, \vec{k}=(%.2f, %.2f), \omega=%.2f)$' %
(name.replace('_', '.'), time, KX, KZ, OMEGA))
fig.subplots_adjust(hspace=0.7, wspace=0.6)
savefig = save_fmt_str % (t_idx)
plt.savefig('%s/%s' % (snapshots_dir, savefig))
logger.info('Saved %s/%s' % (snapshots_dir, savefig))
plt.close()
z = domain.grid(1)
state_vars = defaultdict(list)
with h5py.File(filename, mode='r') as dat:
sim_times = np.array(dat['scales']['sim_time'])
for varname in dyn_vars:
values = dat['tasks'][varname]
state_vars[varname].append(np.copy(values))
# cast to np arrays
for key in state_vars.keys():
state_vars[key] = np.array(state_vars[key])[0]
x = domain.grid(0, scales=1)
z = domain.grid(1, scales=1)
xmesh, zmesh = quad_mesh(x=x[:, 0], y=z[0])
for t_idx, sim_time in list(enumerate(sim_times)):
fig = plt.figure(figsize=(8, 12), dpi=200)
idx = 1
for var in dyn_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)
def plot(name, params):
slice_suffix = '(x=0)'
sum_suffix = '(mean)'
sub_suffix = ' (- mean)'
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 = ['ux']
# c_vars = ['uz_c']
# f_vars = ['uz_f']
f2_vars = ['ux']
z_vars = ['%s%s' % (i, sum_suffix) for i in ['ux']] # sum these over x
# slice_vars = ['%s%s' % (i, slice_suffix) for i in ['ux']]
sub_vars = ['%s%s' % (i, sub_suffix) for i in ['ux']]
c_vars = []
f_vars = []
# f2_vars = []
# z_vars = []
slice_vars = []
# sub_vars = []
n_cols = 4
n_rows = 1
plot_stride = 1
N_X = params['N_X'] * params['INTERP_X']
N_Z = params['N_Z'] * params['INTERP_Z']
# z_b = N_Z // 4
z_b = 0
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'])[:, z_b:]
xmesh, zmesh = quad_mesh(x=x[:, 0], y=z[0])
x2mesh, z2mesh = quad_mesh(x=np.arange(params['N_X'] // 2), y=z[0])
for var in z_vars:
state_vars[var] = np.sum(state_vars[var.replace(sum_suffix, '')],
axis=1) / N_X
for var in slice_vars:
state_vars[var] = np.copy(state_vars[var.replace(slice_suffix,
'')][:, 0, :])
for var in sub_vars:
# can't figure out how to numpy this together
means = state_vars[var.replace(sub_suffix, sum_suffix)]
state_vars[var] = np.copy(state_vars[var.replace(sub_suffix, '')])
for idx, _ in enumerate(state_vars[var]):
mean = state_vars[var.replace(sub_suffix, sum_suffix)][idx]
state_vars[var][idx] -= np.tile(mean, (N_X, 1))
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 + sub_vars:
axes = fig.add_subplot(n_rows, n_cols, idx, title=var)
var_dat = state_vars[var][:, :, z_b:]
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)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
idx += 1
for var in f2_vars:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title='log %s (x-FT)' % var)
var_dat = state_vars[var][:, :, z_b:]
var_dat_t = np.fft.fft(var_dat[t_idx], axis=0)
var_dat_shaped = np.log(
np.abs(2 * var_dat_t.real[:params['N_X'] // 2, :]))
p = axes.pcolormesh(x2mesh, z2mesh, var_dat_shaped.T)
# vmin=var_dat.min(), vmax=var_dat.max())
axes.axis(pad_limits(x2mesh, z2mesh))
cb = fig.colorbar(p, ax=axes)
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_b:]
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[t_idx].min(), var_dat[t_idx].max()]
axes.set_xlim(*xlims)
axes.set_ylim(z_pts.min(), z_pts.max())
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_front(name, params):
''' few plots for front, defined where flux drops below 1/3 of theory '''
N_X = params['N_X'] * params['INTERP_X']
N_Z = params['N_Z'] * params['INTERP_Z']
dyn_vars = ['uz', 'ux', 'rho', 'P', 'ux_z']
snapshots_dir = SNAPSHOTS_DIR % name
logfile = '%s/data.log' % snapshots_dir
sim_times = []
front_pos = []
ri_inv = []
fluxes = []
flux_th = (params['F'] * get_uz_f_ratio(params))**2 / 2 \
* abs(params['KZ'] / params['KX']) \
* params['RHO0'] * np.exp(-params['Z0'] / params['H'])
flux_threshold = flux_th / 3
# load if exists
if not os.path.exists(logfile):
print('log file not found, generating')
sim_times, domain, state_vars = load(name,
params,
dyn_vars,
2,
start=10)
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])
z_pts = (zmesh[1:, 0] + zmesh[:-1, 0]) / 2
ux_z = np.sum(state_vars['ux_z'], axis=1) / N_X
F_px = np.sum(state_vars['F_px'], axis=1) / N_X
for t_idx, sim_time in enumerate(sim_times):
max_pos = len(F_px[t_idx]) - 1
while F_px[t_idx][max_pos] < flux_threshold and max_pos >= 0:
max_pos -= 1
front_pos.append(z_pts[max_pos])
ri_inv.append(ux_z[t_idx][max_pos] / (params['g'] / params['H']))
fluxes.append(F_px[t_idx][max_pos - 10] -
F_px[t_idx][max_pos + 10])
with open(logfile, 'w') as data:
data.write(repr(sim_times.tolist()))
data.write('\n')
data.write(repr(front_pos))
data.write('\n')
data.write(repr(ri_inv))
data.write('\n')
data.write(repr(fluxes))
else:
print('data.log found, loading')
with open(logfile) as data:
sim_times = eval(data.readline())
front_pos = eval(data.readline())
ri_inv = eval(data.readline())
fluxes = eval(data.readline())
flux_anal = flux_th * np.ones(np.shape(fluxes))
velocities_anal = -flux_anal / \
(params['RHO0'] * np.exp(-np.array(front_pos) / params['H'])) * \
params['KX'] / params['OMEGA']
pos_anal = np.cumsum(velocities_anal) * np.gradient(sim_times) # dt
pos_anal += front_pos[-2] - pos_anal[-2]
plt.plot(sim_times, front_pos, label='Data')
plt.plot(sim_times, pos_anal, label='Analytical')
plt.ylabel('Critical Layer Position')
plt.xlabel('Time')
plt.title(name)
plt.legend()
plt.savefig('%s/front.png' % snapshots_dir, dpi=200)
plt.clf()
plt.plot(sim_times,
np.gradient(front_pos) / np.gradient(sim_times),
label='Data')
plt.plot(sim_times, velocities_anal, label='Analytic')
plt.ylabel('Critical Layer Velocity')
plt.xlabel('Time')
plt.legend()
plt.title(name)
plt.savefig('%s/front_v.png' % snapshots_dir, dpi=200)
plt.clf()
plt.plot(sim_times, 1 / np.array(ri_inv)**2)
plt.ylabel('Ri')
plt.xlabel('Time')
plt.title(name)
plt.ylim([0, 10])
plt.savefig('%s/f_ri.png' % snapshots_dir, dpi=200)
plt.clf()
plt.plot(sim_times, np.array(fluxes) * 1e5, label='Data')
plt.plot(sim_times, flux_anal * 1e5, label='Analytical')
plt.ylabel('F_px (1e-5)')
plt.xlabel('Time')
plt.title(name)
plt.legend()
plt.savefig('%s/fluxes.png' % snapshots_dir, dpi=200)
plt.clf()
def plot(name, params):
rank = CW.rank
size = CW.size
slice_suffix = '(x=0)'
mean_suffix = '(mean)'
sub_suffix = ' (- mean)'
snapshots_dir = SNAPSHOTS_DIR % name
path = '{s}/{s}_s1'.format(s=snapshots_dir)
matplotlib.rcParams.update({'font.size': 6})
N_X = params['N_X'] * params['INTERP_X']
N_Z = params['N_Z'] * params['INTERP_Z']
KX = params['KX']
KZ = params['KZ']
V_GZ = get_vgz(params['g'], params['H'], KX, KZ)
# available cfgs:
# plot_vars: 2D plot
# c_vars: horizontally-summed vertical chebyshev components
# f_vars: vertically-summed Fourier components
# f2_vars: 2D plot w/ horizontal Fourier transform
# slice_vars: sliced at x=0
# mean_vars: horizontally averaged
# sub_vars: 2D plot, mean-subtracted
def get_plot_vars(cfg):
''' unpacks above variables from cfg shorthand '''
ret_vars = [
cfg.get('plot_vars', []),
[i + '_c' for i in cfg.get('c_vars', [])],
[i + '_f' for i in cfg.get('f_vars', [])],
cfg.get('f2_vars', []),
[i + mean_suffix for i in cfg.get('mean_vars', [])],
[i + slice_suffix for i in cfg.get('slice_vars', [])],
[i + sub_suffix for i in cfg.get('sub_vars', [])]
]
n_cols = cfg.get('n_cols', sum([len(arr) for arr in ret_vars]))
n_rows = cfg.get('n_rows', 1)
ret = [n_cols, n_rows, cfg['save_fmt_str']]
ret.extend(ret_vars)
return ret
plot_cfgs = [
{
'save_fmt_str': 'p_%03i.png',
'mean_vars': ['F_px', 'ux', 'ux_z'],
'slice_vars': ['uz'],
'sub_vars': ['ux'],
},
{
'save_fmt_str': 'm_%03i.png',
'plot_vars': ['ux', 'uz', 'rho', 'P'],
},
# {
# 'save_fmt_str': 't_%03i.png',
# 'mean_vars': ['F_px', 'ux'],
# 'slice_vars': ['uz'],
# },
]
dyn_vars = ['uz', 'ux', 'rho', 'P', 'ux_z']
sim_times, domain, state_vars = load(name,
params,
dyn_vars,
plot_stride,
start=0)
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])
x2mesh, z2mesh = quad_mesh(x=np.arange(params['N_X'] // 2), y=z[0])
# preprocess
for var in dyn_vars + ['F_px']:
state_vars[var + mean_suffix] = (np.sum(state_vars[var], axis=1) / N_X,
np.min(state_vars[var], axis=1),
np.max(state_vars[var], axis=1))
for var in dyn_vars:
state_vars[var + slice_suffix] = np.copy(state_vars[var][:, 0, :])
for var in dyn_vars:
# can't figure out how to numpy this together
means = state_vars[var + mean_suffix][0]
state_vars[var + sub_suffix] = np.copy(state_vars[var])
for idx, _ in enumerate(state_vars[var + sub_suffix]):
mean = means[idx]
state_vars[var + sub_suffix][idx] -= np.tile(mean, (N_X, 1))
for cfg in plot_cfgs:
n_cols, n_rows, save_fmt_str, plot_vars, c_vars, f_vars,\
f2_vars, mean_vars, slice_vars, sub_vars = get_plot_vars(cfg)
uz_est = params['F'] * get_uz_f_ratio(params)
for t_idx, sim_time in list(enumerate(sim_times))[rank::size]:
fig = plt.figure(dpi=200)
idx = 1
for var in plot_vars + sub_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)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
idx += 1
for var in f2_vars:
axes = fig.add_subplot(n_rows,
n_cols,
idx,
title='log %s (x-FT)' % var)
var_dat = state_vars[var]
var_dat_t = np.fft.fft(var_dat[t_idx], axis=0)
var_dat_shaped = np.log(
np.abs(2 * var_dat_t.real[:params['N_X'] // 2, :]))
p = axes.pcolormesh(x2mesh,
z2mesh,
var_dat_shaped.T,
vmin=var_dat.min(),
vmax=var_dat.max())
axes.axis(pad_limits(x2mesh, z2mesh))
cb = fig.colorbar(p, ax=axes)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
idx += 1
for var in mean_vars + slice_vars:
axes = fig.add_subplot(n_rows, n_cols, idx, title=var)
if var in slice_vars:
var_dat = state_vars[var]
else:
var_dat, var_min, var_max = state_vars[var]
z_pts = (zmesh[1:, 0] + zmesh[:-1, 0]) / 2
p = axes.plot(var_dat[t_idx], z_pts, 'r-', 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)
if var == 'F_px%s' % mean_suffix:
p = axes.plot(
(params['F'] * get_uz_f_ratio(params))**2 / 2 *
abs(params['KZ'] / params['KX']) * params['RHO0'] *
np.exp(-params['Z0'] / params['H']) *
np.ones(np.shape(z_pts)),
z_pts,
'orange',
linewidth=0.5)
if var == 'ux%s' % mean_suffix:
# mean flow = E[ux * uz] / V_GZ
p = axes.plot(uz_est**2 * abs(KZ) / KX / (2 * abs(V_GZ)) *
np.exp(
(z_pts - params['Z0']) / (params['H'])),
z_pts,
'orange',
linewidth=0.5)
# critical = omega / kx
p = axes.plot(params['OMEGA'] / params['KX'] *
np.ones(np.shape(z_pts)),
z_pts,
'green',
linewidth=0.5)
if var in mean_vars:
p = axes.plot(var_min[t_idx], z_pts, 'r:', linewidth=0.2)
p = axes.plot(var_max[t_idx], z_pts, 'r:', linewidth=0.2)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
xlims = [var_dat.min(), var_dat.max()]
axes.set_xlim(*xlims)
axes.set_ylim(z_pts.min(), z_pts.max())
p = axes.plot(xlims,
[params['SPONGE_LOW'] + params['SPONGE_WIDTH']] *
len(xlims),
'r:',
linewidth=0.5)
p = axes.plot(
xlims, [params['SPONGE_HIGH'] - params['SPONGE_WIDTH']] *
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, %.2f), w=%.2f, v_gz=%.2f)' %
(name, sim_time, KX, KZ, params['OMEGA'], V_GZ))
fig.subplots_adjust(hspace=0.7, wspace=0.6)
savefig = save_fmt_str % (t_idx)
plt.savefig('%s/%s' % (snapshots_dir, savefig))
logger.info('Saved %s/%s' % (snapshots_dir, savefig))
plt.close()
# 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')
def main(filename, start, count, output):
"""Save plot of specified tasks for given range of analysis writes."""
# Plot settings
tasks = ['c1', 'c2', 'c3']
scale = 2.5
dpi = 100
savename_func = lambda write: 'write_{:06}.png'.format(write)
fig = plt.figure()
# Plot writes
with h5py.File(filename, mode='r') as file:
# r = file['scales/r/10']
# theta = file['scales/θ/10']
r = file['scales/r/1.0']
theta = file['scales/θ/1.0']
theta, r = plot_tools.quad_mesh(theta, r)
c1 = file['tasks/c1']
c2 = file['tasks/c2']
c3 = file['tasks/c3']
phi = file['tasks/left(φ)']
for index in range(start, start + count):
logger.info("Plotting frame {:d}".format(index))
c1_mod = c1[index]
c1_mod[c1[index] < 0] = 0
c1_mod[c1[index] > 1] = 1
c2_mod = c2[index]
c2_mod[c2[index] < 0] = 0
c2_mod[c2[index] > 1] = 1
c3_mod = c3[index]
c3_mod[c3[index] < 0] = 0
c3_mod[c3[index] > 1] = 1
phi_in = phi[index, 0, 0]
colors = tuple(
np.array([
c1_mod.T.flatten(),
c2_mod.T.flatten(),
c3_mod.T.flatten()
]).transpose().tolist())
ax = fig.add_subplot(111, projection='polar')
ax.pcolormesh(theta, r, c1[0, :, :].T, color=colors)
ax.annotate("",
xy=(phi_in, 0.6),
xytext=(phi_in, 1.),
arrowprops=dict(arrowstyle='-',
connectionstyle='arc3',
linewidth=2))
ax.spines['polar'].set_visible(False)
## removing the tick marks
ax.set_xticks([])
ax.set_yticks([])
# Save figure
savename = savename_func(file['scales/write_number'][index])
savepath = output.joinpath(savename)
fig.savefig(str(savepath), dpi=dpi)
fig.clear()
plt.close(fig)
x_lres_list.append(midpoint_x_lres)
y_lres_list.append(midpoint_y_lres)
r_zeros_hres = np.loadtxt('double_Re1e4_0p25_hres%s/contour_flux.dat' %
ending)
(t, midpoint_x_hres, midpoint_y_hres) = np.loadtxt(
'double_Re1e4_0p25_hres%s/thermal_midpoint_flux.dat' % ending)
r_hres_list.append(r_zeros_hres)
x_hres_list.append(midpoint_x_hres)
y_hres_list.append(midpoint_y_hres)
lw = 1
# plot slices
c_im = []
xm_lres, ym_lres = plot_tools.quad_mesh(x_lres, y_lres)
xm_hres, ym_hres = plot_tools.quad_mesh(x_hres, y_hres)
for i in range(2):
for j in range(5):
if i == 0:
xm = xm_lres
ym = ym_lres
data = rho_lres_list[j]
midpoint_x = x_lres_list[j]
midpoint_y = y_lres_list[j]
r_zeros = r_lres_list[j]
z = z_lres
elif i == 1:
xm = xm_hres
ym = ym_hres
print('main loop time:', main_loop_time)
print(' total time:', total_time)
print(' iterations:', solver.iteration)
print(' loop sec/iter:', main_loop_time / solver.iteration)
print(' average dt:', solver.sim_time / n_steps)
print(
" N_cores, Nx, Ny, Nz, startup main loop, main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid"
)
print(
'scaling:', ' {:d} {:d} {:d} {:d}'.format(N_TOTAL_CPU, Nx, Ny, 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 * Ny * Nz),
N_TOTAL_CPU * main_loop_time / n_steps / (Nx * Ny * Nz)))
print('-' * 40)
xmesh, ymesh = plot_tools.quad_mesh(
x=domain.grid(0, scales=domain.dealias).flatten(),
y=domain.grid(2, scales=domain.dealias).flatten())
# Plot
# plt.figure(figsize=(12, 8))
# print(b_array[2][:][5][:])
# np.savetxt("hola.txt",b_array)
#print(b_array[0][:][1][:])
plt.pcolormesh(xmesh, ymesh, b_array[25][:][1][:].T, cmap='RdBu_r')
#plt.axis(plot_tools.pad_limits(xmesh, ymesh))
plt.colorbar()
plt.xlabel('x')
plt.ylabel('z')
# plt.title('A dispersive shock!')
plt.show()
def plot(setup_problem,
XMAX,
ZMAX,
N_X,
N_Z,
T_F,
DT,
KX,
KZ,
H,
G,
A,
RHO0,
INTERP_X,
INTERP_Z,
name=None,
**_):
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', 'ux', 'rho', 'P', 'rho0']
z_vars = ['F_z', 'E'] # sum these over x
n_cols = 3
n_rows = 3
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(setup_problem,
XMAX=XMAX,
ZMAX=ZMAX,
N_X=N_X,
N_Z=N_Z,
T_F=T_F,
DT=DT,
KX=KX,
KZ=KZ,
H=H,
G=G,
A=A,
RHO0=RHO0,
INTERP_X=INTERP_X,
INTERP_Z=INTERP_Z,
name=name)
x = domain.grid(0, scales=INTERP_X)
z = domain.grid(1, scales=INTERP_Z)
OMEGA = get_omega(G, H, KX, KZ)
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 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))
fig.colorbar(p, ax=axes)
idx += 1
for var in z_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)
axes.set_xlim(var_dat.min(), var_dat.max())
idx += 1
fig.suptitle('Config: %s (t=%.2f, kx=-2pi/H, kz=2pi/H, omega=%.2f)' %
(name, sim_time, OMEGA))
fig.subplots_adjust(hspace=0.4, wspace=0.4)
savefig = SAVE_FMT_STR % (t_idx // plot_stride)
plt.savefig('%s/%s' % (path, savefig))
print('Saved %s/%s' % (path, savefig))
plt.close()
os.system('ffmpeg -y -framerate 6 -i %s/%s %s.mp4' %
(path, SAVE_FMT_STR, name))
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']
# z_vars = ['F_z', 'E'] # sum these over x
z_vars = []
slice_vars = ['%s%s' % (i, slice_suffix) for i in ['uz']]
n_cols = 3
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
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))
report_cadence = 10
dt = 1e-3
plot_cadence = report_cadence
plot = theta_target in theta
var = T['g']
name = 'T_diff'
if plot:
i_theta = np.argmin(np.abs(theta[0, :, 0] - theta_target))
import matplotlib.pyplot as plt
from dedalus.extras import plot_tools
r = r[0, 0, :]
phi = phi[:, 0, 0]
rm, phim = plot_tools.quad_mesh(r, phi)
x = rm * np.cos(phim)
y = rm * np.sin(phim)
fig, ax = plt.subplots()
im = ax.pcolormesh(x, y, var[:, i_theta, :].real)
plt.colorbar(im)
plt.savefig('frames/%s_%04i.png' %
(name, solver.iteration // plot_cadence))
# Main loop
start_time = time.time()
while solver.ok:
if solver.iteration % report_cadence == 0 and solver.iteration > 0:
T0 = np.sum(vol_correction * weight_r * weight_theta * T['g'].real**2)
T0 = T0 * (np.pi) / (Lmax + 1) / L_dealias
T0 = reducer.reduce_scalar(T0, MPI.SUM)
x_list.append(np.array(f['scales/x/1.0'])) y_list.append(np.array(f['scales/z/1.0'])) f.close() lw = 1 # plot slices c_im = [] for i in range(3): x = x_list[i] y = y_list[i] xm, ym = plot_tools.quad_mesh(x, y) T = np.copy(T_list[i]) T[T < 0] += np.nan c_im.append(slice_axes[i].pcolormesh(xm, ym, T.T, cmap=cmaps[0])) T = np.copy(T_list[i]) u = u_list[i] w = w_list[i] om = np.gradient(w, x, y)[0] - np.gradient(u, x, y)[1] om[T > 0] += np.nan c_im.append(slice_axes[i].pcolormesh(xm, ym, om.T, cmap=cmaps[1]))
phi_shape = phi.shape
r_shape = r.shape
phi = phi.reshape((phi_shape[0], 1))
r = r.reshape((1, r_shape[-1]))
phi_1D = np.concatenate((phi[:, 0], [0]))
plot_data = []
plot_data.append(ur) # ur
plot_data.append(uphi) # uphi
plot_data.append(ur * np.cos(phi) - uphi * np.sin(phi)) # ux
plot_data.append(ur * np.sin(phi) + uphi * np.cos(phi)) # uy
labels = [r'$u_r$', r'$u_\phi$', r'$u_x$', r'$u_y$']
r, phi = plot_tools.quad_mesh(r.ravel(), phi.ravel())
r[:, 0] = 0
x = r * np.cos(phi)
y = r * np.sin(phi)
c_im = []
cbars = []
eps = 0.02
for i in range(num):
# c_im.append(plot_axes[i].pcolormesh(r,phi,plot_data[i]))
c_im.append(plot_axes[i].pcolormesh(x, y, plot_data[i], cmap='RdBu'))
plot_axes[i].plot((1 + eps / 2) * np.sin(phi), (1 + eps / 2) * np.cos(phi),
color='k',
phi = file['phi']
theta = file['theta']
time = file['t'][0]
data = file[field]
signs = [0, -1, 1]
count = 0
for axes in axs:
sign = signs[count]
count += 1
# Create plot
#if i == first_frame + comm.rank:
lon = (phi + phi[1] / 2 - np.pi +
sign * Omega * time) * 180 / np.pi
lat = (np.pi / 2 - theta) * 180 / np.pi
xmesh, ymesh = plot_tools.quad_mesh(lon, lat)
image = axes.pcolormesh(xmesh,
ymesh,
data.T,
cmap='RdBu_r',
transform=ccrs.PlateCarree())
if show_time:
title = fig.suptitle('t = %.4f' % time)
if gridlines:
axes.gridlines(xlocs=np.arange(0, 361, 30),
ylocs=np.arange(-60, 61, 30),
color='k')
axes.set_global()
axes.outline_patch.set_edgecolor(edgecolor)
axes.set_title('%i' % (sign))
raise
finally:
end_time = time.time()
logger.info('Iterations: %i' % sh.solver.iteration)
logger.info('Sim end time: %f' % sh.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 * sh.domain.dist.comm_cart.size))
# Create space-time plot
u_array = np.array(sh.u_list)
ux_array = np.array(sh.ux_list)
print(str(math.sqrt(2) * np.linalg.norm(u_array[-1], 2)) + '\n')
print(u_array[-1])
print(np.linalg.norm(u_array[-1] - u_array[0], 1) / 1000)
t_array = np.array(sh.t_list)
xmesh, ymesh = quad_mesh(x=sh.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')
# -------------edited for travelling wave ----------------
plt.title('Swift-Honenberg, r=%g' % (sh.problem.parameters['k']))
# plt.savefig('Swift_Honenberg.pdf',format='pdf')
plt.savefig('Swift_Honenberg.png')
plt.show()
# sh.save_array(u_array[-1], "./")
def plot(get_solver, setup_problem, name, params):
slice_suffix = '(x=0)' # slice suffix
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', 'ux']
z_vars = ['F_z', 'E'] # sum these over x
slice_vars = [
'%s%s' % (i, slice_suffix) for i in ['uz', 'ux', 'rho1', 'P1']
]
n_cols = 3
n_rows = 3
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(get_solver, setup_problem, 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, :]
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)
plt.xticks(rotation=30)
plt.yticks(rotation=30)
# xlims = [var_dat.min(), var_dat.max()]
xlims = [var_dat[t_idx].min(), var_dat[t_idx].max()]
axes.set_xlim(*xlims)
p = axes.plot(xlims, [params['SPONGE_START_LOW']] * len(xlims),
'r--')
p = axes.plot(xlims, [params['SPONGE_START_HIGH']] * len(xlims),
'r--')
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.5, wspace=0.6)
savefig = SAVE_FMT_STR % (t_idx // plot_stride)
plt.savefig('%s/%s' % (path, savefig))
print('Saved %s/%s' % (path, savefig))
plt.close()
os.system('ffmpeg -y -framerate 12 -i %s/%s %s.mp4' %
(path, 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()
