Add tests for isotope config options

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.1, atol=1.e-12)
+    assert np.isclose(abundances.loc[14, 4], 0.2, atol=1.e-12)
+    assert np.isclose(isotope_abundance.loc[(28, 56), 1], 0.1, atol=1.e-12)
+    assert np.isclose(isotope_abundance.loc[(28, 58), 2], 0.1, 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/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)