Merge pull request #778 from yeganer/j_blues_dataframe_master

Fix NLTE

tardis/plasma/base.py

@@ -1,4 +1,5 @@
 import os
+import re
 import logging
 import tempfile
 import fileinput
@@ -145,9 +146,9 @@ def _init_properties(self, plasma_properties, property_kwargs=None, **kwargs):
 
     def store_previous_properties(self):
         for property in self.previous_iteration_properties:
-            base_property = property.outputs[0][9:]
-            self.outputs_dict[property.outputs[0]].set_value(
-                self.get_value(base_property))
+            p = property.outputs[0]
+            self.outputs_dict[p].set_value(
+                self.get_value(re.sub(r'^previous_', '', p)))
 
     def update(self, **kwargs):
         for key in kwargs:
@@ -193,7 +194,7 @@ def _resolve_update_list(self, changed_properties):
 
         descendants_ob.sort(key=lambda val: sort_order.index(val) )
 
-        logger.debug('Updating modules in the following order:'.format(
+        logger.debug('Updating modules in the following order: {}'.format(
             '->'.join(descendants_ob)))
 
         return descendants_ob

tardis/plasma/properties/j_blues.py

@@ -22,7 +22,7 @@ def calculate(lines, nu, t_rad):
         j_blues = intensity_black_body(nu.values[np.newaxis].T, t_rad)
         j_blues = pd.DataFrame(j_blues, index=lines.index,
                                columns=np.arange(len(t_rad)))
-        return np.array(j_blues, copy=False)
+        return j_blues
 
 
 class JBluesDiluteBlackBody(ProcessingPlasmaProperty):
@@ -34,7 +34,7 @@ def calculate(lines, nu, t_rad, w):
         j_blues = w * intensity_black_body(nu.values[np.newaxis].T, t_rad)
         j_blues = pd.DataFrame(j_blues, index=lines.index,
                                columns=np.arange(len(t_rad)))
-        return np.array(j_blues, copy=False)
+        return j_blues
 
 
 class JBluesDetailed(ProcessingPlasmaProperty):
@@ -63,7 +63,7 @@ def calculate(self, lines, nu, t_rad, w, j_blues_norm_factor,
                     self.w_epsilon *
                     intensity_black_body(nu[zero_j_blues].values,
                                          t_rad[i]))
-            return np.array(j_blues, copy=False)
+            return j_blues
 
 
 class JBluesNormFactor(ProcessingPlasmaProperty):

tardis/plasma/properties/nlte.py

@@ -34,12 +34,11 @@ class PreviousBetaSobolev(PreviousIterationProperty):
     outputs = ('previous_beta_sobolev',)
 
     def set_initial_value(self, kwargs):
-        try:
-            lines = len(kwargs['atomic_data'].lines)
-        except:
-            lines = len(kwargs['atomic_data']._lines)
-        initial_value = np.ones((lines,
-            len(kwargs['abundance'].columns)))
+        initial_value = pd.DataFrame(
+                1.,
+                index=kwargs['atomic_data'].lines.index,
+                columns=kwargs['abundance'].columns,
+                )
         self._set_initial_value(initial_value)
 
 class HeliumNLTE(ProcessingPlasmaProperty):

tardis/plasma/properties/partition_function.py

@@ -1,6 +1,7 @@
 import logging
 
 import numpy as np
+from numpy.linalg.linalg import LinAlgError
 import pandas as pd
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
@@ -12,6 +13,7 @@
            'LevelBoltzmannFactorNoNLTE', 'LevelBoltzmannFactorNLTE',
            'PartitionFunction']
 
+
 class LevelBoltzmannFactorLTE(ProcessingPlasmaProperty):
     """
     Attributes
@@ -38,6 +40,7 @@ def calculate(excitation_energy, g, beta_rad, levels):
                                               dtype=np.float64)
         return level_boltzmann_factor
 
+
 class LevelBoltzmannFactorDiluteLTE(ProcessingPlasmaProperty):
     """
     Attributes
@@ -55,13 +58,15 @@ class LevelBoltzmannFactorDiluteLTE(ProcessingPlasmaProperty):
     latex_formula = ('Wg_{i,j,k}e^{\\dfrac{-\\epsilon_{i,j,k}}{k_{\
         \\textrm{B}}T_{\\textrm{rad}}}}',)
 
-    def calculate(self, levels, g, excitation_energy, beta_rad, w,
-        metastability):
+    def calculate(
+            self, levels, g, excitation_energy, beta_rad, w,
+            metastability):
         level_boltzmann_factor = LevelBoltzmannFactorLTE.calculate(
             excitation_energy, g, beta_rad, levels)
         level_boltzmann_factor[~metastability] *= w
         return level_boltzmann_factor
 
+
 class LevelBoltzmannFactorNoNLTE(ProcessingPlasmaProperty):
     """
     Attributes
@@ -76,6 +81,7 @@ class LevelBoltzmannFactorNoNLTE(ProcessingPlasmaProperty):
     def calculate(general_level_boltzmann_factor):
         return general_level_boltzmann_factor
 
+
 class LevelBoltzmannFactorNLTE(ProcessingPlasmaProperty):
     """
     Attributes
@@ -87,64 +93,83 @@ class LevelBoltzmannFactorNLTE(ProcessingPlasmaProperty):
     outputs = ('level_boltzmann_factor',)
 
     def calculate(self):
-        pass
+        raise AttributeError(
+                'This attribute is not defined on the parent class.'
+                'Please use one of the subclasses.')
+
+    @classmethod
+    def from_config(cls, nlte_conf):
+        if nlte_conf.classical_nebular and not nlte_conf.coronal_approximation:
+            return LevelBoltzmannFactorNLTEClassic
+        elif (
+                nlte_conf.coronal_approximation and
+                not nlte_conf.classical_nebular):
+            return LevelBoltzmannFactorNLTECoronal
+        elif nlte_conf.coronal_approximation and nlte_conf.classical_nebular:
+            raise PlasmaConfigError('Both coronal approximation and '
+                                    'classical nebular specified in the '
+                                    'config.')
+        else:
+            return LevelBoltzmannFactorNLTEGeneral
 
-    def __init__(self, plasma_parent, classical_nebular=False,
-        coronal_approximation=False):
+    def __init__(self, plasma_parent):
         """
         Selects appropriate 'calculate' function based on NLTE config
         options selected.
         """
         super(LevelBoltzmannFactorNLTE, self).__init__(plasma_parent)
-        if classical_nebular == True and coronal_approximation == False:
-            self.calculate = self._calculate_classical_nebular
-        elif coronal_approximation == True and classical_nebular == False:
-            self.calculate = self._calculate_coronal_approximation
-        elif coronal_approximation == True and classical_nebular == True:
-            raise PlasmaConfigError('Both coronal approximation and classical'
-                                    'nebular specified in the config.')
-        else:
-            self.calculate = self._calculate_general
+
         self._update_inputs()
 
-    def _main_nlte_calculation(self, atomic_data, nlte_data,
-        t_electrons, j_blues, beta_sobolevs, general_level_boltzmann_factor,
-        previous_electron_densities):
+    def _main_nlte_calculation(
+            self, atomic_data, nlte_data, t_electrons,
+            j_blues, beta_sobolevs, general_level_boltzmann_factor,
+            previous_electron_densities):
         """
         The core of the NLTE calculation, used with all possible config.
         options.
         """
         for species in nlte_data.nlte_species:
-            j_blues = j_blues.values
             logger.info('Calculating rates for species %s', species)
             number_of_levels = atomic_data.levels.energy.ix[species].count()
             lnl = nlte_data.lines_level_number_lower[species]
             lnu = nlte_data.lines_level_number_upper[species]
-            lines_index = nlte_data.lines_idx[species]
+            lines_index, = nlte_data.lines_idx[species]
+
+            try:
+                j_blues_filtered = j_blues.iloc[lines_index]
+            except AttributeError:
+                j_blues_filtered = j_blues
+            try:
+                beta_sobolevs_filtered = beta_sobolevs.iloc[lines_index]
+            except AttributeError:
+                beta_sobolevs_filtered = beta_sobolevs
             A_uls = nlte_data.A_uls[species]
             B_uls = nlte_data.B_uls[species]
             B_lus = nlte_data.B_lus[species]
             r_lu_index = lnu * number_of_levels + lnl
             r_ul_index = lnl * number_of_levels + lnu
-            r_ul_matrix = np.zeros((number_of_levels, number_of_levels,
-                len(t_electrons)), dtype=np.float64)
-            r_ul_matrix_reshaped = r_ul_matrix.reshape((number_of_levels**2,
-                len(t_electrons)))
+            r_ul_matrix = np.zeros(
+                    (number_of_levels, number_of_levels, len(t_electrons)),
+                    dtype=np.float64)
+            r_ul_matrix_reshaped = r_ul_matrix.reshape(
+                    (number_of_levels**2, len(t_electrons)))
             r_ul_matrix_reshaped[r_ul_index] = A_uls[np.newaxis].T + \
-                B_uls[np.newaxis].T * j_blues[lines_index]
-            r_ul_matrix_reshaped[r_ul_index] *= beta_sobolevs[lines_index]
+                B_uls[np.newaxis].T * j_blues_filtered
+            r_ul_matrix_reshaped[r_ul_index] *= beta_sobolevs_filtered
             r_lu_matrix = np.zeros_like(r_ul_matrix)
-            r_lu_matrix_reshaped = r_lu_matrix.reshape((number_of_levels**2,
-                len(t_electrons)))
+            r_lu_matrix_reshaped = r_lu_matrix.reshape(
+                    (number_of_levels**2, len(t_electrons)))
             r_lu_matrix_reshaped[r_lu_index] = B_lus[np.newaxis].T * \
-                j_blues[lines_index] * beta_sobolevs[lines_index]
+                j_blues_filtered * beta_sobolevs_filtered
             if atomic_data.has_collision_data:
                 if previous_electron_densities is None:
                     collision_matrix = r_ul_matrix.copy()
                     collision_matrix.fill(0.0)
                 else:
-                    collision_matrix = nlte_data.get_collision_matrix(species,
-                        t_electrons) * previous_electron_densities.values
+                    collision_matrix = nlte_data.get_collision_matrix(
+                            species, t_electrons
+                            ) * previous_electron_densities.values
             else:
                 collision_matrix = r_ul_matrix.copy()
                 collision_matrix.fill(0.0)
@@ -155,61 +180,88 @@ def _main_nlte_calculation(self, atomic_data, nlte_data,
             x = np.zeros(rates_matrix.shape[0])
             x[0] = 1.0
             for i in xrange(len(t_electrons)):
-                level_boltzmann_factor = \
-                    np.linalg.solve(rates_matrix[:, :, i], x)
+                try:
+                    level_boltzmann_factor = np.linalg.solve(
+                            rates_matrix[:, :, i], x)
+                except LinAlgError as e:
+                    if e.message == 'Singular matrix':
+                        raise ValueError(
+                                'SingularMatrixError during solving of the '
+                                'rate matrix. Does the atomic data contain '
+                                'collision data?')
+                    else:
+                        raise e
                 general_level_boltzmann_factor[i].ix[species] = \
                     level_boltzmann_factor
         return general_level_boltzmann_factor
 
-    def _calculate_classical_nebular(self, t_electrons, lines, atomic_data,
-        nlte_data, general_level_boltzmann_factor, j_blues,
-        previous_electron_densities):
+    def _calculate_classical_nebular(
+            self, t_electrons, lines, atomic_data,
+            nlte_data, general_level_boltzmann_factor, j_blues,
+            previous_electron_densities):
         """
         Performs NLTE calculations using the classical nebular treatment.
         All beta sobolev values taken as 1.
         """
-        beta_sobolevs = np.ones((len(lines), len(t_electrons)))
-        j_blues = pd.DataFrame(j_blues, index=lines.index, columns=
-                               range(len(t_electrons)))
+        beta_sobolevs = 1.0
+
         general_level_boltzmann_factor = self._main_nlte_calculation(
-            atomic_data, nlte_data, t_electrons, j_blues,
-            beta_sobolevs, general_level_boltzmann_factor,
-            previous_electron_densities)
+                atomic_data,
+                nlte_data,
+                t_electrons,
+                j_blues,
+                beta_sobolevs,
+                general_level_boltzmann_factor,
+                previous_electron_densities)
         return general_level_boltzmann_factor
 
-    def _calculate_coronal_approximation(self, t_electrons, lines, atomic_data,
-        nlte_data, general_level_boltzmann_factor,
-        previous_electron_densities):
+    def _calculate_coronal_approximation(
+            self, t_electrons, lines, atomic_data,
+            nlte_data, general_level_boltzmann_factor,
+            previous_electron_densities):
         """
         Performs NLTE calculations using the coronal approximation.
         All beta sobolev values taken as 1 and j_blues taken as 0.
         """
-        beta_sobolevs = np.ones((len(lines), len(t_electrons)))
-        j_blues = np.zeros((len(lines), len(t_electrons)))
+        beta_sobolevs = 1.0
+        j_blues = 0.0
         general_level_boltzmann_factor = self._main_nlte_calculation(
             atomic_data, nlte_data, t_electrons, j_blues,
             beta_sobolevs, general_level_boltzmann_factor,
             previous_electron_densities)
         return general_level_boltzmann_factor
 
-    def _calculate_general(self, t_electrons, lines, atomic_data, nlte_data,
-        general_level_boltzmann_factor, j_blues,
-        previous_beta_sobolev, previous_electron_densities):
+    def _calculate_general(
+            self, t_electrons, lines, atomic_data, nlte_data,
+            general_level_boltzmann_factor, j_blues,
+            previous_beta_sobolev, previous_electron_densities):
         """
         Full NLTE calculation without approximations.
         """
         if previous_beta_sobolev is None:
-            beta_sobolevs = np.ones((len(lines), len(t_electrons)))
+            beta_sobolevs = 1.0
         else:
             beta_sobolevs = previous_beta_sobolev
-        j_blues = pd.DataFrame(j_blues, index=lines.index, columns=
-                               range(len(t_electrons)))
+
         general_level_boltzmann_factor = self._main_nlte_calculation(
             atomic_data, nlte_data, t_electrons, j_blues,
             beta_sobolevs, general_level_boltzmann_factor,
             previous_electron_densities)
         return general_level_boltzmann_factor
 
+
+class LevelBoltzmannFactorNLTECoronal(LevelBoltzmannFactorNLTE):
+    calculate = LevelBoltzmannFactorNLTE._calculate_coronal_approximation
+
+
+class LevelBoltzmannFactorNLTEClassic(LevelBoltzmannFactorNLTE):
+    calculate = LevelBoltzmannFactorNLTE._calculate_classical_nebular
+
+
+class LevelBoltzmannFactorNLTEGeneral(LevelBoltzmannFactorNLTE):
+    calculate = LevelBoltzmannFactorNLTE._calculate_general
+
+
 class PartitionFunction(ProcessingPlasmaProperty):
     """
     Attributes