Merge be51046 into badc89f

docs/conf.py

@@ -49,6 +49,7 @@
     'sphinx.ext.mathjax',
     'sphinx.ext.graphviz',
     'numpydoc',
+    'nbsphinx',
 ]
 
 source_suffix = '.rst'

docs/to_hdf.rst

@@ -0,0 +1,8 @@
+**************************
+Storing simulations to HDF
+**************************
+
+You can ask TARDIS to store the state of each iteration of the simulation you are
+ running. To see an example of how this can be achieved see the following jupyter notebook.
+
+:ref:`notebooks/to_hdf.ipynb`

tardis/io/util.py

@@ -1,5 +1,6 @@
 #Utility functions for the IO part of TARDIS
 
+import os
 import pandas as pd
 import numpy as np
 import collections
@@ -196,3 +197,60 @@ def check_equality(item1, item2):
         return False
     else:
         return True
+
+
+def to_hdf(path_or_buf, path, elements):
+    """
+    A function to uniformly store TARDIS data
+    to an HDF file.
+
+    Scalars will be stored in a Series under path/scalars
+    1D arrays will be stored under path/property_name as distinct Series
+    2D arrays will be stored under path/property_name as distinct DataFrames
+
+    Units will be stored as their SI value
+
+    Parameters
+    ----------
+    path_or_buf:
+        Path or buffer to the HDF store
+    path: str
+        Path inside the HDF store to store the `elements`
+    elements: dict
+        A dict of property names and their values to be
+        stored.
+
+    Returns
+    -------
+
+    """
+    scalars = {}
+    for key, value in elements.iteritems():
+        if hasattr(value, 'si'):
+            value = value.si.value
+        if np.isscalar(value):
+            scalars[key] = value
+        elif hasattr(value, 'shape'):
+            if value.ndim == 1:
+                # This try,except block is only for model.plasma.levels
+                try:
+                    pd.Series(value).to_hdf(path_or_buf,
+                                            os.path.join(path, key))
+                except NotImplementedError:
+                    pd.DataFrame(value).to_hdf(path_or_buf,
+                                               os.path.join(path, key))
+            else:
+                pd.DataFrame(value).to_hdf(path_or_buf, os.path.join(path, key))
+        else:
+            data = pd.DataFrame([value])
+            data.to_hdf(path_or_buf, os.path.join(path, key))
+
+    if scalars:
+        scalars_series = pd.Series(scalars)
+
+        # Unfortunately, with to_hdf we cannot append, so merge beforehand
+        scalars_path = os.path.join(path, 'scalars')
+        with pd.HDFStore(path_or_buf) as store:
+            if scalars_path in store:
+                scalars_series = store[scalars_path].append(scalars_series)
+        scalars_series.to_hdf(path_or_buf, os.path.join(path, 'scalars'))

tardis/montecarlo/base.py

@@ -6,6 +6,7 @@
 from scipy.special import zeta
 
 from tardis.montecarlo import montecarlo, packet_source
+from tardis.io.util import to_hdf
 
 import numpy as np
 import pandas as pd
@@ -223,58 +224,23 @@ def calculate_f_nu(self, frequency):
     def calculate_f_lambda(self, wavelength):
         pass
 
-    def property_to_hdf(self, property):
-        """
-        Convert an attribute of MontecarloRunner to a pandas
-         DataFrame
-
-        Parameters
-        ----------
-        property: str
-            The MontecarloRunner property name to be converted.
-
-        Returns
-        -------
-            : pandas.DataFrame
-
-        """
-        output_value = getattr(self, property)
-        if np.isscalar(output_value):
-            output_value = [output_value]
-
-        return pd.DataFrame(output_value)
-
-    def to_hdf(self, path_or_buf, path='', close_hdf=True):
+    def to_hdf(self, path_or_buf, path=''):
         """
         Store the runner to an HDF structure.
-        Currently only output_nu and output_energy are stored.
 
         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]
-        filter: ~Filter
-            Only properties returned by the filter will be stored
+            Path inside the HDF store to store the runner
         Returns
         -------
             : None
 
         """
-        try:
-            hdf_store = pd.HDFStore(path_or_buf)
-        except TypeError:
-            hdf_store = path_or_buf
-
+        runner_path = os.path.join(path, 'runner')
         properties = ['output_nu', 'output_energy', 'nu_bar_estimator',
                       'j_estimator']
-
-        for prop in properties:
-            hdf_store[os.path.join(path, 'runner', prop)] = \
-                self.property_to_hdf(prop)
-
-        if close_hdf:
-            hdf_store.close()
+        to_hdf(path_or_buf, runner_path, {name: getattr(self, name) for name
+                                          in properties})

tardis/plasma/base.py

@@ -260,7 +260,7 @@ def remove_hidden_properties(self, print_graph):
                 print_graph.remove_node(str(item.name))
         return print_graph
 
-    def to_hdf(self, path_or_buf, path='', close_hdf=True, collection=None):
+    def to_hdf(self, path_or_buf, path='', collection=None):
         """
         Store the plasma to an HDF structure
 
@@ -270,8 +270,6 @@ def to_hdf(self, path_or_buf, path='', close_hdf=True, collection=None):
             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
@@ -285,17 +283,10 @@ def to_hdf(self, path_or_buf, path='', close_hdf=True, collection=None):
             : 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'))
-
-        if close_hdf:
-            hdf_store.close()
+            prop.to_hdf(path_or_buf, os.path.join(path, 'plasma'))

tardis/plasma/properties/base.py

@@ -4,6 +4,8 @@
 import numpy as np
 import pandas as pd
 
+from tardis.io.util import to_hdf
+
 __all__ = ['BasePlasmaProperty', 'BaseAtomicDataProperty',
            'HiddenPlasmaProperty', 'Input', 'ArrayInput', 'DataFrameInput',
            'ProcessingPlasmaProperty', 'PreviousIterationProperty']
@@ -63,37 +65,14 @@ 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):
+    def to_hdf(self, path_or_buf, path):
         """
         Stores plugin value to an HDFStore
 
         Parameters
         ----------
-        hdf_store: ~pandas.HDFStore object
-            an open pandas datastore
-
+        path_or_buf:
+            Path or buffer to the HDF store
         path: ~str
             path to store the modules data under
             - will be joined to <path>/output_name
@@ -103,9 +82,8 @@ def to_hdf(self, hdf_store, path):
             : None
 
         """
-        for output in self.outputs:
-            hdf_store[os.path.join(path, output)] = self._get_output_to_hdf(
-                output)
+        outputs = {name: getattr(self, name) for name in self.outputs}
+        to_hdf(path_or_buf, path, outputs)
 
 
 class ProcessingPlasmaProperty(BasePlasmaProperty):

tardis/plasma/properties/plasma_input.py

@@ -31,7 +31,7 @@ class AtomicData(Input):
     """
     outputs = ('atomic_data',)
 
-    def to_hdf(self, hdf_store, path):
+    def to_hdf(self, path_or_buf, path):
         """
         Will not store anything to hdf as derivatives (lines, levels) are stored
 

tardis/simulation/base.py

@@ -419,7 +419,7 @@ def to_hdf(self, model, path_or_buf, path='', plasma_properties=None):
         -------
         None
         """
-        self.runner.to_hdf(path_or_buf, path, False)
+        self.runner.to_hdf(path_or_buf, path)
         model.to_hdf(path_or_buf, path, plasma_properties)