Tidying/Testing/Documenting Plasma

tardis/plasma/base.py

@@ -4,6 +4,7 @@
 from tardis.plasma.exceptions import PlasmaMissingModule, NotInitializedModule
 
 import tempfile
+import fileinput
 
 
 logger = logging.getLogger(__name__)
@@ -16,8 +17,8 @@ def __init__(self, plasma_properties, **kwargs):
         self.input_properties = []
         self.plasma_properties = self._init_properties(plasma_properties,
                                                        **kwargs)
-
         self._build_graph()
+#        self.write_to_tex('Plasma_Graph', 'Plasma_Formulae')
         self.update(**kwargs)
 
     def __getattr__(self, item):
@@ -58,7 +59,6 @@ def _build_graph(self):
         """
 
         self.graph = nx.DiGraph()
-
         ## Adding all nodes
         self.graph.add_nodes_from([(plasma_property.name, {})
                                    for plasma_property
@@ -79,8 +79,15 @@ def _build_graph(self):
                                               '{1} which has not been added'
                                               ' to this plasma'.format(
                         plasma_property.name, input))
+                try:
+                    position = self.outputs_dict[input].outputs.index(input)
+                    label = self.outputs_dict[input].latex_name[position]
+                    label = '$' + label + '$'
+                    label = label.replace('\\', '\\\\')
+                except:
+                    label = input.replace('_', '-')
                 self.graph.add_edge(self.outputs_dict[input].name,
-                plasma_property.name, label=input)
+                    plasma_property.name, label = label)
 
     def _init_properties(self, plasma_properties, **kwargs):
         """
@@ -175,22 +182,22 @@ def _resolve_update_list(self, changed_properties):
         return descendants_ob
 
     def write_to_dot(self, fname, latex_label=True):
-        self._update_module_type_str()
+#        self._update_module_type_str()
 
         try:
             import pygraphviz
         except ImportError:
             raise ImportError('pygraphviz is needed for method '
                               '\'write_to_dot\'')
+        print_graph = self.graph.copy()
+        print_graph.remove_node('LinesUpperLevelIndex')
+        print_graph.remove_node('LinesLowerLevelIndex')
+        for node in print_graph:
+            print_graph.node[str(node)]['label'] = node
 
-        for node in self.graph:
-            self.graph.node[node]['label'] = self.module_dict[node].get_latex_label()
-            self.graph.node[node]['color'] = 'red'
-            self.graph.node[node]['shape'] = 'box '
-
-        nx.write_dot(self.graph, fname)
+        nx.write_dot(print_graph, fname)
 
-    def write_to_tex(self, fname):
+    def write_to_tex(self, fname_graph, fname_formulae):
         try:
             import dot2tex
         except ImportError:
@@ -202,9 +209,32 @@ def write_to_tex(self, fname):
 
         dot_string = open(temp_fname).read()
 
-        open(fname, 'w').write(dot2tex.dot2tex(dot_string, texmode='raw'))
-
-
+        open(fname_graph, 'w').write(dot2tex.dot2tex(dot_string,
+            texmode='raw'))
+
+        for line in fileinput.input(fname_graph, inplace = 1):
+            print line.replace('\documentclass{article}',
+                '\documentclass[class=minimal,border=20pt]{standalone}'),
+
+        for line in fileinput.input(fname_graph, inplace = 1):
+            print line.replace('\enlargethispage{100cm}', ''),
+
+        formulae = open(fname_formulae, 'w')
+        print>>formulae, '\\documentclass{minimal}', '\n',\
+            '\usepackage{amsmath}', '\n', '\\begin{document}', '\n'
+        for key in self.outputs_dict.keys():
+            output_number = self.outputs_dict[key].outputs.index(key)
+            try:
+                label = self.outputs_dict[key].latex_name[output_number]
+            except:
+                label = key.replace('_', '-')
+            if hasattr(self.outputs_dict[key], 'latex_formula'):
+                print>>formulae, '$' + label + '$' + '\\['
+                print>>formulae, self.outputs_dict[key].latex_formula[
+                    output_number]
+                print>>formulae, '\\]' + '\n'
+        print>>formulae, '\\end{document}'
+        formulae.close()
 
 class StandardPlasma(BasePlasma):
 

tardis/plasma/properties/atomic.py

@@ -21,156 +21,207 @@ def __init__(self, plasma_parent):
 
         super(BaseAtomicDataProperty, self).__init__(plasma_parent)
         self.value = None
-        assert len(self.outputs) == 1
 
     @abstractmethod
-    def _set_index(self, raw_atomic_property, atomic_data):
+    def _set_index(self, raw_atomic_property):
         raise NotImplementedError('Needs to be implemented in subclasses')
 
     @abstractmethod
     def _filter_atomic_property(self, raw_atomic_property):
         raise NotImplementedError('Needs to be implemented in subclasses')
 
-
     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])
+                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), atomic_data)
-
+                    raw_atomic_property, selected_atoms))
 
 class Levels(BaseAtomicDataProperty):
-    outputs = ('levels',)
+    """
+    Outputs:
+    levels : Pandas DataFrame
+        Levels data needed for particular simulation
+    """
+    outputs = ('levels', 'excitation_energy', 'metastability', 'g')
+    latex_name = ('\\textrm{levels}', '\\epsilon_{\\textrm{k}}', '\\textrm{metastability}',
+        'g')
 
     def _filter_atomic_property(self, levels, selected_atoms):
-        return levels[levels.atomic_number.isin([selected_atoms]
-                                                if np.isscalar(selected_atoms)
-                                                else selected_atoms)]
+        return levels[levels.atomic_number.isin(selected_atoms)]
 
-    def _set_index(self, levels, atomic_data):
-        return levels.set_index(['atomic_number', 'ion_number',
-            'level_number'])
+    def _set_index(self, levels):
+        levels = levels.set_index(['atomic_number', 'ion_number',
+                                 'level_number'])
+        return (levels.index, levels['energy'], levels['metastable'],
+            levels['g'])
 
 class Lines(BaseAtomicDataProperty):
-    outputs = ('lines',)
+    """
+    Outputs:
+    lines : Pandas DataFrame
+        Lines data needed for particular simulation
+    """
+    outputs = ('lines', 'nu', 'f_lu', 'wavelength_cm')
 
     def _filter_atomic_property(self, lines, selected_atoms):
         return lines[lines.atomic_number.isin(selected_atoms)]
 
-    def _set_index(self, lines, atomic_data):
+    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(atomic_data.lines.index)
+            reindexed = lines.reindex(lines.index)
         except:
-            reindexed = lines.reindex(atomic_data._lines.index)
-        return reindexed
+            reindexed = lines.reindex(lines.index)
+        lines = reindexed.dropna(subset=['atomic_number'])
+        return lines, lines['nu'], lines['f_lu'], lines['wavelength_cm']
 
 class LinesLowerLevelIndex(ProcessingPlasmaProperty):
+    """
+    Outputs:
+    lines_lower_level_index : One-dimensional Numpy Array
+        Levels data for lower levels of particular lines
+        Usage: levels.ix[lines_lower_level_index]
+    """
     outputs = ('lines_lower_level_index',)
 
     def calculate(self, levels, lines):
         levels_index = pd.Series(np.arange(len(levels), dtype=np.int64),
-                                 index=levels.index)
+                                 index=levels)
         lines_index = lines.set_index(
             ['atomic_number', 'ion_number',
              'level_number_lower']).index
         return np.array(levels_index.ix[lines_index])
 
 class LinesUpperLevelIndex(ProcessingPlasmaProperty):
+    """
+    Outputs:
+    lines_upper_level_index : One-dimensional Numpy Array
+        Levels data for upper levels of particular lines
+        Usage: levels.ix[lines_upper_level_index]
+    """
     outputs = ('lines_upper_level_index',)
 
     def calculate(self, levels, lines):
         levels_index = pd.Series(np.arange(len(levels), dtype=np.int64),
-                                 index=levels.index)
+                                 index=levels)
         lines_index = lines.set_index(
             ['atomic_number', 'ion_number',
              'level_number_upper']).index
         return np.array(levels_index.ix[lines_index])
 
-
 class IonCXData(BaseAtomicDataProperty):
     outputs = ('ion_cx_data',)
 
     def _filter_atomic_property(self, ion_cx_data, selected_atoms):
-        return filtered_ion_cx_data
-
-    def _set_index(self, ion_cx_data, atomic_data):
-        return levels.set_index(['atomic_number', 'ion_number',
-                                 'level_number'])
+        return ion_cx_data[ion_cx_data.atomic_number.isin([selected_atoms]
+                                                          if np.isscalar(
+            selected_atoms)
+                                                          else selected_atoms)]
 
+    def _set_index(self, ion_cx_data):
+        return ion_cx_data.set_index(['atomic_number', 'ion_number',
+                                      'level_number'])
 
 class AtomicMass(ProcessingPlasmaProperty):
+    """
+    Outputs:
+    atomic_mass : Pandas Series
+        Atomic masses of the elements used, indexed by atomic number
+    """
     outputs = ('atomic_mass',)
 
     def calculate(self, atomic_data, selected_atoms):
         if getattr(self, self.outputs[0]) is not None:
-            return (getattr(self, self.outputs[0]),)
+            return getattr(self, self.outputs[0]),
         else:
             return atomic_data.atom_data.ix[selected_atoms].mass
 
 class IonizationData(BaseAtomicDataProperty):
+    """
+    Outputs:
+    ionization_data : Pandas DataFrame
+        Ionization energies of the elements used
+    """
     outputs = ('ionization_data',)
 
     def _filter_atomic_property(self, ionization_data, selected_atoms):
-        ionization_data['atomic_number'] = ionization_data.index.labels[0]+1
-        ionization_data['ion_number'] = ionization_data.index.labels[1]+1
+        ionization_data['atomic_number'] = ionization_data.index.labels[0] + 1
+        ionization_data['ion_number'] = ionization_data.index.labels[1] + 1
         ionization_data = ionization_data[ionization_data.atomic_number.isin(
             selected_atoms)]
         ion_data_check = counter(ionization_data.atomic_number.values)
         keys = np.array(ion_data_check.keys())
         values = np.array(ion_data_check.values())
-        if np.alltrue(keys==values):
+        if np.alltrue(keys == values):
             return ionization_data
         else:
             raise IncompleteAtomicData('ionization data for the ion (' +
-                            str(keys[keys!=values]) +
-                            str(values[keys!=values]) + ')')
+                                       str(keys[keys != values]) +
+                                       str(values[keys != values]) + ')')
 
-    def _set_index(self, ionization_data, atomic_data):
+    def _set_index(self, ionization_data):
         return ionization_data.set_index(['atomic_number', 'ion_number'])
 
 class ZetaData(BaseAtomicDataProperty):
+    """
+    Outputs:
+    zeta_data : Pandas DataFrame
+        Zeta data for the elements used
+        Required for the nebular ionization scheme.
+        The zeta value represents the fraction of recombination events
+        from the ionized state that go directly to the ground state.
+    """
     outputs = ('zeta_data',)
 
     def _filter_atomic_property(self, zeta_data, selected_atoms):
-        zeta_data['atomic_number'] = zeta_data.index.labels[0]+1
-        zeta_data['ion_number'] = zeta_data.index.labels[1]+1
-        zeta_data =  zeta_data[zeta_data.atomic_number.isin(selected_atoms)]
+        zeta_data['atomic_number'] = zeta_data.index.labels[0] + 1
+        zeta_data['ion_number'] = zeta_data.index.labels[1] + 1
+        zeta_data = zeta_data[zeta_data.atomic_number.isin(selected_atoms)]
         zeta_data_check = counter(zeta_data.atomic_number.values)
         keys = np.array(zeta_data_check.keys())
         values = np.array(zeta_data_check.values())
-        if np.alltrue(keys+1==values):
+        if np.alltrue(keys + 1 == values):
             return zeta_data
         else:
 #            raise IncompleteAtomicData('zeta data')
+# This currently replaces missing zeta data with 1, which is necessary with
+# the present atomic data. Will replace with the error above when I have
+# complete atomic data.
             logger.warn('Zeta_data missing - replaced with 1s')
             updated_index = []
             for atom in selected_atoms:
-                for ion in range(1, atom+2):
-                    updated_index.append([atom,ion])
+                for ion in range(1, atom + 2):
+                    updated_index.append([atom, ion])
             updated_index = np.array(updated_index)
             updated_dataframe = pd.DataFrame(index=pd.MultiIndex.from_arrays(
                 updated_index.transpose().astype(int)),
-                columns = zeta_data.columns)
+                columns=zeta_data.columns)
             for value in range(len(zeta_data)):
                 updated_dataframe.ix[zeta_data.atomic_number.values[value]].ix[
                     zeta_data.ion_number.values[value]] = \
                     zeta_data.ix[zeta_data.atomic_number.values[value]].ix[
-                    zeta_data.ion_number.values[value]]
+                        zeta_data.ion_number.values[value]]
             updated_dataframe = updated_dataframe.astype(float)
             updated_index = pd.DataFrame(updated_index)
             updated_dataframe['atomic_number'] = np.array(updated_index[0])
             updated_dataframe['ion_number'] = np.array(updated_index[1])
             updated_dataframe.fillna(1.0, inplace=True)
             return updated_dataframe
 
-    def _set_index(self, zeta_data, atomic_data):
+    def _set_index(self, zeta_data):
         return zeta_data.set_index(['atomic_number', 'ion_number'])
 
 class NLTEData(ProcessingPlasmaProperty):
@@ -181,4 +232,3 @@ def calculate(self, atomic_data):
             return (getattr(self, self.outputs[0]),)
         else:
             return atomic_data.nlte_data
-

tardis/plasma/properties/base.py

@@ -38,9 +38,9 @@ def get_latex_label(self):
         outputs = self.outputs.replace('_', r'\_')
         latex_name = getattr(self, 'latex_name', '')
         if latex_name != '':
-            complete_name = '{0} [{1}]'.format(name, self.latex_name)
+            complete_name = '{0} [{1}]'.format(latex_name, self.latex_name)
         else:
-            complete_name = name
+            complete_name = latex_name
 
 
         latex_label = latex_template.format(name=complete_name,