In [11]:
from bqplot import *
import numpy as np
import pandas as pd
from collections import OrderedDict
from pyne import nucname
from traitlets import link
from ipywidgets import ToggleButtons, VBox, HTML
from bqplot.interacts import (
FastIntervalSelector, IndexSelector, BrushIntervalSelector,
BrushSelector, MultiSelector, LassoSelector, PanZoom, HandDraw
)
import tardis
from tardis.io.atom_data.base import AtomData
In [31]:
class IonizationAnalysis(object):
"""A plotter generating various diagnostics plots for analyzing
the ionization state of a TARDIS simulation.
Parameters
----------
simdict : OrderedDict
Dictionary where each {key : value} pair is {label : sim}
where sim is the HDFStore output of a TARDIS simulation
generated by the TARDIS to_hdf() method, and label is a
string descriptor of the corresponding sim HDFStore. The
label is used in legends for plots.
"""
def __init__(self, simdict):
self.simdict = simdict
return
def ioniz_plot(self, simlabel, x_val, ion_atomic_num, ions, frac=True, yscale='Linear'):
"""
Parameters
----------
simlabel : string
Key in self.simdict specifying which sim HDFStore to use.
x_val : string
Either 'velocity' or 'temp' determines the physical quantity
plotted on the x-axis.
ion_atomic_num : int
Atomic number of the ion to analyze.
ions : list
List of which ionization levels to include in the plot.
frac : Boolean
Normalizes by the total Ca number density in each shell
when set to True.
yscale : string
Default to 'Linear' but can be set to 'Log'.
Returns
-------
fig : bqplot.figure
"""
sim = self.simdict[simlabel]
ion_name = nucname.name(ion_atomic_num)
temp = sim['/simulation/model/t_radiative'].values
vel = sim['/simulation/model/v_outer'].values
sim_num_dens = sim['/simulation/plasma/ion_number_density']
cols = np.arange(len(ions))
lines = []
x_sc = LinearScale()
y_sc = LinearScale()
col_line = ColorScale(scheme='gist_rainbow', min=0, max=len(ions))
col_ax = ColorAxis(label='', scale=col_line)
def_tt = Tooltip(fields=['name'], formats=[''], labels=['id'])
if frac:
ion_tots = np.sum(sim_num_dens.loc[ion_atomic_num].values, axis=0)
if x_val == 'temp':
x_axis_vals = temp
else:
x_axis_vals = vel
for i,ion in enumerate(ions):
if frac:
ioniz = sim_num_dens.loc[ion_atomic_num].T.iloc[:,ion] / ion_tots
else:
ioniz = sim_num_dens.loc[ion_atomic_num].T.iloc[:,ion]
line = Lines(x=x_axis_vals, y=ioniz, scales={'x':x_sc, 'y':y_sc, 'color':col_line},
display_legend=True, labels=['Ion '+str(ion)], color=[cols[i]],
tooltip=def_tt)
#scat = Scatter(x=vel, y=ioniz, scales={'x':x_sc, 'y':y_sc, 'color':col_line},
# colors=['red'])
lines.append(line)
if x_val == 'temp':
xlab = 'Shell Temperature (K)'
else:
xlab = 'Shell Outer Velocity (km/s)'
ax_x = Axis(scale=x_sc, grid_lines='solid', label=xlab)
if frac:
ylab = ion_name + ' Fractional Number Density'
else:
ylab = ion_name + ' Number Density'
ax_y = Axis(scale=y_sc, orientation='vertical', grid_lines='solid',
label=ylab)
fig = Figure(marks=lines, axes=[ax_x, ax_y], title='', legend_location='bottom-right')
pz = PanZoom(scales={'x':[x_sc]})
deb = HTML()
deb.value = '[]'
selection_interacts = ToggleButtons(options=OrderedDict([('PanZoom', pz),('None', None)]))
link((selection_interacts, 'value'), (fig, 'interaction'))
fig=VBox([deb, fig, selection_interacts], align_self='stretch')
return fig
def electron_dens_plot(self):
"""
Plots the electron number density curves as a function
of shell velocity for each TARDIS simulation HDFStore
in self.simdict
Returns
-------
fig : bqplot.figure
"""
x_sc = LinearScale()
y_sc = LogScale()
col_line = ColorScale(scheme='gist_rainbow', min=0, max=len(self.simdict.keys()))
col_ax = ColorAxis(label='', scale=col_line)
def_tt = Tooltip(fields=['name'], formats=[''], labels=['id'])
cols = np.arange(len(self.simdict.keys()))
lines = []
for i,simkey in enumerate(self.simdict):
sim = self.simdict[simkey]
n_e = sim['/simulation/plasma/electron_densities']
v = sim['/simulation/model/v_outer'] / 1e5
line = Lines(x=v, y=n_e, scales={'x':x_sc, 'y':y_sc, 'color':col_line},
display_legend=True, labels=[simkey], color=[cols[i]],
tooltip=def_tt)
lines.append(line)
ax_x = Axis(scale=x_sc, grid_lines='solid', label='Shell Velocity (km/s)')
ax_y = Axis(scale=y_sc, orientation='vertical', grid_lines='solid',
label='Electron Density (1/cm^3)')
fig = Figure(marks=lines, axes=[ax_x, ax_y], title='', legend_location='top-right')
pz = PanZoom(scales={'x':[x_sc]})
deb = HTML()
deb.value = '[]'
selection_interacts = ToggleButtons(options=OrderedDict([('PanZoom', pz),('None', None)]))
link((selection_interacts, 'value'), (fig, 'interaction'))
fig=VBox([deb, fig, selection_interacts], align_self='stretch')
return fig
def avg_ioniz_level(self, ion_atomic_num):
"""
Plots the average ionization level of the atom specified
by ion_atomic_num in each shell.
Parameters
----------
ion_atomic_num : int
Atomic number of ion to plot
Returns
-------
fig : bqplot.figure
"""
x_sc = LinearScale()
y_sc = LinearScale()
col_line = ColorScale(scheme='gist_rainbow', min=0, max=len(self.simdict.keys()))
col_ax = ColorAxis(label='', scale=col_line)
def_tt = Tooltip(fields=['name'], formats=[''], labels=['id'])
cols = np.arange(len(self.simdict.keys()))
lines = []
for i,simkey in enumerate(self.simdict):
sim = self.simdict[simkey]
ion_dens = sim['/simulation/plasma/ion_number_density'].loc[ion_atomic_num]
ion_tots = np.sum(ion_dens, axis=0)
norm_ion_dens = ion_dens / ion_tots
norm_ion_dens = norm_ion_dens.T
print(norm_ion_dens.shape)
lvls = np.arange(norm_ion_dens.shape[1])
print(lvls.shape)
avg_ion = np.dot(norm_ion_dens, lvls)
v = sim['/simulation/model/v_outer'] / 1e5
line = Lines(x=v, y=avg_ion, scales={'x':x_sc, 'y':y_sc, 'color':col_line},
display_legend=True, labels=[simkey], color=[cols[i]],
tooltip=def_tt)
lines.append(line)
ax_x = Axis(scale=x_sc, grid_lines='solid', label='Shell Velocity (km/s)')
ax_y = Axis(scale=y_sc, orientation='vertical', grid_lines='solid',
label=nucname.name(ion_atomic_num)+' Avg Ionization Level')
fig = Figure(marks=lines, axes=[ax_x, ax_y], title='', legend_location='bottom-left')
pz = PanZoom(scales={'x':[x_sc]})
deb = HTML()
deb.value = '[]'
selection_interacts = ToggleButtons(options=OrderedDict([('PanZoom', pz),('None', None)]))
link((selection_interacts, 'value'), (fig, 'interaction'))
fig=VBox([deb, fig, selection_interacts], align_self='stretch')
return fig
def ion_level_pop_dist(self, ion_atomic_num, shell_num, ion_lvl, atom_data_fname):
"""
Plots the Ion level population as a function of level energy
in a single shell.
Parameters
----------
ion_atomic_num : int
Atomic number of ion to plot.
shell_num : int
Zero indexed shell number.
ion_lvl : int
Which ionization level, where 0 is neutral.
Returns
-------
fig : bqplot.figure
"""
atom_data = AtomData.from_hdf(atom_data_fname)
energy_lvls = atom_data.levels.loc[(ion_atomic_num,ion_lvl)]['energy'].values
ionname = nucname.name(ion_atomic_num)+str(ion_lvl)
x_sc = LinearScale()
y_sc = LogScale()
col_line = ColorScale(scheme='gist_rainbow', min=0, max=len(self.simdict.keys()))
col_ax = ColorAxis(label='', scale=col_line)
def_tt = Tooltip(fields=['name'], formats=[''], labels=['id'])
cols = np.arange(len(self.simdict.keys()))
lines = []
for i,simkey in enumerate(self.simdict):
sim = self.simdict[simkey]
ion_lvls = sim['/simulation/plasma/level_number_density'].loc\
[ion_atomic_num].loc[ion_lvl].iloc[:,shell_num]
line = Lines(x=energy_lvls, y=ion_lvls, scales={'x':x_sc, 'y':y_sc, 'color':col_line},
display_legend=True, labels=[simkey], color=[cols[i]],
tooltip=def_tt)
lines.append(line)
ax_x = Axis(scale=x_sc, grid_lines='solid', label=ionname+' Level Energy (eV)')
ax_y = Axis(scale=y_sc, orientation='vertical', grid_lines='solid',
label=ionname + ' Shell '+str(shell_num) + ' Level Pop')
fig = Figure(marks=lines, axes=[ax_x, ax_y], title='', legend_location='bottom-left')
pz = PanZoom(scales={'x':[x_sc]})
deb = HTML()
deb.value = '[]'
selection_interacts = ToggleButtons(options=OrderedDict([('PanZoom', pz),('None', None)]))
link((selection_interacts, 'value'), (fig, 'interaction'))
fig=VBox([deb, fig, selection_interacts], align_self='stretch')
return fig
In [25]:
sim16 = pd.HDFStore('hach16d_sim.hdf')
sim22 = pd.HDFStore('hach22d_sim.hdf')
sim30 = pd.HDFStore('hach30d_sim.hdf')
sim40 = pd.HDFStore('hach40d_sim.hdf')
In [26]:
simdict = OrderedDict({'16days':sim16, '22days':sim22, '30days':sim30, '40days':sim40})
ion_analysis = IonizationAnalysis(simdict)
In [27]:
ion_analysis.ioniz_plot('40days', 'velocity', 20, [0,1,2,3,4], frac=True, yscale='Linear')
In [28]:
ion_analysis.electron_dens_plot()
In [29]:
ion_analysis.avg_ioniz_level(ion_atomic_num=20)
In [35]:
atom = '/Users/marcwilliamson/Research/TARDIS/tardis-refdata/atom_data/kurucz_cd23_chianti_H_He.h5'
ion_analysis.ion_level_pop_dist(ion_atomic_num=20, shell_num=0, ion_lvl=1, atom_data_fname=atom)
In [ ]: