Initial steps.

tardis/plasma/base.py

@@ -12,22 +12,22 @@
 class BasePlasma(object):
 
     def __init__(self, plasma_modules, **kwargs):
-        self.module_dict = {}
+        self.outputs_dict = {}
+        self.plasma_modules = plasma_modules
+
         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:
+        if item in self.outputs_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:
+        if key != 'outputs_dict' and key in self.outputs_dict:
             raise AttributeError('Plasma inputs can only be updated using '
                                  'the \'update\' method')
         else:
@@ -56,56 +56,74 @@ def _build_graph(self):
         self.graph = nx.DiGraph()
 
         ## Adding all nodes
-        self.graph.add_nodes_from([(key, {})
-                                   for key, value in self.module_dict.items()])
+
+        for item in self.outputs_dict: self.graph.add_node(item, {})
 
         #Flagging all input modules
-        self.input_modules = [key for key, item in self.module_dict.items()
-                              if not hasattr(item, 'inputs')]
+        self.input_modules = [item for item in self.outputs_dict.keys()
+                                if hasattr(self.outputs_dict[item],
+                                'set_value')]
 
-        for plasma_module in self.module_dict.values():
+        edge_labels = {}
+        for plasma_module in self.outputs_dict.keys():
             #Skipping any module that is an input module
-            if plasma_module.name in self.input_modules:
+            if plasma_module in self.input_modules:
                 continue
 
-            for input in plasma_module.inputs:
-                if input not in self.graph:
+            for input in self.outputs_dict[plasma_module].inputs:
+                if input not in self.outputs_dict:
                     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)
+                self.graph.add_edge(plasma_module, input,
+                    label=self.outputs_dict[plasma_module].name)
+                edge_labels[(plasma_module, input)] = \
+                    self.outputs_dict[plasma_module].name
 
     def _init_modules(self, plasma_modules, **kwargs):
         """
         Builds a dictionary with the plasma module names as keys
-        :param plasma_modules:
-        :return:
+
+        Parameters
+        ----------
+
+        plasma_modules: ~list
+            list of Plasma properties
+        kwargs: dictionary
+            input values for input properties. For example, t_rad=[5000, 6000,],
+            j_blues=[..]
+
         """
 
-        self.module_dict = {}
+        self.outputs_dict = {}
         for module in plasma_modules:
             if hasattr(module, 'set_value'):
-                if module.name not in kwargs:
+                #duck-typing for PlasmaInputProperty
+                #that means if it is an input property from model
+                if not set(kwargs.keys()).issuperset(module.outputs):
+                    missing_input_values = (set(module.outputs) -
+                    set(kwargs.keys()))
                     raise NotInitializedModule('Input {0} required for '
                                                'plasma but not given when '
                                                'instantiating the '
-                                               'plasma'.format(module.name))
+                                               'plasma'.format(
+                                               missing_input_values))
                 current_module_object = module()
             else:
                 current_module_object = module(self)
-
-            self.module_dict[module.name] = current_module_object
+            for output in module.outputs:
+                self.outputs_dict[output] = current_module_object
 
     def update(self, **kwargs):
         for key in kwargs:
-            if key not in self.module_dict:
+            if key not in self.outputs_dict:
                 raise PlasmaMissingModule('Trying to update property {0}'
                                           ' that is unavailable'.format(key))
-            self.module_dict[key].set_value(kwargs[key])
+            self.outputs_dict[key].set_value(kwargs[key])
 
         for module_name in self._resolve_update_list(kwargs.keys()):
-            self.module_dict[module_name].update()
+            self.outputs_dict[module_name].update()
 
     def _update_module_type_str(self):
         for node in self.graph:
@@ -120,8 +138,6 @@ def _resolve_update_list(self, changed_modules):
         Parameters
         ----------
 
-        graph: ~networkx.Graph
-            the plasma graph as
         changed_modules: ~list
             all modules changed in the plasma
 
@@ -147,7 +163,7 @@ def _resolve_update_list(self, changed_modules):
 
         return descendants_ob
 
-    def write_to_dot(self, fname):
+    def write_to_dot(self, fname, latex_label=True):
         self._update_module_type_str()
 
         try:
@@ -157,7 +173,7 @@ def write_to_dot(self, fname):
                               '\'write_to_dot\'')
 
         for node in self.graph:
-            self.graph.node[node]['label'] = self.module_dict[node].get_latex_label()
+            self.graph.node[node]['label'] = self.outputs_dict[node].name
             self.graph.node[node]['color'] = 'red'
             self.graph.node[node]['shape'] = 'box '
 
@@ -187,4 +203,3 @@ def __init__(self, number_densities, atom_data, time_explosion,
                  link_t_rad_t_electron=0.9):
 
         pass
-

tardis/plasma/properties/atomic.py

@@ -35,17 +35,17 @@ def calculate(self, atomic_data, selected_atoms):
             return self.value
         else:
             try:
-                raw_atomic_property = getattr(atomic_data, '_' + self.name)
+                raw_atomic_property = getattr(atomic_data, '_' + self.outputs)
                 return self._set_index(self._filter_atomic_property(
                     raw_atomic_property, selected_atoms), atomic_data)
             except:
-                raw_atomic_property = getattr(atomic_data, self.name)
+                raw_atomic_property = getattr(atomic_data, self.outputs)
                 return self._set_index(self._filter_atomic_property(
                     raw_atomic_property, selected_atoms), atomic_data)
 
 
 class Levels(BaseAtomicDataProperty):
-    name = 'levels'
+    outputs = ('levels',)
 
     def _filter_atomic_property(self, levels, selected_atoms):
         return levels[levels.atomic_number.isin(selected_atoms)]
@@ -55,7 +55,7 @@ def _set_index(self, levels, atomic_data):
             'level_number'])
 
 class Lines(BaseAtomicDataProperty):
-    name = 'lines'
+    outputs = ('lines',)
 
     def _filter_atomic_property(self, lines, selected_atoms):
         return lines[lines.atomic_number.isin(selected_atoms)]
@@ -68,7 +68,7 @@ def _set_index(self, lines, atomic_data):
         return reindexed
 
 class LinesLowerLevelIndex(ProcessingPlasmaProperty):
-    name = '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),
@@ -79,7 +79,7 @@ def calculate(self, levels, lines):
         return np.array(levels_index.ix[lines_index])
 
 class LinesUpperLevelIndex(ProcessingPlasmaProperty):
-    name = '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),
@@ -91,7 +91,7 @@ def calculate(self, levels, lines):
 
 
 class IonCXData(BaseAtomicDataProperty):
-    name = 'ion_cx_data'
+    outputs = ('ion_cx_data',)
 
     def _filter_atomic_property(self, ion_cx_data, selected_atoms):
         return filtered_ion_cx_data
@@ -102,7 +102,7 @@ def _set_index(self, ion_cx_data, atomic_data):
 
 
 class AtomicMass(ProcessingPlasmaProperty):
-    name = 'atomic_mass'
+    outputs = ('atomic_mass',)
 
     def calculate(self, atomic_data, selected_atoms):
         if self.value is not None:
@@ -111,7 +111,7 @@ def calculate(self, atomic_data, selected_atoms):
             return atomic_data.atom_data.ix[selected_atoms].mass
 
 class IonizationData(BaseAtomicDataProperty):
-    name = 'ionization_data'
+    outputs = ('ionization_data',)
 
     def _filter_atomic_property(self, ionization_data, selected_atoms):
         ionization_data['atomic_number'] = ionization_data.index.labels[0]+1
@@ -132,7 +132,7 @@ def _set_index(self, ionization_data, atomic_data):
         return ionization_data.set_index(['atomic_number', 'ion_number'])
 
 class ZetaData(BaseAtomicDataProperty):
-    name = 'zeta_data'
+    outputs = ('zeta_data',)
 
     def _filter_atomic_property(self, zeta_data, selected_atoms):
         zeta_data['atomic_number'] = zeta_data.index.labels[0]+1
@@ -144,27 +144,27 @@ def _filter_atomic_property(self, zeta_data, selected_atoms):
         if np.alltrue(keys+1==values):
             return zeta_data
         else:
-            raise IncompleteAtomicData('zeta 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])
-#            updated_index = np.array(updated_index)
-#            updated_dataframe = pd.DataFrame(index=pd.MultiIndex.from_arrays(
-#                updated_index.transpose().astype(int)),
-#                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]]
-#            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
+#            raise IncompleteAtomicData('zeta 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])
+            updated_index = np.array(updated_index)
+            updated_dataframe = pd.DataFrame(index=pd.MultiIndex.from_arrays(
+                updated_index.transpose().astype(int)),
+                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]]
+            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):
         return zeta_data.set_index(['atomic_number', 'ion_number'])

tardis/plasma/properties/base.py

@@ -7,9 +7,13 @@ class BasePlasmaProperty(object):
     __metaclass__ = ABCMeta
 
     @abstractproperty
-    def name(self):
+    def outputs(self):
         pass
 
+    @property
+    def name(self):
+        return self.__class__.__name__
+
     def __init__(self):
         self.value = None
 
@@ -30,7 +34,7 @@ def get_latex_label(self):
 \textbf{{Formula}} {formula}
 {description}
 """
-        name = self.name.replace('_', r'\_')
+        outputs = self.outputs.replace('_', r'\_')
         latex_name = getattr(self, 'latex_name', '')
         if latex_name != '':
             complete_name = '{0} [{1}]'.format(name, self.latex_name)
@@ -81,8 +85,3 @@ def update(self):
     def calculate(self, *args, **kwargs):
         raise NotImplementedError('This method needs to be implemented by '
                                   'processing plasma modules')
-
-
-
-
-

tardis/plasma/properties/general.py

@@ -11,7 +11,7 @@
            'ElectronTemperature', 'BetaElectron']
 
 class BetaRadiation(ProcessingPlasmaProperty):
-    name = 'beta_rad'
+    outputs = ('beta_rad',)
     latex_name = r'$\beta_\textrm{rad}$'
     latex_formula = r'$\frac{1}{K_B T_\textrm{rad}}$'
 
@@ -25,7 +25,7 @@ def calculate(self, t_rad):
 
 class GElectron(ProcessingPlasmaProperty):
 
-    name = 'g_electron'
+    outputs = ('g_electron',)
     latex_name = r'$g_\textrm{electron}$'
 
     latex_formula = (r'$\left(\frac{2\pi m_\textrm{e} '
@@ -38,26 +38,26 @@ def calculate(beta_rad):
 
 
 class NumberDensity(ProcessingPlasmaProperty):
-    name = 'number_density'
+    outputs = ('number_density',)
 
     def calculate(self, atomic_mass, abundance, density):
         number_densities = (abundance * density)
         return number_densities.div(atomic_mass.ix[abundance.index], axis=0)
 
 class SelectedAtoms(ProcessingPlasmaProperty):
-    name = 'selected_atoms'
+    outputs = ('selected_atoms',)
 
     def calculate(self, abundance):
         return abundance.index
 
 class ElectronTemperature(ProcessingPlasmaProperty):
-    name = 't_electron'
+    outputs = ('t_electron',)
 
     def calculate(self, t_rad, link_t_rad_t_electron):
         return t_rad * link_t_rad_t_electron
 
 class BetaElectron(ProcessingPlasmaProperty):
-    name = 'beta_electron'
+    outputs = ('beta_electron',)
 
     def __init__(self, plasma_parent):
         super(BetaElectron, self).__init__(plasma_parent)