[TEP002] Add to_hdf methods to the plasma

tardis/gui/widgets.py

@@ -750,7 +750,7 @@ def on_atom_header_double_clicked(self, index):
         """Called when a header in the first column is clicked to show 
         ion populations."""
         self.current_atom_index = self.table1_data.index.values.tolist()[index]
-        self.table2_data = self.parent.model.plasma_array.ion_number_density[
+        self.table2_data = self.parent.model.plasma.ion_number_density[
             self.shell_index].ix[self.current_atom_index]
         self.ionsdata = self.createTable([['Ion: '], 
             ['Count (Z = %d)' % self.current_atom_index]], 
@@ -776,7 +776,7 @@ def on_atom_header_double_clicked(self, index):
     def on_ion_header_double_clicked(self, index):
         """Called on double click of ion headers to show level populations."""
         self.current_ion_index = self.table2_data.index.values.tolist()[index]
-        self.table3_data = self.parent.model.plasma_array.level_number_density[
+        self.table3_data = self.parent.model.plasma.level_number_density[
             self.shell_index].ix[self.current_atom_index, self.current_ion_index]
         self.levelsdata = self.createTable([['Level: '], 
             ['Count (Ion %d)' % self.current_ion_index]], 
@@ -795,7 +795,7 @@ def on_ion_header_double_clicked(self, index):
 
     def update_tables(self):
         """Update table data for shell info viewer."""
-        self.table1_data = self.parent.model.plasma_array.number_density[
+        self.table1_data = self.parent.model.plasma.number_density[
             self.shell_index]
         self.atomsdata.index_info=self.table1_data.index.values.tolist()
         self.atomsdata.arraydata = []

tardis/model.py

@@ -6,6 +6,7 @@
 import numpy as np
 import pandas as pd
 from astropy import constants, units as u
+from astropy.utils.decorators import deprecated
 
 from util import intensity_black_body
 
@@ -103,7 +104,7 @@ def __init__(self, tardis_config):
         heating_rate_data_file = getattr(
             tardis_config.plasma, 'heating_rate_data_file', None)
 
-        self.plasma_array = LegacyPlasmaArray(
+        self.plasma = LegacyPlasmaArray(
             tardis_config.number_densities, tardis_config.atom_data,
             tardis_config.supernova.time_explosion.to('s').value,
             nlte_config=tardis_config.plasma.nlte,
@@ -129,6 +130,16 @@ def __init__(self, tardis_config):
         self.update_plasmas(initialize_nlte=True)
 
 
+    @property
+    @deprecated('v1.5',
+                'plasma_array has been renamed to plasma and will be removed in the future. Please use model.plasma instead.')
+    def plasma_array(self):
+        return self.plasma
+
+    @plasma_array.setter
+    def plasma_array(self, value):
+        self.plasma = value
+
     @property
     def line_interaction_type(self):
         return self._line_interaction_type
@@ -225,15 +236,15 @@ def calculate_j_blues(self, init_detailed_j_blues=False):
 
     def update_plasmas(self, initialize_nlte=False):
 
-        self.plasma_array.update_radiationfield(
+        self.plasma.update_radiationfield(
             self.t_rads.value, self.ws, self.j_blues,
             self.tardis_config.plasma.nlte, initialize_nlte=initialize_nlte,
             n_e_convergence_threshold=0.05)
 
         if self.tardis_config.plasma.line_interaction_type in ('downbranch',
                                                                'macroatom'):
             self.transition_probabilities = (
-                self.plasma_array.transition_probabilities)
+                self.plasma.transition_probabilities)
 
     def save_spectra(self, fname):
         self.spectrum.to_ascii(fname)

tardis/montecarlo/base.py

@@ -95,7 +95,7 @@ def run(self, model, no_of_packets, no_of_virtual_packets=0, nthreads=1):
         self.time_of_simulation = model.time_of_simulation
         self.volume = model.tardis_config.structure.volumes
         self._initialize_estimator_arrays(self.volume.shape[0],
-                                          model.plasma_array.tau_sobolevs.shape)
+                                          model.plasma.tau_sobolevs.shape)
         self._initialize_geometry_arrays(model.tardis_config.structure)
 
         self._initialize_packets(model.t_inner.value,

tardis/montecarlo/montecarlo.pyx

@@ -119,10 +119,10 @@ cdef initialize_storage_model(model, runner, storage_model_t *storage):
     storage.inverse_time_explosion = 1.0 / storage.time_explosion
     #electron density
     storage.electron_densities = <double*> PyArray_DATA(
-        model.plasma_array.electron_densities.values)
+        model.plasma.electron_densities.values)
 
     runner.inverse_electron_densities = (
-        1.0 / model.plasma_array.electron_densities.values)
+        1.0 / model.plasma.electron_densities.values)
     storage.inverse_electron_densities = <double*> PyArray_DATA(
         runner.inverse_electron_densities)
     # Switch for continuum processes
@@ -148,7 +148,7 @@ cdef initialize_storage_model(model, runner, storage_model_t *storage):
     storage.no_of_lines = model.atom_data.lines.nu.values.size
     storage.line_list_nu = <double*> PyArray_DATA(model.atom_data.lines.nu.values)
     runner.line_lists_tau_sobolevs = (
-            model.plasma_array.tau_sobolevs.values.flatten(order='F')
+            model.plasma.tau_sobolevs.values.flatten(order='F')
             )
     storage.line_lists_tau_sobolevs = <double*> PyArray_DATA(
             runner.line_lists_tau_sobolevs
@@ -162,13 +162,13 @@ cdef initialize_storage_model(model, runner, storage_model_t *storage):
     # macro atom & downbranch
     if storage.line_interaction_id >= 1:
         runner.transition_probabilities = (
-                model.plasma_array.transition_probabilities.values.flatten(order='F')
+                model.plasma.transition_probabilities.values.flatten(order='F')
         )
         storage.transition_probabilities = <double*> PyArray_DATA(
                 runner.transition_probabilities
                 )
         storage.transition_probabilities_nd = (
-        model.plasma_array.transition_probabilities.values.shape[0])
+        model.plasma.transition_probabilities.values.shape[0])
         storage.line2macro_level_upper = <int_type_t*> PyArray_DATA(
             model.atom_data.lines_upper2macro_reference_idx)
         storage.macro_block_references = <int_type_t*> PyArray_DATA(
@@ -214,7 +214,7 @@ cdef initialize_storage_model(model, runner, storage_model_t *storage):
     storage.reflective_inner_boundary = model.tardis_config.montecarlo.enable_reflective_inner_boundary
     storage.inner_boundary_albedo = model.tardis_config.montecarlo.inner_boundary_albedo
     # Data for continuum implementation
-    cdef np.ndarray[double, ndim=1] t_electrons = model.plasma_array.t_electrons
+    cdef np.ndarray[double, ndim=1] t_electrons = model.plasma.t_electrons
     storage.t_electrons = <double*> t_electrons.data
 
 def montecarlo_radial1d(model, runner, int_type_t virtual_packet_flag=0,

tardis/plasma/base.py

@@ -1,8 +1,10 @@
+import os
 import logging
 import tempfile
 import fileinput
 
 import networkx as nx
+import pandas as pd
 
 from tardis.plasma.exceptions import PlasmaMissingModule, NotInitializedModule
 from tardis.plasma.properties.base import *
@@ -39,7 +41,7 @@ def __dir__(self):
                  if not item.startswith('_')]
         attrs += [item for item in self.__class__.__dict__
                  if not item.startswith('_')]
-        attrs += self.module_dict.keys()
+        attrs += self.outputs_dict.keys()
         return attrs
 
     @property
@@ -258,10 +260,42 @@ def remove_hidden_properties(self, print_graph):
                 print_graph.remove_node(str(item.name))
         return print_graph
 
-class StandardPlasma(BasePlasma):
+    def to_hdf(self, path_or_buf, path='', close_hdf=True, collection=None):
+        """
+        Store the plasma to an HDF structure
+
+        Parameters
+        ----------
+        path_or_buf:
+            Path or buffer to the HDF store
+        path:
+            Path inside the HDF store to store the plasma
+        close_hdf : bool
+            if  True close the HDFStore [default=True]
+        collection:
+            `None` or a `PlasmaPropertyCollection` of which members are
+            the property types which will be stored. If `None` then
+            all types of properties will be stored.
+
+            This acts like a filter, for example if a value of
+            `property_collections.basic_inputs` is given, only
+            those input parameters will be stored to the HDF store.
+        Returns
+        -------
+            : None
+
+        """
+        try:
+            hdf_store = pd.HDFStore(path_or_buf)
+        except TypeError:
+            hdf_store = path_or_buf
+        if collection:
+            properties = [prop for prop in self.plasma_properties if
+                          isinstance(prop, tuple(collection))]
+        else:
+            properties = self.plasma_properties
+        for prop in properties:
+            prop.to_hdf(hdf_store, os.path.join(path, 'plasma'))
 
-    def __init__(self, number_densities, atom_data, time_explosion,
-                 nlte_config=None, ionization_mode='lte',
-                 excitation_mode='lte', w=None,
-                 link_t_rad_t_electron=0.9):
-        pass
+        if close_hdf:
+            hdf_store.close()

tardis/plasma/properties/base.py

@@ -10,6 +10,7 @@
 
 logger = logging.getLogger(__name__)
 
+import os
 
 class BasePlasmaProperty(object):
     """
@@ -62,6 +63,50 @@ def get_latex_label(self):
                     ''))
         return latex_label.replace('\\', r'\\')
 
+    def _get_output_to_hdf(self, output):
+        """
+        Convert output into a pandas DataFrame to enable storing in an HDFStore
+
+        Parameters
+        ----------
+        output: ~str
+            output string name
+
+        Returns
+        -------
+            : ~pandas.DataFrame
+
+        """
+
+        output_value = getattr(self, output)
+
+        if np.isscalar(output_value):
+            output_value = [output_value]
+
+        return pd.DataFrame(output_value)
+
+    def to_hdf(self, hdf_store, path):
+        """
+        Stores plugin value to an HDFStore
+
+        Parameters
+        ----------
+        hdf_store: ~pandas.HDFStore object
+            an open pandas datastore
+
+        path: ~str
+            path to store the modules data under
+            - will be joined to <path>/output_name
+
+        Returns
+        -------
+            : None
+
+        """
+        for output in self.outputs:
+            hdf_store[os.path.join(path, output)] = self._get_output_to_hdf(
+                output)
+
 
 class ProcessingPlasmaProperty(BasePlasmaProperty):
     """

tardis/plasma/properties/plasma_input.py

@@ -31,6 +31,17 @@ class AtomicData(Input):
     """
     outputs = ('atomic_data',)
 
+    def to_hdf(self, hdf_store, path):
+        """
+        Will not store anything to hdf as derivatives (lines, levels) are stored
+
+        Returns
+        -------
+            : None
+
+        """
+        pass
+
 class Abundance(Input):
     """
     Attributes

tardis/plasma/tests/test_plasmas_full.py

@@ -48,7 +48,7 @@ def test_lte_plasma(self, plasma_compare_data):
 
         new_plasma_t_rads = self.lte_model.t_rads / u.Unit('K')
         new_plasma_levels = \
-            self.lte_model.plasma_array.get_value(
+            self.lte_model.plasma.get_value(
             'level_number_density').ix[8].ix[1][10].values
         np.testing.assert_allclose(
             new_plasma_t_rads, old_plasma_t_rads, atol=100)
@@ -60,7 +60,7 @@ def test_nlte_plasma(self, plasma_compare_data):
         old_plasma_levels = plasma_compare_data['test_nlte1/levels']
         new_plasma_t_rads = self.nlte_model.t_rads / u.Unit('K')
         new_plasma_levels = \
-            self.nlte_model.plasma_array.get_value(
+            self.nlte_model.plasma.get_value(
             'level_number_density').ix[2].ix[1][10].values
         np.testing.assert_allclose(
             new_plasma_t_rads, old_plasma_t_rads, atol=150)

tardis/util.py

@@ -192,7 +192,7 @@ def create_synpp_yaml(radial1d_mdl, fname, shell_no=0, lines_db=None):
     for key, value in radial1d_mdl.atom_data.synpp_refs.iterrows():
         try:
             radial1d_mdl.atom_data.synpp_refs['ref_log_tau'].ix[key] = np.log10(
-                radial1d_mdl.plasma_array.tau_sobolevs[0].ix[value['line_id']])
+                radial1d_mdl.plasma.tau_sobolevs[0].ix[value['line_id']])
         except KeyError:
             pass