Merge pull request #382 from aoifeboyle/plasma/graph

Improvements to plasma graph

tardis/data/tardis_config_definition.yml

@@ -132,7 +132,6 @@ plasma:
             help: sets all beta_sobolevs to 1
 
 
-
 model:
     structure:
         property_type : container-property

tardis/io/config_reader.py

@@ -976,6 +976,7 @@ def from_config_dict(cls, config_dict, atom_data=None, test_parser=False,
 
 
 
+
         ##### NLTE subsection of Plasma start
         nlte_validated_config_dict = {}
         nlte_species = []

tardis/plasma/base.py

@@ -2,6 +2,7 @@
 
 import networkx as nx
 from tardis.plasma.exceptions import PlasmaMissingModule, NotInitializedModule
+from tardis.plasma.properties.base import HiddenPlasmaProperty
 
 import tempfile
 import fileinput
@@ -18,7 +19,7 @@ def __init__(self, plasma_properties, **kwargs):
         self.plasma_properties = self._init_properties(plasma_properties,
                                                        **kwargs)
         self._build_graph()
-#        self.write_to_tex('Plasma_Graph', 'Plasma_Formulae')
+        self.write_to_tex('Plasma_Graph', 'Plasma_Formulae')
         self.update(**kwargs)
 
     def __getattr__(self, item):
@@ -186,12 +187,12 @@ def write_to_dot(self, fname, latex_label=True):
 
         try:
             import pygraphviz
-        except ImportError:
-            raise ImportError('pygraphviz is needed for method '
-                              '\'write_to_dot\'')
+        except:
+            logger.warn('pygraphviz missing. Plasma graph will not be '
+                        'generated.')
+            return
         print_graph = self.graph.copy()
-        print_graph.remove_node('LinesUpperLevelIndex')
-        print_graph.remove_node('LinesLowerLevelIndex')
+        print_graph = self.remove_hidden_properties(print_graph)
         for node in print_graph:
             print_graph.node[str(node)]['label'] = node
             if hasattr(self.plasma_properties_dict[node],
@@ -209,8 +210,10 @@ def write_to_dot(self, fname, latex_label=True):
     def write_to_tex(self, fname_graph, fname_formulae):
         try:
             import dot2tex
-        except ImportError:
-            raise ImportError('dot2tex is needed for method\'write_to_tex\'')
+        except:
+            logger.warn('dot2tex missing. Plasma graph will not be '
+                        'generated.')
+            return
 
         temp_fname = tempfile.NamedTemporaryFile().name
 
@@ -228,6 +231,31 @@ def write_to_tex(self, fname_graph, fname_formulae):
         for line in fileinput.input(fname_graph, inplace = 1):
             print line.replace('\enlargethispage{100cm}', ''),
 
+    def remove_hidden_properties(self, print_graph):
+        for item in self.plasma_properties_dict.values():
+            module = self.plasma_properties_dict[item.name].__class__
+            if (issubclass(module, HiddenPlasmaProperty)):
+                output = module.outputs[0]
+                inputs = self.plasma_properties_dict[item.name].inputs
+                for value in self.plasma_properties_dict.keys():
+                    if output in getattr(self.plasma_properties_dict[value],
+                        'inputs', []):
+                        for input in self.plasma_properties_dict[
+                            item.name].inputs:
+                            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,
+                                value, label = label)
+                print_graph.remove_node(str(item.name))
+        return print_graph
+
 class StandardPlasma(BasePlasma):
 
     def __init__(self, number_densities, atom_data, time_explosion,

tardis/plasma/properties/atomic.py

@@ -4,7 +4,8 @@
 from collections import Counter as counter
 import logging
 
-from tardis.plasma.properties.base import ProcessingPlasmaProperty
+from tardis.plasma.properties.base import (ProcessingPlasmaProperty,
+    HiddenPlasmaProperty)
 from tardis.plasma.exceptions import IncompleteAtomicData
 
 logger = logging.getLogger(__name__)
@@ -88,15 +89,14 @@ def _set_index(self, lines):
         lines = reindexed.dropna(subset=['atomic_number'])
         return lines, lines['nu'], lines['f_lu'], lines['wavelength_cm']
 
-class LinesLowerLevelIndex(ProcessingPlasmaProperty):
+class LinesLowerLevelIndex(HiddenPlasmaProperty):
     """
     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)
@@ -105,7 +105,7 @@ def calculate(self, levels, lines):
              'level_number_lower']).index
         return np.array(levels_index.ix[lines_index])
 
-class LinesUpperLevelIndex(ProcessingPlasmaProperty):
+class LinesUpperLevelIndex(HiddenPlasmaProperty):
     """
     Outputs:
     lines_upper_level_index : One-dimensional Numpy Array

tardis/plasma/properties/base.py

@@ -94,3 +94,9 @@ def update(self):
     def calculate(self, *args, **kwargs):
         raise NotImplementedError('This method needs to be implemented by '
                                   'processing plasma modules')
+
+class HiddenPlasmaProperty(ProcessingPlasmaProperty):
+    __metaclass__ = ABCMeta
+
+    def __init__(self, plasma_parent):
+        super(HiddenPlasmaProperty, self).__init__(plasma_parent)

tardis/plasma/properties/nlte.py

@@ -14,4 +14,5 @@ class PreviousElectronDensities(PreviousIterationProperty):
     outputs = ('previous_electron_densities',)
 
 class PreviousBetaSobolevs(PreviousIterationProperty):
-    outputs = ('previous_beta_sobolevs',)
+    outputs = ('previous_beta_sobolevs',)
+

tardis/plasma/properties/property_collections.py

@@ -34,4 +34,4 @@ class PlasmaPropertyCollection(list):
     LevelBoltzmannFactorDiluteLTE])
 non_nlte_properties = PlasmaPropertyCollection([LevelBoltzmannFactorNoNLTE])
 nlte_properties = PlasmaPropertyCollection([
-    LevelBoltzmannFactorNLTE, NLTEData, NLTESpecies, LTEJBlues])
+    LevelBoltzmannFactorNLTE, NLTEData, NLTESpecies, LTEJBlues])