WIP: Helium approx.

docs/physics/plasma/helium_nlte.rst

@@ -0,0 +1,45 @@
+Helium NLTE
+============
+
+The `helium_treatment` setting in the Tardis config. file will accept one of three options:
+ * `none`: The default setting. Populate helium in the same way as the other elements.
+ * `recomb-nlte`: Treats helium in NLTE using the analytical approximation outlined in an upcoming paper. 
+ * `numerical-nlte`: To be implemented. Will allow the use of a separate module (not distributed with Tardis) to perform helium NLTE calculations numerically.
+
+`recomb-NLTE` (Paper version)
+-----------------------------
+This section will summarise the equations used in the calculation of the helium state for the `recomb-NLTE` approximation in Tardis. A full physical justification for these equations will be provided in an upcoming paper. All of the level populations are given as a function of the He II ground state population (:math:`n_{2,1,0}`), and the values are then normalised using the helium number density to give the correct level number densities.
+
+Symbols/Indexing:
+ * :math:`N_{i,j}`: Ion Number Density
+ * :math:`n_{i,j,k}`: Level Number Density
+
+  * i: atomic number
+  * j: ion number
+  * k: level number
+
+.. math::
+    n_{2,0,0} = 0
+
+.. math::
+    n_{2,0,k}~(k\geq1) = n_{2,1,0}\times n_{e}\times\frac{1}{W}\times\frac{g_{2,0,k}}{2g_{2,1,0}}\times\left(\frac{h^{2}}{2\pi m_{e}kT_{r}}\right)^{3/2}\times\exp{\left(\frac{\chi_{2,1}-\epsilon_{2,0,k}}{kT_{r}}}\right)\times\left(\frac{T_{r}}{T_{e}}\right)^{1/2}
+
+(Note: An extra factor of :math:`\frac{1}{W}` is included for the metastable states of He I.)
+
+.. math::
+    n_{2,1,0} = 1
+
+.. math::
+    n_{2,1,k}~(k\geq1) = W\times\frac{g_{2,0,k}}{g_{2,1,0}}\times n_{2,1,0}\times\exp{\left(-\frac{\epsilon_{2,1,k}}{kT_{r}}\right)}
+
+.. math::
+    n_{2,2,0} = \frac{n_{2,1,0}}{n_{e}}\times[W(\delta_{2,2}\times\zeta_{2,2}+W(1-\zeta_{2,2})]\left(\frac{T_{e}}{T_{r}}\right)^{1/2}\times\frac{2g_{2,2,0}}{g_{2,1,0}}\times\left(\frac{2\pi m_{e}kT_{r}}{h^{2}}\right)^{3/2}\times\exp{\left(-\frac{\chi_{2,1}}{kT_{r}}\right)}
+
+`recomb-NLTE` (Code version)
+-----------------------------
+
+In the Tardis plasma, some of these equations are re-written slightly to make use of existing property methods (e.g. `PhiSahaLTE`, `PhiSahaNebular`) often using the relation:
+
+.. math::
+    \frac{N_{i,j}}{Z_{i,j}} = \frac{n_{i,j,k}}{g_{i,j,k}}
+

docs/physics/plasma/index.rst

@@ -8,3 +8,4 @@ Overview of plasma calculations
     nebular_plasma
     macroatom
     nlte
+    helium_nlte

tardis/data/tardis_config_definition.yml

@@ -131,6 +131,12 @@ plasma:
             mandatory: False
             help: sets all beta_sobolevs to 1
 
+    helium_treatment:
+        property_type: string
+        mandatory: False
+        default: none
+        allowed_value: none recomb-nlte
+        help: none to treat He as the other elements. recomb-nlte to treat with NLTE approximation.
 
 model:
     structure:

tardis/io/config_reader.py

@@ -973,6 +973,11 @@ def from_config_dict(cls, config_dict, atom_data=None, test_parser=False,
             logger.warn('Disabling electron scattering - this is not physical')
             validated_config_dict['montecarlo']['sigma_thomson'] = 1e-200 / (u.cm ** 2)
 
+        if plasma_section['helium_treatment'] == 'recomb-nlte':
+            validated_config_dict['plasma']['helium_treatment'] == 'recomb-nlte'
+        else:
+            validated_config_dict['plasma']['helium_treatment'] == 'dilute-lte'
+
 
 
 

tardis/model.py

@@ -129,7 +129,7 @@ def __init__(self, tardis_config):
                                                          ionization_mode=tardis_config.plasma.ionization,
                                                          excitation_mode=tardis_config.plasma.excitation,
                                                          line_interaction_type=tardis_config.plasma.line_interaction_type,
-                                                         link_t_rad_t_electron=0.9)
+                                                         link_t_rad_t_electron=0.9, helium_treatment=tardis_config.plasma.helium_treatment)
 
         self.spectrum = TARDISSpectrum(tardis_config.spectrum.frequency, tardis_config.supernova.distance)
         self.spectrum_virtual = TARDISSpectrum(tardis_config.spectrum.frequency, tardis_config.supernova.distance)

tardis/plasma/properties/ion_population.py

@@ -53,11 +53,12 @@ class PhiSahaNebular(ProcessingPlasmaProperty):
                      \\dfrac{T_{\\textrm{electron}}}{T_{\\textrm{rad}}}\
                      \\right)^{1/2}',)
 
-    def calculate(self, t_rad, w, zeta_data, t_electrons, delta,
+    @staticmethod
+    def calculate(t_rad, w, zeta_data, t_electrons, delta,
             g_electron, beta_rad, partition_function, ionization_data):
         phi_lte = PhiSahaLTE.calculate(g_electron, beta_rad,
             partition_function, ionization_data)
-        zeta = self.get_zeta_values(zeta_data, phi_lte, t_rad)
+        zeta = PhiSahaNebular.get_zeta_values(zeta_data, phi_lte, t_rad)
         phis = phi_lte * w * ((zeta * delta) + w * (1 - zeta)) * \
                (t_electrons/t_rad) ** .5
         return phis
@@ -141,6 +142,12 @@ def __init__(self, plasma_parent, ion_zero_threshold=1e-20):
         super(IonNumberDensity, self).__init__(plasma_parent)
         self.ion_zero_threshold = ion_zero_threshold
 
+    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):
         ion_populations = pd.DataFrame(data=0.0,
@@ -166,6 +173,12 @@ def calculate(self, phi, partition_function, number_density):
         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_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(

tardis/plasma/properties/level_population.py

@@ -15,12 +15,34 @@ class LevelNumberDensity(ProcessingPlasmaProperty):
     latex_name = ('N_{i,j,k}',)
     latex_formula = ('N_{i,j}\\dfrac{bf_{i,j,k}}{Z_{i,j}}',)
 
-    def calculate(self, level_boltzmann_factor, ion_number_density,
+    def calculate():
+        pass
+
+    def __init__(self, plasma_parent, helium_treatment='dilute-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_recomb_nlte
+        elif helium_treatment=='dilute-lte':
+            self.calculate = self._calculate_dilute_lte
+        self._update_inputs()
+
+    def _calculate_dilute_lte(self, level_boltzmann_factor, ion_number_density,
         levels, partition_function):
         partition_function_broadcast = partition_function.ix[
             levels.droplevel(2)].values
         level_population_fraction = level_boltzmann_factor /\
             partition_function_broadcast
         ion_number_density_broadcast = ion_number_density.ix[
             level_population_fraction.index.droplevel(2)].values
-        return level_population_fraction * ion_number_density_broadcast
+        return level_population_fraction * ion_number_density_broadcast
+
+    def _calculate_helium_recomb_nlte(self, level_boltzmann_factor,
+        ion_number_density, levels, partition_function, helium_population):
+        level_number_density = self._calculate_dilute_lte(
+            level_boltzmann_factor, ion_number_density, levels,
+            partition_function)
+        level_number_density.ix[2].update(helium_population)

tardis/plasma/properties/nlte.py

@@ -1,6 +1,13 @@
-from tardis.plasma.properties.base import BasePlasmaProperty
+import numpy as np
+import pandas as pd
+from scipy import interpolate
 
-__all__ = ['PreviousElectronDensities', 'PreviousBetaSobolevs']
+from tardis.plasma.properties.base import (BasePlasmaProperty,
+                                           ProcessingPlasmaProperty)
+from tardis.plasma.properties import PhiSahaNebular, PhiSahaLTE
+
+__all__ = ['PreviousElectronDensities', 'PreviousBetaSobolevs',
+           'HeliumNLTE']
 
 class PreviousIterationProperty(BasePlasmaProperty):
     def _set_output_value(self, output, value):
@@ -16,3 +23,83 @@ class PreviousElectronDensities(PreviousIterationProperty):
 class PreviousBetaSobolevs(PreviousIterationProperty):
     outputs = ('previous_beta_sobolevs',)
 
+class HeliumNLTE(ProcessingPlasmaProperty):
+    outputs = ('helium_population',)
+
+    def calculate(self, level_boltzmann_factor, electron_densities,
+        ionization_data, beta_rad, g, g_electron, w, t_rad, t_electrons,
+        delta, zeta_data, number_density, partition_function):
+        helium_population = level_boltzmann_factor.ix[2].copy()
+        # He I excited states
+        he_one_population = self.calculate_helium_one(g_electron, beta_rad,
+            partition_function, ionization_data, level_boltzmann_factor,
+            electron_densities, g, w, t_rad, t_electrons)
+        helium_population.ix[0].update(he_one_population)
+        #He I metastable states
+        helium_population.ix[0].ix[1] *= (1 / w)
+        helium_population.ix[0].ix[2] *= (1 / w)
+        #He I ground state
+        helium_population.ix[0].ix[0] = 0.0
+        #He II excited states
+        he_two_population = level_boltzmann_factor.ix[2,1].mul(
+            (g.ix[2,1].ix[0]**(-1)) * (t_electrons / t_rad)**0.5)
+        helium_population.ix[1].update(he_two_population)
+        #He II ground state
+        helium_population.ix[1].ix[0] = 1.0
+        #He III states
+        helium_population.ix[2].ix[0] = self.calculate_helium_three(t_rad, w,
+            zeta_data, t_electrons, delta, g_electron, beta_rad,
+            partition_function, ionization_data, electron_densities)
+        unnormalised = helium_population.sum()
+        normalised = helium_population.mul(number_density.ix[2] / unnormalised)
+        helium_population.update(normalised)
+        return helium_population
+
+    def calculate_helium_one(self, g_electron, beta_rad, partition_function,
+            ionization_data, level_boltzmann_factor, electron_densities, g,
+            w, t_rad, t_electron):
+        (partition_function_index, ionization_data_index, partition_function,
+            ionization_data) = self.filter_with_helium_index(2, 1,
+            partition_function, ionization_data)
+        phis = (1 / PhiSahaLTE.calculate(g_electron, beta_rad,
+            partition_function, ionization_data)) * electron_densities * \
+            (1.0/g.ix[2,1,0]) * (1/w) * (t_rad/t_electron)**(0.5)
+        return level_boltzmann_factor.ix[2].ix[0].mul(
+            pd.DataFrame(phis.ix[2].ix[1].values)[0].transpose())
+
+    def calculate_helium_three(self, t_rad, w, zeta_data, t_electrons, delta,
+        g_electron, beta_rad, partition_function, ionization_data,
+        electron_densities):
+        (partition_function_index, ionization_data_index, partition_function,
+            ionization_data) = self.filter_with_helium_index(2, 2,
+            partition_function, ionization_data)
+        zeta_data = pd.DataFrame(zeta_data.ix[2].ix[2].values,
+            columns=ionization_data_index, index=zeta_data.columns).transpose()
+        delta = pd.DataFrame(delta.ix[2].ix[2].values,
+            columns=ionization_data_index, index=delta.columns).transpose()
+        phis = PhiSahaNebular.calculate(t_rad, w,
+            zeta_data, t_electrons, delta, g_electron,
+            beta_rad, partition_function, ionization_data)
+        return (phis * (partition_function.ix[2].ix[1] /
+            partition_function.ix[2].ix[2]) * (1 /
+            electron_densities)).ix[2].ix[2]
+
+    @staticmethod
+    def filter_with_helium_index(atomic_number, ion_number, partition_function,
+        ionization_data):
+        partition_function_index = pd.MultiIndex.from_tuples([(atomic_number,
+            ion_number-1), (atomic_number, ion_number)],
+            names=['atomic_number', 'ion_number'])
+        ionization_data_index = pd.MultiIndex.from_tuples([(atomic_number,
+            ion_number)],
+            names=['atomic_number', 'ion_number'])
+        partition_function = pd.DataFrame(
+            partition_function.ix[atomic_number].ix[
+            ion_number-1:ion_number].values,
+            index=partition_function_index, columns=partition_function.columns)
+        ionization_data = pd.DataFrame(
+            ionization_data.ix[atomic_number].ix[ion_number][
+            'ionization_energy'], index=ionization_data_index,
+            columns=['ionization_energy'])
+        return partition_function_index, ionization_data_index,\
+               partition_function, ionization_data

tardis/plasma/properties/property_collections.py

@@ -10,7 +10,7 @@
     RadiationFieldCorrection, RadiationFieldCorrectionInput,
     LevelBoltzmannFactorNoNLTE, LevelBoltzmannFactorNLTE, NLTEData,
     NLTESpecies, PreviousBetaSobolevs, LTEJBlues,
-    PreviousElectronDensities, Chi0)
+    PreviousElectronDensities, Chi0, HeliumNLTE)
 
 class PlasmaPropertyCollection(list):
     pass
@@ -34,4 +34,7 @@ class PlasmaPropertyCollection(list):
     LevelBoltzmannFactorDiluteLTE])
 non_nlte_properties = PlasmaPropertyCollection([LevelBoltzmannFactorNoNLTE])
 nlte_properties = PlasmaPropertyCollection([
-    LevelBoltzmannFactorNLTE, NLTEData, NLTESpecies, LTEJBlues])
+    LevelBoltzmannFactorNLTE, NLTEData, NLTESpecies, LTEJBlues])
+helium_nlte_properties = PlasmaPropertyCollection([HeliumNLTE,
+    RadiationFieldCorrection, ZetaData,
+    BetaElectron, Chi0])

tardis/plasma/standard_plasmas.py

@@ -7,7 +7,7 @@
     basic_properties, lte_excitation_properties, lte_ionization_properties,
     macro_atom_properties, dilute_lte_excitation_properties,
     nebular_ionization_properties, non_nlte_properties,
-    nlte_properties)
+    nlte_properties, helium_nlte_properties)
 from tardis.io.util import parse_abundance_dict_to_dataframe
 
 logger = logging.getLogger(__name__)
@@ -55,7 +55,8 @@ def update_radiationfield(self, t_rad, ws, j_blues, nlte_config,
     def __init__(self, number_densities, atomic_data, time_explosion,
         t_rad=None, delta_treatment=None, nlte_config=None,
         ionization_mode='lte', excitation_mode='lte',
-        line_interaction_type='scatter', link_t_rad_t_electron=0.9):
+        line_interaction_type='scatter', link_t_rad_t_electron=0.9,
+        helium_treatment='lte'):
 
         plasma_modules = basic_inputs + basic_properties
 
@@ -104,6 +105,9 @@ def __init__(self, number_densities, atomic_data, time_explosion,
         else:
             nlte_species = None
 
+        if helium_treatment=='recomb-nlte':
+            plasma_modules += helium_nlte_properties
+
         super(LegacyPlasmaArray, self).__init__(plasma_properties=plasma_modules,
             t_rad=t_rad, abundance=abundance, density=density,
             atomic_data=atomic_data, time_explosion=time_explosion,

tardis/plasma/tests/test_property_nlte.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+from tardis.plasma.standard_plasmas import LegacyPlasmaArray
+
+def test_he_nlte_plasma(number_density, atomic_data, time_explosion,
+    t_rad, w, j_blues):
+    he_nlte_plasma = LegacyPlasmaArray(number_densities=number_density,
+        atomic_data=atomic_data, time_explosion=time_explosion,
+        helium_treatment='recomb-nlte')
+    he_nlte_plasma.update_radiationfield(t_rad, w, j_blues, nlte_config=None)
+    assert np.allclose(he_nlte_plasma.get_value(
+        'ion_number_density').ix[2].ix[1], number_density.ix[2])
+    assert np.all(he_nlte_plasma.get_value(
+        'level_number_density').ix[2].ix[0].ix[0])==0.0
+    assert np.allclose(he_nlte_plasma.get_value(
+        'level_number_density').ix[2].sum(), number_density.ix[2])