diff --git a/tardis/io/tests/test_config_reader.py b/tardis/io/tests/test_config_reader.py index 7574db68893..ec869a51bb8 100644 --- a/tardis/io/tests/test_config_reader.py +++ b/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): diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 4cf2304a8f4..2d7b67ca4d6 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/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",) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 006083b5585..9716bda572b 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -17,6 +17,7 @@ class PlasmaPropertyCollection(list): HeliumTreatment, ContinuumInteractionSpecies, NLTEIonizationSpecies, + NLTEExcitationSpecies, ] ) basic_properties = PlasmaPropertyCollection( diff --git a/tardis/plasma/properties/rate_matrix_index.py b/tardis/plasma/properties/rate_matrix_index.py index e6a154d926b..ab16121b57f 100644 --- a/tardis/plasma/properties/rate_matrix_index.py +++ b/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. diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 081396f8147..4961d0c6dd6 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/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 diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index 0283ba4f61a..7062a1d01cb 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/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 diff --git a/tardis/tests/fixtures/atom_data.py b/tardis/tests/fixtures/atom_data.py index 520ede4b980..a1cdd8456d6 100644 --- a/tardis/tests/fixtures/atom_data.py +++ b/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)