Merge pull request #471 from yeganer/quickfix_455

Quickfix to solve Issue #455 and #464

tardis/atomic.py

@@ -17,6 +17,9 @@
 import pandas as pd
 
 
+class AtomDataNotPreparedError(Exception):
+    pass
+
 
 logger = logging.getLogger(__name__)
 
@@ -358,6 +361,8 @@ def from_hdf5(cls, fname=None):
     def __init__(self, atom_data, ionization_data, levels_data, lines_data, macro_atom_data=None, zeta_data=None,
                  collision_data=None, synpp_refs=None, ion_cx_data=None):
 
+        self.prepared = False
+
         if levels_data is not None:
             self.has_levels = True
         else:
@@ -419,13 +424,12 @@ def __init__(self, atom_data, ionization_data, levels_data, lines_data, macro_at
         self.ionization_data.ionization_energy = units.Unit('eV').to('erg',
                                                                      self.ionization_data.ionization_energy.values)
 
-        self._levels = DataFrame(levels_data.__array__())
-        self._levels.energy = units.Unit('eV').to('erg', self._levels.energy.values)
+        self.levels = DataFrame(levels_data.__array__())
+        self.levels.energy = units.Unit('eV').to('erg', self.levels.energy.values)
 
-        self._lines = DataFrame(lines_data.__array__())
-        self._lines.set_index('line_id', inplace=True)
-        self._lines['nu'] = units.Unit('angstrom').to('Hz', self._lines['wavelength'], units.spectral())
-        self._lines['wavelength_cm'] = units.Unit('angstrom').to('cm', self._lines['wavelength'])
+        self.lines = DataFrame(lines_data.__array__())
+        self.lines['nu'] = units.Unit('angstrom').to('Hz', self.lines['wavelength'], units.spectral())
+        self.lines['wavelength_cm'] = units.Unit('angstrom').to('cm', self.lines['wavelength'])
 
 
 
@@ -463,11 +467,14 @@ def prepare_atom_data(self, selected_atomic_numbers, line_interaction_type='scat
             maximum ion number to be included in the calculation
 
         """
+        if not self.prepared:
+            self.prepared = True
+        else:
+            raise AtomDataNotPreparedError("AtomData was already prepared")
         self.selected_atomic_numbers = selected_atomic_numbers
 
         self.nlte_species = nlte_species
-        self._levels = self._levels.reset_index(drop=True)
-        self.levels = self._levels.copy()
+        self.levels = self.levels.reset_index(drop=True)
 
         self.levels = self.levels[self.levels['atomic_number'].isin(self.selected_atomic_numbers)]
 
@@ -479,15 +486,16 @@ def prepare_atom_data(self, selected_atomic_numbers, line_interaction_type='scat
 
         self.levels_index = pd.Series(np.arange(len(self.levels), dtype=int), index=self.levels.index)
         #cutting levels_lines
-        self.lines = self._lines.copy()
         self.lines = self.lines[self.lines['atomic_number'].isin(self.selected_atomic_numbers)]
         if max_ion_number is not None:
             self.lines = self.lines[self.lines['ion_number'] <= max_ion_number]
 
-        self.lines.sort('wavelength', inplace=True)
+        # self.lines.sort(['wavelength', 'line_id'], inplace=True)
+        self.lines.sort(['wavelength'], inplace=True)
+        self.lines.set_index('line_id', inplace=True)
+
 
 
-
         self.lines_index = pd.Series(np.arange(len(self.lines), dtype=int), index=self.lines.index)
 
         tmp_lines_lower2level_idx = pd.MultiIndex.from_arrays([self.lines['atomic_number'], self.lines['ion_number'],
@@ -593,7 +601,6 @@ def __init__(self, atom_data, nlte_species):
             if atom_data.has_collision_data:
                 self._create_collision_coefficient_matrix()
         else:
-
             self._create_nlte_mask()
 
     def _init_indices(self):

tardis/conftest.py

@@ -39,7 +39,7 @@ def pytest_report_header(config):
     s = "\n"
     if six.PY2:
         args = [x.decode('utf-8') for x in config.args]
-    elif six.PY3:
+    else:
         args = config.args
     s += "Running tests in {0}.\n\n".format(" ".join(args))
 
@@ -182,6 +182,3 @@ def included_he_atomic_data(test_data_path):
 def tardis_config_verysimple():
     return yaml.load(
         open('tardis/io/tests/data/tardis_configv1_verysimple.yml'))
-
-
-

tardis/io/tests/data/tardis_configv1_verysimple.yml

@@ -5,7 +5,6 @@ supernova:
     time_explosion: 13 day
 
 atom_data: kurucz_atom_pure_simple.h5
-
 model:
 
     structure:

tardis/plasma/properties/atomic.py

@@ -31,11 +31,12 @@ class Levels(BaseAtomicDataProperty):
         'g')
 
     def _filter_atomic_property(self, levels, selected_atoms):
-        return levels[levels.atomic_number.isin(selected_atoms)]
+        return levels
+        # return levels[levels.atomic_number.isin(selected_atoms)]
 
     def _set_index(self, levels):
-        levels = levels.set_index(['atomic_number', 'ion_number',
-                                 'level_number'])
+        # levels = levels.set_index(['atomic_number', 'ion_number',
+        #                          'level_number'])
         return (levels.index, levels['energy'], levels['metastable'],
             levels['g'])
 
@@ -53,23 +54,15 @@ class Lines(BaseAtomicDataProperty):
     wavelength_cm: Pandas DataFrame (index=line_id), dtype float
             Line wavelengths in cm
     """
-# Would like for lines to just be the line_id values, but this is necessary because of the filtering
-# error noted below.
+# Would like for lines to just be the line_id values
     outputs = ('lines', 'nu', 'f_lu', 'wavelength_cm')
 
     def _filter_atomic_property(self, lines, selected_atoms):
-        return lines[lines.atomic_number.isin(selected_atoms)]
+        # return lines[lines.atomic_number.isin(selected_atoms)]
+        return lines
 
     def _set_index(self, lines):
-# Filtering process re-arranges the index. This re-orders it.
-# It seems to be important that the lines stay indexed in the correct order
-# so that the tau_sobolevs values are in the right order for the montecarlo
-# code, or it returns the wrong answer.
-        try:
-            reindexed = lines.reindex(lines.index)
-        except:
-            reindexed = lines.reindex(lines.index)
-        lines = reindexed.dropna(subset=['atomic_number'])
+        # lines = lines.set_index('line_id')
         return lines, lines['nu'], lines['f_lu'], lines['wavelength_cm']
 
 class LinesLowerLevelIndex(HiddenPlasmaProperty):

tardis/plasma/properties/base.py

@@ -10,6 +10,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 class BasePlasmaProperty(object):
     """
     Attributes
@@ -50,14 +51,18 @@ def get_latex_label(self):
         else:
             complete_name = latex_name
 
-        latex_label = latex_template.format(name=complete_name,
-                                     formula=getattr(self,
-                                                     'latex_formula', '--'),
-                                     description=getattr(self,
-                                                         'latex_description',
-                                                         ''))
+        latex_label = latex_template.format(
+                name=complete_name,
+                formula=getattr(
+                    self,
+                    'latex_formula', '--'),
+                description=getattr(
+                    self,
+                    'latex_description',
+                    ''))
         return latex_label.replace('\\', r'\\')
 
+
 class ProcessingPlasmaProperty(BasePlasmaProperty):
     """
     Attributes
@@ -79,8 +84,9 @@ def _update_inputs(self):
         """
         calculate_call_signature = self.calculate.func_code.co_varnames[
                                    :self.calculate.func_code.co_argcount]
-        self.inputs = [item for item in calculate_call_signature if
-                      item != 'self']
+        self.inputs = [
+                item for item in calculate_call_signature
+                if item != 'self']
 
     def _get_input_values(self):
         return (self.plasma_parent.get_value(item) for item in self.inputs)
@@ -105,6 +111,7 @@ def calculate(self, *args, **kwargs):
         raise NotImplementedError('This method needs to be implemented by '
                                   'processing plasma modules')
 
+
 class HiddenPlasmaProperty(ProcessingPlasmaProperty):
     """
     Used for plasma properties that should not be displayed in the final graph (e.g. lines_lower_level_index).
@@ -116,6 +123,7 @@ class HiddenPlasmaProperty(ProcessingPlasmaProperty):
     def __init__(self, plasma_parent):
         super(HiddenPlasmaProperty, self).__init__(plasma_parent)
 
+
 class BaseAtomicDataProperty(ProcessingPlasmaProperty):
     """
     Used for atomic data properties. Main feature is the ability to filter atomic data by the elements required for
@@ -143,21 +151,17 @@ def calculate(self, atomic_data, selected_atoms):
         if getattr(self, self.outputs[0]) is not None:
             return getattr(self, self.outputs[0])
         else:
-# Atomic Data Issue: Some atomic property names in the h5 files are preceded
-# by an underscore, e.g. _levels, _lines.
-            try:
-                raw_atomic_property = getattr(atomic_data, '_'
-                                              + self.outputs[0])
-            except AttributeError:
-                raw_atomic_property = getattr(atomic_data, self.outputs[0])
-            finally:
-                return self._set_index(self._filter_atomic_property(
-                    raw_atomic_property, selected_atoms))
+            raw_atomic_property = getattr(atomic_data, self.outputs[0])
+            return self._set_index(
+                    self._filter_atomic_property(
+                        raw_atomic_property, selected_atoms)
+                    )
+
 
 class Input(BasePlasmaProperty):
     """
-    The plasma property class for properties that are input directly from model and not calculated within the plasma
-    module, e.g. t_rad.
+    The plasma property class for properties that are input directly from model
+    and not calculated within the plasma module, e.g. t_rad.
     """
     def _set_output_value(self, output, value):
         setattr(self, output, value)
@@ -166,14 +170,17 @@ def set_value(self, value):
         assert len(self.outputs) == 1
         self._set_output_value(self.outputs[0], value)
 
+
 class ArrayInput(Input):
     def _set_output_value(self, output, value):
         setattr(self, output, np.array(value, copy=False))
 
+
 class DataFrameInput(Input):
     def _set_output_value(self, output, value):
         setattr(self, output, np.array(pd.DataFrame(value), copy=False))
 
+
 class PreviousIterationProperty(BasePlasmaProperty):
     """
     This class is used for properties where, to prevent a property calculation loop, the property values from the

tardis/plasma/tests/conftest.py

@@ -13,10 +13,12 @@
 # INPUTS
 
 @pytest.fixture
-def atomic_data():
+def atomic_data(selected_atoms):
     atomic_db_fname = os.path.join(tardis.__path__[0], 'tests', 'data',
                                    'chianti_he_db.h5')
-    return AtomData.from_hdf5(atomic_db_fname)
+    atom_data = AtomData.from_hdf5(atomic_db_fname)
+    atom_data.prepare_atom_data(selected_atoms)
+    return atom_data
 
 @pytest.fixture
 def number_of_cells():
@@ -256,4 +258,4 @@ def transition_probabilities(atomic_data, beta_sobolev, j_blues,
     return transition_probabilities_module.calculate(atomic_data, beta_sobolev,
                                                      j_blues,
                                                      stimulated_emission_factor,
-                                                     tau_sobolev)
+                                                     tau_sobolev)

tardis/plasma/tests/test_plasmas_full.py

@@ -66,4 +66,4 @@ def test_nlte_plasma(self, plasma_compare_data):
         np.testing.assert_allclose(
             new_plasma_t_rads, old_plasma_t_rads, atol=150)
         np.testing.assert_allclose(
-            new_plasma_levels, old_plasma_levels, rtol=0.1)
+            new_plasma_levels, old_plasma_levels, rtol=0.1)

tardis/simulation/tests/test_simulation.py

@@ -7,9 +7,6 @@
 from tardis.simulation import Simulation
 
 
-
-import pandas as pd
-
 @pytest.fixture
 def tardis_config(kurucz_atomic_data, tardis_config_verysimple):
     return config_reader.Configuration.from_config_dict(
@@ -28,30 +25,37 @@ def simulation_one_loop(raw_model, tardis_config):
 
     return sim
 
+
 @pytest.fixture()
 def simulation_compare_data_fname():
     return 'tardis/simulation/tests/data/test_data.h5'
 
+
 @pytest.fixture()
 def simulation_compare_data(simulation_compare_data_fname):
     return h5py.File(simulation_compare_data_fname, mode='r')
 
 
-
 def test_plasma_estimates(simulation_one_loop, simulation_compare_data):
-    npt.assert_allclose(simulation_one_loop.runner.j_estimator,
-                        simulation_compare_data['test1/j_estimators'], atol=0.0)
+    t_rad, w = simulation_one_loop.runner.calculate_radiationfield_properties()
+
     npt.assert_allclose(simulation_one_loop.runner.nu_bar_estimator,
                         simulation_compare_data['test1/nubar_estimators'],
                         atol=0.0)
+    npt.assert_allclose(simulation_one_loop.runner.j_estimator,
+                        simulation_compare_data['test1/j_estimators'],
+                        atol=0.0)
 
-    t_rad, w = simulation_one_loop.runner.calculate_radiationfield_properties()
-    npt.assert_allclose(t_rad, simulation_compare_data['test1/t_rad'], atol=0.0)
+    npt.assert_allclose(
+            t_rad, simulation_compare_data['test1/t_rad'], atol=0.0)
     npt.assert_allclose(w, simulation_compare_data['test1/w'], atol=0.0)
 
+
 def test_packet_output(simulation_one_loop, simulation_compare_data):
-    npt.assert_allclose(simulation_one_loop.runner.output_nu,
-                    simulation_compare_data['test1/output_nu'], atol=0.0)
+    npt.assert_allclose(
+            simulation_one_loop.runner.output_nu,
+            simulation_compare_data['test1/output_nu'],
+            atol=0.0)
 
     npt.assert_allclose(simulation_one_loop.runner.output_energy,
                         simulation_compare_data['test1/output_energy'],

tardis/tests/test_atomic.py

@@ -50,6 +50,10 @@ def test_atomic_reprepare():
     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