[TEP007] Isotope stratified file support

tardis/io/model_reader.py

@@ -85,17 +85,25 @@ 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,
+                    'tardis_model': read_simple_isotope_abundances}
+
+    isotope_abundance = pd.DataFrame()
+    if abundance_filetype == 'tardis_model':
+        index, abundances, isotope_abundance = read_simple_isotope_abundances(
+            abundance_filename)
+    else:
+        index, abundances = 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_uniform_abundances(abundances_section, no_of_shells):
@@ -256,3 +264,31 @@ def read_simple_ascii_abundances(fname):
     abundances = pd.DataFrame(data[1:,1:].transpose(), index=np.arange(1, data.shape[1]))
 
     return index, abundances
+
+
+def read_simple_isotope_abundances(fname, delimiter='\s+'):
+    df = pd.read_csv(fname, comment='#', delimiter=delimiter)
+    df = df.transpose()
+
+    abundance = pd.DataFrame(columns=np.arange(df.shape[1]),
+                             index=pd.Index([],
+                                            name='atomic_number'),
+                             dtype=np.float64)
+
+    isotope_index = pd.MultiIndex(
+        [[]] * 2, [[]] * 2, names=['atomic_number', 'mass_number'])
+    isotope_abundance = pd.DataFrame(columns=np.arange(df.shape[1]),
+                                     index=isotope_index,
+                                     dtype=np.float64)
+
+    for element_symbol_string in df.index:
+        if element_symbol_string in nucname.name_zz:
+            z = nucname.name_zz[element_symbol_string]
+            abundance.loc[z, :] = df.loc[element_symbol_string].tolist()
+        else:
+            z = nucname.znum(element_symbol_string)
+            mass_no = nucname.anum(element_symbol_string)
+            isotope_abundance.loc[(
+                z, mass_no), :] = df.loc[element_symbol_string].tolist()
+
+    return abundance.index, abundance, isotope_abundance

tardis/io/setup_package.py

@@ -1,4 +1,4 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 def get_package_data():
-    return {'tardis.io.tests':['data/*.dat', 'data/*.yml']}
+    return {'tardis.io.tests': ['data/*.dat', 'data/*.yml', 'data/*.csv']}

tardis/io/tests/data/tardis_configv1_isotope_uniabund.yml

@@ -0,0 +1,49 @@
+tardis_config_version: v1.0
+
+supernova:
+    luminosity_requested: 2.8e9 solLum
+    time_explosion: 13 day
+
+atom_data: kurucz_atom_pure_simple.h5
+model:
+
+    structure:
+        type: specific
+
+        velocity:
+            start: 1.1e4 km/s
+            stop: 2.0e4 km/s
+            num: 20
+
+        density:
+            type: branch85_w7
+
+    abundances:
+        type: uniform
+        O: 0.19
+        Mg: 0.03
+        Si: 0.42
+        S: 0.19
+        Ar: 0.04
+        Ca: 0.03
+        Ni56: 0.05
+        Ni58: 0.05
+
+
+plasma:
+    ionization: lte
+    excitation: lte
+    radiative_rates_type: dilute-blackbody
+    line_interaction_type: macroatom
+
+montecarlo:
+    seed: 23111963
+    no_of_packets : 2.0e+5
+    iterations: 5
+    last_no_of_packets: 5.0e+5
+    no_of_virtual_packets: 5
+
+spectrum:
+    start: 500 angstrom
+    stop: 20000 angstrom
+    num: 10000

tardis/io/tests/data/tardis_model_abund.csv

@@ -0,0 +1,11 @@
+C O Mg Si Ni56 Ni58
+0 0 0 0.6 0.4 0
+0 0 0 0.1 0.5 0.4
+0 0 0 0.3 0 0.7
+0 0.2 0.8 0 0 0
+0 0.3 0.7 0 0 0
+0 0.2 0.8 0 0 0
+0 0.2 0.8 0 0 0
+0 0.2 0.8 0 0 0
+0.5 0.5 0 0 0 0
+0.5 0.5 0 0 0 0

tardis/io/tests/test_model_reader.py

@@ -5,8 +5,9 @@
 import pytest
 
 import tardis
-from tardis.io.model_reader import (read_artis_density, 
-read_simple_ascii_abundances)
+from tardis.io.config_reader import Configuration
+from tardis.io.model_reader import (
+    read_artis_density, read_simple_ascii_abundances, read_simple_isotope_abundances, read_uniform_abundances)
 
 data_path = os.path.join(tardis.__path__[0], 'io', 'tests', 'data')
 
@@ -18,6 +19,19 @@ def artis_density_fname():
 def artis_abundances_fname():
     return os.path.join(data_path, 'artis_abundances.dat')
 
+
+@pytest.fixture
+def tardis_model_abundance_fname():
+    return os.path.join(data_path, 'tardis_model_abund.csv')
+
+
+@pytest.fixture
+def isotope_uniform_abundance():
+    config_path = os.path.join(
+        data_path, 'tardis_configv1_isotope_uniabund.yml')
+    config = Configuration.from_yaml(config_path)
+    return config.model.abundances
+
 def test_simple_read_artis_density(artis_density_fname):
     time_of_model, velocity, mean_density = read_artis_density(artis_density_fname)
 
@@ -32,3 +46,22 @@ def test_read_simple_ascii_abundances(artis_abundances_fname):
     index, abundances = read_simple_ascii_abundances(artis_abundances_fname)
     assert len(abundances.columns) == 69
     assert np.isclose(abundances[23].ix[2], 2.672351e-08 , atol=1.e-12)
+
+
+def test_read_simple_isotope_abundances(tardis_model_abundance_fname):
+    index, abundances, isotope_abundance = read_simple_isotope_abundances(
+        tardis_model_abundance_fname)
+    assert np.isclose(abundances.loc[6, 9], 0.5, atol=1.e-12)
+    assert np.isclose(abundances.loc[12, 6], 0.8, atol=1.e-12)
+    assert np.isclose(abundances.loc[14, 2], 0.3, atol=1.e-12)
+    assert np.isclose(isotope_abundance.loc[(28, 56), 1], 0.5, atol=1.e-12)
+    assert np.isclose(isotope_abundance.loc[(28, 58), 2], 0.7, atol=1.e-12)
+
+
+def test_read_uniform_abundances(isotope_uniform_abundance):
+    abundances, isotope_abundance = read_uniform_abundances(
+        isotope_uniform_abundance, 20)
+    assert np.isclose(abundances.loc[8, 2], 0.19, atol=1.e-12)
+    assert np.isclose(abundances.loc[20, 5], 0.03, atol=1.e-12)
+    assert np.isclose(isotope_abundance.loc[(28, 56), 15], 0.05, atol=1.e-12)
+    assert np.isclose(isotope_abundance.loc[(28, 58), 2], 0.05, atol=1.e-12)

tardis/model/base.py

@@ -333,8 +333,8 @@ 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, isotope_abundance = read_abundances_file(abundances_fname,
+                                                                       abundances_section.filetype)
 
         abundance = abundance.replace(np.nan, 0.0)
         abundance = abundance[abundance.sum(axis=1) > 0]

tardis/tests/test_util.py

@@ -1,5 +1,6 @@
 import pytest
 import numpy as np
+import os
 from astropy import units as u
 from io import StringIO
 
@@ -8,8 +9,14 @@
                          calculate_luminosity, intensity_black_body, savitzky_golay,
                          species_tuple_to_string, species_string_to_tuple, parse_quantity,
                          element_symbol2atomic_number, atomic_number2element_symbol,
-                         reformat_element_symbol, quantity_linspace)
+                         reformat_element_symbol, quantity_linspace, convert_abundances_format)
 
+data_path = os.path.join('tardis', 'io', 'tests', 'data')
+
+
+@pytest.fixture
+def artis_abundances_fname():
+    return os.path.join(data_path, 'artis_abundances.dat')
 
 def test_malformed_species_error():
     malformed_species_error = MalformedSpeciesError('He')
@@ -203,3 +210,11 @@ def test_quantity_linspace(start, stop, num, expected):
 
     with pytest.raises(ValueError):
         quantity_linspace(u.Quantity(0.5, 'eV'), '0.6 eV', 3)
+
+
+def test_convert_abundances_format(artis_abundances_fname):
+    abundances = convert_abundances_format(artis_abundances_fname)
+    assert np.isclose(abundances.loc[3, 'O'], 1.240199e-08, atol=1.e-12)
+    assert np.isclose(abundances.loc[1, 'Co'], 2.306023e-05, atol=1.e-12)
+    assert np.isclose(abundances.loc[69, 'Ni'], 1.029928e-17, atol=1.e-12)
+    assert np.isclose(abundances.loc[2, 'C'], 4.425876e-09, atol=1.e-12)

tardis/util.py

@@ -1,8 +1,10 @@
 # Utilities for TARDIS
 
 from astropy import units as u, constants
+from pyne import nucname
 import numexpr as ne
 import numpy as np
+import pandas as pd
 import os
 import yaml
 import re
@@ -428,4 +430,14 @@ def quantity_linspace(start, stop, num, **kwargs):
         raise ValueError('Both start and stop need to be quantities with a '
                          'unit attribute')
 
-    return np.linspace(start.value, stop.to(start.unit).value, num, **kwargs) * start.unit
+    return np.linspace(start.value, stop.to(start.unit).value, num, **kwargs) * start.unit
+
+
+def convert_abundances_format(fname, delimiter='\s+'):
+    df = pd.read_csv(fname, delimiter=delimiter, comment='#', header=None)
+    #Drop shell index column
+    df.drop(df.columns[0], axis=1, inplace=True)
+    #Assign header row
+    df.columns = [nucname.name(i)
+                  for i in range(1, df.shape[1] + 1)]
+    return df