Adding rate matrix index (#2132)

* adding NlteIndexHelper class

* adding nlte_ionization_species to io/schemas/plasma.yml

* adding nlte_excitation_species to io/schemas/plasma.yml

* added nlte_ionization_species in assemble_plasma

* fixed the issue with getting nlte_ionization_species from config

* fixed transforming string from config to tuple for nlte ionization species

* ran black

* attempt of writing tests 1

* ran black

* test attempt 2

* black on test_hdf_plasma

* experimenting with revert

* attempt 3

* reverted tardis/plasma/tests/test_complete_plasmas.py

* reverted tardis/plasma/tests/test_hdf_plasma.py

* adding cofig for nlte

* adding fixture for tardis_cofig_verysimple_nlte

* added test_plasma_nlte_section_config

* ran black

* fixed the issue with the test, ran it locally

* ran black on necessary files

* updated description in schemas

* switching from nlte_ionization to be used from self

* ran black

* changed nlte_properties_new to nlte_solver_properties

* adding nlte_rate_equation_matrix.py

* reverting to previous commit

tardis/conftest.py

@@ -145,6 +145,13 @@ def tardis_config_verysimple():
     )
 
 
+@pytest.fixture(scope="function")
+def tardis_config_verysimple_nlte():
+    return yaml_load_file(
+        "tardis/io/tests/data/tardis_configv1_nlte.yml", YAMLLoader
+    )
+
+
 ###
 # HDF Fixtures
 ###

tardis/io/schemas/plasma.yml

@@ -111,6 +111,14 @@ properties:
     type: number
     default: 0.9
     description: Value used for estimating the electron temperature from radiation temperature.
+  nlte_ionization_species:
+    type: array
+    default: []
+    description: List of species treated with nlte ionization. In the format ["H I", "He II"] etc.
+  nlte_excitation_species:
+    type: array
+    default: []
+    description: List of species treated with nlte excitation. In the format ["H I", "He II"] etc.
 required:
 - ionization
 - excitation

tardis/io/tests/data/tardis_configv1_nlte.yml

@@ -0,0 +1,54 @@
+tardis_config_version: v1.0
+
+supernova:
+    luminosity_requested: 2.8e9 solLum
+    time_explosion: 13 day
+
+atom_data: kurucz_atom_pure_simple.h5
+
+model:
+    structure:
+        type: specific
+        velocity:
+            start: 1.1e4 km/s
+            stop: 2.0e4 km/s
+            num: 20
+        density:
+            type: branch85_w7
+    abundances:
+        type: uniform
+        H: 0.1
+        O: 0.09
+        Mg: 0.03
+        Si: 0.52
+        S: 0.19
+        Ar: 0.04
+        Ca: 0.03
+
+plasma:
+    ionization: lte
+    excitation: lte
+    radiative_rates_type: dilute-blackbody
+    line_interaction_type: macroatom
+    continuum_interaction:
+      species:
+        - H I
+    nlte_ionization_species: [H I]
+
+montecarlo:
+    seed: 23111963
+    no_of_packets : 2.0e+5
+    iterations: 5
+    last_no_of_packets: 5.0e+5
+    no_of_virtual_packets: 5
+    convergence_strategy:
+        type: damped
+        damping_constant: 0.5
+        threshold: 0.05
+        lock_t_inner_cycles: 1
+        t_inner_update_exponent: -0.5
+
+spectrum:
+    start: 500 angstrom
+    stop: 20000 angstrom
+    num: 10000

tardis/io/tests/test_config_reader.py

@@ -181,6 +181,46 @@ def test_plasma_section_config(tardis_config_verysimple):
     assert ve.type is ValueError
 
 
+def test_plasma_nlte_section_config(tardis_config_verysimple_nlte):
+    """
+    Configuration Validation Test for Plasma Section of the Tardis Config YAML File.
+
+    Validates:
+        nlte_ionization_species: should be included in continuum_interaction
+
+    Parameter
+    ---------
+        `tardis_config_verysimple` : 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_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
+
+
 def test_spectrum_section_config(tardis_config_verysimple):
     """
     Configuration Validation Test for Plasma Section of the Tardis Config YAML File

tardis/plasma/properties/__init__.py

@@ -10,3 +10,4 @@
 from tardis.plasma.properties.continuum_processes import *
 from tardis.plasma.properties.transition_probabilities import *
 from tardis.plasma.properties.helium_nlte import *
+from tardis.plasma.properties.rate_matrix_index import *

tardis/plasma/properties/plasma_input.py

@@ -15,6 +15,7 @@
     "TInner",
     "Volume",
     "ContinuumInteractionSpecies",
+    "NLTEIonizationSpecies",
 ]
 
 
@@ -148,3 +149,8 @@ class ContinuumInteractionSpecies(Input):
     """
 
     outputs = ("continuum_interaction_species",)
+
+
+class NLTEIonizationSpecies(Input):
+
+    outputs = ("nlte_ionization_species",)

tardis/plasma/properties/property_collections.py

@@ -16,6 +16,7 @@ class PlasmaPropertyCollection(list):
         LinkTRadTElectron,
         HeliumTreatment,
         ContinuumInteractionSpecies,
+        NLTEIonizationSpecies,
     ]
 )
 basic_properties = PlasmaPropertyCollection(
@@ -57,6 +58,7 @@ class PlasmaPropertyCollection(list):
         BetaSobolev,
     ]
 )
+nlte_solver_properties = PlasmaPropertyCollection([NLTEIndexHelper])
 helium_nlte_properties = PlasmaPropertyCollection(
     [
         HeliumNLTE,

tardis/plasma/properties/rate_matrix_index.py

@@ -0,0 +1,40 @@
+import pandas as pd
+from tardis.plasma.properties.base import ProcessingPlasmaProperty
+
+__all__ = [
+    "NLTEIndexHelper",
+]
+
+
+class NLTEIndexHelper(ProcessingPlasmaProperty):
+    outputs = ("rate_matrix_index",)
+
+    def __init__(self, plasma_parent, nlte_ionization_species=0):
+        super().__init__(plasma_parent)
+        self.nlte_ionization_species = nlte_ionization_species
+
+    def calculate(self, levels, nlte_ionization_species):
+        nlte_excitation_species = []  # not yet implemented
+        rate_matrix_index = pd.MultiIndex.from_tuples(
+            list(
+                self.calculate_rate_matrix_index(
+                    levels,
+                    self.nlte_ionization_species,
+                    nlte_excitation_species,
+                )
+            ),
+            names=levels.names,
+        ).drop_duplicates()
+        return rate_matrix_index
+
+    def calculate_rate_matrix_index(self, levels, nlte_excitation_species=[]):
+        for level in levels:
+            if level[:2] in self.nlte_ionization_species:
+                yield (*level[:2], "nlte_ion")
+            elif (level[:2] not in self.nlte_ionization_species) and (
+                level[:2] not in nlte_excitation_species
+            ):
+                yield (*level[:2], "lte_ion")
+            else:
+                yield level
+        yield ("n_e", "n_e", "n_e")

tardis/plasma/standard_plasmas.py

@@ -5,6 +5,7 @@
 import pandas as pd
 
 from tardis.io.atom_data import AtomData
+from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper
 from tardis.util.base import species_string_to_tuple
 from tardis.plasma import BasePlasma
 from tardis.plasma.properties.base import TransitionProbabilitiesProperty
@@ -28,6 +29,7 @@
     adiabatic_cooling_properties,
     two_photon_properties,
     isotope_properties,
+    nlte_solver_properties,
 )
 from tardis.plasma.exceptions import PlasmaConfigError
 
@@ -123,6 +125,11 @@ def assemble_plasma(config, model, atom_data=None):
             "in the configuration."
         )
 
+    nlte_ionization_species = [
+        species_string_to_tuple(s)
+        for s in config.plasma.nlte_ionization_species
+    ]
+
     kwargs = dict(
         t_rad=model.t_radiative,
         abundance=model.abundance,
@@ -132,6 +139,7 @@ def assemble_plasma(config, model, atom_data=None):
         w=model.dilution_factor,
         link_t_rad_t_electron=config.plasma.link_t_rad_t_electron,
         continuum_interaction_species=continuum_interaction_species,
+        nlte_ionization_species=nlte_ionization_species,
     )
 
     plasma_modules = basic_inputs + basic_properties
@@ -167,6 +175,17 @@ 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."
+                    )
+            property_kwargs[NLTEIndexHelper] = {
+                "nlte_ionization_species": nlte_ionization_species
+            }
+            plasma_modules += nlte_solver_properties
 
         kwargs.update(
             gamma_estimator=None,