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',)

tardis/plasma/properties/partition_function.py

@@ -8,18 +8,17 @@
 logger = logging.getLogger(__name__)
 
 __all__ = ['LevelBoltzmannFactorLTE', 'LevelBoltzmannFactorDiluteLTE',
+           'LevelBoltzmannFactorNoNLTE', 'LevelBoltzmannFactorNLTE',
            'PartitionFunction']
 
 class LevelBoltzmannFactorLTE(ProcessingPlasmaProperty):
     """
     Calculate the level population Boltzmann factor
-
     .. math:
         {latex_formula}
-
     """
 
-    outputs = ('level_boltzmann_factor',)
+    outputs = ('general_level_boltzmann_factor',)
     latex_formula = r'$g_{i, j, k} e^{E_{i, j, k} \times \beta_\textrm{rad}}$'
 
     @staticmethod
@@ -36,7 +35,7 @@ def calculate(levels, beta_rad):
 
 class LevelBoltzmannFactorDiluteLTE(ProcessingPlasmaProperty):
 
-    outputs = ('level_boltzmann_factor',)
+    outputs = ('general_level_boltzmann_factor',)
 
     @staticmethod
     def calculate(levels, beta_rad, w):
@@ -51,6 +50,124 @@ def calculate(levels, beta_rad, w):
         level_boltzmann_factor[~levels.metastable] *= w
         return level_boltzmann_factor
 
+class LevelBoltzmannFactorNoNLTE(ProcessingPlasmaProperty):
+
+    outputs = ('level_boltzmann_factor',)
+
+    @staticmethod
+    def calculate(general_level_boltzmann_factor):
+        return general_level_boltzmann_factor
+
+class LevelBoltzmannFactorNLTE(ProcessingPlasmaProperty):
+    outputs = ('level_boltzmann_factor',)
+
+    def calculate(self):
+        pass
+
+    def __init__(self, plasma_parent, classical_nebular=False,
+        coronal_approximation=False):
+        super(LevelBoltzmannFactorNLTE, self).__init__(plasma_parent)
+        if classical_nebular == True:
+            self.calculate = self._calculate_classical_nebular
+        elif coronal_approximation == True:
+            self.calculate = self._calculate_coronal_approximation
+        else:
+            self.calculate = self._calculate_general
+        self._update_inputs()
+
+    def _main_nlte_calculation(self, nlte_species, atomic_data, nlte_data,
+        t_electron, j_blues, beta_sobolevs, general_level_boltzmann_factor,
+        previous_electron_densities):
+        for species in nlte_species:
+            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]
+            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_electron)), dtype=np.float64)
+            r_ul_matrix_reshaped = r_ul_matrix.reshape((number_of_levels**2,
+                len(t_electron)))
+            r_ul_matrix_reshaped[r_ul_index] = A_uls[np.newaxis].T + \
+                B_uls[np.newaxis].T * j_blues.ix[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(t_electron)))
+            r_lu_matrix_reshaped[r_lu_index] = B_lus[np.newaxis].T * \
+                j_blues.ix[lines_index] * beta_sobolevs[lines_index]
+            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_electron) * previous_electron_densities.values
+            else:
+                collision_matrix = r_ul_matrix.copy()
+                collision_matrix.fill(0.0)
+            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(t_electron)):
+                level_boltzmann_factor = \
+                    np.linalg.solve(rates_matrix[:, :, i], x)
+                general_level_boltzmann_factor[i].ix[species] = \
+                    level_boltzmann_factor
+        return general_level_boltzmann_factor
+
+    def _calculate_classical_nebular(self, t_electron, lines, atomic_data,
+        nlte_data, general_level_boltzmann_factor, nlte_species, j_blues,
+        previous_beta_sobolevs, lte_j_blues, previous_electron_densities):
+        beta_sobolevs = np.ones((len(lines), len(t_electron)))
+        if len(j_blues)==0:
+            j_blues = lte_j_blues
+        else:
+            j_blues = pd.DataFrame(j_blues, index=lines.index, columns =
+                range(len(t_electron)))
+        general_level_boltzmann_factor = self._main_nlte_calculation(
+            nlte_species, atomic_data, nlte_data, t_electron, j_blues,
+            beta_sobolevs, general_level_boltzmann_factor,
+            previous_electron_densities)
+        return general_level_boltzmann_factor
+
+    def _calculate_coronal_approximation(self, t_electron, lines, atomic_data,
+        nlte_data, general_level_boltzmann_factor, nlte_species,
+        previous_electron_densities):
+        beta_sobolevs = np.ones((len(lines), len(t_electron)))
+        j_blues = np.zeros((len(lines), len(t_electron)))
+        general_level_boltzmann_factor = self._main_nlte_calculation(
+            nlte_species, atomic_data, nlte_data, t_electron, j_blues,
+            beta_sobolevs, general_level_boltzmann_factor,
+            previous_electron_densities)
+        return general_level_boltzmann_factor
+
+    def _calculate_general(self, t_electron, lines, atomic_data, nlte_data,
+        general_level_boltzmann_factor, nlte_species, j_blues,
+        previous_beta_sobolevs, lte_j_blues, previous_electron_densities):
+        if previous_beta_sobolevs is None:
+            beta_sobolevs = np.ones((len(lines), len(t_electron)))
+        else:
+            beta_sobolevs = previous_beta_sobolevs
+        if len(j_blues)==0:
+            j_blues = lte_j_blues
+        else:
+            j_blues = pd.DataFrame(j_blues, index=lines.index, columns =
+                range(len(t_electron)))
+        general_level_boltzmann_factor = self._main_nlte_calculation(
+            nlte_species, atomic_data, nlte_data, t_electron, j_blues,
+            beta_sobolevs, general_level_boltzmann_factor,
+            previous_electron_densities)
+        return general_level_boltzmann_factor
+
 class PartitionFunction(ProcessingPlasmaProperty):
     outputs = ('partition_function',)
     latex_outputs = '$Z_{i, j}$'
@@ -61,4 +178,4 @@ class PartitionFunction(ProcessingPlasmaProperty):
     @staticmethod
     def calculate(levels, level_boltzmann_factor):
         return level_boltzmann_factor.groupby(
-            level=['atomic_number', 'ion_number']).sum()
+            level=['atomic_number', 'ion_number']).sum()

tardis/plasma/properties/plasma_input.py

@@ -4,7 +4,9 @@
 
 __all__ = ['TRadiative', 'DilutionFactor', 'AtomicData', 'Abundance', 'Density',
            'TimeExplosion', 'JBlues', 'LinkTRadTElectron',
-           'RadiationFieldCorrectionInput']
+           'RadiationFieldCorrectionInput', 'NLTESpecies',
+           'PreviousIterationBetaSobolevs',
+           'PreviousIterationElectronDensities']
 
 
 class Input(BasePlasmaProperty):
@@ -65,3 +67,12 @@ class JBlues(DataFrameInput):
 
 class LinkTRadTElectron(StaticInput):
     outputs = ('link_t_rad_t_electron',)
+
+class NLTESpecies(StaticInput):
+    outputs = ('nlte_species',)
+
+class PreviousIterationBetaSobolevs(DynamicInput):
+    outputs = ('previous_beta_sobolevs',)
+
+class PreviousIterationElectronDensities(DynamicInput):
+    outputs = ('previous_electron_densities',)

tardis/plasma/properties/property_collections.py

@@ -7,14 +7,18 @@
     TransitionProbabilities, StimulatedEmissionFactor, SelectedAtoms,
     PhiGeneral, PhiSahaNebular, LevelBoltzmannFactorDiluteLTE, DilutionFactor,
     ZetaData, ElectronTemperature, LinkTRadTElectron, BetaElectron,
-    RadiationFieldCorrection, RadiationFieldCorrectionInput)
+    RadiationFieldCorrection, RadiationFieldCorrectionInput,
+    LevelBoltzmannFactorNoNLTE, LevelBoltzmannFactorNLTE, NLTEData,
+    NLTESpecies, PreviousBetaSobolevs, LTEJBlues,
+    PreviousElectronDensities)
 
 class PlasmaPropertyCollection(list):
     pass
 
 basic_inputs = PlasmaPropertyCollection([TRadiative, Abundance, Density,
     TimeExplosion, AtomicData, JBlues, DilutionFactor, LinkTRadTElectron,
-    RadiationFieldCorrectionInput])
+    RadiationFieldCorrectionInput, NLTESpecies, PreviousBetaSobolevs,
+    PreviousElectronDensities])
 basic_properties = PlasmaPropertyCollection([BetaRadiation,
     Levels, Lines, AtomicMass, LevelPopulation, PartitionFunction,
     GElectron, IonizationData, NumberDensity, LinesLowerLevelIndex,
@@ -28,3 +32,6 @@ class PlasmaPropertyCollection(list):
     ZetaData, ElectronTemperature, BetaElectron, RadiationFieldCorrection])
 dilute_lte_excitation_properties = PlasmaPropertyCollection([
     LevelBoltzmannFactorDiluteLTE])
+non_nlte_properties = PlasmaPropertyCollection([LevelBoltzmannFactorNoNLTE])
+nlte_properties = PlasmaPropertyCollection([
+    LevelBoltzmannFactorNLTE, NLTEData, NLTESpecies, LTEJBlues])