def getRayleighVelocities(fname, direction=[]):
"""Returns positions of the shock, and both inner an outer rayleigh line velocities joining
both sides of the shocked cell.
(xin, xout, cjin, cjout, float(ray.ds.current_time), ray.ds.parameters['x_match'])
"""
# fields = ['sound_speed', 'density', 'pressure']
# data, _ = reader.getLineout(fname, fields=fields, species=False, geom=geom)
# time, params, _, _, _ = reader.getMeta(fname)
# rad, cs, dens, pres = data[0], data[1], data[2], data[3]
fields = ['sound_speed', 'density', 'pressure']
time, pars, _, _, paths = directMeta(fname)
if len(direction)>(pars['dimensionality']-1):
print("Direction doesn't match dimensionality: {}".format(pars['dimensionality']))
return None
data, _ = getLineout(fname, fields=fields, species=False, direction=direction, geom=pars['geometry'])
#time, params, _, _, _ = directMeta(fname)
rad, cs, dens, pres = data[0], data[1], data[2], data[3]
# this fails for x_match = y_match = z_match = 0.0
linepos = ut.estimateMatch(direction, pars, vvv=False)
# print(linepos)
shockin, shockout = ut.locateShock(rad, cs, linepos, vvv=False)
# shockin, shockout = ut.locateShock(rad, cs, pars['x_match'], vvv=False)
xin, xout = rad[shockin], rad[shockout]
cjin = ut.roughCJ(dens, pres, shockin)
cjout = ut.roughCJ(dens, pres, shockout)
return xin, xout, cjin, cjout, time, pars['x_match']
def getShockConditions(fname, inward=False, addvar='temp', direction=[]):
"""(1D) Returns bulk conditions at both sides of shock.
Conditions are sorted so that output has the form: [ash, shock, fuel]
Args:
fname(str): filepath
inward(bool): toggle for inward/outward bound shock.
addvar(float): extra variable to get from lineout.
**kwargs: arguments for lineout (direction, geometry, etc.)
Returns:
(list): radii of states.
(list): densities of states.
(list): pressures of states.
(list): addvar at each state.
(float): direct Rayleigh speed for the ash state.
(float): match head position
(float): timestamp of file.
"""
fields = ['sound_speed', 'density', 'pressure', addvar]
time, pars, _, _, _ = directMeta(fname)
if len(direction) > (pars['dimensionality'] - 1):
print("Direction doesn't match dimensionality: "
"{}".format(pars['dimensionality']))
return None
data, _ = getLineout(fname,
fields=fields,
species=False,
direction=direction,
geom=pars['geometry'])
# time, params, _, _, _ = directMeta(fname)
rad, cs, dens, pres, var = data[0], data[1], data[2], data[3], data[4]
# this fails for x_match = y_match = z_match = 0.0
linepos = ut.estimateMatch(direction, pars, vvv=False)
# print(linepos)
shockin, shockout = ut.locateShock(rad, cs, linepos, vvv=False)
# print('check shockConditions')
if inward:
ind = shockin
offset = -1
else:
ind = shockout
offset = 1
# directly calculate the rayleigh speed joining both points
cjest = rayleighSpeed(pres[ind + offset], 1.0 / dens[ind + offset],
pres[ind - offset], 1.0 / dens[ind - offset])
pos = [rad[ind - offset], rad[ind], rad[ind + offset]]
condd = [dens[ind - offset], dens[ind], dens[ind + offset]]
condp = [pres[ind - offset], pres[ind], pres[ind + offset]]
xvar = [var[ind - offset], var[ind], var[ind + offset]]
return pos, condd, condp, xvar, cjest, linepos, time
def buildHelmTrojan(fname, offset=1, geom='spherical'):
"""Frankensteinian bridge between flash checkpoints and
J.Schwab's Helmholtz python module.
Joined inward/outward and spewing cj data to avoid
calling yt more than once.
Args:
fname(str): filename.
offset(int): zone offset from shock.
"""
props = ['dens', 'temp', 'pres', 'eint']
nprops = len(props)
data, species = hdf5yt.getLineout(fname,
fields=props,
species=True,
geom=geom)
nspecs = len(species)
xin, xout, cjin, cjout, time, xmatch = getRayleighVelocities(fname)
# get fuel and ash for outward shock
# xin/xout == ray @ len([x for x in ray['r'][rsort] if x<xout])
inw = len([x for x in data[0] if x < xout]) + offset
ouw = inw - 2 * offset
inv, ouv = [], []
for i in range(1, nprops + 1):
inv.append(data[i][inw])
ouv.append(data[i][ouw])
xmin, xmou = [], []
for i in range(nprops + 1, nprops + 1 + nspecs):
xmin.append(data[i][inw])
xmou.append(data[i][ouw])
_, abar, zbar = convXmass2Abun(species, xmin)
# [pressure, eint], [rho], [temp], [abar], [zbar]]
fuelo = [[inv[-2], inv[-1]], [inv[0]], [inv[1]], [abar], [zbar]]
_, abar, zbar = convXmass2Abun(species, xmou)
asho = [[ouv[-2], ouv[-1]], [ouv[0]], [ouv[1]], [abar], [zbar]]
# get fuel and ash for inward shock
inw = len([x for x in data[0] if x < xin]) - offset
ouw = inw + 2 * offset
inv, ouv = [], []
for i in range(1, nprops + 1):
inv.append(data[i][inw])
ouv.append(data[i][ouw])
xmin, xmou = [], []
for i in range(nprops + 1, nprops + 1 + nspecs):
xmin.append(data[i][inw])
xmou.append(data[i][ouw])
_, abar, zbar = convXmass2Abun(species, xmin)
fueli = [[inv[-2], inv[-1]], [inv[0]], [inv[1]], [abar], [zbar]]
_, abar, zbar = convXmass2Abun(species, xmou)
ashi = [[ouv[-2], ouv[-1]], [ouv[0]], [ouv[1]], [abar], [zbar]]
return fueli, ashi, fuelo, asho, [xin, xout, cjin, cjout, time, xmatch]
def flashProfile(fname, thresh=1e-6, xrange=[0.0, 0.0],
filetag='prof', batch=False, byM=True, direction=[]):
"""Plot bulk properties and species in a chekpoint through a ray.
Args:
fname(str): path of file.
thresh(float): threshold for species fraction.
xrange(float list): abscissa range.
filetag(str): prefix for batch mode.
batch(bool): skips returning figure, saving it to a structured directory instead.
byM(bool): plot by mass instead of radius.
direction(float list): list of spherical angles (alt, azimuth), empty for 1D.
Returns:
(mpl figure) or (None)
"""
time, pars, _, _, paths = directMeta(fname)
if len(direction)>(pars['dimensionality']-1):
print("Direction doesn't match dimensionality: {}".format(pars['dimensionality']))
return None
ad, allsp = getLineout(fname, geom=pars['geometry'],
direction=direction, srcnames=False)
keys = ['radius', 'density', 'temperature', 'pressure']
keys +=allsp
prof = dataMatrix([keys, ad.transpose()])
fig = plotDMatMerged(prof, byM=byM,
thresh=thresh, xrange=xrange)
ax = plt.gca()
a = ax.annotate("{:.5f} s".format(time),
xy=(0.0, 0.0), xytext=(0.82, 0.10), size=12,
textcoords='figure fraction', xycoords='figure fraction',
bbox=dict(boxstyle='round', fc='w', ec='k'))
if not batch:
return fig
else:
num = paths[1][-5:] # checkpoint number 'flash_hdf5_chk_0001'
dest = os.path.join(os.path.dirname(paths[0]), filetag)
name = os.path.join(dest, '{}{}.png'.format(filetag, num))
os.makedirs(dest, exist_ok=True) # bless you, p3
plt.savefig(name, format='png')
plt.close(fig)
print("Wrote: {}".format(name))
def radialSpeeds(fname,
elevation=5,
depth=5,
geom='cartesian',
dimension=2,
ref='x',
antipode=False):
"""radius vs radial speed for a wedge centered at the equator.
"""
if dimension == 2:
rawd = wedge2d(fname,
elevation,
depth,
fields=['x', 'y', 'z', 'velx', 'vely', 'velz'])
rs, vrs = [], []
for x, y, z, vx, vy, vz in zip(*rawd):
vec = np.array([x, y, z])
vel = np.array([vx, vy, vz])
r = np.sqrt(vec.dot(vec))
v = np.sqrt(vel.dot(vel))
rxv = np.cross(vec, vel)
normrxv = np.sqrt(rxv.dot(rxv))
angle = np.arcsin(normrxv / r / v)
vrad = v * np.cos(angle)
rs.append(r)
vrs.append(vrad)
elif dimension == 3:
rawd = wedge3d(
fname,
elevation,
depth,
fields=['spherical_radius', 'velocity_spherical_radius'],
reference=ref,
antipode=antipode)
rs, vrs = rawd
else:
rawd, _ = getLineout(fname, fields=['velx'], geom=geom, species=False)
rs, vrs = rawd
return rs, vrs
def speedHisto(fname,
resolution=4e7,
velrange=[1e9, 5e9],
elevation=5,
depth=5,
geom='cartesian',
dimension=2,
ref='x',
antipode=False,
cylvel=False,
species=_alphas):
"""Calculate speeds within a wedge and return masses.
histogram bin range is fixed so that one can mix wedges.
# recommended resolution: 100 km/s (Fink, 2010)
# default resolution: 400 km/s
# default max velocity: a sixth of c.
Args:
fname(str): filename.
resolution(float): bin size in cm/s.
velrange(float list): historgram range.
elevation(float): equator-north pole degree.
depth(float): equator-south pole degree.
geom(str): specify geometry for 1d file.
dimension(int): specify file dimension.
ref(str): reference axis (3D only, see datahaul.hdf5yt.wedge3d).
antipode(bool): antipodal wedge (see datahaul.hdf5yt.wedge3d).
cylvel(bool): take velx (axial velocity in cylindrical) as speed.
species(str list): specify which species data to get.
Returns:
np.array list, np.array: Masses for species queries, bin limits.
"""
# split into wedges process each, then read outputs...
# wedge2d call without 'fields' calculates cylindrical volumes
if dimension == 2:
rawd, ispecies = wedge2d(fname, elevation, depth, cylvel=cylvel)
elif dimension == 3:
rawd, ispecies = wedge3d(fname,
elevation,
depth,
reference=ref,
antipode=antipode)
else:
rawd, ispecies = getLineout(
fname, fields=['velx', 'vely', 'velz', 'cell_mass'], geom=geom)
rawd = rawd[1:] # remove the radius column
offset = len(rawd) - len(ispecies)
# see datahaul.hdf5yt.wedge2d for positions of data in rawd
speeds = rawd[0][:]
celltpls = []
for i in range(len(rawd[0])):
# multiply each X_i by the cell_masses
masses = [
rawd[j][i] * rawd[1][i] for j in range(offset,
len(ispecies) + offset)
]
celltpls.append((speeds[i], masses))
sortedcells = sorted(celltpls)
speeds, massgrid = zip(*sortedcells)
# massgrid = [[Mh1_0, Mhe4_0,...], [Mh1_1, Mhe4_1,...]]
# respecting increasing speed order
# get ranges for histogram bins
vmin, vmax = velrange
binnum = int((vmax - vmin) / resolution)
print("post.speedHisto: Velocity range: {:e} {:e}".format(vmin, vmax))
print("post.speedHisto: Bins: {}".format(binnum))
try:
counts, bins = np.histogram(speeds, bins=int(binnum))
except ValueError:
print('post.speedHisto: zero speed range.')
return None
# sort by species
species = [s.strip() for s in species]
buckets = []
for s in species:
buckets.append(np.zeros(len(bins)))
psp = dict(zip(species, range(len(species))))
for i in range(len(ispecies)):
if ispecies[i] in species: # main species filter
weights = [0]
start = 0
for c in counts:
# since the data is already sorted by speed
# counts follows the data in each bin
totalspmass = sum([m[i] for m in massgrid[start:start + c]])
# force the summed data value as histo "weight"
weights.append(totalspmass)
start += c
weights = np.array(weights)
# finally sum the array of masses to the corresponding histo
buckets[psp[ispecies[i]]] += weights
else:
continue
return buckets, bins
def PARsimProfile(fname, simfolder='', thresh=1e-4, xrange=[0.0, 0.0],
filetag='metaprof', batch=False, byM=False, direction=[], meta=False):
"""bloated method for IPP.
WARN: slowest pos will not match for non x-axis directions.
Args:
fname(str): path to checkpoint file.
simfolder(str): simulation run metadata folder.
thresh(float): species threshold.
xrange(float list): abscissa range for all plots.
filetag(str): prefix for output files (if any).
batch(bool): write file instead of returning figure.
byM(bool): plot by mass coordinate instead of radius.
direction(str list): altitude and azimuthal direction of profile ray.
Return:
(mpl figure) or (None)
"""
time, pars, _, _, paths = directMeta(fname)
if len(direction)>(pars['dimensionality']-1):
print("Direction doesn't match dimensionality: {}".format(pars['dimensionality']))
return None
ad, allsp = getLineout(fname, geom=pars['geometry'],
direction=direction, srcnames=False)
sim = simulation(simfolder)
n, t = sim.time2step(time)
step = sim.steps[int(n)]
xm, ym, zm = step.slowx, step.slowy, step.slowz
rm = np.sqrt(xm**2+ym**2+zm**2)
keys = ['radius', 'density', 'temperature', 'pressure']
keys +=allsp
prof = dataMatrix([keys, ad.transpose()])
fig = plotDMatMerged(prof, byM=byM,
thresh=thresh, xrange=xrange)
# add extra marks from simulation obj
ax = plt.gca()
a = ax.annotate("{:.5f} s".format(time),
xy=(0.0, 0.0), xytext=(0.84, 0.10), size=12,
textcoords='figure fraction', xycoords='figure fraction',
bbox=dict(boxstyle='round', fc='w', ec='k'))
axlist = fig.axes
for ax in axlist[1:]:
p = rm if not byM else prof.masses[np.where(prof.radius>rm)[0][0]]
ax.axvline(p, ls=':', color='brown', alpha=0.5, label='slowest')
if sim.CJ:
times, xins, cjins, xouts, cjouts = np.genfromtxt(sim.CJ[0], unpack=True)
if byM:
xo, xi = xouts[pars['checkpointfilenumber']], xins[pars['checkpointfilenumber']]
no = np.where(prof.radius>xo)[0][0]
ni = np.where(prof.radius>xi)[0][0]
po, pi = prof.masses[no], prof.masses[ni]
else:
po, pi = xouts[pars['checkpointfilenumber']], xins[pars['checkpointfilenumber']]
ax.axvline(po, ls=':', color='green', alpha=0.5, label='shock')
ax.axvline(pi, ls=':', color='red', alpha=0.5, label='shock')
lgd = ax.legend(ncol=1, loc='upper left', bbox_to_anchor=(1.00, 0.50),
columnspacing=0.0, labelspacing=0.0, markerfirst=False,
numpoints=3, handletextpad=0.0, edgecolor='k')
if meta:
markings = [po, pi, p, paths]
return fig, markings
#I/O
if not batch:
return fig
else:
num = paths[1][-5:] # checkpoint number 'flash_hdf5_chk_0001'
dest = os.path.join(os.path.dirname(paths[0]), filetag)
name = os.path.join(dest, '{}{}.png'.format(filetag, num))
os.makedirs(dest, exist_ok=True) # bless you, p3
plt.savefig(name, format='png')
plt.close(fig)
print("Wrote: {}".format(name))
def flashSpecies(fname, thresh=1e-6, filetag='spec', batch=False,
byM=True, direction=[], plotall=False):
"""Plot species and aggregated masses in a chekpoint through a ray.
Args:
fname(str): path of file.
thresh(float): threshold for species fraction and mass yields.
filetag(str): prefix for batch mode.
batch(bool): skips returning figure, saving it to a structured directory instead.
byM(bool): plot by mass instead of radius.
direction(float list): list of spherical angles (alt, azimuth), empty for 1D.
plotall(bool): force plotting every species found.
Returns:
(mpl figure) or (None)
"""
fields = ['density', 'temperature', 'pressure', 'velx']
time, pars, _, _, paths = directMeta(fname)
if len(direction)>(pars['dimensionality']-1):
print("Direction doesn't match dimensionality: {}".format(pars['dimensionality']))
return None
ad, allsp = getLineout(fname, geom=pars['geometry'],
direction=direction, fields=fields, srcnames=False)
keys = ['radius'] + fields
keys +=allsp
prof = dataMatrix([keys, ad.transpose()])
fig = plt.figure(figsize=(10, 8))
layout = (3,3)
# plot species
ax1 = plt.subplot2grid(layout, (0, 0), colspan=2)
skip = plotSpecies(ax1, prof, byMass=byM, thresh=thresh, plotall=plotall)
ax1.set_xlabel('Mass ($M_{\odot}$)')
# timestamp and legend
a = ax1.annotate("{:.5f} s".format(time),
xy=(0.0, 0.0), xytext=(0.65, 0.1), size=12,
textcoords='figure fraction', xycoords='figure fraction',
bbox=dict(boxstyle='round', fc='w', ec='k'))
lgd = ax1.legend(ncol=4, loc='upper left', bbox_to_anchor=(1.05, 1.0),
columnspacing=0.0, labelspacing=0.0, markerfirst=False,
numpoints=3, handletextpad=0.0)
lgd.get_frame().set_edgecolor('k')
# write otp files for yields
massfiletuples = []
# plot masses
ax2 = plt.subplot2grid(layout, (1,0), colspan=2)
for i, sp in enumerate(prof.species):
props = next(ax2._get_lines.prop_cycler)
spmass = trapz(getattr(prof, sp), x=prof.masses)
massfiletuples.append((sp, spmass))
if i in skip:
continue
a = elemSplit(sp)[1]
spmass = trapz(getattr(prof, sp), x=prof.masses)
ax2.scatter(a, spmass, label=sp, color=props['color'])
ax2.set_ylabel('Total Mass ($M_\odot$)', rotation=90, labelpad=5)
ax2.set_xlabel('Mass Number (A)')
ax2.set_ylim([thresh, 1.5])
ax2.set_yscale("log", nonposy='clip')
#ax2.yaxis.set_major_formatter(StrMethodFormatter('{x:.2f}'))
ax2.yaxis.set_minor_formatter(StrMethodFormatter(''))
ax2.xaxis.set_major_formatter(StrMethodFormatter('{x:.0f}'))
ax2.xaxis.set_minor_formatter(StrMethodFormatter(''))
# decay yield
names, decmasses = decayYield(*zip(*massfiletuples))
# plot mass vs speed
ax3 = plt.subplot2grid(layout, (2, 0), colspan=2)
nbins = 60
for i, sp in enumerate(prof.species):
props = next(ax2._get_lines.prop_cycler)
if i in skip:
continue
counts, bins = np.histogram(abs(prof.velx), bins=nbins)
joined = sorted(zip(abs(prof.velx), getattr(prof, sp)))
speeds, masses = zip(*joined)
start = 0
weights = []
for c in counts:
massf = sum(masses[start:start+c])
weights.append(massf/c)
start+=c
weights = np.array(weights)
bins = bins[1:]
mpln, mplbins, patches = ax3.hist(bins, bins=len(bins), weights=weights, histtype='step', log=True, label=sp)
ax3.set_ylim([1e-6, 2])
ax3.yaxis.set_minor_formatter(StrMethodFormatter(''))
ax3.xaxis.set_major_formatter(customFormatter(10))
ax3.xaxis.set_minor_formatter(StrMethodFormatter(''))
ax3.set_xlabel('Speed ($10^{10}$ cm/s)')
fig.subplots_adjust(hspace=0.4)
if not batch:
return fig
else:
num = paths[1][-5:] # checkpoint number 'flash_hdf5_chk_0001'
dest = os.path.join(os.path.dirname(paths[0]), filetag)
name = os.path.join(dest, '{}{}.png'.format(filetag, num))
os.makedirs(dest, exist_ok=True) # bless you, p3
plt.savefig(name, format='png')
plt.close(fig)
print("Wrote: {}".format(name))
yieldfile = os.path.join(dest, '{}{}.yield'.format(filetag, num))
with open(yieldfile, 'w') as f:
f.write('# {}\n'.format(paths[1]))
f.write('# time: {}\n'.format(time))
f.write('\n'.join(['{} {}'.format(*a) for a in massfiletuples]))
f.write('\n')
print("Wrote: {}".format(yieldfile))
decayfile = os.path.join(dest, '{}{}.decayed'.format(filetag, num))
with open(decayfile, 'w') as f:
f.write('# {}\n'.format(paths[1]))
f.write('# time: {}\n'.format(time))
f.write('\n'.join(['{} {}'.format(*a) for a in zip(names, decmasses)]))
f.write('\n')
print("Wrote: {}".format(decayfile))