Reading nlte_excitation_species from config (#2195)

* implementing nlte_excitation

* ran black

* fixed a typo

* got rid of unnecessary lines

* rewrote tests

* made a change on assigning the config values to plasma properties

* fixed the tests

* Update tardis/io/tests/test_config_reader.py

Co-authored-by: Christian Vogl <[email protected]>

---------

Co-authored-by: Christian Vogl <[email protected]>

tardis/io/tests/test_config_reader.py

@@ -8,6 +8,8 @@
 
 from tardis.io import config_reader
 from tardis.io.config_reader import Configuration
+from tardis.plasma.exceptions import PlasmaConfigError
+from tardis.plasma.standard_plasmas import assemble_plasma
 
 
 def data_path(filename):
@@ -181,7 +183,11 @@ def test_plasma_section_config(tardis_config_verysimple):
     assert ve.type is ValueError
 
 
-def test_plasma_nlte_section_config(tardis_config_verysimple_nlte):
+def test_plasma_nlte_section_config(
+    tardis_config_verysimple_nlte,
+    nlte_raw_model,
+    nlte_atom_data,
+):
     """
     Configuration Validation Test for Plasma Section of the Tardis Config YAML File.
 
@@ -190,35 +196,53 @@ def test_plasma_nlte_section_config(tardis_config_verysimple_nlte):
 
     Parameter
     ---------
-        `tardis_config_verysimple` : YAML File
+        `tardis_config_verysimple_nlte` : YAML File
+        `nlte_raw_model` : A simple model
+        `nlte_atom_data` : An example atomic dataset
 
     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_ionization_species"] = ["H I"]
-    with pytest.raises(ValueError) as ve:
-        nlte_ionization_species = tardis_config_verysimple_nlte["plasma"][
-            "nlte_ionization_species"
-        ]
-
-        for species in nlte_ionization_species:
-            if not (
-                species
-                in tardis_config_verysimple_nlte["plasma"][
-                    "continuum_interaction"
-                ]["species"]
-            ):
-                raise ValueError("Nlte ionization species not in continuum.")
-    assert ve.type is ValueError
+    config = Configuration.from_config_dict(tardis_config_verysimple_nlte)
+    with pytest.raises(PlasmaConfigError) as ve:
+        assemble_plasma(config, nlte_raw_model, nlte_atom_data)
+
+
+def test_plasma_nlte_exc_section_config(
+    tardis_config_verysimple_nlte, nlte_raw_model, nlte_atom_data
+):
+    """
+    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
+        `nlte_raw_model` : A simple model
+        `nlte_atom_data` : An example atomic dataset
+
+    Result
+    ------
+        Assertion based on validation for specified values
+    """
+    tardis_config_verysimple_nlte["plasma"]["continuum_interaction"][
+        "species"
+    ] = [
+        "He I",
+    ]
+    tardis_config_verysimple_nlte["plasma"]["nlte_excitation_species"] = ["H I"]
+    config = Configuration.from_config_dict(tardis_config_verysimple_nlte)
+    with pytest.raises(PlasmaConfigError):
+        plasma = assemble_plasma(config, nlte_raw_model, nlte_atom_data)
 
 
 def test_spectrum_section_config(tardis_config_verysimple):

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,31 @@ 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
+        ):
+            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": config.plasma.nlte_ionization_species,
+                "nlte_excitation_species": config.plasma.nlte_excitation_species,
             }
             plasma_modules += nlte_solver_properties
 

tardis/plasma/tests/test_nlte_solver.py

@@ -228,55 +228,6 @@ def test_jacobian_matrix(
     assert_almost_equal(actual_jacobian_matrix, desired_jacobian_matrix)
 
 
-@pytest.fixture(scope="session")
-def nlte_atomic_data_fname(tardis_ref_path):
-    """
-    File name for the atomic data file used in NTLE ionization solver tests.
-    """
-    atomic_data_fname = os.path.join(
-        tardis_ref_path, "nlte_atom_data", "TestNLTE_He_Ti.h5"
-    )
-
-    atom_data_missing_str = (
-        f"{atomic_data_fname} atomic datafiles " f"does not seem to exist"
-    )
-
-    if not os.path.exists(atomic_data_fname):
-        pytest.exit(atom_data_missing_str)
-
-    return atomic_data_fname
-
-
-@pytest.fixture(scope="session")
-def nlte_atomic_dataset(nlte_atomic_data_fname):
-    """
-    Atomic dataset used for NLTE ionization solver tests.
-    """
-    nlte_atomic_data = AtomData.from_hdf(nlte_atomic_data_fname)
-    return nlte_atomic_data
-
-
-@pytest.fixture
-def nlte_atom_data(nlte_atomic_dataset):
-
-    atomic_data = deepcopy(nlte_atomic_dataset)
-    return atomic_data
-
-
-data_path = os.path.join("tardis", "io", "tests", "data")
-
-
-@pytest.fixture
-def tardis_model_config_nlte():
-    filename = "tardis_configv1_nlte.yml"
-    return Configuration.from_yaml(os.path.join(data_path, filename))
-
-
-@pytest.fixture
-def nlte_raw_model(tardis_model_config_nlte):
-    return Radial1DModel.from_config(tardis_model_config_nlte)
-
-
 @pytest.fixture
 def nlte_raw_plasma_w1(
     tardis_model_config_nlte, nlte_raw_model, nlte_atom_data

tardis/tests/fixtures/atom_data.py

@@ -4,6 +4,8 @@
 import pytest
 
 from tardis.io.atom_data.base import AtomData
+from tardis.io.config_reader import Configuration
+from tardis.model.base import Radial1DModel
 
 DEFAULT_ATOM_DATA_UUID = "864f1753714343c41f99cb065710cace"
 
@@ -40,3 +42,52 @@ def atomic_dataset(atomic_data_fname):
 def kurucz_atomic_data(atomic_dataset):
     atomic_data = deepcopy(atomic_dataset)
     return atomic_data
+
+
+@pytest.fixture  # (scope="session")
+def nlte_atomic_data_fname(tardis_ref_path):
+    """
+    File name for the atomic data file used in NTLE ionization solver tests.
+    """
+    atomic_data_fname = os.path.join(
+        tardis_ref_path, "nlte_atom_data", "TestNLTE_He_Ti.h5"
+    )
+
+    atom_data_missing_str = (
+        f"{atomic_data_fname} atomic datafiles " f"does not seem to exist"
+    )
+
+    if not os.path.exists(atomic_data_fname):
+        pytest.exit(atom_data_missing_str)
+
+    return atomic_data_fname
+
+
+@pytest.fixture  # (scope="session")
+def nlte_atomic_dataset(nlte_atomic_data_fname):
+    """
+    Atomic dataset used for NLTE ionization solver tests.
+    """
+    nlte_atomic_data = AtomData.from_hdf(nlte_atomic_data_fname)
+    return nlte_atomic_data
+
+
+@pytest.fixture  # (scope="session")
+def nlte_atom_data(nlte_atomic_dataset):
+
+    atomic_data = deepcopy(nlte_atomic_dataset)
+    return atomic_data
+
+
+data_path = os.path.join("tardis", "io", "tests", "data")
+
+
+@pytest.fixture  # (scope="session")
+def tardis_model_config_nlte():
+    filename = "tardis_configv1_nlte.yml"
+    return Configuration.from_yaml(os.path.join(data_path, filename))
+
+
+@pytest.fixture  # (scope="session")
+def nlte_raw_model(tardis_model_config_nlte):
+    return Radial1DModel.from_config(tardis_model_config_nlte)