Enable nlte ionization as plasma component (#2458)

* Rename nlte ion and electron

* Enable NLTE

* Fix typo

* Add docs

* Fix transitionprob dtype

* dtype conversion

* dtype for sparse matrix

* Typing in numba

* Fix test variable name

* Add missing electron density case

* Fix typo

docs/io/configuration/components/plasma.rst

@@ -44,4 +44,15 @@ NLTE
 The NLTE configuration currently allows setting ``coronal_approximation``, which sets all :math:`J_\textrm{blue}` to 0.
 This is useful for debugging with :term:`chianti` for example. Furthermore, one can enable 'classical_nebular' to set all
 :math:`\beta_\textrm{Sobolev}` to 1. Both options are used for checking with other codes and should not be enabled in
-normal operations.
+normal operations.
+
+NLTE Ionization
+^^^^^^^^^^^^^^^
+
+.. code-block:: yaml
+
+    plasma:
+        nlte_ionization_species: [H I, H II, He I, He II]
+    
+This option allows the user to specify which species should be included in the NLTE ionization treatment. Note that the
+species must be present in the continuum interaction species as well.

docs/physics/setup/plasma/index.rst

@@ -73,6 +73,9 @@ The next more complex class is `LTEPlasma` which will calculate the ionization b
 
 TARDIS also allows for NLTE treatments of specified species, as well as special NLTE treatments for Helium.
 
+.. note::
+    The NLTE treatment of specified species is currently incompatible with the NLTE treatment for helium and cannot be used simulataneously.
+
 .. toctree::
     :maxdepth: 2
 

tardis/montecarlo/montecarlo_numba/numba_interface.py

@@ -217,7 +217,9 @@ def opacity_state_initialize(plasma, line_interaction_type):
         ].values
 
         phot_nus = phot_nus.values
-        ff_opacity_factor = plasma.ff_cooling_factor / np.sqrt(t_electrons)
+        ff_opacity_factor = (
+            plasma.ff_cooling_factor / np.sqrt(t_electrons)
+        ).astype(np.float64)
         emissivities = plasma.fb_emission_cdf.loc[
             plasma.level2continuum_idx.index
         ].values

tardis/plasma/properties/nlte_rate_equation_solver.py

@@ -11,7 +11,7 @@
 
 
 class NLTERateEquationSolver(ProcessingPlasmaProperty):
-    outputs = ("ion_number_density_nlte", "electron_densities_nlte")
+    outputs = ("ion_number_density", "electron_densities")
 
     def calculate(
         self,

tardis/plasma/properties/transition_probabilities.py

@@ -41,10 +41,16 @@ def normalize_trans_probs(p):
         all probabilites with the same source_level_idx sum to one.
         Indexed by source_level_idx, destination_level_idx.
     """
-    p_summed = p.groupby(level=0).sum()
+    # Dtype conversion is needed for pandas to return nan instead of
+    # a ZeroDivisionError in cases where the sum is zero.
+    p = p.astype(np.float64)
+    p_summed = p.groupby(level=0).sum().astype(np.float64)
     index = p.index.get_level_values("source_level_idx")
     p_norm = p / p_summed.loc[index].values
     p_norm = p_norm.fillna(0.0)
+    # Convert back to original dtypes to avoid typing problems later on
+    # in the numba code.
+    p_norm = p_norm.convert_dtypes()
     return p_norm
 
 
@@ -78,7 +84,9 @@ def series2matrix(series, idx2reduced_idx):
             idx2reduced_idx.loc[q_indices[1]].values,
         )
         max_idx = idx2reduced_idx.max() + 1
-        matrix = sp.coo_matrix((series, q_indices), shape=(max_idx, max_idx))
+        matrix = sp.coo_matrix(
+            (series.astype(np.float64), q_indices), shape=(max_idx, max_idx)
+        )
         return matrix
 
     @staticmethod

tardis/plasma/standard_plasmas.py

@@ -6,6 +6,10 @@
 import pandas as pd
 
 from tardis.io.atom_data import AtomData
+from tardis.plasma.properties.level_population import LevelNumberDensity
+from tardis.plasma.properties.nlte_rate_equation_solver import (
+    NLTERateEquationSolver,
+)
 from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper
 from tardis.util.base import species_string_to_tuple
 from tardis.plasma import BasePlasma
@@ -265,6 +269,26 @@ def assemble_plasma(config, simulation_state, atom_data=None):
             delta_treatment=config.plasma.delta_treatment
         )
 
+    if (
+        config.plasma.helium_treatment == "recomb-nlte"
+        or config.plasma.helium_treatment == "numerical-nlte"
+    ) and (
+        config.plasma.nlte_ionization_species
+        or config.plasma.nlte_excitation_species
+    ):
+        # Prevent the user from using helium NLTE treatment with
+        # NLTE ionization and excitation treatment. This is because
+        # the helium_nlte_properties could overwrite the NLTE ionization
+        # and excitation ion number and electron densities.
+        # helium_numerical_nlte_properties is also included here because
+        # it is currently in the same if else block, and thus may block
+        # the addition of the components from the else block.
+        raise PlasmaConfigError(
+            "Helium NLTE treatment is incompatible with the NLTE eonization and excitation treatment."
+        )
+
+    # TODO: Disentangle these if else block such that compatible components
+    # can be added independently.
     if config.plasma.helium_treatment == "recomb-nlte":
         plasma_modules += helium_nlte_properties
     elif config.plasma.helium_treatment == "numerical-nlte":
@@ -277,7 +301,16 @@ def assemble_plasma(config, simulation_state, atom_data=None):
                 heating_rate_data_file=config.plasma.heating_rate_data_file
             )
     else:
-        plasma_modules += helium_lte_properties
+        # If nlte ionization species are present, we don't want to add the
+        # IonNumberDensity from helium_lte_properties, since we want
+        # to use the IonNumberDensity provided by the NLTE solver.
+        if (
+            config.plasma.nlte_ionization_species
+            or config.plasma.nlte_excitation_species
+        ):
+            plasma_modules += [LevelNumberDensity]
+        else:
+            plasma_modules += helium_lte_properties
 
     if simulation_state._electron_densities is not None:
         electron_densities = pd.Series(
@@ -287,6 +320,13 @@ def assemble_plasma(config, simulation_state, atom_data=None):
             property_kwargs[IonNumberDensityHeNLTE] = dict(
                 electron_densities=electron_densities
             )
+        elif (
+            config.plasma.nlte_ionization_species
+            or config.plasma.nlte_excitation_species
+        ):
+            property_kwargs[NLTERateEquationSolver] = dict(
+                electron_densities=electron_densities
+            )
         else:
             property_kwargs[IonNumberDensity] = dict(
                 electron_densities=electron_densities

tardis/plasma/tests/test_nlte_solver.py

@@ -260,7 +260,7 @@ def nlte_raw_plasma_w0(
 
 def test_critical_case_w1(nlte_raw_plasma_w1):
     """Check that the LTE and NLTE solution agree for w=1.0."""
-    ion_number_density_nlte = nlte_raw_plasma_w1.ion_number_density_nlte.values
+    ion_number_density_nlte = nlte_raw_plasma_w1.ion_number_density.values
     ion_number_density_nlte[ion_number_density_nlte < 1e-10] = 0.0
 
     ind = IonNumberDensity(nlte_raw_plasma_w1)