def plot_diff_xz(ax, var, rel=False, lim=None):
if rel:
if lim is not None:
bplt.plot_xz(ax, geom, np.abs((dump1[var] - dump2[var])/dump1[var]), vmin=0, vmax=lim, label=var, cbar=False, arrayspace=USEARRSPACE)
else:
bplt.plot_xz(ax, geom, np.abs((dump1[var] - dump2[var])/dump1[var]), label=var, cbar=False, arrayspace=USEARRSPACE)
else:
if lim is not None:
bplt.plot_xz(ax, geom, np.log10(np.abs(dump1[var] - dump2[var])), vmin=log_floor, vmax=lim, label=var, cbar=False, arrayspace=USEARRSPACE)
else:
bplt.plot_xz(ax, geom, np.log10(np.abs(dump1[var] - dump2[var])), vmin=log_floor, vmax=0, label=var, cbar=False, arrayspace=USEARRSPACE)
def plot(args):
n = args
print '%08d / ' % (n+1) + '%08d' % len(files)
dump = io.load_dump(files[n], geom)
fig = plt.figure(figsize=(FIGX, FIGY))
ax = plt.subplot(2,2,1)
#ax.pcolormesh(geom['x'][:,:,0], geom['z'][:,:,0],
# np.log10(dump['RHO'][:,:,0]), vmin=-4, vmax=0)
bplt.plot_xz(ax, geom, np.log10(dump['RHO']), dump,
vmin=-4, vmax = 0, label='RHO')
ax.set_xlim([0, SIZE]); ax.set_ylim([-SIZE, SIZE])
ax = plt.subplot(2,2,2)
bplt.plot_xz(ax, geom, np.log10(dump['Thetae']), dump,
vmin=-3, vmax = 3, label='Thetae', cmap='RdBu_r')
#bplt.plot_xz(ax, geom, np.log10(dump['Thetae']), dump,
# vmin=-2, vmax=2, label='Thetae', cmap='RdBu_r')
ax.set_xlim([0, SIZE]); ax.set_ylim([-SIZE, SIZE])
ax = plt.subplot(2,2,3)
bplt.plot_xz(ax, geom, np.log10(dump['beta']), dump,
vmin=-3, vmax = 3, label='beta', cmap='RdBu_r')
#bplt.plot_xy(ax, geom, np.log10(dump['RHO']), dump,
# vmin=-4, vmax=0, label='RHO')
ax.set_xlim([0, SIZE]); ax.set_ylim([-SIZE, SIZE])
ax = plt.subplot(2,2,4)
bplt.plot_xz(ax, geom, np.log10(dump['TpTe']), dump,
vmin=-3, vmax = 3, label='TpTe', cmap='RdBu_r')
#bplt.plot_xy(ax, geom, np.log10(dump['Thetae']), dump,
# vmin=-2, vmax=2, label='Thetae', cmap='RdBu_r')
ax.set_xlim([0, SIZE]); ax.set_ylim([-SIZE, SIZE])
#ax.pcolormesh(dump['X1'][:,:,0], dump['X2'][:,:,0], dump['RHO'][:,:,0])
plt.savefig(os.path.join(FRAMEDIR, 'frame_%08d.png' % n),
bbox_inches='tight', dpi=100)
plt.close(fig)
def plot(n):
imname = os.path.join(frame_dir, 'frame_%08d.png' % n)
tdump = io.get_dump_time(files[n])
if (tstart is not None and tdump < tstart) or (tend is not None and tdump > tend):
return
print("{} / {}".format((n+1),len(files)))
fig = plt.figure(figsize=(FIGX, FIGY))
if movie_type not in ["simplest", "simpler", "simple"]:
dump = io.load_dump(files[n], hdr, geom, derived_vars=True, extras=False)
#fig.suptitle("t = %d"%dump['t']) # TODO put this at the bottom somehow?
else:
# Simple movies don't need derived vars
dump = io.load_dump(files[n], hdr, geom, derived_vars=False, extras=False)
# Put the somewhat crazy rho values from KORAL dumps back in plottable range
if np.max(dump['RHO']) < 1e-10:
dump['RHO'] *= 1e15
# Zoom in for small problems
if geom['r'][-1,0,0] > 100:
window = [-100,100,-100,100]
nlines = 20
rho_l, rho_h = -3, 2
iBZ = i_of(geom,100) # most MADs
rBZ = 100
else:
window = [-50,50,-50,50]
nlines = 5
rho_l, rho_h = -4, 1
iBZ = i_of(geom,40) # most SANEs
rBZ = 40
if movie_type == "simplest":
# Simplest movie: just RHO
ax_slc = [plt.subplot(1,2,1), plt.subplot(1,2,2)]
bplt.plot_xz(ax_slc[0], geom, np.log10(dump['RHO']),
label="", vmin=rho_l, vmax=rho_h, window=window,
xlabel=False, ylabel=False, xticks=False, yticks=False,
cbar=False, cmap='jet')
bplt.plot_xy(ax_slc[1], geom, np.log10(dump['RHO']),
label="", vmin=rho_l-0.5, vmax=rho_h-0.5, window=window,
xlabel=False, ylabel=False, xticks=False, yticks=False,
cbar=False, cmap='jet')
pad = 0.0
plt.subplots_adjust(hspace=0, wspace=0, left=pad, right=1-pad, bottom=pad, top=1-pad)
elif movie_type == "simpler":
# Simpler movie: RHO and phi
gs = gridspec.GridSpec(2, 2, height_ratios=[6, 1], width_ratios=[16,17])
ax_slc = [fig.subplot(gs[0,0]), fig.subplot(gs[0,1])]
ax_flux = [fig.subplot(gs[1,:])]
bplt.plot_slices(ax_slc[0], ax_slc[1], geom, dump, np.log10(dump['RHO']),
label=r"$\log_{10}(\rho)$", vmin=rho_l, vmax=rho_h, window=window,
overlay_field=False, cmap='jet')
bplt.diag_plot(ax_flux[0], diag, 'phi_b', dump['t'], ylabel=r"$\phi_{BH}$", logy=LOG_PHI, xlabel=False)
elif movie_type == "simple":
# Simple movie: RHO mdot phi
gs = gridspec.GridSpec(3, 2, height_ratios=[4, 1, 1])
ax_slc = [fig.subplot(gs[0,0]), fig.subplot(gs[0,1])]
ax_flux = [fig.subplot(gs[1,:]), fig.subplot(gs[2,:])]
bplt.plot_slices(ax_slc[0], ax_slc[1], geom, dump, np.log10(dump['RHO']),
label=r"$\log_{10}(\rho)$", vmin=rho_l, vmax=rho_h, window=window, cmap='jet')
bplt.diag_plot(ax_flux[0], diag, 'mdot', dump['t'], ylabel=r"$\dot{M}$", logy=LOG_MDOT)
bplt.diag_plot(ax_flux[1], diag, 'phi_b', dump['t'], ylabel=r"$\phi_{BH}$", logy=LOG_PHI)
elif movie_type == "radial":
# TODO just record these in analysis output...
rho_r = eht_profile(geom, dump['RHO'], jmin, jmax)
B_r = eht_profile(geom, np.sqrt(dump['bsq']), jmin, jmax)
uphi_r = eht_profile(geom, dump['ucon'][:,:,:,3], jmin, jmax)
Pg = (hdr['gam']-1.)*dump['UU']
Pb = dump['bsq']/2
Pg_r = eht_profile(geom, Pg, jmin, jmax)
Ptot_r = eht_profile(geom, Pg + Pb, jmin, jmax)
betainv_r = eht_profile(geom, Pb/Pg, jmin, jmax)
ax_slc = lambda i: plt.subplot(2, 3, i)
bplt.radial_plot(ax_slc(1), geom, rho_r, ylabel=r"$<\rho>$", logy=True, ylim=[1.e-2, 1.e0])
bplt.radial_plot(ax_slc(2), geom, Pg_r, ylabel=r"$<P_g>$", logy=True, ylim=[1.e-6, 1.e-2])
bplt.radial_plot(ax_slc(3), geom, B_r, ylabel=r"$<|B|>$", logy=True, ylim=[1.e-4, 1.e-1])
bplt.radial_plot(ax_slc(4), geom, uphi_r, ylabel=r"$<u^{\phi}>$", logy=True, ylim=[1.e-3, 1.e1])
bplt.radial_plot(ax_slc(5), geom, Ptot_r, ylabel=r"$<P_{tot}>$", logy=True, ylim=[1.e-6, 1.e-2])
bplt.radial_plot(ax_slc(6), geom, betainv_r, ylabel=r"$<\beta^{-1}>$", logy=True, ylim=[1.e-2, 1.e1])
elif movie_type == "fluxes_cap":
axes = [plt.subplot(2, 2, i) for i in range(1,5)]
bplt.plot_thphi(axes[0], geom, np.log10(d_fns['FE'](dump)[iBZ,:,:]), iBZ, vmin=-8, vmax=-4, label =r"FE $\theta-\phi$ slice")
bplt.plot_thphi(axes[1], geom, np.log10(d_fns['FM'](dump)[iBZ,:,:]), iBZ, vmin=-8, vmax=-4, label =r"FM $\theta-\phi$ slice")
bplt.plot_thphi(axes[2], geom, np.log10(d_fns['FL'](dump)[iBZ,:,:]), iBZ, vmin=-8, vmax=-4, label =r"FL $\theta-\phi$ slice")
bplt.plot_thphi(axes[3], geom, np.log10(dump['RHO'][iBZ,:,:]), iBZ, vmin=-4, vmax=1, label =r"\rho $\theta-\phi$ slice")
for i,axis in enumerate(axes):
if i == 0:
overlay_thphi_contours(axis, geom, diag, legend=True)
else:
overlay_thphi_contours(axis, geom, diag)
max_th = geom['n2']//2
x = bplt.loop_phi(geom['x'][iBZ,:max_th,:])
y = bplt.loop_phi(geom['y'][iBZ,:max_th,:])
prep = lambda var : bplt.loop_phi(var[:max_th,:])
axis.contour(x,y, prep(geom['th'][iBZ]), [1.0], colors='k')
axis.contour(x,y, prep(d_fns['betagamma'](dump)[iBZ]), [1.0], colors='k')
axis.contour(x,y, prep(d_fns['sigma'](dump)[iBZ]), [1.0], colors='xkcd:green')
axis.contour(x,y, prep(d_fns['FE'](dump)[iBZ]), [0.0], colors='xkcd:pink')
axis.contour(x,y, prep(d_fns['Be_nob'](dump)[iBZ]), [0.02], colors='xkcd:red')
axis.contour(x,y, prep(d_fns['mu'](dump)[iBZ]), [2.0], colors='xkcd:blue')
elif movie_type == "rho_cap":
# Note cmaps are different between left 2 and right plot, due to the latter being far away from EH
bplt.plot_slices(plt.subplot(1,3,1), plt.subplot(1,3,2), geom, dump, np.log10(dump['RHO']),
label=r"$\log_{10}(\rho)$", vmin=-3, vmax=2, cmap='jet')
bplt.overlay_contours(plt.subplot(1,3,1), geom, geom['r'], [rBZ], color='k')
bplt.plot_thphi(plt.subplot(1,3,3), geom, np.log10(dump['RHO'][iBZ,:,:]), iBZ, vmin=-4, vmax=1, label=r"$\log_{10}(\rho)$ $\theta-\phi$ slice r="+str(rBZ))
elif movie_type == "funnel_wall":
rKH = 20
iKH = i_of(geom, rKH)
win=[0,rBZ/2,0,rBZ]
gs = gridspec.GridSpec(1, 3, width_ratios=[1,1,1])
axes = [plt.subplot(gs[0,i]) for i in range(3)]
bplt.plot_xz(axes[0], geom, np.log10(dump['RHO']),
label=r"$\log_{10}(\rho)$", vmin=-3, vmax=2, cmap='jet', window=win, shading='flat')
bplt.plot_xz(axes[1], geom, np.log10(dump['ucon'][:,:,:,3]),
label=r"$\log_{10}(u^{\phi})$", vmin=-3, vmax=0, cmap='Reds', window=win, cbar=False, shading='flat')
bplt.plot_xz(axes[1], geom, np.log10(-dump['ucon'][:,:,:,3]),
label=r"$\log_{10}(u^{\phi})$", vmin=-3, vmax=0, cmap='Blues', window=win, cbar=False, shading='flat')
bplt.plot_xz(axes[2], geom, np.log10(dump['beta'][:,:,:,3]),
label=r"$\log_{10}(u_{\phi})$", vmin=-3, vmax=3, window=win, shading='flat')
for axis in axes:
bplt.overlay_field(axis, geom, dump, nlines=nlines*4)
# bplt.plot_thphi(axes[2], geom, np.log10(dump['RHO'][iKH,:,:]), iKH,
# label=r"$\log_{10}(\rho)$ $\theta-\phi$ slice r="+str(rKH), vmin=-4, vmax=1, cmap='jet', shading='flat')
elif movie_type == "kh_radii":
if True: # Half-theta (one jet) switch
awindow = [0,1,0.5,1]
bwindow = [0,rBZ/2,0,rBZ]
else:
awindow = [0,1,0,1]
bwindow = [0,rBZ/2,-rBZ/2,rBZ/2]
rlevels = [10, 20, 40, 80]
axes = [plt.subplot(2,3,1), plt.subplot(2,3,2), plt.subplot(2,3,4), plt.subplot(2,3,5)]
bigaxis = plt.subplot(1,3,3)
for ax,rlevel in zip(axes, rlevels):
bplt.plot_thphi(ax, geom, np.log10(dump['RHO'][i_of(geom, rlevel),:,:]), i_of(geom, rlevel),
label=r"$\log_{10}(\rho) (r = "+str(rlevel)+")$", vmin=-3, vmax=2, cmap='jet', shading='flat',
arrayspace=True, window=awindow)
# bplt.plot_xz(bigaxis, geom, np.log10(dump['RHO']), label=r"$\log_{10}(\rho) (\phi slice)$",
# vmin=-3, vmax=2, cmap='jet', shading='flat', window=bwindow)
bplt.plot_xz(bigaxis, geom, np.log10(dump['ucon'][:,:,:,3]),
label="", vmin=-3, vmax=0, cmap='Reds', window=bwindow, cbar=False, shading='flat')
bplt.plot_xz(bigaxis, geom, np.log10(-dump['ucon'][:,:,:,3]),
label=r"$\log_{10}(u^{\phi})$", vmin=-3, vmax=0, cmap='Blues', window=bwindow, shading='flat')
bplt.overlay_field(bigaxis, geom, dump)
bplt.overlay_contours(bigaxis, geom, geom['r'][:,:,0], levels=rlevels, color='r')
else: # All other movie types share a layout
ax_slc = lambda i: plt.subplot(2, 4, i)
ax_flux = lambda i: plt.subplot(4, 2, i)
if movie_type == "traditional":
# Usual movie: RHO beta fluxes
# CUTS
bplt.plot_slices(ax_slc(1), ax_slc(2), geom, dump, np.log10(dump['RHO']),
label=r"$\log_{10}(\rho)$", vmin=-3, vmax=2, cmap='jet')
bplt.plot_slices(ax_slc(5), ax_slc(6), geom, dump, np.log10(dump['beta']),
label=r"$\beta$", vmin=-2, vmax=2, cmap='RdBu_r')
# FLUXES
bplt.diag_plot(ax_flux(2), diag, 'mdot', dump['t'], ylabel=r"$\dot{M}$", logy=LOG_MDOT)
bplt.diag_plot(ax_flux(4), diag, 'phi_b', dump['t'], ylabel=r"$\phi_{BH}$", logy=LOG_PHI)
# Mixins:
# Zoomed in RHO
bplt.plot_slices(ax_slc(7), ax_slc(8), geom, dump, np.log10(dump['RHO']),
label=r"$\log_{10}(\rho)$", vmin=-3, vmax=2, window=[-10,10,-10,10], field_overlay=False)
# Bsq
# bplt.plot_slices(ax_slc[6], ax_slc[7], geom, dump, np.log10(dump['bsq']),
# label=r"$b^2$", vmin=-5, vmax=0, cmap='Blues')
# Failures: all failed zones, one per nonzero pflag
# bplt.plot_slices(ax_slc[6], ax_slc[7], geom, dump, dump['fail'] != 0,
# label="Failed zones", vmin=0, vmax=20, cmap='Reds', int=True) #, arrspace=True)
# 2D histograms
# bplt.hist_2d(ax_slc[6], np.log10(dump['RHO']), np.log10(dump['UU']),r"$\log_{10}(\rho)$", r"$\log_{10}(U)$", logcolor=True)
# bplt.hist_2d(ax_slc[7], np.log10(dump['UU']), np.log10(dump['bsq']),r"$\log_{10}(U)$", r"$b^2$", logcolor=True)
# Extra fluxes:
# bplt.diag_plot(ax_flux[1], diag, dump, 'edot', r"\dot{E}", logy=LOG_PHI)
elif movie_type == "e_ratio":
# Energy ratios: difficult places to integrate, with failures
bplt.plot_slices(ax_slc(0), ax_slc(1), geom, dump, np.log10(dump['UU']/dump['RHO']),
label=r"$\log_{10}(U / \rho)$", vmin=-3, vmax=3, average=True)
bplt.plot_slices(ax_slc(2), ax_slc(3), geom, dump, np.log10(dump['bsq']/dump['RHO']),
label=r"$\log_{10}(b^2 / \rho)$", vmin=-3, vmax=3, average=True)
bplt.plot_slices(ax_slc(4), ax_slc(5), geom, dump, np.log10(1/dump['beta']),
label=r"$\beta^{-1}$", vmin=-3, vmax=3, average=True)
bplt.plot_slices(ax_slc(6), ax_slc(7), geom, dump, dump['fail'] != 0,
label="Failures", vmin=0, vmax=20, cmap='Reds', int=True) #, arrspace=True)
elif movie_type == "conservation":
# Continuity plots to verify local conservation of energy, angular + linear momentum
# Integrated T01: continuity for momentum conservation
bplt.plot_slices(ax_slc[0], ax_slc[1], geom, dump, Tmixed(dump, 1, 0),
label=r"$T^1_0$ Integrated", vmin=0, vmax=600, arrspace=True, integrate=True)
# integrated T00: continuity plot for energy conservation
bplt.plot_slices(ax_slc[4], ax_slc[5], geom, dump, np.abs(Tmixed(dump, 0, 0)),
label=r"$T^0_0$ Integrated", vmin=0, vmax=3000, arrspace=True, integrate=True)
# Usual fluxes for reference
bplt.diag_plot(ax_flux[1], diag, 'mdot', dump['t'], ylabel=r"$\dot{M}$", logy=LOG_MDOT)
#bplt.diag_plot(ax_flux[3], diag, 'phi_b', dump['t'], ylabel=r"$\phi_{BH}$", logy=LOG_PHI)
# Radial conservation plots
E_r = sum_shell(geom,Tmixed(geom, dump, 0,0))
Ang_r = sum_shell(geom,Tmixed(geom, dump, 0,3))
mass_r = sum_shell(dump['ucon'][:,:,:,0]*dump['RHO'])
# TODO arrange legend better -- add labels when radial/diag plotting
bplt.radial_plot(ax_flux[3], geom, np.abs(E_r), 'Conserved vars at R', ylim=(0,1000), rlim=(0,20), arrayspace=True)
bplt.radial_plot(ax_flux[3], geom, np.abs(Ang_r)/10, '', ylim=(0,1000), rlim=(0,20), col='r', arrayspace=True)
bplt.radial_plot(ax_flux[3], geom, np.abs(mass_r), '', ylim=(0,1000), rlim=(0,20), col='b', arrayspace=True)
# Radial energy accretion rate
Edot_r = sum_shell(geom, Tmixed(geom, dump,1,0))
bplt.radial_plot(ax_flux[5], geom, np.abs(Edot_r), 'Edot at R', ylim=(0,200), rlim=(0,20), arrayspace=True)
# Radial integrated failures
bplt.radial_plot(ax_flux[7], geom, (dump['fail'] != 0).sum(axis=(1,2)), 'Fails at R', arrayspace=True, rlim=[0,50], ylim=[0,1000])
elif movie_type == "floors":
# TODO add measures of all floors' efficacy. Record ceilings in header or extras?
bplt.plot_slices(ax_flux[0], ax_flux[1], geom, dump['bsq']/dump['RHO'] - 100,
vmin=-100, vmax=100, cmap='RdBu_r')
bplt.diag_plot(ax, diag, dump, 'sigma_max', 'sigma_max')
elif movie_type in d_fns: # Hail mary for plotting new functions one at a time
axes = [plt.subplot(1,2,1), plt.subplot(1,2,2)]
win=[l*2 for l in window]
var = d_fns[movie_type](dump)
bplt.plot_slices(axes[0], axes[1], geom, dump, np.log10(var), vmin=-3, vmax=3, cmap='Reds', window=win)
bplt.plot_slices(axes[0], axes[1], geom, dump, np.log10(-var), vmin=-3, vmax=3, cmap='Blues', window=win)
else:
print("Movie type not known!")
exit(1)
# Extra padding for crowded 4x2 plots
pad = 0.03
plt.subplots_adjust(left=pad, right=1-pad, bottom=pad, top=1-pad)
plt.savefig(imname, dpi=1920/FIGX) # TODO the group projector is like 4:3 man
plt.close(fig)
dump.clear()
del dump
def make_snap(dfnam,
vnam,
coords,
size,
cmap,
logplot,
savefig,
label,
vmin,
vmax,
index=None,
geom=None):
if not os.path.exists(dfnam):
print('ERROR File ' + dfnam + ' does not exist!')
sys.exit()
import matplotlib
if savefig is not None and savefig:
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import plot as bplt
import hdf5_to_dict as io
import numpy as np
font = {'size': 16}
matplotlib.rc('font', **font)
hdr = io.load_hdr(dfnam)
if geom is None:
geom = io.load_geom(hdr)
dump = io.load_dump(dfnam, geom=geom)
if not vnam in dump.keys():
print('ERROR Variable ' + vnam + ' is not contained in dump file!')
print('Available variables are:')
for key in dump.keys():
if type(dump[key]) is np.ndarray:
if (len(dump[key].shape) == 3
and dump[key].shape[0] == hdr['N1']
and dump[key].shape[1] == hdr['N2']
and dump[key].shape[2] == hdr['N3']):
print(key, end=' ')
print('')
sys.exit()
IS_3D = hdr['N3'] > 1
var = dump[vnam]
if index is not None:
var = var[..., index]
if logplot:
var = np.log10(var)
if label is None:
label = vnam
if index is not None:
label += '[{}]'.format(index)
if IS_3D:
if coords == 'mks':
fig, (a0, a1) = plt.subplots(1,
2,
gridspec_kw={'width_ratios': [1, 1]},
figsize=(12, 6))
elif coords == 'cart':
fig, (a0, a1) = plt.subplots(1,
2,
gridspec_kw={'width_ratios': [1, 1]},
figsize=(13, 7))
ax = a0
if coords == 'mks':
bplt.plot_X1X2(ax,
geom,
var,
dump,
cmap=cmap,
vmin=vmin,
vmax=vmax,
cbar=False,
label=label,
ticks=None,
shading='gouraud')
elif coords == 'cart':
bplt.plot_xz(ax,
geom,
var,
dump,
cmap=cmap,
vmin=vmin,
vmax=vmax,
cbar=False,
label=label,
ticks=None,
shading='gouraud')
ax.set_xlim([-size, size])
ax.set_ylim([-size, size])
ax = a1
if coords == 'mks':
bplt.plot_X1X3(ax,
geom,
var,
dump,
cmap=cmap,
vmin=vmin,
vmax=vmax,
cbar=True,
label=label,
ticks=None,
shading='gouraud')
elif coords == 'cart':
bplt.plot_xy(ax,
geom,
var,
dump,
cmap=cmap,
vmin=vmin,
vmax=vmax,
cbar=True,
label=label,
ticks=None,
shading='gouraud')
ax.set_xlim([-size, size])
ax.set_ylim([-size, size])
else:
if coords == 'mks':
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
bplt.plot_X1X2(ax,
geom,
var,
dump,
cmap=cmap,
vmin=vmin,
vmax=vmax,
cbar=True,
label=label,
ticks=None,
shading='gouraud')
elif coords == 'cart':
fig, ax = plt.subplots(1, 1, figsize=(7, 10))
bplt.plot_xz(ax,
geom,
var,
dump,
cmap=cmap,
vmin=vmin,
vmax=vmax,
cbar=True,
label=label,
ticks=None,
shading='gouraud')
ax.set_xlim([0, size])
ax.set_ylim([-size, size])
if savefig == False:
plt.show()
else:
plt.savefig(savefig, bbox_inches='tight')
plt.cla()
plt.clf()
plt.close()
def plot(args):
n = args
print '%08d / ' % (n + 1) + '%08d' % len(files)
diag = io.load_diag(path)
dump = io.load_dump(files[n], geom)
#hdr = dump['hdr']
fig = plt.figure(figsize=(FIGX, FIGY))
# GET SHELL AVERAGES
Thetae_sadw = np.zeros(hdr['N1'])
sigma = np.zeros(hdr['N1'])
mdot = np.zeros(hdr['N1'])
for i in xrange(hdr['N1']):
vol = 0.
for j in xrange(hdr['N2']):
for k in xrange(hdr['N3']):
Thetae_sadw[i] += dump['Thetae'][i, j, k] * dump['RHO'][
i, j,
k] * geom['gdet'][i,
j] * hdr['dx1'] * hdr['dx2'] * hdr['dx3']
sigma[i] += dump['RHO'][i, j, k] * hdr['dx2'] * geom['gdet'][i,
j]
mdot[i] += dump['RHO'][
i, j, k] * hdr['dx2'] * hdr['dx3'] * geom['gdet'][i, j]
vol += dump['RHO'][i, j, k] * geom['gdet'][
i, j] * hdr['dx1'] * hdr['dx2'] * hdr['dx3']
sigma[i] /= (hdr['N2'] * hdr['N3'])
Thetae_sadw[i] /= vol
ax = plt.subplot(3, 3, 1)
bplt.plot_xz(ax,
geom,
np.log10(dump['RHO']),
dump,
vmin=-4,
vmax=0,
label='RHO',
ticks=[-4, -3, -2, -1, 0])
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
ax.set_xticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax.set_yticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax = plt.subplot(3, 3, 2)
bplt.plot_xz(ax,
geom,
np.log10(dump['Thetae']),
dump,
vmin=-2,
vmax=2,
label='Thetae',
cmap='RdBu_r',
ticks=[-2, -1, 0, 1, 2])
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
ax.set_xticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax.set_yticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax = plt.subplot(3, 3, 3)
ax.plot(geom['r'][:, 0, 0], sigma, color='b', label='Sigma')
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlim([1, SIZE])
ax.set_ylim([1.e-2, 1.e2])
ax.set_aspect(np.log(SIZE / 1) / np.log(1.e2 / 1.e-2))
ax.plot(geom['r'][:, 0, 0], Thetae_sadw, color='r', label='SADW Thetae')
ax.legend(loc=2, fontsize=10)
ax.set_xlabel('r/M')
ax = plt.subplot(3, 3, 4)
bplt.plot_xy(ax,
geom,
np.log10(dump['RHO']),
dump,
vmin=-4,
vmax=0,
label='RHO',
ticks=[-4, -3, -2, -1, 0])
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
ax.set_xticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax.set_yticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax = plt.subplot(3, 3, 5)
bplt.plot_xy(ax,
geom,
np.log10(dump['Thetae']),
dump,
vmin=-2,
vmax=2,
label='Thetae',
cmap='RdBu_r',
ticks=[-2, -1, 0, 1, 2])
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
ax.set_xticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax.set_yticks([-SIZE, -SIZE / 2, 0, SIZE / 2, SIZE])
ax = plt.subplot(3, 1, 3)
ax.plot(diag['t'], diag['mdot'], color='k')
ax.axvline(dump['t'], color='k', linestyle='--')
ax.set_xlim([0, dump['hdr']['tf']])
ax.set_xlabel('t/M')
ax.set_ylabel('mdot')
ax2 = ax.twinx()
ax2.set_ylabel('phi')
plt.subplots_adjust(hspace=0.25)
#ax.pcolormesh(dump['X1'][:,:,0], dump['X2'][:,:,0], dump['RHO'][:,:,0])
plt.savefig(os.path.join(FRAMEDIR, 'frame_%08d.png' % n),
bbox_inches='tight',
dpi=100)
plt.close(fig)
flux_in[0] = np.sum(dump['B2'][0, N2 // 2, :] *
geom['gdet'][0, N2 // 2, None] * hdr['dx1'] * hdr['dx3'])
for n in range(1, N1):
flux_in[n] = flux_in[n - 1] + np.sum(
dump['B2'][n, N2 // 2, :] * geom['gdet'][n, N2 // 2, None] *
hdr['dx1'] * hdr['dx3'])
ax = plt.subplot(NPLOTSY, NPLOTSX, 2)
bplt.radial_plot(ax, geom, flux_in, ylabel=r"Flux in r", rlim=[0, SIZE])
# Density 2D
ax = plt.subplot(NPLOTSY, NPLOTSX, 3)
bplt.plot_xz(ax,
geom,
np.log10(dump['RHO']),
vmin=-4,
vmax=0,
label=r"$\log_{10}(\rho)$",
window=[0, SIZE, -SIZE / 2, SIZE / 2])
# Beta 2D
ax = plt.subplot(NPLOTSY, NPLOTSX, 4)
bplt.plot_xz(ax,
geom,
np.log10(dump['beta']),
label=r"$\beta$",
cmap='RdBu_r',
vmin=1,
vmax=4,
window=[0, SIZE, -SIZE / 2, SIZE / 2])
bplt.overlay_field(ax, geom, dump, nlines=NLINES)
def plot(args):
n = args
imname = 'frame_%08d.png' % n
imname = os.path.join(FRAMEDIR, imname)
print '%08d / ' % (n + 1) + '%08d' % len(files)
print imname
# Ignore if frame already exists
if os.path.isfile(imname):
return
dump = io.load_dump(files[n], geom)
fig = plt.figure(figsize=(FIGX, FIGY))
fig.suptitle('t = %05.2g' % dump['t'])
ax = plt.subplot(2, 2, 1)
bplt.plot_xz(ax,
geom,
np.log10(dump['RHO']),
dump,
vmin=-4,
vmax=0,
label='RHO')
bplt.overlay_field(ax, geom, dump, NLEV=10)
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
ax = plt.subplot(2, 2, 2)
bplt.plot_xz(ax,
geom,
np.log10(dump['beta']),
dump,
vmin=-2,
vmax=2,
label='beta',
cmap='RdBu_r')
bplt.overlay_field(ax, geom, dump, NLEV=10)
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
ax = plt.subplot(2, 2, 3)
bplt.plot_xy(ax,
geom,
np.log10(dump['RHO']),
dump,
vmin=-4,
vmax=0,
label='RHO')
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
ax = plt.subplot(2, 2, 4)
bplt.plot_xy(ax,
geom,
np.log10(dump['beta']),
dump,
vmin=-2,
vmax=2,
label='beta',
cmap='RdBu_r')
ax.set_xlim([-SIZE, SIZE])
ax.set_ylim([-SIZE, SIZE])
#ax.pcolormesh(dump['X1'][:,:,0], dump['X2'][:,:,0], dump['RHO'][:,:,0])
plt.savefig(imname, bbox_inches='tight', dpi=100)
plt.close(fig)
dumpfile = os.path.join(
"/scratch/03002/bprather/pharm_dumps/M87SimulationLibrary/GRMHD", run_name,
"dumps/dump_00001500.h5")
hdr, geom, dump = io.load_all(dumpfile)
plotfile = os.path.join("/work/03002/bprather/stampede2/movies", run_name,
"eht_out.p")
avg = pickle.load(open(plotfile, "rb"))
fig = plt.figure(figsize=(FIGX, FIGY))
gs = gridspec.GridSpec(2, 2, width_ratios=[1, 2])
ax = plt.subplot(gs[0, 0])
bplt.plot_xz(ax,
geom,
np.log10(d_fns['FE_EM'](dump)),
arrayspace=USEARRSPACE,
average=True,
window=window)
ax.set_title(r"$\log_{10}( -{{T_{EM}}^r}_t )$")
bplt.overlay_contours(ax, geom, geom['r'], [AT_R], color='k')
overlay_rth_contours(ax, geom, avg, legend=True)
ax = plt.subplot(gs[1, 0])
bplt.plot_xz(ax,
geom,
np.log10(d_fns['FE'](dump)),
arrayspace=USEARRSPACE,
average=True,
window=window)
ax.set_title(r"$\log_{10}( -{T^r}_t - \rho u^r )$")
iBZ,
project=False,
label=r"FL $\theta-\phi$ slice")
overlay_thphi_contours(ax[1, 0], geom, avg)
bplt.plot_thphi(ax[1, 1],
geom,
np.log10(avg['rho_100_thphi']),
iBZ,
project=False,
label=r"$\rho$ $\theta-\phi$ slice")
overlay_thphi_contours(ax[1, 1], geom, avg)
plt.savefig(run_name.replace("/", "_") + '_L_100_thphi.png')
plt.close(fig)
fig, ax = plt.subplots(2, 2, figsize=(FIGX, FIGY))
bplt.plot_xz(ax[0, 0], geom, np.log10(avg['FE_rth']), label="FE X-Z Slice")
overlay_rth_contours(ax[0, 0], geom, avg, legend=True)
bplt.plot_xz(ax[0, 1], geom, np.log10(avg['FM_rth']), label="FM X-Z Slice")
overlay_rth_contours(ax[0, 1], geom, avg)
bplt.plot_xz(ax[1, 0], geom, np.log10(avg['FL_rth']), label="FL X-Z Slice")
overlay_rth_contours(ax[1, 0], geom, avg)
bplt.plot_xz(ax[1, 1],
geom,
np.log10(avg['rho_rth']),
label="RHO X-Z Slice")
overlay_rth_contours(ax[1, 1], geom, avg)
plt.savefig(run_name.replace("/", "_") + '_L_rth.png')
plt.close(fig)
def plot(n):
imname = os.path.join(frame_dir, 'frame_%08d.png' % n)
tdump = io.get_dump_time(files1[n])
if (tstart is not None and tdump < tstart) or (tend is not None and tdump > tend):
return
# Don't calculate b/ucon/cov/e- stuff unless we need it below
dump1 = io.load_dump(files1[n], hdr1, geom1, derived_vars = False, extras = False)
dump2 = io.load_dump(files2[n], hdr2, geom2, derived_vars = False, extras = False)
fig = plt.figure(figsize=(FIGX, FIGY))
# Keep same parameters betwen plots, even of SANE/MAD
rho_l, rho_h = -3, 2
window = [-20,20,-20,20]
nlines1 = 20
nlines2 = 5
# But BZ stuff is done individually
if hdr1['r_out'] < 100:
iBZ1 = i_of(geom1,40) # most SANEs
rBZ1 = 40
else:
iBZ1 = i_of(geom1,100) # most MADs
rBZ1 = 100
if hdr2['r_out'] < 100:
iBZ2 = i_of(geom2,40)
rBZ2 = 40
else:
iBZ2 = i_of(geom2,100)
rBZ2 = 100
if movie_type == "simplest":
# Simplest movie: just RHO
gs = gridspec.GridSpec(1, 2)
ax_slc = [plt.subplot(gs[0]), plt.subplot(gs[1])]
bplt.plot_xz(ax_slc[0], geom, np.log10(dump1['RHO']), label=r"$\log_{10}(\rho)$, MAD",
ylabel=False, vmin=rho_l, vmax=rho_h, window=window, half_cut=True, cmap='jet')
bplt.overlay_field(ax_slc[0], geom, dump1, nlines1)
bplt.plot_xz(ax_slc[1], geom, np.log10(dump2['RHO']), label=r"$\log_{10}(\rho)$, SANE",
ylabel=False, vmin=rho_l, vmax=rho_h, window=window, half_cut=True, cmap='jet')
bplt.overlay_field(ax_slc[1], geom, dump2, nlines2)
elif movie_type == "simpler":
# Simpler movie: RHO and phi
gs = gridspec.GridSpec(2, 2, height_ratios=[5, 1])
ax_slc = [plt.subplot(gs[0,0]), plt.subplot(gs[0,1])]
ax_flux = [plt.subplot(gs[1,:])]
bplt.plot_xz(ax_slc[0], geom, np.log10(dump1['RHO']), label=r"$\log_{10}(\rho)$, MAD",
ylabel=False, vmin=rho_l, vmax=rho_h, window=window, cmap='jet')
bplt.overlay_field(ax_slc[0], geom, dump1, nlines1)
bplt.plot_xz(ax_slc[1], geom, np.log10(dump2['RHO']), label=r"$\log_{10}(\rho)$, SANE",
ylabel=False, vmin=rho_l, vmax=rho_h, window=window, cmap='jet')
bplt.overlay_field(ax_slc[1], geom, dump2, nlines2)
# This is way too custom
ax = ax_flux[0]; ylim=[0,80]
slc1 = np.where((diag1['phi'] > ylim[0]) & (diag1['phi'] < ylim[1]))
slc2 = np.where((diag2['phi'] > ylim[0]) & (diag2['phi'] < ylim[1]))
ax.plot(diag1['t'][slc1], diag1['phi'][slc1], 'r', label="MAD")
ax.plot(diag2['t'][slc2], diag2['phi'][slc2], 'b', label="SANE")
ax.set_xlim([diag1['t'][0], diag1['t'][-1]])
ax.axvline(dump1['t'], color='r')
ax.set_ylim(ylim)
ax.set_ylabel(r"$\phi_{BH}$")
ax.legend(loc=2)
elif movie_type == "rho_cap":
axes = [plt.subplot(2,3,i) for i in range(1,7)]
bplt.plot_slices(axes[0], axes[1], geom1, dump1, np.log10(dump1['RHO']),
label=r"$\log_{10}(\rho) (1)$", vmin=-3, vmax=2, cmap='jet')
bplt.overlay_contours(axes[0], geom1, geom1['r'], [rBZ1], color='k')
bplt.plot_thphi(axes[2], geom1, np.log10(dump1['RHO'][iBZ1,:,:]), iBZ1, vmin=-4, vmax=1,
label=r"$\log_{10}(\rho)$ $\theta-\phi$ slice r="+str(rBZ1)+" (1)")
bplt.plot_slices(axes[3], axes[4], geom2, dump2, np.log10(dump2['RHO']),
label=r"$\log_{10}(\rho) (2)$", vmin=-3, vmax=2, cmap='jet')
bplt.overlay_contours(axes[3], geom2, geom2['r'], [rBZ2], color='k')
bplt.plot_thphi(axes[5], geom2, np.log10(dump2['RHO'][iBZ2,:,:]), iBZ2, vmin=-4, vmax=1,
label=r"$\log_{10}(\rho)$ $\theta-\phi$ slice r="+str(rBZ2)+" (2)")
pad = 0.05
plt.subplots_adjust(left=pad, right=1-pad, bottom=2*pad, top=1-pad)
plt.savefig(imname, dpi=1920/FIGX)
plt.close(fig)
vmax=1e12)
plt.tight_layout()
plt.savefig(name + "_xy.png", dpi=100)
plt.close(fig)
fig = plt.figure(figsize=(FIGX, FIGY))
# Plot XZ
if var in ['jcon', 'ucon', 'ucov', 'bcon', 'bcov']:
axes = [plt.subplot(2, 2, i) for i in range(1, 5)]
for n in range(4):
bplt.plot_xz(axes[n],
geom,
np.log10(dump[var][:, :, :, n]),
arrayspace=USEARRSPACE,
window=window)
elif var in ['divB2D', 'divE2D', 'divE2D_face', 'divB3D']:
ax = plt.subplot(1, 1, 1)
bplt.plot_xz(ax,
geom,
np.log10(np.abs(dump[var])),
arrayspace=USEARRSPACE,
window=window,
vmin=-6,
vmax=0)
if var in ['divE2D', 'divE2D_face']:
#JE1 = -T_mixed(dump, 1,0)
#JE2 = -T_mixed(dump, 2,0)
