Updates to He NLTE treatment.

tardis/plasma/properties/ion_population.py

@@ -7,16 +7,14 @@
 
 from scipy import interpolate
 
-
-
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
 from tardis.plasma.exceptions import PlasmaIonizationError
 
 
 logger = logging.getLogger(__name__)
 
 __all__ = ['PhiSahaNebular', 'PhiSahaLTE', 'RadiationFieldCorrection',
-           'IonNumberDensity']
+           'IonNumberDensity', 'IonNumberDensityHeNLTE']
 
 
 def calculate_block_ids_from_dataframe(dataframe):
@@ -195,23 +193,19 @@ def __init__(self, plasma_parent, ion_zero_threshold=1e-20):
         self.ion_zero_threshold = ion_zero_threshold
         self.block_ids = None
 
-    def update_helium_nlte(self, ion_number_density, number_density):
-        ion_number_density.ix[2].ix[0] = 0.0
-        ion_number_density.ix[2].ix[2] = 0.0
-        ion_number_density.ix[2].ix[1].update(number_density.ix[2])
-        return ion_number_density
-
-    def calculate_with_n_electron(self, phi, partition_function,
-                                  number_density, n_electron):
-        if self.block_ids is None:
-            self.block_ids = self._calculate_block_ids(phi)
+    @staticmethod
+    def calculate_with_n_electron(phi, partition_function,
+                                  number_density, n_electron, block_ids,
+                                  ion_zero_threshold):
+        if block_ids is None:
+            block_ids = IonNumberDensity._calculate_block_ids(phi)
 
         ion_populations = np.empty_like(partition_function.values)
 
         phi_electron = np.nan_to_num(phi.values / n_electron.values)
 
-        for i, start_id in enumerate(self.block_ids[:-1]):
-            end_id = self.block_ids[i + 1]
+        for i, start_id in enumerate(block_ids[:-1]):
+            end_id = block_ids[i + 1]
             current_phis = phi_electron[start_id:end_id]
             phis_product = np.cumprod(current_phis, 0)
 
@@ -223,11 +217,10 @@ def calculate_with_n_electron(self, phi, partition_function,
 
             ion_populations[start_id + i:end_id + 1 + i] = tmp_ion_populations
 
-        ion_populations[ion_populations < self.ion_zero_threshold] = 0.0
+        ion_populations[ion_populations < ion_zero_threshold] = 0.0
 
         return pd.DataFrame(data = ion_populations,
-                            index=partition_function.index)
-
+                            index=partition_function.index), block_ids
 
     @staticmethod
     def _calculate_block_ids(phi):
@@ -239,14 +232,10 @@ def calculate(self, phi, partition_function, number_density):
         n_electron_iterations = 0
 
         while True:
-            ion_number_density = self.calculate_with_n_electron(
-                phi, partition_function, number_density, n_electron)
-            if hasattr(self.plasma_parent, 'plasma_properties_dict'):
-                if 'HeliumNLTE' in \
-                    self.plasma_parent.plasma_properties_dict.keys():
-                    ion_number_density = \
-                        self.update_helium_nlte(ion_number_density,
-                        number_density)
+            ion_number_density, self.block_ids = \
+                self.calculate_with_n_electron(
+                phi, partition_function, number_density, n_electron,
+                    self.block_ids, self.ion_zero_threshold)
             ion_numbers = ion_number_density.index.get_level_values(1).values
             ion_numbers = ion_numbers.reshape((ion_numbers.shape[0], 1))
             new_n_electron = (ion_number_density.values * ion_numbers).sum(
@@ -264,3 +253,78 @@ def calculate(self, phi, partition_function, number_density):
                 break
             n_electron = 0.5 * (new_n_electron + n_electron)
         return ion_number_density, n_electron
+
+class IonNumberDensityHeNLTE(ProcessingPlasmaProperty):
+    """
+    Attributes:
+    ion_number_density : Pandas DataFrame, dtype float
+                         Index atom number, ion number. Columns zones.
+    electron_densities : Numpy Array, dtype float
+
+    Convergence process to find the correct solution. A trial value for
+    the electron density is initiated in a particular zone. The ion
+    number densities are then calculated using the Saha equation. The
+    electron density is then re-calculated by using the ion number
+    densities to sum over the number of free electrons. If the two values
+    for the electron densities are not similar to within the threshold
+    value, a new guess for the value of the electron density is chosen
+    and the process is repeated.
+    """
+    outputs = ('ion_number_density', 'electron_densities',
+               'helium_population_updated')
+    latex_name = ('N_{i,j}','n_{e}',)
+
+    def __init__(self, plasma_parent, ion_zero_threshold=1e-20):
+        super(IonNumberDensityHeNLTE, self).__init__(plasma_parent)
+        self.ion_zero_threshold = ion_zero_threshold
+        self.block_ids = None
+
+    def update_he_population(self, helium_population, n_electron,
+                             number_density):
+        helium_population_updated = helium_population.copy()
+        he_one_population = helium_population_updated.ix[0].mul(n_electron)
+        he_three_population = helium_population_updated.ix[2].mul(
+            1./n_electron)
+        helium_population_updated.ix[0].update(he_one_population)
+        helium_population_updated.ix[2].update(he_three_population)
+        unnormalised = helium_population_updated.sum()
+        normalised = helium_population_updated.mul(number_density.ix[2] /
+            unnormalised)
+        helium_population_updated.update(normalised)
+        return helium_population_updated
+
+    def calculate(self, phi, partition_function, number_density,
+                  helium_population):
+        n_e_convergence_threshold = 0.05
+        n_electron = number_density.sum(axis=0)
+        n_electron_iterations = 0
+        while True:
+            ion_number_density, self.block_ids = \
+                IonNumberDensity.calculate_with_n_electron(
+                phi, partition_function, number_density, n_electron,
+                    self.block_ids, self.ion_zero_threshold)
+            helium_population_updated = self.update_he_population(
+                helium_population, n_electron, number_density)
+            ion_number_density.ix[2].ix[0].update(helium_population_updated.ix[
+                                                      0].sum(axis=0))
+            ion_number_density.ix[2].ix[1].update(helium_population_updated.ix[
+                                                      1].sum(axis=0))
+            ion_number_density.ix[2].ix[2].update(helium_population_updated.ix[
+                                                      2].ix[0])
+            ion_numbers = ion_number_density.index.get_level_values(1).values
+            ion_numbers = ion_numbers.reshape((ion_numbers.shape[0], 1))
+            new_n_electron = (ion_number_density.values * ion_numbers).sum(
+                axis=0)
+            if np.any(np.isnan(new_n_electron)):
+                raise PlasmaIonizationError('n_electron just turned "nan" -'
+                                            ' aborting')
+            n_electron_iterations += 1
+            if n_electron_iterations > 100:
+                logger.warn('n_electron iterations above 100 ({0}) -'
+                            ' something is probably wrong'.format(
+                    n_electron_iterations))
+            if np.all(np.abs(new_n_electron - n_electron)
+                              / n_electron < n_e_convergence_threshold):
+                break
+            n_electron = 0.5 * (new_n_electron + n_electron)
+        return ion_number_density, n_electron, helium_population_updated

tardis/plasma/properties/level_population.py

@@ -6,7 +6,7 @@
 
 logger = logging.getLogger(__name__)
 
-__all__ = ['LevelNumberDensity']
+__all__ = ['LevelNumberDensity', 'LevelNumberDensityHeNLTE']
 
 class LevelNumberDensity(ProcessingPlasmaProperty):
     """
@@ -21,21 +21,59 @@ class LevelNumberDensity(ProcessingPlasmaProperty):
     def calculate(self):
         pass
 
-    def __init__(self, plasma_parent, helium_treatment='dilute-lte'):
+    def __init__(self, plasma_parent):
         """
         Calculates the level populations with the Boltzmann equation in LTE.
         """
         super(LevelNumberDensity, self).__init__(plasma_parent)
-        if hasattr(self.plasma_parent, 'plasma_properties_dict'):
-            if 'HeliumNLTE' in \
-                self.plasma_parent.plasma_properties_dict.keys():
-                    helium_treatment='recomb-nlte'
-        if helium_treatment=='recomb-nlte':
-            self.calculate = self._calculate_helium_nlte
-        elif helium_treatment=='dilute-lte':
-            self.calculate = self._calculate_dilute_lte
+        self.calculate = self._calculate_dilute_lte
         self._update_inputs()
+        self.initialize_indices = True
+
+    def _initialize_indices(self, levels, partition_function):
+        indexer = pd.Series(np.arange(partition_function.shape[0]),
+                            index=partition_function.index)
+        self._ion2level_idx = indexer.ix[levels.droplevel(2)].values
 
+    def _calculate_dilute_lte(self, level_boltzmann_factor, ion_number_density,
+        levels, partition_function):
+        """
+        Reduces non-metastable level populations by a factor of W compared to LTE in the case of dilute-lte excitation.
+        """
+        if self.initialize_indices:
+            self._initialize_indices(levels, partition_function)
+            self.initialize_indices = False
+        partition_function_broadcast = partition_function.values[
+            self._ion2level_idx]
+        level_population_fraction = (level_boltzmann_factor.values
+                                     / partition_function_broadcast)
+        ion_number_density_broadcast = ion_number_density.values[
+            self._ion2level_idx]
+        level_number_density = (level_population_fraction *
+                                ion_number_density_broadcast)
+        return pd.DataFrame(level_number_density,
+                            index=level_boltzmann_factor.index)
+
+class LevelNumberDensityHeNLTE(ProcessingPlasmaProperty):
+    """
+    Attributes:
+    level_number_density : Pandas DataFrame, dtype float
+                           Index atom number, ion number, level number. Columns are zones.
+    """
+    outputs = ('level_number_density',)
+    latex_name = ('N_{i,j,k}',)
+    latex_formula = ('N_{i,j}\\dfrac{bf_{i,j,k}}{Z_{i,j}}',)
+
+    def calculate(self):
+        pass
+
+    def __init__(self, plasma_parent):
+        """
+        Calculates the level populations with the Boltzmann equation in LTE.
+        """
+        super(LevelNumberDensityHeNLTE, self).__init__(plasma_parent)
+        self.calculate = self._calculate_helium_nlte
+        self._update_inputs()
         self.initialize_indices = True
 
 
@@ -64,14 +102,15 @@ def _calculate_dilute_lte(self, level_boltzmann_factor, ion_number_density,
                             index=level_boltzmann_factor.index)
 
     def _calculate_helium_nlte(self, level_boltzmann_factor,
-        ion_number_density, levels, partition_function, helium_population):
+        ion_number_density, levels, partition_function,
+                               helium_population_updated):
         """
         If one of the two helium NLTE methods is used, this updates the helium level populations to the appropriate
         values.
         """
         level_number_density = self._calculate_dilute_lte(
             level_boltzmann_factor, ion_number_density, levels,
             partition_function)
-        if helium_population is not None:
-            level_number_density.ix[2].update(helium_population)
+        if helium_population_updated is not None:
+            level_number_density.ix[2].update(helium_population_updated)
         return level_number_density

tardis/plasma/properties/nlte.py

@@ -6,7 +6,7 @@
 
 from tardis.plasma.properties.base import (PreviousIterationProperty,
                                            ProcessingPlasmaProperty)
-from tardis.plasma.properties import PhiSahaNebular, PhiSahaLTE
+from tardis.plasma.properties.ion_population import PhiSahaNebular
 
 __all__ = ['PreviousElectronDensities', 'PreviousBetaSobolev',
            'HeliumNLTE', 'HeliumNumericalNLTE']
@@ -45,20 +45,18 @@ def set_initial_value(self, kwargs):
 class HeliumNLTE(ProcessingPlasmaProperty):
     outputs = ('helium_population',)
 
-    def calculate(self, level_boltzmann_factor, electron_densities,
+    @staticmethod
+    def calculate(level_boltzmann_factor,
         ionization_data, beta_rad, g, g_electron, w, t_rad, t_electrons,
         delta, zeta_data, number_density, partition_function):
         """
         Updates all of the helium level populations according to the helium NLTE recomb approximation.
         """
         helium_population = level_boltzmann_factor.ix[2].copy()
         # He I excited states
-        he_one_population = self.calculate_helium_one(g_electron, beta_rad,
-            ionization_data, level_boltzmann_factor, electron_densities, g, w)
+        he_one_population = HeliumNLTE.calculate_helium_one(g_electron, beta_rad,
+            ionization_data, level_boltzmann_factor, g, w)
         helium_population.ix[0].update(he_one_population)
-        #He I metastable states
-        helium_population.ix[0,1] *= (1 / w)
-        helium_population.ix[0,2] *= (1 / w)
         #He I ground state
         helium_population.ix[0,0] = 0.0
         #He II excited states
@@ -68,38 +66,38 @@ def calculate(self, level_boltzmann_factor, electron_densities,
         #He II ground state
         helium_population.ix[1,0] = 1.0
         #He III states
-        helium_population.ix[2,0] = self.calculate_helium_three(t_rad, w,
+        helium_population.ix[2,0] = HeliumNLTE.calculate_helium_three(t_rad, w,
             zeta_data, t_electrons, delta, g_electron, beta_rad,
-            ionization_data, electron_densities, g)
-        unnormalised = helium_population.sum()
-        normalised = helium_population.mul(number_density.ix[2] / unnormalised)
-        helium_population.update(normalised)
+            ionization_data, g)
+#        unnormalised = helium_population.sum()
+#        normalised = helium_population.mul(number_density.ix[2] /
+# unnormalised)
+#        helium_population.update(normalised)
         return helium_population
 
     @staticmethod
     def calculate_helium_one(g_electron, beta_rad, ionization_data,
-        level_boltzmann_factor, electron_densities, g, w):
+        level_boltzmann_factor, g, w):
         """
         Calculates the He I level population values, in equilibrium with the He II ground state.
         """
-        return level_boltzmann_factor.ix[2,0].mul(
-            g.ix[2,0], axis=0) * (1./(2*g.ix[2,1,0])) * \
-            (1/g_electron) * (1/(w**2)) * np.exp(
-            ionization_data.ionization_energy.ix[2,1] * beta_rad) * \
-            electron_densities
+        return level_boltzmann_factor.ix[2,0] * (1./(2*g.ix[2,1,0])) * \
+            (1/g_electron) * (1/(w**2.)) * np.exp(
+            ionization_data.ionization_energy.ix[2,1] * beta_rad)
 
     @staticmethod
     def calculate_helium_three(t_rad, w, zeta_data, t_electrons, delta,
-        g_electron, beta_rad, ionization_data, electron_densities, g):
+        g_electron, beta_rad, ionization_data, g):
         """
         Calculates the He III level population values.
         """
         zeta = PhiSahaNebular.get_zeta_values(zeta_data, 2, t_rad)[1]
-        he_three_population = (2 / electron_densities) * \
+        he_three_population = 2 * \
             (float(g.ix[2,2,0])/g.ix[2,1,0]) * g_electron * \
             np.exp(-ionization_data.ionization_energy.ix[2,2] * beta_rad) \
             * w * (delta.ix[2,2] * zeta + w * (1. - zeta)) * \
             (t_electrons / t_rad) ** 0.5
+        return he_three_population
 
 class HeliumNumericalNLTE(ProcessingPlasmaProperty):
     '''

tardis/plasma/properties/plasma_input.py

@@ -1,7 +1,7 @@
 from tardis.plasma.properties.base import (Input, ArrayInput, DataFrameInput)
 
 __all__ = ['TRadiative', 'DilutionFactor', 'AtomicData', 'Abundance', 'Density',
-           'TimeExplosion', 'JBlues', 'LinkTRadTElectron']
+           'TimeExplosion', 'JBlues', 'LinkTRadTElectron', 'HeliumTreatment']
 
 class TRadiative(ArrayInput):
     """
@@ -80,3 +80,6 @@ class LinkTRadTElectron(Input):
     """
     outputs = ('link_t_rad_t_electron',)
     latex_name = ('T_{\\textrm{electron}}/T_{\\textrm{rad}}',)
+
+class HeliumTreatment(Input):
+    outputs = ('helium_treatment',)

tardis/plasma/properties/property_collections.py

@@ -4,11 +4,12 @@ class PlasmaPropertyCollection(list):
     pass
 
 basic_inputs = PlasmaPropertyCollection([TRadiative, Abundance, Density,
-    TimeExplosion, AtomicData, JBlues, DilutionFactor, LinkTRadTElectron])
+    TimeExplosion, AtomicData, JBlues, DilutionFactor, LinkTRadTElectron,
+    HeliumTreatment])
 basic_properties = PlasmaPropertyCollection([BetaRadiation,
     Levels, Lines, AtomicMass, PartitionFunction,
     GElectron, IonizationData, NumberDensity, LinesLowerLevelIndex,
-    LinesUpperLevelIndex, TauSobolev, LevelNumberDensity, IonNumberDensity,
+    LinesUpperLevelIndex, TauSobolev,
     StimulatedEmissionFactor, SelectedAtoms, ElectronTemperature])
 lte_ionization_properties = PlasmaPropertyCollection([PhiSahaLTE])
 lte_excitation_properties = PlasmaPropertyCollection([LevelBoltzmannFactorLTE])
@@ -24,6 +25,8 @@ class PlasmaPropertyCollection(list):
     PreviousBetaSobolev, BetaSobolev])
 helium_nlte_properties = PlasmaPropertyCollection([HeliumNLTE,
     RadiationFieldCorrection, ZetaData,
-    BetaElectron])
+    BetaElectron, LevelNumberDensityHeNLTE, IonNumberDensityHeNLTE])
+helium_lte_properties = PlasmaPropertyCollection([LevelNumberDensity,
+                                                  IonNumberDensity])
 helium_numerical_nlte_properties = PlasmaPropertyCollection([
     HeliumNumericalNLTE])