def test_rho_plot(self):
name = 'cat_R1000_M15_Ni007'
path = join(dirname(abspath(__file__)), 'data', 'stella')
rho_file = os.path.join(path, name + '.rho')
presn = sn_eve.load_rho(rho_file)
ax = presn.plot_chem(ylim=(1e-12, 1.))
plt.show()
def plot_swd_chem(dic_axes, argsrho, path_relativ, alpha=0.25):
"""Add chemical data to the plot_shock_details plot
:type dic_axes: dict, dic_axes['r' or 'm'].append({'itime': i, 't': t, 'rho': axrho, 'par': axpar})
:param argsrho: rho-file + Elements
:param path_relativ:
:param alpha: the transparanc
"""
from pystella.model import sn_eve
from pystella.util.phys_var import phys
colors = sn_eve.eve_colors
s_elements = None
if '+' in argsrho:
frho, s_elements = argsrho.split('+')
else:
frho = argsrho
frho = os.path.join(path_relativ, frho)
pathrho, namerho = os.path.split(frho)
eve = sn_eve.load_rho(os.path.join(pathrho, namerho))
elements = eve.Elements
if s_elements is not None:
elements = s_elements.split(':')
print('Add chem [{}] from {}'.format(':'.join(elements), frho))
x = eve.m / phys.M_sun
if min(x) > 0.1:
print('Chem is shifted at {} Msun'.format(-min(x)))
x = x - min(x)
def two_y(x, lims, up=0.85):
return (lims[1] * up - lims[0]) * x + lims[0]
for i, d in enumerate(dic_axes['m']):
ax = d['par']
xlim = ax.get_xlim()
ylim = ax.get_ylim()
# print(ylim)
xx = np.linspace(xlim[1] * 0.8, xlim[0], len(elements))
yy_tot = np.zeros_like(x) # + ylim[1]
yy_prev = np.zeros_like(x) # + ylim[0]
for ie, el in enumerate(elements):
# for el in reversed(elements):
y = eve.el(el)
# yy += y
# yy = y
# yyy = yy
yy_tot += y
# yyy = np.log10(yy) + ylim[1]-1
ax.fill_between(x,
y1=two_y(yy_tot, ylim),
y2=two_y(yy_prev, ylim),
color=colors[el],
alpha=alpha)
# ax.plot(x, yyy, color=colors[el], alpha=alpha)
# Print the element name
# xp = np.average(x, weights=y)
# yp = np.average(yyy, weights=y) * 0.9
xp = xx[ie]
ax.text(xp, ylim[1] * 0.9, el, color=colors[el], fontsize=11)
yy_prev = yy_tot.copy()
def test_rho_el(self):
name = 'cat_R1000_M15_Ni007'
path = join(dirname(abspath(__file__)), 'data', 'stella')
rho_file = os.path.join(path, name + '.rho')
presn = sn_eve.load_rho(rho_file)
h = presn.lg_el('H')
self.assertTrue(len(h) == len(presn.m), "No data for el: %s" % 'H')
def test_rho_load(self):
name = 'cat_R1000_M15_Ni007'
path = join(dirname(abspath(__file__)), 'data', 'stella')
rho_file = os.path.join(path, name + '.rho')
presn = sn_eve.load_rho(rho_file)
self.assertTrue(presn.nzon > 0,
"Rho-file have not been loaded: %s" % rho_file)
def get_eve(self, name=None, path=None, is_hyd_abn=False, **kwargs):
from pystella.model import sn_eve
if name is None:
name = self.Name
if path is None:
path = self.Path
if is_hyd_abn:
eve = sn_eve.load_hyd_abn(name, path, **kwargs)
else:
eve = sn_eve.load_rho(os.path.join(path, name + '.rho'))
return eve
def test_rho_write_hyd_abn(self):
name = 'cat_R1000_M15_Ni007'
path = join(dirname(abspath(__file__)), 'data', 'stella')
rho_file = os.path.join(path, name + '.rho')
presn = sn_eve.load_rho(rho_file)
fname = 'tmprho.hyd'
res = presn.write_hyd(fname)
self.assertTrue(res, "Fail to write in %s" % fname)
fname = 'tmprho.abn'
res = presn.write_abn(fname)
self.assertTrue(res, "Fail to write in %s" % fname)
def test_mass_tot(self):
name = 'cat_R1000_M15_Ni007'
path = join(dirname(abspath(__file__)), 'data', 'stella')
rho_file = os.path.join(path, name + '.rho')
presn = sn_eve.load_rho(rho_file)
s, s1 = 0., 0.
for el, m in presn.mass_tot_el().items():
m = m / ps.phys.M_sun
m1 = presn.mass_tot_el(el) / ps.phys.M_sun
print("{:4s}: {:.3e} df: {:.3e}".format(el, m, m1))
s += m
s1 += m1
print(
f"Total: {s} dm: {s1} {presn.m_tot / ps.phys.M_sun - presn.m_core / ps.phys.M_sun}"
)
def test_presn_cut(self):
name = 'cat_R1000_M15_Ni007'
path = join(dirname(abspath(__file__)), 'data', 'stella')
rho_file = os.path.join(path, name + '.rho')
eve = sn_eve.load_rho(rho_file)
# for same size must be the same
s, e = int(eve.nzon / 10), int(eve.nzon / 2)
evenew = eve.cut(start=s, end=e)
self.assertEqual(
evenew.nzon, e - s,
"Cut PreSn: you have {} zone, but it should be {}".format(
evenew.nzon, e - s))
self.assertEqual(
eve.m[s], evenew.m[0],
"Mass PreSn: you have the first zone where old mass {} = new {}".
format(eve.m[s], evenew.m[0]))
self.assertEqual(
eve.m[e - 1], evenew.m[-1],
"Mass PreSn: you have the last zone where old mass {} = new {}".
format(eve.m[e], evenew.m[-1]))
def main():
import os
import sys
from itertools import cycle
def get(arr, i, default):
if i < len(arr):
if a[i] != '*':
return a[i]
return default
parser = get_parser()
args, unknownargs = parser.parse_known_args()
eve_prev = None
markersize = 6
fig = None
if args.path:
pathDef = os.path.expanduser(args.path)
else:
pathDef = os.getcwd()
# if args.elements:
if '_' in args.elements:
elements = list(sneve.eve_elements)
excluded = args.elements.split(':')
for e in excluded:
if not e.startswith('_'):
logger.error(
'For excluded mode all elements should be starts from _. Even element: '
+ e)
sys.exit(2)
e = e[1:]
if e not in sneve.eve_elements:
logger.error('No such element: ' + e)
sys.exit(2)
elements.remove(e)
else:
elements = args.elements.split(':')
for e in elements:
if e not in sneve.eve_elements:
logger.error('No such element: ' + e)
sys.exit(2)
# Set model names
names = []
if args.input:
for nm in args.input:
names.append(nm[0]) # remove extension
else:
if len(unknownargs) > 0:
names.append(unknownargs[0])
if len(names) == 0:
# logger.error(" No data. Use key '-i' ")
parser.print_help()
sys.exit(2)
if len(names) > 1: # special case
markers_cycler = cycle(markers_style)
lines_cycler = cycle(lines_style)
else:
markers_cycler = cycle([None])
lines_cycler = cycle(['-'])
ax = None
ax2 = None
handles_nm = []
for nm in names:
# print("Run eve-model %s" % nm)
path, fullname = os.path.split(nm)
if len(path) == 0:
path = pathDef
# print("Run eve-model %s in %s" % (name, path))
if fullname.endswith('hyd') or fullname.endswith('abn'):
name = fullname.replace('.hyd', '') # remove extension
name = name.replace('.abn', '') # remove extension
try:
# With header
eve = sneve.load_hyd_abn(name=name,
path=path,
is_dm=False,
is_dum=args.is_dum)
except ValueError:
# No header
eve = sneve.load_hyd_abn(name=name,
path=path,
is_dm=False,
is_dum=args.is_dum,
skiprows=0)
else:
name = fullname.replace('.rho', '') # remove extension
rho_file = os.path.join(path, name + '.rho')
eve = sneve.load_rho(rho_file)
if args.reshape is not None:
a = args.reshape.split(':')
nz, axis, xmode = get(a, 0,
eve.nzon), get(a, 1,
'M'), get(a, 2,
'resize') # rlog
start, end = get(a, 3, 0), get(a, 4, None)
kind = get(a, 5, 'np')
start = int(start)
if end is None or end.upper() in 'NONE':
end = None
if end is not None:
end = int(end)
nz = int(nz)
print(f'Resize: before Nzon={eve.nzon}')
print(
f'Resize parameters: nznew= {nz} axis={axis} xmode={xmode} '
f'start= {start} end= {end} kind= {kind}')
print("The element masses: before Resize")
print_masses(eve)
eve = eve.reshape(nz=nz,
axis=axis,
xmode=xmode,
start=start,
end=end,
kind=kind)
eve.chem_norm()
# eve = eve_resize
print(f'Resize: after Nzon={eve.nzon}')
print("The element masses: after Resize")
print_masses(eve)
# Boxcar
if args.box is not None:
is_info = False
s = args.box.split(':')
dm, n = float(s[0]), int(s[1])
if len(s) == 3:
is_info = bool(s[2])
print(f'Running boxcar average: dm= {dm} Msun Repeats= {n}')
print("The element masses: Before boxcar")
print_masses(eve)
eve_box = eve.boxcar(box_dm=dm,
n=n,
el_included=elements,
is_info=is_info)
print("The element masses: After boxcar")
print_masses(eve_box)
eve, eve_prev = eve_box, eve
# Smooth
if args.smooth is not None:
is_info = False
s = args.smooth.split(':')
window_length, polyorder = int(s[0]), int(s[1])
mode = 'interp'
if len(s) == 3:
mode = s[2]
if len(s) == 4:
is_info = True
print(
f'Running Savitzky-Golay filter to Rho: '
f'window_length= {window_length} Msun polyorder= {polyorder} mode= {mode}'
)
print("The element masses: Before smoothing")
print_masses(eve, is_el=is_info)
eve_smooth = eve.smooth(window_length=window_length,
polyorder=polyorder,
mode=mode,
is_info=is_info)
print("The element masses: After smoothing")
print_masses(eve_smooth, is_el=is_info)
eve, eve_prev = eve_smooth, eve
if args.write_to:
fname = os.path.expanduser(args.write_to)
# fname = os.path.join(path, name)
# f = fname + '.eve.abn'
fname = fname.replace('.rho', '')
f = fname + '.abn'
if eve.write_abn(f, is_header=True):
print(" abn has been saved to {}".format(f))
else:
print("Error with abn saving to {}".format(f))
# f = fname + '.eve.hyd'
f = fname + '.hyd'
if eve.write_hyd(f):
print(" hyd has been saved to {}".format(f))
else:
print("Error with hyd saving to {}".format(f))
continue
marker = next(markers_cycler)
ls = next(lines_cycler)
if args.is_structure:
fig = eve.plot_structure(elements=elements,
title=name,
ylimChem=(1e-8, 1.))
else:
if args.is_chem:
# print "Plot eve-model %s" % name
ax = eve.plot_chem(elements=elements,
ax=ax,
x=args.x,
ylim=(1e-8, 1.),
marker=marker,
markersize=markersize,
leg_loc='lower center')
if eve_prev is not None:
eve_prev.plot_chem(elements=elements,
ax=ax,
x=args.x,
ylim=(1e-8, 1.),
marker=marker,
markersize=max(1, markersize - 2),
alpha=0.5,
leg_loc='lower center')
# ax.set_title('{}: before boxcar'.format(eve_prev.Name))
if args.rho:
if args.is_chem:
if ax2 is None:
ax2 = ax.twinx()
ax2.set_ylabel(r'$\rho, [g/cm^3]$ ')
else:
ax2 = ax
ax2 = eve.plot_rho(x=args.x, ax=ax2, ls=ls, marker=marker)
if eve_prev is not None:
eve_prev.plot_rho(x=args.x,
ax=ax2,
ls=ls,
markersize=max(1, markersize - 2),
alpha=0.5)
else:
ls = 'None'
handle = mlines.Line2D([], [],
color='black',
marker=marker,
markersize=markersize,
label=name,
linestyle=ls)
handles_nm.append(handle)
if len(names) > 1:
if ax2 is None:
ax2 = ax.twinx()
ax2.legend(handles=handles_nm,
loc=4,
fancybox=False,
frameon=False)
if args.is_verb:
m_tot = 0.
m_ni56 = 0.
print('{:22s}: '.format(eve.Name))
for n, m in eve.mass_tot_el().items():
m_tot += m
print(f'{n} {m/phys.M_sun:.3e}')
print(' m_tot= {:6.3f} m_tot(El)= {:6.3f} '.format(
eve.m_tot / phys.M_sun, m_tot / phys.M_sun))
# eve.chem_norm()
# print('{:22s} after Norm '.format(eve.Name))
# for n, m in eve.mass_tot_el().items():
# m_tot += m
# print(f'{n} {m/phys.M_sun:.3e}')
# print(' m_tot= {:6.3f} m_tot(El)= {:6.3f} '.format(eve.m_tot/phys.M_sun, m_tot/phys.M_sun))
if not args.write_to:
if args.save_plot:
fsave = os.path.expanduser(args.save_plot)
# if args.rho:
# fsave = os.path.join(os.path.expanduser('~/'), 'rho_%s.pdf' % names[0])
# else:
# fsave = os.path.join(os.path.expanduser('~/'), 'chem_%s.pdf' % names[0])
logger.info(" Save plot to %s " % fsave)
if fig is None:
fig = ax.get_figure()
fig.savefig(fsave, bbox_inches='tight')
else:
plt.show()