Reimplement tests for the atomic module

tardis/conftest.py

@@ -88,7 +88,7 @@ def atomic_data_fname():
 
 @pytest.fixture
 def kurucz_atomic_data(atomic_data_fname):
-    atomic_data = AtomData.from_hdfstore(atomic_data_fname)
+    atomic_data = AtomData.from_hdf(atomic_data_fname)
 
     if atomic_data.md5 != '06709f04d369f07c58e01a7d0f34bd10':
         pytest.skip('Need default Kurucz atomic dataset '
@@ -103,9 +103,9 @@ def test_data_path():
 
 
 @pytest.fixture
-def included_he_atomic_data(test_data_path):
+def atom_data_he(test_data_path):
     atomic_db_fname = os.path.join(test_data_path, 'chianti_he_db.h5')
-    return AtomData.from_hdf5(atomic_db_fname)
+    return AtomData.from_hdf(atomic_db_fname)
 
 
 @pytest.fixture

tardis/io/tests/test_atomic.py

@@ -1,155 +1,67 @@
-import os
 import pytest
-from numpy import testing
-from tardis.io import atomic
 
+from astropy.tests.helper import assert_quantity_allclose
+from astropy import units as u
+from astropy import constants as const
+from tardis.io.atomic import atomic_number2symbol, symbol2atomic_number
 
-@pytest.fixture(scope="module")
-def default_atom_h5_path():
-    return atomic.data_path('atom_data.h5')
 
+@pytest.fixture
+def basic_atom_data_he(atom_data_he):
+    return atom_data_he.atom_data
 
-@pytest.fixture(scope="module")
-def chianti_he_db_h5_path():
-    return atomic.tests_data_path('chianti_he_db.h5')
 
+@pytest.fixture
+def ionization_data_he(atom_data_he):
+    return atom_data_he.ionization_data
 
-def test_data_path():
-    data_path = atomic.data_path('test')
-    assert data_path.split('/')[-3:] == ['tardis', 'data', 'test']
 
+@pytest.fixture
+def levels_he(atom_data_he):
+    return atom_data_he.levels
 
-def test_read_basic_atom_data(default_atom_h5_path):
-    data = atomic.read_basic_atom_data(default_atom_h5_path)
-    assert data['atomic_number'][13] == 14
-    assert data['symbol'][13] == "Si"
-    testing.assert_almost_equal(data['mass'][13], 28.085, decimal=4)
 
+@pytest.fixture
+def lines_he(atom_data_he):
+    return atom_data_he.lines
 
-def test_read_ionization_data(default_atom_h5_path):
-    data = atomic.read_ionization_data(default_atom_h5_path)
-    assert data['atomic_number'][0] == 1
-    assert data['ion_number'][0] == 1
-    testing.assert_almost_equal(data['ionization_energy'][0], 13.59844, decimal=4)
 
+def test_atom_data_basic_atom_data(basic_atom_data_he):
+    assert basic_atom_data_he.loc[2, 'symbol'] == 'He'
+    assert_quantity_allclose(
+        basic_atom_data_he.loc[2, 'mass'] * u.Unit('g'),
+        4.002602 * const.u.cgs
+     )
 
-def test_read_levels_data(default_atom_h5_path):
-    data = atomic.read_levels_data(default_atom_h5_path)
-    assert data['atomic_number'][4] == 14
-    assert data['ion_number'][4] == 0
-    assert data['level_number'][4] == 4
-    testing.assert_almost_equal(data['energy'][4], 1.90865, decimal=4)
-    assert data['g'][4] == 1
-    assert data['metastable'][4] == False
 
+def test_atom_data_ionization_data(ionization_data_he):
+    assert_quantity_allclose(
+        ionization_data_he.loc[(2, 1), 'ionization_energy'] * u.Unit('erg'),
+        24.587387936 * u.Unit('eV')
+    )
 
-def test_read_lines_data(default_atom_h5_path):
-    data = atomic.read_lines_data(default_atom_h5_path)
-    assert data['line_id'][0] == 8
-    assert data['atomic_number'][0] == 14
-    assert data['ion_number'][0] == 5
-    testing.assert_almost_equal(data['wavelength'][0], 66.772, decimal=4)
-    testing.assert_almost_equal(data['f_ul'][0], 0.02703, decimal=4)
-    testing.assert_almost_equal(data['f_lu'][0], 0.04054, decimal=4)
-    assert data['level_number_lower'][0] == 0.0
-    assert data['level_number_upper'][0] == 36.0
 
+def test_atom_data_levels(levels_he):
+    levels_he = levels_he.set_index(['atomic_number', 'ion_number', 'level_number'])
+    assert_quantity_allclose(
+        u.Quantity(levels_he.loc[(2, 0, 2), 'energy'], u.Unit('erg')).to(u.Unit('cm-1'), equivalencies=u.spectral()),
+        166277.542 * u.Unit('cm-1')
+    )
 
-def test_read_synpp_refs(chianti_he_db_h5_path):
-    data = atomic.read_synpp_refs(chianti_he_db_h5_path)
-    assert data['atomic_number'][0] == 1
-    assert data['ion_number'][0] == 0
-    testing.assert_almost_equal(data['wavelength'][0], 6562.7973633, decimal=4)
-    assert data['line_id'][0] == 564995
 
+def test_atom_data_lines(lines_he):
+    lines_he = lines_he.set_index(['atomic_number', 'ion_number',
+                                   'level_number_lower', 'level_number_upper'])
 
-def test_read_zeta_data(default_atom_h5_path, chianti_he_db_h5_path):
-    data = atomic.read_zeta_data(chianti_he_db_h5_path)
-    testing.assert_almost_equal(data[2000][1][1], 0.339, decimal=4)
-    testing.assert_almost_equal(data[2000][1][2], 0.000, decimal=4)
+    assert_quantity_allclose(
+        lines_he.loc[(2,0,0,6), 'wavelength_cm'] * u.Unit('cm'),
+        584.335 * u.Unit('Angstrom')
+    )
 
-    with pytest.raises(ValueError):
-        atomic.read_zeta_data(None)
-
-    with pytest.raises(IOError):
-        atomic.read_zeta_data('fakepath')
-
-    with pytest.raises(ValueError):
-        atomic.read_zeta_data(default_atom_h5_path)
-
-
-def test_read_collision_data(default_atom_h5_path, chianti_he_db_h5_path):
-    data = atomic.read_collision_data(chianti_he_db_h5_path)
-    assert data[0]['atomic_number'][0] == 2
-    assert data[0]['ion_number'][0] == 0
-    assert data[0]['level_number_upper'][0] == 18
-    assert data[0]['level_number_lower'][0] == 2
-    assert data[0]['g_ratio'][0] == 1.0
-    testing.assert_almost_equal(data[0]['delta_e'][0], 35484.251143, decimal=4)
-    assert data[1][0] == 2000.0
-    assert data[1][1] == 4000.0
-
-    with pytest.raises(ValueError):
-        atomic.read_zeta_data(None)
-
-    with pytest.raises(IOError):
-        atomic.read_zeta_data('fakepath')
-
-    with pytest.raises(ValueError):
-        atomic.read_zeta_data(default_atom_h5_path)
-
-
-def test_read_macro_atom_data(default_atom_h5_path, chianti_he_db_h5_path):
-    data = atomic.read_macro_atom_data(chianti_he_db_h5_path)
-    assert data[0]['atomic_number'][0] == 2
-    assert data[0]['ion_number'][0] == 0
-    assert data[0]['source_level_number'][0] == 0.0
-    assert data[0]['destination_level_number'][0] == 48.0
-    assert data[0]['transition_type'][0] == 1
-    assert data[0]['transition_probability'][0] == 0.0
-    assert data[0]['transition_line_id'][0] == 564957
-
-    assert data[1]['count_down'][0] == 0
-    assert data[1]['count_up'][0] == 7
-    assert data[1]['count_total'][0] == 7
-
-    with pytest.raises(ValueError):
-        atomic.read_macro_atom_data(None)
-
-    with pytest.raises(IOError):
-        atomic.read_macro_atom_data('fakepath')
-
-    with pytest.raises(ValueError):
-        atomic.read_macro_atom_data(default_atom_h5_path)
-
-
-def test_atom_levels():
-    atom_data = atomic.AtomData.from_hdf5(atomic.default_atom_h5_path)
-    with pytest.raises(Exception):
-        raise Exception('test the atom_data thoroughly')
 
 def test_atomic_symbol():
-    assert atomic.atomic_number2symbol[14] == 'Si'
+    assert atomic_number2symbol[14] == 'Si'
 
-def test_atomic_symbol_reverse():
-    assert atomic.symbol2atomic_number['Si'] == 14
-
-@pytest.mark.skipif(not pytest.config.getvalue("atomic-dataset"),
-                    reason='--atomic_database was not specified')
-def test_atomic_reprepare():
-    atom_data_filename = os.path.expanduser(os.path.expandvars(
-        pytest.config.getvalue('atomic-dataset')))
-    assert os.path.exists(atom_data_filename), ("{0} atomic datafiles "
-                                                         "does not seem to "
-                                                         "exist".format(
-        atom_data_filename))
-    atom_data = atomic.AtomData.from_hdf5(atom_data_filename)
-    atom_data.prepare_atom_data([14])
-    assert len(atom_data.lines) > 0
-    # Fix for new behavior of prepare_atom_data
-    # Consider doing only one prepare_atom_data and check
-    # len(atom_data.lines) == N where N is known
-    atom_data = atomic.AtomData.from_hdf5(atom_data_filename)
-    atom_data.prepare_atom_data([20])
-    assert len(atom_data.lines) > 0
 
+def test_atomic_symbol_reverse():
+    assert symbol2atomic_number['Si'] == 14