Merge pull request #312 from aoifeboyle/plasma/testing_level_population

Added tests for plasma level population properties.

tardis/plasma/tests/conftest.py

@@ -9,8 +9,10 @@
 NumberDensity)
 from tardis.plasma.properties.partition_function import (LevelBoltzmannFactor,
 LTEPartitionFunction)
-from tardis.plasma.properties.atomic import Levels, AtomicMass, IonizationData
+from tardis.plasma.properties.atomic import (Levels, Lines, AtomicMass,
+IonizationData)
 from tardis.plasma.standard_plasmas import LTEPlasma
+from tardis.plasma.properties.level_population import LevelPopulationLTE
 
 @pytest.fixture
 def number_of_cells():
@@ -64,6 +66,11 @@ def levels(included_he_atomic_data, selected_atoms):
     selected_atoms = [2]
     return levels_module.calculate(included_he_atomic_data, selected_atoms)
 
+@pytest.fixture
+def lines(included_he_atomic_data, selected_atoms):
+    lines_module = Lines(None)
+    return lines_module.calculate(included_he_atomic_data, selected_atoms)
+
 @pytest.fixture
 def level_boltzmann_factor(levels, beta_rad):
     level_boltzmann_factor_module = LevelBoltzmannFactor(None)
@@ -101,4 +108,10 @@ def number_density(atomic_mass, abundance, density):
 def ion_number_density(phi_saha_lte, partition_function, number_density):
     ion_number_density_module = IonNumberDensity(None)
     return ion_number_density_module.calculate(phi_saha_lte,
-        partition_function, number_density)
+        partition_function, number_density)
+
+@pytest.fixture
+def level_population_lte(levels, partition_function, level_boltzmann_factor):
+    level_population_lte_module = LevelPopulationLTE(None)
+    return level_population_lte_module.calculate(levels,
+        partition_function, level_boltzmann_factor)

tardis/plasma/tests/test_property_atomic.py

@@ -1,33 +1,23 @@
 import numpy as np
 
-import tardis
-from tardis.plasma.properties.atomic import (Levels, Lines, 
-LinesLowerLevelIndex, LinesUpperLevelIndex, AtomicMass, IonizationData)
+from tardis.plasma.properties.atomic import (LinesLowerLevelIndex,
+LinesUpperLevelIndex)
 
-def test_levels_property(levels, selected_atoms, included_he_atomic_data):
+def test_levels_property(levels):
     assert np.isclose(levels.ix[2].ix[0].ix[1]['energy'], 3.17545416e-11)
 
-def test_lines_property(included_he_atomic_data, selected_atoms):
-    lines_module = Lines(None)
-    lines = lines_module.calculate(included_he_atomic_data, selected_atoms)
+def test_lines_property(lines):
     assert np.isclose(lines.ix[564954]['wavelength'], 10833.307)
 
-def test_lines_lower_level_index_property(included_he_atomic_data,
-        selected_atoms, levels):
-    lines_module = Lines(None)
-    lines = lines_module.calculate(included_he_atomic_data, selected_atoms)
+def test_lines_lower_level_index_property(levels, lines):
     lines_lower_module = LinesLowerLevelIndex(None)
     lines_lower_level_index = lines_lower_module.calculate(levels, lines)
     assert lines_lower_level_index[9] == 0
 
-def test_lines_upper_level_index_property(included_he_atomic_data,
-        selected_atoms, levels):
-    lines_module = Lines(None)
-    lines = lines_module.calculate(included_he_atomic_data, selected_atoms)
+def test_lines_upper_level_index_property(levels, lines):
     lines_upper_module = LinesUpperLevelIndex(None)
     lines_upper_level_index = lines_upper_module.calculate(levels, lines)
     assert lines_upper_level_index[9] == 30
 
-def test_ionization_data_property(included_he_atomic_data,
-        ionization_data):
+def test_ionization_data_property(included_he_atomic_data, ionization_data):
     assert np.isclose(float(ionization_data.ix[2].ix[1]), 3.9393336e-11)

tardis/plasma/tests/test_property_ion_population.py

@@ -1,15 +1,5 @@
-import pytest
-
 import numpy as np
 
-from tardis.plasma.properties.ion_population import (PhiSahaLTE,
-IonNumberDensity)
-from tardis.plasma.properties.general import (BetaRadiation, GElectron,
-NumberDensity)
-from tardis.plasma.properties.partition_function import (LevelBoltzmannFactor,
-LTEPartitionFunction)
-from tardis.plasma.properties.atomic import Levels, AtomicMass, IonizationData
-
 def test_phi_saha_lte(t_rad, beta_rad, g_electron, ionization_data,
         phi_saha_lte):
     assert(phi_saha_lte.shape == (2,20))

tardis/plasma/tests/test_property_level_population.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+from tardis.plasma.properties.level_population import (LevelPopulationLTE,
+LevelNumberDensity)
+
+def test_level_population_lte(level_population_lte):
+    assert np.allclose(level_population_lte.ix[2].ix[0].ix[0], 1.0)
+    assert np.allclose(level_population_lte.ix[2].ix[0].sum(), 1.0)
+
+def test_level_number_density(level_population_lte, ion_number_density):
+    level_number_density_module = LevelNumberDensity(None)
+    level_number_density = level_number_density_module.calculate(
+        level_population_lte, ion_number_density)
+    assert np.allclose(level_number_density.ix[2].ix[0].ix[0],
+        ion_number_density.ix[2].ix[0])
+    assert np.allclose(level_number_density.sum(), ion_number_density.sum())