[TEP007] Isotope stratified file support

tardis/io/model_reader.py

@@ -4,7 +4,7 @@
 from numpy import recfromtxt, genfromtxt
 import pandas as pd
 from astropy import units as u
-
+from pyne import nucname
 import logging
 # Adding logging support
 logger = logging.getLogger(__name__)
@@ -85,17 +85,51 @@ def read_abundances_file(abundance_filename, abundance_filetype,
     """
 
     file_parsers = {'simple_ascii': read_simple_ascii_abundances,
-                    'artis': read_simple_ascii_abundances}
+                    'artis': read_simple_ascii_abundances,
+                    'isotopes': read_simple_ascii_isotopes}
+
+    isotope_abundance = None
+    if abundance_filetype != 'isotopes':
+        index, abundances = file_parsers[abundance_filetype](
+            abundance_filename)
+    else:
+        index, abundances, isotope_abundance = file_parsers[abundance_filetype](
+            abundance_filename)
 
-    index, abundances = file_parsers[abundance_filetype](abundance_filename)
     if outer_boundary_index is not None:
         outer_boundary_index_m1 = outer_boundary_index - 1
     else:
         outer_boundary_index_m1 = None
     index = index[inner_boundary_index:outer_boundary_index]
     abundances = abundances.ix[:, slice(inner_boundary_index, outer_boundary_index_m1)]
     abundances.columns = np.arange(len(abundances.columns))
-    return index, abundances
+    return index, abundances, isotope_abundance
+
+
+def read_simple_ascii_isotopes(fname):
+
+    data = np.genfromtxt(fname, names=True, dtype=np.float64)
+    isotopes_name = data.dtype.names[31:]
+
+    #Reshape Numpy Structured arrays
+    #File rows are read as tuples here, this line converts it to a Matrix form
+    #First column is ignored
+    data = data.view(np.float64).reshape(-1, len(data.dtype))[:, 1:]
+
+    #Parse isotopes
+    mass_no = [nucname.anum(name) for name in isotopes_name]
+    z = [nucname.znum(name) for name in isotopes_name]
+    isotope_index = pd.MultiIndex.from_tuples(
+        zip(z, mass_no), names=['atomic_number', 'mass_number'])
+
+    isotope_abundance = pd.DataFrame(data[:, 30:].transpose(),
+                                     index=isotope_index,
+                                     dtype=np.float64)
+
+    index = pd.Index(np.arange(1, 30), name='atomic_number')
+    abundances = pd.DataFrame(data[:, :29].transpose(), index=index)
+
+    return index, abundances, isotope_abundance
 
 
 def read_simple_ascii_density(fname):

tardis/model/base.py

@@ -2,9 +2,10 @@
 import logging
 import numpy as np
 import pandas as pd
+from pyne import nucname
 from astropy import constants, units as u
 
-from tardis.util import quantity_linspace, element_symbol2atomic_number
+from tardis.util import quantity_linspace, element_symbol2atomic_number, MalformedElementSymbolError
 from tardis.io.model_reader import read_density_file, read_abundances_file
 from tardis.io.util import HDFWriterMixin
 from density import HomologousDensity
@@ -59,11 +60,10 @@ class Radial1DModel(HDFWriterMixin):
     """
     hdf_properties = ['t_inner', 'w', 't_radiative', 'v_inner', 'v_outer', 'homologous_density']
     hdf_name = 'model'
-
-    def __init__(self, velocity, homologous_density, abundance, time_explosion,
-                 t_inner, luminosity_requested=None, t_radiative=None,
-                 dilution_factor=None, v_boundary_inner=None,
-                 v_boundary_outer=None):
+
+    def __init__(self, velocity, homologous_density, abundance, isotope_abundance,
+                 time_explosion, t_inner, luminosity_requested=None, t_radiative=None,
+                 dilution_factor=None, v_boundary_inner=None, v_boundary_outer=None):
         self._v_boundary_inner = None
         self._v_boundary_outer = None
         self._velocity = None
@@ -73,6 +73,10 @@ def __init__(self, velocity, homologous_density, abundance, time_explosion,
         self.homologous_density = homologous_density
         self._abundance = abundance
         self.time_explosion = time_explosion
+
+        self.raw_abundance = self._abundance
+        self.raw_isotope_abundance = isotope_abundance
+
         if t_inner is None:
             if luminosity_requested is not None:
                 self.t_inner = ((luminosity_requested /
@@ -317,17 +321,29 @@ def from_config(cls, config):
             t_inner = config.plasma.initial_t_inner
 
         abundances_section = config.model.abundances
+        isotope_index = pd.MultiIndex(
+            [[]] * 2, [[]] * 2, names=['atomic_number', 'mass_number'])
+        isotope_abundance = pd.DataFrame(columns=np.arange(no_of_shells),
+                                         index=isotope_index,
+                                         dtype=np.float64)
+
         if abundances_section.type == 'uniform':
             abundance = pd.DataFrame(columns=np.arange(no_of_shells),
                                      index=pd.Index(np.arange(1, 120),
                                                     name='atomic_number'),
                                      dtype=np.float64)
-
             for element_symbol_string in abundances_section:
                 if element_symbol_string == 'type':
                     continue
-                z = element_symbol2atomic_number(element_symbol_string)
-                abundance.ix[z] = float(abundances_section[element_symbol_string])
+                try:
+                    z = element_symbol2atomic_number(element_symbol_string)
+                    abundance.ix[z] = float(
+                        abundances_section[element_symbol_string])
+                except MalformedElementSymbolError:
+                    mass_no = nucname.anum(element_symbol_string)
+                    z = nucname.znum(element_symbol_string)
+                    isotope_abundance.loc[(z, mass_no), :] = float(
+                        abundances_section[element_symbol_string])
 
         elif abundances_section.type == 'file':
             if os.path.isabs(abundances_section.filename):
@@ -336,23 +352,26 @@ def from_config(cls, config):
                 abundances_fname = os.path.join(config.config_dirname,
                                                 abundances_section.filename)
 
-            index, abundance = read_abundances_file(abundances_fname,
-                                                    abundances_section.filetype)
+            index, abundance, isotopes = read_abundances_file(abundances_fname,
+                                                              abundances_section.filetype)
+            if isotopes is not None:
+                isotope_abundance = isotopes
 
         abundance = abundance.replace(np.nan, 0.0)
         abundance = abundance[abundance.sum(axis=1) > 0]
 
-        norm_factor = abundance.sum(axis=0)
+        norm_factor = abundance.sum(axis=0) + isotope_abundance.sum(axis=0)
 
         if np.any(np.abs(norm_factor - 1) > 1e-12):
             logger.warning("Abundances have not been normalized to 1."
                            " - normalizing")
             abundance /= norm_factor
-
+            isotope_abundance /= norm_factor
 
         return cls(velocity=velocity,
                    homologous_density=homologous_density,
                    abundance=abundance,
+                   isotope_abundance=isotope_abundance,
                    time_explosion=time_explosion,
                    t_radiative=t_radiative,
                    t_inner=t_inner,