Reading nlte_excitation_species from config

tardis/io/tests/test_config_reader.py

@@ -190,7 +190,7 @@ def test_plasma_nlte_section_config(tardis_config_verysimple_nlte):
 
     Parameter
     ---------
-        `tardis_config_verysimple` : YAML File
+        `tardis_config_verysimple_nlte` : YAML File
 
     Result
     ------
@@ -221,6 +221,46 @@ def test_plasma_nlte_section_config(tardis_config_verysimple_nlte):
     assert ve.type is ValueError
 
 
+def test_plasma_nlte_exc_section_config(tardis_config_verysimple_nlte):
+    """
+    Configuration Validation Test for Plasma Section of the Tardis Config YAML File.
+
+    Validates:
+        nlte_excitation_species: should be included in continuum_interaction
+
+    Parameter
+    ---------
+        `tardis_config_verysimple_nlte` : YAML File
+
+    Result
+    ------
+        Assertion based on validation for specified values
+    """
+    conf = Configuration.from_config_dict(
+        tardis_config_verysimple_nlte, validate=True, config_dirname="test"
+    )
+    tardis_config_verysimple_nlte["plasma"]["continuum_interaction"][
+        "species"
+    ] = [
+        "He I",
+    ]
+    tardis_config_verysimple_nlte["plasma"]["nlte_excitation_species"] = ["H I"]
+    with pytest.raises(ValueError) as ve:
+        nlte_excitation_species = tardis_config_verysimple_nlte["plasma"][
+            "nlte_excitation_species"
+        ]
+
+        for species in nlte_excitation_species:
+            if not (
+                species
+                in tardis_config_verysimple_nlte["plasma"][
+                    "continuum_interaction"
+                ]["species"]
+            ):
+                raise ValueError("Nlte excitation species not in continuum.")
+    assert ve.type is ValueError
+
+
 def test_spectrum_section_config(tardis_config_verysimple):
     """
     Configuration Validation Test for Plasma Section of the Tardis Config YAML File

tardis/plasma/properties/plasma_input.py

@@ -16,6 +16,7 @@
     "Volume",
     "ContinuumInteractionSpecies",
     "NLTEIonizationSpecies",
+    "NLTEExcitationSpecies",
 ]
 
 
@@ -154,3 +155,8 @@ class ContinuumInteractionSpecies(Input):
 class NLTEIonizationSpecies(Input):
 
     outputs = ("nlte_ionization_species",)
+
+
+class NLTEExcitationSpecies(Input):
+
+    outputs = ("nlte_excitation_species",)

tardis/plasma/properties/property_collections.py

@@ -17,6 +17,7 @@ class PlasmaPropertyCollection(list):
         HeliumTreatment,
         ContinuumInteractionSpecies,
         NLTEIonizationSpecies,
+        NLTEExcitationSpecies,
     ]
 )
 basic_properties = PlasmaPropertyCollection(

tardis/plasma/properties/rate_matrix_index.py

@@ -9,11 +9,19 @@
 class NLTEIndexHelper(ProcessingPlasmaProperty):
     outputs = ("rate_matrix_index",)
 
-    def __init__(self, plasma_parent, nlte_ionization_species=0):
+    def __init__(
+        self,
+        plasma_parent,
+        nlte_ionization_species=0,
+        nlte_excitation_species=0,
+    ):
         super().__init__(plasma_parent)
         self.nlte_ionization_species = nlte_ionization_species
+        self.nlte_excitation_species = nlte_excitation_species
 
-    def calculate(self, levels, nlte_ionization_species):
+    def calculate(
+        self, levels, nlte_ionization_species, nlte_excitation_species
+    ):
         """Generates rate_matrix_index using levels and changing the last index(level) to
         "lte_ion" if that ion_number is treated in LTE or nebular, "nlte_ion" for NLTE ionization and
         keeps the levels for the rest.

tardis/plasma/standard_plasmas.py

@@ -129,6 +129,10 @@ def assemble_plasma(config, model, atom_data=None):
         species_string_to_tuple(s)
         for s in config.plasma.nlte_ionization_species
     ]
+    nlte_excitation_species = [
+        species_string_to_tuple(s)
+        for s in config.plasma.nlte_excitation_species
+    ]
 
     kwargs = dict(
         t_rad=model.t_radiative,
@@ -140,6 +144,7 @@ def assemble_plasma(config, model, atom_data=None):
         link_t_rad_t_electron=config.plasma.link_t_rad_t_electron,
         continuum_interaction_species=continuum_interaction_species,
         nlte_ionization_species=nlte_ionization_species,
+        nlte_excitation_species=nlte_excitation_species,
     )
 
     plasma_modules = basic_inputs + basic_properties
@@ -175,15 +180,33 @@ def assemble_plasma(config, model, atom_data=None):
         property_kwargs[MarkovChainTransProbsCollector] = {
             "inputs": transition_probabilities_outputs
         }
-        if config.plasma.nlte_ionization_species:
-            nlte_ionization_species = config.plasma.nlte_ionization_species
-            for species in nlte_ionization_species:
-                if not (species in config.plasma.continuum_interaction.species):
-                    raise PlasmaConfigError(
-                        f"NLTE ionization species {species} not in continuum species."
-                    )
+        if (
+            config.plasma.nlte_ionization_species
+            or config.plasma.nlte_excitation_species
+        ):
+            nlte_ionization_species = []
+            nlte_excitation_species = []
+            if config.plasma.nlte_ionization_species:
+                nlte_ionization_species = config.plasma.nlte_ionization_species
+                for species in nlte_ionization_species:
+                    if not (
+                        species in config.plasma.continuum_interaction.species
+                    ):
+                        raise PlasmaConfigError(
+                            f"NLTE ionization species {species} not in continuum species."
+                        )
+            if config.plasma.nlte_excitation_species:
+                nlte_excitation_species = config.plasma.nlte_excitation_species
+                for species in nlte_excitation_species:
+                    if not (
+                        species in config.plasma.continuum_interaction.species
+                    ):
+                        raise PlasmaConfigError(
+                            f"NLTE excitation species {species} not in continuum species."
+                        )
             property_kwargs[NLTEIndexHelper] = {
-                "nlte_ionization_species": nlte_ionization_species
+                "nlte_ionization_species": nlte_ionization_species,
+                "nlte_excitation_species": nlte_excitation_species,
             }
             plasma_modules += nlte_solver_properties