NLTE

tardis/model.py

@@ -235,9 +235,8 @@ def calculate_updated_radiationfield(self, nubar_estimator, j_estimator):
 
     def update_plasmas(self, initialize_nlte=False):
 
-        self.plasma_array.update_radiationfield(self.t_rads.value, self.ws, j_blues=self.j_blues,
-                                        initialize_nlte=initialize_nlte, n_e_convergence_threshold=0.05)
-
+        self.plasma_array.update_radiationfield(self.t_rads.value, self.ws, self.j_blues,
+            self.tardis_config.plasma.nlte, initialize_nlte=initialize_nlte, n_e_convergence_threshold=0.05)
 
         if self.tardis_config.plasma.line_interaction_type in ('downbranch', 'macroatom'):
             self.transition_probabilities = self.plasma_array.transition_probabilities

tardis/plasma/base.py

@@ -119,6 +119,12 @@ def _init_properties(self, plasma_properties, **kwargs):
                 plasma_property_objects.append(current_property_object)
         return plasma_property_objects
 
+    def store_previous_properties(self):
+        self.outputs_dict['previous_electron_densities'].set_value(
+            self.get_value('electron_densities'))
+        self.outputs_dict['previous_beta_sobolevs'].set_value(
+            self.get_value('beta_sobolev'))
+
     def update(self, **kwargs):
         for key in kwargs:
             if key not in self.outputs_dict:
@@ -205,6 +211,8 @@ class StandardPlasma(BasePlasma):
     def __init__(self, number_densities, atom_data, time_explosion,
                  delta_treatment=None, nlte_config=None, ionization_mode='lte',
                  excitation_mode='lte', w=None,
-                 link_t_rad_t_electron=0.9):
+                 link_t_rad_t_electron=0.9, nlte_species=None,
+                 previous_beta_sobolevs=None,
+                 previous_electron_densities=None):
 
         pass

tardis/plasma/properties/__init__.py

@@ -5,4 +5,5 @@
 from tardis.plasma.properties.partition_function import *
 from tardis.plasma.properties.plasma_input import *
 from tardis.plasma.properties.radiative_properties import *
+from tardis.plasma.properties.nlte import *
 

tardis/plasma/properties/atomic.py

@@ -10,7 +10,7 @@
 logger = logging.getLogger(__name__)
 
 __all__ = ['Levels', 'Lines', 'LinesLowerLevelIndex', 'LinesUpperLevelIndex',
-           'AtomicMass', 'IonizationData', 'ZetaData']
+           'AtomicMass', 'IonizationData', 'ZetaData', 'NLTEData']
 
 class BaseAtomicDataProperty(ProcessingPlasmaProperty):
     __metaclass__ = ABCMeta
@@ -173,3 +173,12 @@ def _filter_atomic_property(self, zeta_data, selected_atoms):
     def _set_index(self, zeta_data, atomic_data):
         return zeta_data.set_index(['atomic_number', 'ion_number'])
 
+class NLTEData(ProcessingPlasmaProperty):
+    outputs = ('nlte_data',)
+
+    def calculate(self, atomic_data):
+        if getattr(self, self.outputs[0]) is not None:
+            return (getattr(self, self.outputs[0]),)
+        else:
+            return atomic_data.nlte_data
+

tardis/plasma/properties/level_population.py

@@ -30,73 +30,4 @@ class LevelNumberDensity(ProcessingPlasmaProperty):
     def calculate(level_population_fraction, ion_number_density):
         ion_number_density_broadcast = ion_number_density.ix[
             level_population_fraction.index.droplevel(2)].values
-        return level_population_fraction * ion_number_density_broadcast
-
-class LevelPopulationNLTE(ProcessingPlasmaProperty):
-    @staticmethod
-    def calculate(self):
-        """
-        Calculating the NLTE level populations for specific ions
-
-        """
-
-        if not hasattr(self, 'beta_sobolevs'):
-            self.beta_sobolevs = np.zeros_like(self.tau_sobolevs.values)
-
-        macro_atom.calculate_beta_sobolev(self.tau_sobolevs.values.ravel(order='F'),
-                                          self.beta_sobolevs.ravel(order='F'))
-        self.beta_sobolevs_precalculated = True
-
-        if self.nlte_config.get('coronal_approximation', False):
-            beta_sobolevs = np.ones_like(self.beta_sobolevs)
-            j_blues = np.zeros_like(self.j_blues)
-            logger.info('using coronal approximation = setting beta_sobolevs to 1 AND j_blues to 0')
-        else:
-            beta_sobolevs = self.beta_sobolevs
-            j_blues = self.j_blues.values
-
-        if self.nlte_config.get('classical_nebular', False):
-            logger.info('using Classical Nebular = setting beta_sobolevs to 1')
-            beta_sobolevs = np.ones_like(self.beta_sobolevs)
-
-        for species in self.nlte_config.species:
-            logger.info('Calculating rates for species %s', species)
-            number_of_levels = self.atom_data.levels.energy.ix[species].count()
-
-            level_population_proportionalitys = self.level_population_proportionalitys.ix[species].values
-            lnl = self.atom_data.nlte_data.lines_level_number_lower[species]
-            lnu = self.atom_data.nlte_data.lines_level_number_upper[species]
-
-            lines_index = self.atom_data.nlte_data.lines_idx[species]
-            A_uls = self.atom_data.nlte_data.A_uls[species]
-            B_uls = self.atom_data.nlte_data.B_uls[species]
-            B_lus = self.atom_data.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(self.t_rads)), dtype=np.float64)
-            r_ul_matrix_reshaped = r_ul_matrix.reshape((number_of_levels**2, len(self.t_rads)))
-            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]
-
-            r_lu_matrix = np.zeros_like(r_ul_matrix)
-            r_lu_matrix_reshaped = r_lu_matrix.reshape((number_of_levels**2, len(self.t_rads)))
-            r_lu_matrix_reshaped[r_lu_index] = B_lus[np.newaxis].T * j_blues[lines_index] * beta_sobolevs[lines_index]
-
-            collision_matrix = self.atom_data.nlte_data.get_collision_matrix(species, self.t_electrons) * \
-                               self.electron_densities.values
-
-
-            rates_matrix = r_lu_matrix + r_ul_matrix + collision_matrix
-
-            for i in xrange(number_of_levels):
-                rates_matrix[i, i] = -rates_matrix[:, i].sum(axis=0)
-
-            rates_matrix[0, :, :] = 1.0
-
-            x = np.zeros(rates_matrix.shape[0])
-            x[0] = 1.0
-            for i in xrange(len(self.t_rads)):
-                relative_level_population_proportionalitys = np.linalg.solve(rates_matrix[:, :, i], x)
-                self.level_population_proportionalitys[i].ix[species] = relative_level_population_proportionalitys * self.ion_populations[i].ix[species]
+        return level_population_fraction * ion_number_density_broadcast

tardis/plasma/properties/nlte.py

@@ -0,0 +1,17 @@
+from tardis.plasma.properties.base import BasePlasmaProperty
+
+__all__ = ['PreviousElectronDensities', 'PreviousBetaSobolevs']
+
+class PreviousIterationProperty(BasePlasmaProperty):
+    def _set_output_value(self, output, value):
+        setattr(self, output, value)
+
+    def set_value(self, value):
+        assert len(self.outputs) == 1
+        self._set_output_value(self.outputs[0], value)
+
+class PreviousElectronDensities(PreviousIterationProperty):
+    outputs = ('previous_electron_densities',)
+
+class PreviousBetaSobolevs(PreviousIterationProperty):
+    outputs = ('previous_beta_sobolevs',)