New Plasma restructure using networkx

tardis/atomic.py

@@ -358,12 +358,20 @@ 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):
 
+        if levels_data is not None:
+            self.has_levels = True
+        else:
+            self.has_levels = False
+
+        if lines_data is not None:
+            self.has_lines = True
+        else:
+            self.has_lines = False
 
         if macro_atom_data is not None:
             self.has_macro_atom = True
             self.macro_atom_data_all = DataFrame(macro_atom_data[0].__array__())
             self.macro_atom_references_all = DataFrame(macro_atom_data[1].__array__())
-
         else:
             self.has_macro_atom = False
 
@@ -431,7 +439,7 @@ def __init__(self, atom_data, ionization_data, levels_data, lines_data, macro_at
 
 
 
-
+        self.ion_cx_data = ion_cx_data
 
 
 

tardis/conftest.py

@@ -139,4 +139,35 @@ def pytest_report_header(config):
     if not six.PY3:
         s = s.encode(stdoutencoding, 'replace')
 
-    return s
+    return s
+
+
+import os
+
+@pytest.fixture
+def atomic_data_fname():
+    atomic_data_fname = pytest.config.getvalue("atomic-dataset")
+    if atomic_data_fname is None:
+        pytest.skip('--atomic_database was not specified')
+    else:
+        return os.path.expandvars(os.path.expanduser(atomic_data_fname))
+
+from tardis.atomic import AtomData
+
+@pytest.fixture
+def kurucz_atomic_data(atomic_data_fname):
+    atomic_data = AtomData.from_hdf5(atomic_data_fname)
+
+    if atomic_data.md5 != '21095dd25faa1683f4c90c911a00c3f8':
+        pytest.skip('Need default Kurucz atomic dataset '
+                    '(md5="21095dd25faa1683f4c90c911a00c3f8"')
+    else:
+        return atomic_data
+
+@pytest.fixture
+def included_he_atomic_data():
+    import os, tardis
+    atomic_db_fname = os.path.join(tardis.__path__[0], 'tests', 'data',
+                                   'chianti_he_db.h5')
+    return AtomData.from_hdf5(atomic_db_fname)
+

tardis/plasma/__init__.py

@@ -0,0 +1,2 @@
+from tardis.plasma.base import BasePlasma
+from tardis.plasma.standard_plasmas import LTEPlasma

tardis/plasma/base.py

@@ -0,0 +1,189 @@
+import logging
+
+import networkx as nx
+from tardis.plasma.exceptions import PlasmaMissingModule, NotInitializedModule
+
+import tempfile
+
+
+logger = logging.getLogger(__name__)
+
+
+class BasePlasma(object):
+
+    def __init__(self, plasma_modules, **kwargs):
+        self.module_dict = {}
+        self.input_modules = []
+        self._init_modules(plasma_modules, **kwargs)
+        self._build_graph()
+        self.update(**kwargs)
+
+
+    def __getattr__(self, item):
+        if item in self.module_dict:
+            return self.get_value(item)
+        else:
+            super(BasePlasma, self).__getattribute__(item)
+
+
+    def __setattr__(self, key, value):
+        if key != 'module_dict' and key in self.module_dict:
+            raise AttributeError('Plasma inputs can only be updated using '
+                                 'the \'update\' method')
+        else:
+            super(BasePlasma, self).__setattr__(key, value)
+
+    def __dir__(self):
+        attrs = [item for item in self.__dict__
+                 if not item.startswith('_')]
+        attrs += [item for item in self.__class__.__dict__
+                 if not item.startswith('_')]
+        attrs += self.module_dict.keys()
+
+        return attrs
+
+    def get_value(self, item):
+        return self.module_dict[item].value
+
+    def _build_graph(self):
+        """
+        Builds the directed Graph using network X
+
+        :param plasma_modules:
+        :return:
+        """
+
+        self.graph = nx.DiGraph()
+
+        ## Adding all nodes
+        self.graph.add_nodes_from([(key, {})
+                                   for key, value in self.module_dict.items()])
+
+        #Flagging all input modules
+        self.input_modules = [key for key, item in self.module_dict.items()
+                              if not hasattr(item, 'inputs')]
+
+        for plasma_module in self.module_dict.values():
+            #Skipping any module that is an input module
+            if plasma_module.name in self.input_modules:
+                continue
+
+            for input in plasma_module.inputs:
+                if input not in self.graph:
+                    raise PlasmaMissingModule('Module {0} requires input '
+                                              '{1} which has not been added'
+                                              ' to this plasma'.format(
+                        plasma_module.name, input))
+                self.graph.add_edge(input, plasma_module.name)
+
+    def _init_modules(self, plasma_modules, **kwargs):
+        """
+        Builds a dictionary with the plasma module names as keys
+        :param plasma_modules:
+        :return:
+        """
+
+        self.module_dict = {}
+        for module in plasma_modules:
+            if hasattr(module, 'set_value'):
+                if module.name not in kwargs:
+                    raise NotInitializedModule('Input {0} required for '
+                                               'plasma but not given when '
+                                               'instantiating the '
+                                               'plasma'.format(module.name))
+                current_module_object = module()
+            else:
+                current_module_object = module(self)
+
+            self.module_dict[module.name] = current_module_object
+
+    def update(self, **kwargs):
+        for key in kwargs:
+            if key not in self.module_dict:
+                raise PlasmaMissingModule('Trying to update property {0}'
+                                          ' that is unavailable'.format(key))
+            self.module_dict[key].set_value(kwargs[key])
+
+        for module_name in self._resolve_update_list(kwargs.keys()):
+            self.module_dict[module_name].update()
+
+    def _update_module_type_str(self):
+        for node in self.graph:
+            self.module_dict[node]._update_type_str()
+
+    def _resolve_update_list(self, changed_modules):
+        """
+        Returns a list of all plasma models which are affected by the
+        changed_modules due to there dependency in the
+        the plasma_graph.
+
+        Parameters
+        ----------
+
+        graph: ~networkx.Graph
+            the plasma graph as
+        changed_modules: ~list
+            all modules changed in the plasma
+
+        Returns
+        -------
+
+            : ~list
+            all affected modules.
+        """
+
+        descendants_ob = []
+
+        for module in changed_modules:
+            descendants_ob += nx.descendants(self.graph, module)
+
+        descendants_ob = list(set(descendants_ob))
+        sort_order = nx.topological_sort(self.graph)
+
+        descendants_ob.sort(key=lambda val: sort_order.index(val) )
+
+        logger.debug('Updating modules in the following order:'.format(
+            '->'.join(descendants_ob)))
+
+        return descendants_ob
+
+    def write_to_dot(self, fname):
+        self._update_module_type_str()
+
+        try:
+            import pygraphviz
+        except ImportError:
+            raise ImportError('pygraphviz is needed for method '
+                              '\'write_to_dot\'')
+
+        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)
+
+    def write_to_tex(self, fname):
+        try:
+            import dot2tex
+        except ImportError:
+            raise ImportError('dot2tex is needed for method\'write_to_tex\'')
+
+        temp_fname = tempfile.NamedTemporaryFile().name
+
+        self.write_to_dot(temp_fname)
+
+        dot_string = open(temp_fname).read()
+
+        open(fname, 'w').write(dot2tex.dot2tex(dot_string, texmode='raw'))
+
+
+
+class StandardPlasma(BasePlasma):
+
+    def __init__(self, number_densities, atom_data, time_explosion,
+                 delta_treatment=None, nlte_config=None, ionization_mode='lte',
+                 excitation_mode='lte'):
+
+        pass
+

tardis/plasma/exceptions.py

@@ -0,0 +1,22 @@
+class PlasmaException(Exception):
+    pass
+
+class IncompleteAtomicData(PlasmaException):
+    def __init__(self, atomic_data_name):
+        message = ('The current plasma calculation requires {0}, '
+                   'which is not provided by the given atomic data'.format(
+            atomic_data_name))
+        super(PlasmaException, self).__init__(message)
+
+
+class PlasmaMissingModule(PlasmaException):
+    pass
+
+class PlasmaIsolatedModule(PlasmaException):
+    pass
+
+class NotInitializedModule(PlasmaException):
+    pass
+
+class PlasmaIonizationError(PlasmaException):
+    pass

tardis/plasma/properties/__init__.py

@@ -0,0 +1,8 @@
+from tardis.plasma.properties.atomic import *
+from tardis.plasma.properties.general import *
+from tardis.plasma.properties.ion_population import *
+from tardis.plasma.properties.level_population import *
+from tardis.plasma.properties.partition_function import *
+from tardis.plasma.properties.plasma_input import *
+from tardis.plasma.properties.radiative_properties import *
+