def avg_dump_file(infile,
rmin,
rmax,
zmin,
zmax,
fix_ghosts=False,
N2CPU=2,
N3CPU=11,
geom=None):
"Average dump data in infile and save to new file."
dump = io.load_dump(infile, geom=geom)
if 'dump_id' in dump.keys():
count = dump['dump_id']
else:
numbers = [int(n) for n in re.findall('\d+', infile)]
if numbers:
count = max(numbers)
else:
count = None
outfile = new_dumpname(count)
outfile = os.path.join(os.path.dirname(infile), outfile)
with h5py.File(infile, 'r') as src:
with h5py.File(outfile, 'w') as dest:
copy_hdr(src, dest)
avg_dump(src, dest, dump, rmin, rmax, zmin, zmax, fix_ghosts,
N2CPU, N3CPU)
def get_from_name(self, name, hdr, geom):
dump = io.load_dump(name, geom=geom)
t = dump['t']
print("t = ", t)
if self.finder is not None:
mean, std, equil, equil_std, diff = self.get_from_dump(
dump, hdr, geom)
else:
mean, std = self.get_from_dump(dump, hdr, geom)
print("\tt = ", t, " complete")
if self.finder is not None:
return t, mean, std, equil, equil_std, diff
else:
return t, mean, std
def get_ravg(n):
print 'n = %i' % n
dump = io.load_dump(files[n], geom)
t[n] = dump['t']
rho_av = (gdet * dx2 * dx3 * dump['RHO']).sum(axis=-1).sum(axis=-1)
mdot[n, :] = -(gdet * dx2 * dx3 * dump['ucon'][:, :, :, 1] *
dump['RHO']).sum(axis=-1).sum(axis=-1)
Phi[n, :] = 1. / r * (gdet[:, N2 / 2, :] * dx1 * dx3 *
np.fabs(dump['B2'][:, N2 / 2, :])).sum(axis=-1)
Sigma[n, :] = 1. / (2. * np.pi * r[:]) * (
gdet * dx2 * dx3 * dump['RHO']).sum(axis=-1).sum(axis=-1)
b2[n, :] = dump['bsq'][:, N2 / 2, :].mean(axis=-1)
Thetae_sadw[n, :] = (gdet * dx2 * dx3 * dump['Thetae'] *
dump['RHO']).sum(axis=-1).sum(axis=-1)
Thetap_sadw[n, :] = (gdet * dx2 * dx3 * dump['Thetap'] *
dump['RHO']).sum(axis=-1).sum(axis=-1)
b2_sadw[n, :] = (gdet * dx2 * dx3 * dump['bsq'] *
dump['RHO']).sum(axis=-1).sum(axis=-1)
ucon, ucov, bcon, bcov = io.get_state(dump, geom)
bcon[:, :, :, 0] = dump['B2'][:, :, :] * ucov[:, :, :, 2]
bcon[:, :, :, 1] = bcon[:, :, :, 0] * ucon[:, :, :, 1] / ucon[:, :, :, 0]
bcon[:, :, :,
2] = dump['B2'][:, :, :] + bcon[:, :, :, 0] * ucon[:, :, :,
2] / ucon[:, :, :,
0]
bcon[:, :, :, 3] = bcon[:, :, :, 0] * ucon[:, :, :, 3] / ucon[:, :, :, 0]
for mu in xrange(4):
bcov[:, :, :, mu] = (bcon[:, :, :, :] *
geom['gcov'][:, :, None, mu, :]).sum(axis=-1)
b2z = ((bcov * bcon).mean(axis=-1)).mean(axis=-1)
P = ((hdr['gam'] - 1.) * dump['UU']).mean(axis=-1)
betaz[n, :] = 2. * P[:, N2 / 2] / b2z[:, N2 / 2]
Thetae_sadw[n, :] /= rho_av
Thetap_sadw[n, :] /= rho_av
b2_sadw[n, :] /= rho_av
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)
matplotlib.rc('font', **font)
if len(sys.argv) < 2:
print('ERROR Format is ')
print(' spec.py [filename] freq=None')
sys.exit()
fnam = sys.argv[1]
dfolder = fnam.rsplit('/', 1)[0]
freq = None
if len(sys.argv) >= 3:
freq = float(sys.argv[2][5:])
hdr = io.load_hdr(fnam)
dump = io.load_dump(fnam)
nu = hdr['nu_spec']
if freq is not None:
print(dump['nuLnu'].shape)
nuLnu = dump['nuLnu'].sum(axis=0).sum(axis=-2)
for n in range(hdr['nth']):
print(
np.interp(freq, nu, nuLnu[n, :]) * 4 * np.pi /
hdr['dOmega'][n, :].sum(axis=-1))
print(n)
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
nuLnu = dump['nuLnu'].sum(axis=0).sum(axis=0).sum(axis=0)
ax.step(nu, nuLnu, where='mid', color='k')
default=None,
help='Label for y axis')
parser.add_argument('--rmin',
dest='rmin',
type=float,
help='Min radius',
default = None)
parser.add_argument('--rmax'
,dest='rmax',
type=float,
help='Max radius',
default=None)
args = parser.parse_args()
mpl.rcParams.update({'font.size':18})
dump = io.load_dump(args.dump)
geom = dump['geom']
hdr = dump['hdr']
varname = args.qname
var = dump[varname]
fig,ax = plt.subplots()
if args.ylabel is not None:
ylabel = args.ylabel
else:
ylabel = varname
plot_density_weighted_profile(ax,geom,var,dump,ylabel=ylabel)
if args.rmin is None:
rmin = geom['r'].min()
else:
rmin = args.rmin
if args.rmax is None:
def get_phis(i):
dump = io.load_dump(dfiles[i], geom)
phi0 = dump['var0'].reshape(N)
phi1 = dump['var1'].reshape(N)
return phi0, phi1
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))
ax = plt.subplot(2, 2, 1)
ax.pcolormesh(np.log10(dump['RHO'][:, :, 0]).transpose(),
vmin=-4,
vmax=0,
cmap='jet')
ax.set_xlabel('i')
ax.set_ylabel('j')
#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)
ax.pcolormesh(np.log10(dump['beta'][:, :, 0]).transpose(),
vmin=-2,
vmax=2,
cmap='RdBu_r')
ax.set_xlabel('i')
ax.set_ylabel('j')
#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)
ax.pcolormesh(np.log10(dump['RHO'][:, hdr['N2'] / 2, :]).transpose(),
vmin=-4,
vmax=0,
cmap='jet')
ax.set_xlabel('i')
ax.set_ylabel('k')
#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)
ax.pcolormesh(np.log10(dump['beta'][:, hdr['N2'] / 2, :]).transpose(),
vmin=-2,
vmax=2,
cmap='RdBu_r')
ax.set_xlabel('i')
ax.set_ylabel('k')
#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])
fig.suptitle('t = %05.2g' % dump['t'])
#ax.pcolormesh(dump['X1'][:,:,0], dump['X2'][:,:,0], dump['RHO'][:,:,0])
plt.savefig(imname, bbox_inches='tight', dpi=100)
plt.close(fig)
import numpy as np
import matplotlib.pyplot as plt
RES = [32, 64, 128, 256]
NVAR = 8
L1 = np.zeros(len(RES))
# RUN PROBLEM FOR EACH RESOLUTION AND ANALYZE RESULT
for m in range(len(RES)):
os.chdir('../dumps_' + str(RES[m]))
dfiles = io.get_dumps_list(".")
hdr, geom, dump0 = io.load_all(dfiles[0])
dump1 = io.load_dump(dfiles[-1], hdr, geom)
r = geom['r'][:, hdr['n2'] // 2, 0]
# print("r_eh is {}".format(hdr['r_eh']))
imin = 0
while r[imin] < hdr['r_eh']:
imin += 1
rho0 = np.mean(dump0['RHO'][imin:, :, 0], axis=1)
rho1 = np.mean(dump1['RHO'][imin:, :, 0], axis=1)
L1[m] = np.mean(np.fabs(rho1 - rho0))
# MEASURE CONVERGENCE
os.environ['OMP_NUM_THREADS'] = '%d' % NTHREADS
bcall(run_args, int(NPROC), stdin=DEVNULL)
else:
bcall(run_args)
os.chdir('../')
# load simulation output
dumps = sorted(glob.glob(os.path.join(TMP_DIR, '') + '/dumps/dump*.h5'))
hdr = io.load_hdr(dumps[0])
geom = io.load_geom(hdr)
N1, N2, N3 = [hdr[i] for i in ['N1', 'N2', 'N3']]
times = np.empty(len(dumps))
Yes = np.empty((len(dumps), N1, N2))
dump = io.load_dump(dumps[0])
x, y = geom['xf'][:, :, 0], geom['yf'][:, :, 0]
#STOPPED EDITING HERE
for i, d in enumerate(dumps):
dump = io.load_dump(d)
times[i] = dump['t']
Yes[i] = dump['Ye'].reshape(N1, N2)
times *= hdr['T_unit']
Ymaxs = Yes.max(axis=(-2, -1))
Ymins = Yes.min(axis=(-2, -1))
Yavgs = Yes.mean(axis=(-2, -1))
Ydiffs = 100. * (Yes - 0.225)
dYmax = np.abs(Ymaxs - Yavgs) / Yavgs
NVAR = 8
L1 = np.zeros(len(RES))
os.chdir(TMP_DIR)
# RUN PROBLEM FOR EACH RESOLUTION AND ANALYZE RESULT
for m in xrange(len(RES)):
call_string = ['./bhlight_' + str(RES[m]), '-p', 'param.dat']
print(call_string)
bcall(call_string)
dfiles = np.sort(glob.glob('dumps/dump*.h5'))
hdr = io.load_hdr(dfiles[0])
geom = io.load_geom(hdr)
dump0 = io.load_dump(dfiles[0], geom)
dump1 = io.load_dump(dfiles[-1], geom)
r = geom['r'][:, 0, 0]
imin = 0
while r[imin] < hdr['Reh']:
imin += 1
rho0 = np.mean(dump0['RHO'][imin:, :, 0], axis=1)
rho1 = np.mean(dump1['RHO'][imin:, :, 0], axis=1)
L1[m] = np.mean(np.fabs(rho1 - rho0))
files = glob.glob('dumps/*')
for f in files:
os.remove(f)
files = glob.glob('restarts/*')
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)
def get_luminosity_from_fnam(fnam,hdr,geom):
print(fnam)
dump = io.load_dump(fnam,geom=geom)
t = dump['t']
nu,rate,lum = get_luminosity_from_dump(dump,hdr)
return t,nu,rate,lum
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()
import matplotlib
import matplotlib.pyplot as plt
SADW = True
if len(sys.argv) < 3:
print('ERROR Format is')
print(' savg.py [/path/to/dumpfile] [variable] [SADW=True]')
sys.exit()
path = sys.argv[1]
vnam = sys.argv[2]
hdr = io.load_hdr(path)
geom = io.load_geom(hdr)
dump = io.load_dump(path, geom)
N1 = hdr['N1']
N2 = hdr['N2']
N3 = hdr['N3']
dx1 = hdr['dx'][1]
dx2 = hdr['dx'][2]
dx3 = hdr['dx'][3]
savg = np.zeros(N1)
for i in range(N1):
vol = 0
vol = (dump['RHO'][i, :, :] * dx2 * dx3 *
geom['gdet'][i, :, :]).sum(axis=-1).sum(axis=-1)
val = (dump[vnam][i, :, :] *
(dump['RHO'][i, :, :] * dx2 * dx3 * geom['gdet'][i, :, :])).sum(
else:
bcall(run_args)
os.chdir('../')
# READ SIMULATION OUTPUT
dfiles = sorted(glob.glob(os.path.join(TMP_DIR,'')+'/dumps/dump*.h5'))
Nd = len(dfiles)
hdr = io.load_hdr(dfiles[0])
geom = io.load_geom(hdr)
times = np.empty(Nd)
yes = np.empty(Nd)
us = np.empty(Nd)
rhos = np.empty(Nd)
press = np.empty(Nd)
for i,d in enumerate(dfiles):
dump = io.load_dump(d)
times[i] = dump['t']
us[i] = dump['UU'].mean()
yes[i] = dump['Ye'].mean()
rhos[i] = dump['RHO'].mean()
press[i] = dump['PRESS'].mean()
times *= hdr['T_unit']
rhos *= hdr['RHO_unit']
us *= hdr['U_unit']
press *= hdr['U_unit']
ucon0 = dump['ucon'][0,0,0,0]
ucov0 = dump['ucov'][0,0,0,0]
# GET ANALYTIC SOLUTION
#emissivity
cnu_flat = 1
def avg_dump(n):
out = {}
out['t'] = io.get_dump_time(dumps[n])
# When we don't know times, fudge
if out['t'] == 0 and n != 0:
out['t'] = n
if out['t'] < tstart or out['t'] > tend:
#print("Loaded {} / {}: {} (SKIPPED)".format((n+1), len(dumps), out['t']))
# Still return the time
return out
else:
print("Loaded {} / {}: {}".format((n + 1), len(dumps), out['t']))
dump = io.load_dump(dumps[n], hdr, geom, extras=False)
# EHT Radial profiles: special fn for profile, averaged over phi, 1/3 theta, time
if calc_ravgs:
for var in [
'rho', 'Theta', 'B', 'Pg', 'Ptot', 'beta', 'u^phi', 'u_phi',
'sigma', 'FM'
]:
out[var + '_rt'] = eht_profile(geom, d_fns[var](dump), jmin, jmax)
out[var + '_jet_rt'] = eht_profile(
geom, d_fns[var](dump), 0, jmin) + eht_profile(
geom, d_fns[var](dump), jmax, geom['n2'])
if out['t'] >= tavg_start and out['t'] <= tavg_end:
out[var + '_r'] = out[var + '_rt']
out[var + '_jet_r'] = out[var + '_jet_rt']
if out['t'] >= tavg_start and out['t'] <= tavg_end:
# CORRELATION FUNCTION
for var in ['rho', 'betainv']:
Rvar = corr_midplane(geom, d_fns[var](dump))
out[var + '_cf_rphi'] = Rvar
out[var + '_cf_10_phi'] = Rvar[i_of(geom, 10), :]
# THETA AVERAGES
for var in ['betainv', 'sigma']:
out[var + '_25_th'] = theta_av(geom,
d_fns[var](dump),
i_of(geom, 25),
5,
fold=False)
# These are divided averages, not average of division, so not amenable to d_fns
Fcov01, Fcov13 = Fcov(geom, dump, 0, 1), Fcov(geom, dump, 1, 3)
out['omega_hth'] = theta_av(geom, Fcov01, iEH, 1) / theta_av(
geom, Fcov13, iEH, 1)
out['omega_av_hth'] = theta_av(geom, Fcov01, iEH, 5) / theta_av(
geom, Fcov13, iEH, 5)
# This produces much worse results
#out['omega_alt_hth'] = theta_av(Fcov(dump, 0, 2), iEH, 1) / theta_av(Fcov(dump, 2, 3), iEH, 1)
#out['omega_alt_av_hth'] = theta_av(Fcov(dump, 0, 2), iEH-2, 5) / theta_av(Fcov(dump, 2, 3), iEH-2, 5)
if calc_basic:
# FIELD STRENGTHS
# The HARM B_unit is sqrt(4pi)*c*sqrt(rho) which has caused issues:
#norm = np.sqrt(4*np.pi) # This is what I believe matches T,N,M '11 and Narayan '12
norm = 1 # This is what the EHT comparison uses?
if geom['mixed_metrics']:
# When different, B1 will be in the _vector_ coordinates. Must perform the integral in those instead of zone coords
# Some gymnastics were done to keep in-memory size small
dxEH = np.einsum(
"i,...ij->...j",
np.array([0, geom['dx1'], geom['dx2'], geom['dx3']]),
np.linalg.inv(geom['vec_to_grid'][iEH, :, :, :]))
out['Phi_b'] = 0.5 * norm * np.sum(
np.fabs(dump['B1'][iEH, :, :]) * geom['gdet_vec'][iEH, :, None]
* dxEH[:, None, 2] * dxEH[:, None, 3],
axis=(0, 1))
else:
out['Phi_sph_r'] = 0.5 * norm * sum_shell(geom, np.fabs(
dump['B1']))
out['Phi_b'] = out['Phi_sph_r'][iEH]
out['Phi_mid_r'] = np.zeros_like(out['Phi_sph_r'])
for i in range(geom['n1']):
out['Phi_mid_r'][i] = norm * sum_plane(
geom, -dump['B2'], within=i)
# FLUXES
# Radial profiles of Mdot and Edot, and their particular values
# EHT code-comparison normalization has all these values positive
for var, flux in [['Edot', 'FE'], ['Mdot', 'FM'], ['Ldot', 'FL']]:
if out['t'] >= tavg_start and out['t'] <= tavg_end:
out[flux + '_r'] = sum_shell(geom, d_fns[flux](dump))
out[var] = sum_shell(geom, d_fns[flux](dump), at_zone=iF)
# Mdot and Edot are defined inward
out['Mdot'] *= -1
out['Edot'] *= -1
# Maxima (for gauging floors)
for var in ['sigma', 'betainv', 'Theta']:
out[var + '_max'] = np.max(d_fns[var](dump))
# Minima
for var in ['rho', 'U']:
out[var + '_min'] = np.min(d_fns[var](dump))
# TODO KEL? plot in "floor space"? Full set of energy ratios?
# Profiles of different fluxes to gauge jet power calculations
if calc_jet_profile:
for var in [
'rho', 'bsq', 'FM', 'FE', 'FE_EM', 'FE_Fl', 'FL', 'FL_EM',
'FL_Fl', 'betagamma', 'Be_nob', 'Be_b', 'mu'
]:
out[var + '_100_tht'] = np.sum(d_fns[var](dump)[iBZ], axis=-1)
if out['t'] >= tavg_start and out['t'] <= tavg_end:
out[var + '_100_th'] = out[var + '_100_tht']
out[var + '_100_thphi'] = d_fns[var](dump)[iBZ, :, :]
out[var + '_rth'] = d_fns[var](dump).mean(axis=-1)
# Blandford-Znajek Luminosity L_BZ
# This is a lot of luminosities!
if calc_jet_cuts:
# TODO cut on phi/t averages? -- needs 2-pass cut...
cuts = {
'sigma1': lambda dump: (d_fns['sigma'](dump) > 1),
#'sigma10' : lambda dump : (d_fns['sigma'](dump) > 10),
'Be_b0': lambda dump: (d_fns['Be_b'](dump) > 0.02),
'Be_b1': lambda dump: (d_fns['Be_b'](dump) > 1),
'Be_nob0': lambda dump: (d_fns['Be_nob'](dump) > 0.02),
'Be_nob1': lambda dump: (d_fns['Be_nob'](dump) > 1),
#'mu1' : lambda dump : (d_fns['mu'](dump) > 1),
#'mu2' : lambda dump : (d_fns['mu'](dump) > 2),
#'mu3' : lambda dump : (d_fns['mu'](dump) > 3),
'bg1': lambda dump: (d_fns['betagamma'](dump) > 1.0),
'bg05': lambda dump: (d_fns['betagamma'](dump) > 0.5),
'allp': lambda dump: (d_fns['FE'](dump) > 0)
}
# Terminology:
# LBZ = E&M energy only, any cut
# Lj = full E flux, any cut
# Ltot = Lj_allp = full luminosity wherever it is positive
for lum, flux in [['LBZ', 'FE_EM'], ['Lj', 'FE']]:
for cut in cuts.keys():
out[lum + '_' + cut + '_rt'] = sum_shell(geom,
d_fns[flux](dump),
mask=cuts[cut](dump))
if out['t'] >= tavg_start and out['t'] <= tavg_end:
out[lum + '_' + cut + '_r'] = out[lum + '_' + cut + '_rt']
out[lum + '_' + cut] = out[lum + '_' + cut + '_rt'][iBZ]
if calc_lumproxy:
rho, Pg, B = d_fns['rho'](dump), d_fns['Pg'](dump), d_fns['B'](dump)
# See EHT code comparison paper
j = rho**3 / Pg**2 * np.exp(-0.2 * (rho**2 / (B * Pg**2))**(1. / 3.))
out['Lum_rt'] = eht_profile(geom, j, jmin, jmax)
if calc_etot:
# Total energy and current, summed by shells to allow cuts on radius
for tot_name, var_name in [['Etot', 'JE0']]:
out[tot_name + '_rt'] = sum_shell(geom, d_fns[var_name](dump))
for tot_name, var_name in [['Jsqtot', 'jsq']]:
out[tot_name + '_rt'] = sum_shell(geom, d_fns[var_name](geom,
dump))
if calc_efluxes:
# Conserved (maybe; in steady state) 2D energy flux
for var in ['JE0', 'JE1', 'JE2']:
out[var + '_rt'] = sum_shell(geom, d_fns[var](dump))
if out['t'] >= tavg_start and out['t'] <= tavg_end:
out[var + '_rth'] = d_fns[var](dump).mean(axis=-1)
# Total outflowing portions of variables
if calc_outfluxes:
for name, var in [['outflow', 'FM'], ['outEflow', 'FE']]:
var_temp = d_fns[var](dump)
out[name + '_rt'] = sum_shell(geom, var_temp, mask=(var_temp > 0))
if out['t'] >= tavg_start and out['t'] <= tavg_end:
out[name + '_r'] = out[name + '_rt']
if calc_pdfs:
for var in ['betainv', 'rho']:
out[var + '_pdf'], _ = np.histogram(
np.log10(d_fns[var](dump)),
bins=pdf_nbins,
range=(-3.5, 3.5),
weights=np.repeat(geom['gdet'], geom['n3']).reshape(
(geom['n1'], geom['n2'], geom['n3'])),
density=True)
dump.clear()
del dump
return out
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)
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), frameon=False)
ax = fig.gca()
#ax = plt.subplot(1,1,1)
x = geom['x']
y = geom['y']
z = geom['z']
x = bplt.flatten_xz(x, hdr, flip=True)
y = bplt.flatten_xz(y, hdr, flip=True)
z = bplt.flatten_xz(z, hdr)
# rcylindrical
rcyl = np.sqrt(x**2 + y**2)
lrho = bplt.flatten_xz(np.log10(dump['RHO']), hdr)
mesh = ax.pcolormesh(rcyl[0:N1, 0:N2],
z[0:N1, 0:N2],
lrho[0:N1, 0:N2],
cmap='jet',
vmin=-4,
vmax=0,
shading='gouraud')
cax = fig.add_axes([0.135, 0.15, 0.025, 0.7])
cb = fig.colorbar(mesh, cax=cax, ticks=[-4, -2, 0])
cb.set_label(label='Density', fontsize=20)
cax.tick_params(axis='both', which='major', labelsize=20)
lbeta = bplt.flatten_xz(np.log10(dump['beta']), hdr)
mesh = ax.pcolormesh(rcyl[N1:, 0:N2],
z[N1:, 0:N2],
lbeta[N1:, 0:N2],
cmap='inferno',
vmin=-2,
vmax=2,
shading='gouraud')
cax = fig.add_axes([0.865, 0.15, 0.025, 0.7])
cb = fig.colorbar(mesh, cax=cax, ticks=[-2, 0, 2])
cb.set_label(label='Plasma beta', fontsize=20)
cb.ax.yaxis.set_label_position('left')
cax.yaxis.set_ticks_position('left')
cax.tick_params(axis='both', which='major', labelsize=20)
ax.set_xlim([-4. / 3. * SIZE, 4. / 3. * SIZE])
ax.set_ylim([-SIZE, SIZE])
ax.set_xticks([])
ax.set_yticks([])
circle1 = plt.Circle((0, 0),
1. + np.sqrt(1. - hdr['a']**2),
color='k',
zorder=2)
ax.add_artist(circle1)
#ax.axvline(hdr['Risco'], color='k', linestyle='--', linewidth=2)
#ax.axvline(-hdr['Risco'], color='k', linestyle='--', linewidth=2)
bplt.overlay_field(ax,
geom,
dump,
linestyle='-',
linewidth=1,
linecolor='k',
NLEV=10)
plt.savefig(os.path.join(FRAMEDIR, 'frame_%08d.png' % n),
bbox_inches='tight',
pad_inches=0,
dpi=100)
plt.close(fig)
dvar[4] = -0.224080007379
dvar[5] = 0.406380545676
dvar[6] = -0.203190272838
dvar[7] = -0.203190272838
dvar *= amp
# USE DUMPS IN FOLDERS OF GIVEN FORMAT
for m in xrange(len(RES)):
print '../dumps_' + str(RES[m]) + '_' + str(MODES[n])
os.chdir('../dumps_' + str(RES[m]) + '_' + str(MODES[n]))
dfile = np.sort(glob.glob('dump*.h5'))[-1]
hdr = io.load_hdr(dfile)
geom = io.load_geom(hdr, dfile)
dump = io.load_dump(hdr, geom, dfile)
X1 = dump['x']
X2 = dump['y']
X3 = dump['z']
dvar_code = []
dvar_code.append(dump['RHO'] - var0[0])
dvar_code.append(dump['UU'] - var0[1])
dvar_code.append(dump['U1'] - var0[2])
dvar_code.append(dump['U2'] - var0[3])
dvar_code.append(dump['U3'] - var0[4])
dvar_code.append(dump['B1'] - var0[5])
dvar_code.append(dump['B2'] - var0[6])
dvar_code.append(dump['B3'] - var0[7])
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
gridfile = sys.argv[2]
var = sys.argv[3]
elif len(sys.argv) > 2:
gridfile = None
var = sys.argv[2]
# Optionally take extra name
name = sys.argv[-1]
if UNITS and var not in ['Tp']:
M_unit = float(sys.argv[-1])
if gridfile is not None:
hdr = io.load_hdr(dumpfile)
geom = io.load_geom(hdr, gridfile)
dump = io.load_dump(dumpfile, hdr, geom)
else:
# Assumes gridfile in same directory
hdr, geom, dump = io.load_all(dumpfile)
# If we're plotting a derived variable, calculate + add it
if var in ['jcov', 'jsq']:
dump['jcov'] = np.einsum("...i,...ij->...j", dump['jcon'],
geom['gcov'][:, :, None, :, n])
dump['jsq'] = np.sum(dump['jcon'] * dump['jcov'], axis=-1)
elif var in ['divE2D']:
JE1g, JE2g = T_mixed(dump, 1, 0).mean(axis=-1) * geom['gdet'], T_mixed(
dump, 2, 0).mean(axis=-1) * geom['gdet']
face_JE1 = 0.5 * (JE1g[:-1, :] + JE1g[1:, :])
face_JE2 = 0.5 * (JE2g[:, :-1] + JE2g[:, 1:])
divJE = (face_JE1[1:, 1:-1] - face_JE1[:-1, 1:-1]) / geom['dx1'] + (
# RUN EXECUTABLE
os.chdir(TMP_DIR)
#if FAST:
# util.change_rparm('tf',str(TF),'param_template.dat')
call(['./bhlight', '-p', 'param_template.dat'])
os.chdir('../')
# READ SIMULATION OUTPUT
dfiles = np.sort(glob.glob(os.path.join(TMP_DIR, '') + '/dumps/dump*.h5'))
Nd = len(dfiles)
hdr = io.load_hdr(dfiles[0])
geom = io.load_geom(hdr)
t_code = np.zeros(Nd)
Te_code = np.zeros(Nd)
for n in range(Nd):
dump = io.load_dump(dfiles[n], geom)
t_code[n] = dump['t'] * hdr['T_unit']
Te_code[
n] = dump['Thetae'][0][0][0] * cgs['ME'] * cgs['CL']**2 / cgs['KBOL']
# GET ANALYTIC SOLUTION
tf = 3 / 4 * 1.e8
Te0 = 1.e8
dump = io.load_dump(dfiles[0], geom)
ne = dump['RHO'].mean() * hdr['Ne_unit']
gam = 5. / 3.
#N = 5.4e-39 # cm^3 K^1/2 s^-1 Sr^-1 Hz^-1
t_sol = np.linspace(0, tf, 1024)
#Te0 = Te_code[0]
from scipy.integrate import odeint
def avg_dump(args):
global t
n = args
print '%08d / ' % (n + 1) + '%08d' % len(dumps)
dump = io.load_dump(dumps[n], geom)
t[n] = dump['t']
print dump['t']
rho_SADW[n, :] = WAVG(dump['RHO'], dump['RHO'])
# SHELL AVERAGES
#vol = (dx2*dx3*geom['gdet'][:,:]).sum(axis=-1)
rho_r[n, :] = (dx2 * dx3 * geom['gdet'][:, jmin:jmax, None] *
dump['RHO'][:, jmin:jmax, :]).sum(axis=-1).sum(axis=-1)
Theta = (hdr['gam'] - 1.) * dump['UU'][:, :, :] / dump['RHO'][:, :, :]
Theta_r[n, :] = (dx2 * dx3 * geom['gdet'][:, jmin:jmax, None] *
Theta[:, jmin:jmax, :]).sum(axis=-1).sum(axis=-1)
B = np.sqrt(dump['bsq'][:, :, :])
B_r[n, :] = (dx2 * dx3 * geom['gdet'][:, jmin:jmax, None] *
B[:, jmin:jmax, :]).sum(axis=-1).sum(axis=-1)
Pg = (hdr['gam'] - 1.) * dump['UU'][:, :, :]
Pg_r[n, :] = (dx2 * dx3 * geom['gdet'][:, jmin:jmax, None] *
Pg[:, jmin:jmax, :]).sum(axis=-1).sum(axis=-1)
Ptot = Pg + dump['bsq'][:, :, :] / 2
Ptot_r[n, :] = (dx2 * dx3 * geom['gdet'][:, jmin:jmax, None] *
Ptot[:, jmin:jmax, :]).sum(axis=-1).sum(axis=-1)
betainv = (dump['bsq'][:, :, :] / 2) / Pg
betainv_r[n, :] = (dx2 * dx3 * geom['gdet'][:, jmin:jmax, None] *
betainv[:, jmin:jmax, :]).sum(axis=-1).sum(axis=-1)
uphi = (dump['ucon'][:, :, :, 3])
uphi_r[n, :] = (dx2 * dx3 * geom['gdet'][:, jmin:jmax, None] *
uphi[:, jmin:jmax, :]).sum(axis=-1).sum(axis=-1)
rho_r[n, :] /= vol
Theta_r[n, :] /= vol
B_r[n, :] /= vol
Pg_r[n, :] /= vol
Ptot_r[n, :] /= vol
betainv_r[n, :] /= vol
uphi_r[n, :] /= vol
# FLUXES
iF = 5
Mdot[n] = (dump['RHO'][iF, :, :] * dump['ucon'][iF, :, :, 1] *
geom['gdet'][iF, :, None] * dx2 * dx3).sum()
#Phi[n] = 0.5*(np.fabs(dump['bcon'][iF,:,:,1])*geom['gdet'][iF,:,None]*dx2*dx3).sum()
Phi[n] = 0.5 * (np.fabs(dump['B1'][iF, :, :]) * geom['gdet'][iF, :, None] *
dx2 * dx3).sum()
Trphi = (dump['RHO'] + dump['UU'] + (hdr['gam'] - 1.) * dump['UU'] +
dump['bsq']) * dump['ucon'][:, :, :, 1] * dump['ucov'][:, :, :, 3]
Trphi -= dump['bcon'][:, :, :, 1] * dump['bcov'][:, :, :, 3]
Trt = (dump['RHO'] + dump['UU'] + (hdr['gam'] - 1.) * dump['UU'] +
dump['bsq']) * dump['ucon'][:, :, :, 1] * dump['ucov'][:, :, :, 0]
Trt -= dump['bcon'][:, :, :, 1] * dump['bcov'][:, :, :, 0]
Ldot[n] = (Trphi[iF, :, :] * geom['gdet'][iF, :, None] * dx2 * dx3).sum()
Edot[n] = -(Trt[iF, :, :] * geom['gdet'][iF, :, None] * dx2 * dx3).sum()
rho = dump['RHO']
P = (hdr['gam'] - 1.) * dump['UU']
B = np.sqrt(dump['bsq'])
C = 0.2
j = rho**3 * P**(-2) * np.exp(-C * (rho**2 / (B * P**2))**(1. / 3.))
Lum[n] = (j * geom['gdet'][:, :, None] * dx1 * dx2 * dx3).sum()
dvar[2] = -0.253320198552
if MODES[n] == 2: # ALFVEN
dvar[3] = 0.480384461415
dvar[6] = 0.877058019307
if MODES[n] == 3: # FAST
dvar[4] = 0.480384461415
dvar[7] = 0.877058019307
dvar *= amp
# RUN PROBLEM FOR EACH RESOLUTION AND ANALYZE RESULT
for m in xrange(len(RES)):
print['./bhlight_' + str(RES[m]), '-p', 'param.dat']
call(['./bhlight_' + str(RES[m]), '-p', 'param.dat'])
dfiles = np.sort(glob.glob('dumps/dump*.h5'))
dump = io.load_dump(dfiles[-1])
X1 = dump['X1'][:, 0, 0]
X2 = dump['X2'][:, 0, 0]
dvar_code = []
dvar_code.append(dump['RHO'][:, 0, 0] - var0[0])
dvar_code.append(dump['UU'][:, 0, 0] - var0[1])
dvar_code.append(dump['U1'][:, 0, 0] - var0[2])
dvar_code.append(dump['U2'][:, 0, 0] - var0[3])
dvar_code.append(dump['U3'][:, 0, 0] - var0[4])
dvar_code.append(dump['B1'][:, 0, 0] - var0[5])
dvar_code.append(dump['B2'][:, 0, 0] - var0[6])
dvar_code.append(dump['B3'][:, 0, 0] - var0[7])
dvar_sol = []
for k in xrange(NVAR):
dvar_sol.append(np.real(dvar[k]) * np.cos(k1 * X1))
def get_phi(i):
dump = io.load_dump(dfiles[i], geom)
phi = dump['var0'].reshape(mshape)
return phi
dvar[1] = 0.635193903263
dvar[2] = -0.102965815319
dvar[3] = -0.316873207561
dvar[5] = 0.359559114174
dvar[6] = -0.359559114174
dvar *= amp
# RUN PROBLEM FOR EACH RESOLUTION AND ANALYZE RESULT
for m in xrange(len(RES)):
print ['./bhlight_' + str(RES[m]), '-p', 'param_template.dat']
call(['./bhlight_' + str(RES[m]), '-p', 'param_template.dat'])
dfiles = np.sort(glob.glob('dumps/dump*.h5'))
hdr = io.load_hdr(dfiles[-1])
geom = io.load_geom(hdr, recalc=True)
dump = io.load_dump(dfiles[-1],geom)
X1 = geom['x'][:,:,0]
X2 = geom['y'][:,:,0]
dvar_code = []
dvar_code.append(dump['RHO'][:,:,0] - var0[0])
dvar_code.append(dump['UU'][:,:,0] - var0[1])
dvar_code.append(dump['U1'][:,:,0] - var0[2])
dvar_code.append(dump['U2'][:,:,0] - var0[3])
dvar_code.append(dump['U3'][:,:,0] - var0[4])
dvar_code.append(dump['B1'][:,:,0] - var0[5])
dvar_code.append(dump['B2'][:,:,0] - var0[6])
dvar_code.append(dump['B3'][:,:,0] - var0[7])
#dvar_sol = []
dvar_sol = np.zeros([RES[m], RES[m]])
for k in xrange(NVAR):
USEARRSPACE=True
NLINES = 20
SIZE = 600
FIGX = 20
FIGY = 16
dump1file = sys.argv[1]
dump2file = sys.argv[2]
imname = sys.argv[3]
hdr, geom, dump1 = io.load_all(dump1file, derived_vars=False)
#Hopefully this fails for dumps that shouldn't be compared
dump2 = io.load_dump(dump2file, hdr, geom, derived_vars=False)
N1 = hdr['n1']; N2 = hdr['n2']; N3 = hdr['n3']
log_floor = -60
# TODO properly option log, rel, lim
def plot_diff_xy(ax, var, rel=False, lim=None):
if rel:
if lim is not None:
bplt.plot_xy(ax, geom, np.abs((dump1[var] - dump2[var])/dump1[var]), vmin=0, vmax=lim, label=var, cbar=False, arrayspace=USEARRSPACE)
else:
bplt.plot_xy(ax, geom, np.abs((dump1[var] - dump2[var])/dump1[var]), label=var, cbar=False, arrayspace=USEARRSPACE)
else:
if lim is not None:
bplt.plot_xy(ax, geom, np.log10(np.abs(dump1[var] - dump2[var])), vmin=log_floor, vmax=lim, label=var, cbar=False, arrayspace=USEARRSPACE)