Adding nlte solver

tardis/io/tests/data/tardis_configv1_nlte.yml

@@ -1,54 +1,71 @@
 tardis_config_version: v1.0
 
 supernova:
-    luminosity_requested: 2.8e9 solLum
-    time_explosion: 13 day
+  luminosity_requested: 9.05 log_lsun
+  time_explosion: 16 day
 
-atom_data: kurucz_atom_pure_simple.h5
+atom_data: TestNLTE_Ti.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
+  structure:
+    type: specific
+    velocity:
+      start: 6395 km/s
+      stop: 12500 km/s
+      num: 20
+    density:
+      type : power_law
+      time_0: 16.0 day
+      rho_0: 1.948e-14 g/cm^3
+      v_0: 8000 km/s
+      exponent: -10
+
+  abundances:
+    type: uniform
+    H: 0.9
+    He: 0.099 #0.3
+    Ti: 1e-3
+#     Fe: 1e-3
 
 plasma:
-    ionization: lte
-    excitation: lte
-    radiative_rates_type: dilute-blackbody
-    line_interaction_type: macroatom
-    continuum_interaction:
-      species:
-        - H I
-    nlte_ionization_species: [H I]
+  disable_electron_scattering: no
+  ionization: lte
+  excitation: lte
+  radiative_rates_type: dilute-blackbody
+  line_interaction_type: macroatom
+  initial_t_inner: 12000 K
+  link_t_rad_t_electron: 1.0
+  nlte_ionization_species: [H I, He II, Ti II]
+
+  continuum_interaction:
+    species:
+      - H I
+      - Ti II
+      - He II
+
+    enable_adiabatic_cooling: True
+
+
 
 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
+  seed: 23111963
+  no_of_packets: 100000
+  iterations: 1
+  nthreads: 24
+
+  last_no_of_packets: 1000
+  no_of_virtual_packets: 0
+
+  convergence_strategy:
+    type: damped
+    damping_constant: 0.5
+    threshold: 0.05
+    fraction: 0.8
+    hold_iterations: 3
 
 spectrum:
-    start: 500 angstrom
-    stop: 20000 angstrom
-    num: 10000
+  start: 250 angstrom
+  stop: 10000 angstrom
+  num: 10000
+  method: integrated
+  compute: Automatic

tardis/plasma/properties/nlte_rate_equation_solver.py

@@ -1,5 +1,6 @@
 import pandas as pd
 import numpy as np
+from scipy.optimize import root
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
 
@@ -81,40 +82,49 @@ def calculate(
         )
 
         # >>>TODO:initial electron density should be included in the initial guess, added in a future PR
-        initial_electron_density = number_density.sum(axis=0)
+        initial_electron_densities = number_density.sum(axis=0)
         # <<<
         atomic_numbers = (
             rate_matrix_index.get_level_values("atomic_number")
             .unique()
             .drop("n_e")
         )  # dropping the n_e index, as rate_matrix_index's first index is (atomic_numbers, "n_e")
-        rate_matrix = self.calculate_rate_matrix(
-            atomic_numbers,
-            phi[0],
-            initial_electron_density[0],
-            rate_matrix_index,
-            total_photo_ion_coefficients[0],
-            total_rad_recomb_coefficients[0],
-            total_coll_ion_coefficients[0],
-            total_coll_recomb_coefficients[0],
-        )
-        initial_guess = self.prepare_first_guess(
-            atomic_numbers, number_density[0], initial_electron_density[0]
-        )
-        jacobian_matrix = self.jacobian_matrix(
-            atomic_numbers,
-            initial_guess,
-            rate_matrix,
-            rate_matrix_index,
-            total_rad_recomb_coefficients[0],
-            total_coll_ion_coefficients[0],
-            total_coll_recomb_coefficients[0],
-        )
-        # TODO: change the jacobian and rate matrix to use shell id and get coefficients from the attribute of the class.
 
-        raise NotImplementedError(
-            "NLTE ionization hasn't been fully implemented yet!"
+        index = rate_matrix_index.droplevel("level_number").drop("n_e")
+        ion_number_density_nlte = pd.DataFrame(
+            0.0, index=index, columns=phi.columns
         )
+        electron_densities_nlte = pd.Series(0.0, index=phi.columns)
+
+        for i, shell in enumerate(phi.columns):
+            solution_vector = self.prepare_solution_vector(
+                number_density[shell]
+            )
+            first_guess = self.prepare_first_guess(
+                atomic_numbers,
+                number_density[shell],
+                initial_electron_densities[shell],
+            )
+            solution = root(
+                self.population_objective_function,
+                first_guess,
+                args=(
+                    atomic_numbers,
+                    phi[shell],
+                    solution_vector,
+                    rate_matrix_index,
+                    total_photo_ion_coefficients[shell],
+                    total_rad_recomb_coefficients[shell],
+                    total_coll_ion_coefficients[shell],
+                    total_coll_recomb_coefficients[shell],
+                ),
+                jac=True,
+            )
+            assert solution.success
+            ion_number_density_nlte[i] = solution.x[:-1]
+            electron_densities_nlte[i] = solution.x[-1]
+        # TODO: change the jacobian and rate matrix to use shell id and get coefficients from the attribute of the class.
+        return ion_number_density_nlte, electron_densities_nlte
 
     @staticmethod
     def calculate_rate_matrix(
@@ -549,6 +559,22 @@ def deriv_matrix_block(
     def prepare_first_guess(
         self, atomic_numbers, number_density, electron_density
     ):
+        """Constructs a first guess for ion number densities and electron density, where all species are fully once ionized.
+
+        Parameters
+        ----------
+        atomic_numbers : numpy.array
+            All atomic numbers present in the plasma.
+        number_density : DataFrame
+            Number density of present species.
+        electron_density : float
+            Current value of electron density.
+
+        Returns
+        -------
+        numpy.array
+            Guess for ion number densities and electron density.
+        """
         # TODO needs to be changed for excitation
         array_size = (number_density.index.values + 1).sum() + 1
         first_guess = np.zeros(array_size)
@@ -558,3 +584,114 @@ def prepare_first_guess(
             index += atomic_number + 1
         first_guess[-1] = electron_density
         return first_guess
+
+    def population_objective_function(
+        self,
+        populations,
+        atomic_numbers,
+        phi,
+        solution_vector,
+        rate_matrix_index,
+        total_photo_ion_coefficients,
+        total_rad_recomb_coefficients,
+        total_coll_ion_coefficients,
+        total_coll_recomb_coefficients,
+    ):
+        """Main set of equations for the NLTE ionization solver. A*x - B = 0, where x is the populations,
+        A is the matrix of rates and B is the solution vector.
+
+        Parameters
+        ----------
+        populations : numpy.array
+            Current values of ion number densities and electron density.
+        atomic_numbers : numpy.array
+            All atomic numbers present in the plasma.
+        phi :DataFrame
+            Saha Factors of the current shell.
+        solution_vector : numpy.array
+            Solution vector for the set of equations.
+        rate_matrix_index : MultiIndex
+            (atomic_number, ion_number, treatment type)
+            If ion is treated in LTE or nebular ionization, 3rd index is "lte_ion",
+            if treated in NLTE ionization, 3rd index is "nlte_ion".
+        total_photo_ion_coefficients : DataFrame
+            Photo ion. coefficients for current atomic number
+        total_rad_recomb_coefficients : DataFrame
+            Rad. recomb. coefficients for current atomic number
+        total_coll_ion_coefficients : DataFrame
+            Coll. ion. coefficients for current atomic number
+        total_coll_recomb_coefficients : DataFrame
+            Coll. recomb. coefficients for current atomic number
+        Returns
+        -------
+        _type_
+            _description_
+        """
+        electron_density = populations[-1]
+        rate_matrix = self.calculate_rate_matrix(
+            atomic_numbers,
+            phi,
+            electron_density,
+            rate_matrix_index,
+            total_photo_ion_coefficients,
+            total_rad_recomb_coefficients,
+            total_coll_ion_coefficients,
+            total_coll_recomb_coefficients,
+        )
+        jacobian_matrix = self.jacobian_matrix(
+            atomic_numbers,
+            populations,
+            rate_matrix,
+            rate_matrix_index,
+            total_rad_recomb_coefficients,
+            total_coll_ion_coefficients,
+            total_coll_recomb_coefficients,
+        )
+        return (
+            np.dot(rate_matrix.values, populations) - solution_vector,
+            jacobian_matrix,
+        )
+
+    def solution_vector_block(self, atomic_number, number_density):
+        """Block of the solution vector corresponding to the current atomic number. (0, 0, ..., number_density). Length is equal to atomic_number+1.
+
+        Parameters
+        ----------
+        atomic_number : int
+            Current atomic number.
+        number_density : float
+            Number density of the current atomic number.
+
+        Returns
+        -------
+        numpy.array
+            Block of the solution vector corresponding to the current atomic number.
+        """
+        solution_vector = np.zeros(atomic_number + 1)
+        solution_vector[-1] = number_density
+        return solution_vector
+
+    def prepare_solution_vector(self, number_density):
+        """Constructs the solution vector for the NLTE ionization solver set of equations by combining
+        all solution verctor blocks.
+
+        Parameters
+        ----------
+        number_density : DataFrame
+            Number densities of all present species.
+
+        Returns
+        -------
+        numpy.array
+            Solution vector for the NLTE ionization solver.
+        """
+        atomic_numbers = number_density.index
+        solution_array = []
+        for atomic_number in atomic_numbers:
+            solution_array.append(
+                self.solution_vector_block(
+                    atomic_number, number_density.loc[atomic_number]
+                )
+            )
+        solution_vector = np.hstack(solution_array + [0])
+        return solution_vector

tardis/plasma/properties/transition_probabilities.py

@@ -357,7 +357,9 @@ def calculate(
         deactivation_channel_probs = p_deactivation.copy()
         deactivation_channel_probs = pd.concat(
             [
-                level_idxs2transition_idx.loc[deactivation_channel_probs.index],
+                level_idxs2transition_idx.reindex(
+                    deactivation_channel_probs.index
+                ),
                 deactivation_channel_probs,
             ],
             axis=1,