diff --git a/docs/io/configuration/components/plasma.rst b/docs/io/configuration/components/plasma.rst index bf0b5b3579d..448d0a48545 100644 --- a/docs/io/configuration/components/plasma.rst +++ b/docs/io/configuration/components/plasma.rst @@ -44,4 +44,15 @@ NLTE The NLTE configuration currently allows setting ``coronal_approximation``, which sets all :math:`J_\textrm{blue}` to 0. This is useful for debugging with :term:`chianti` for example. Furthermore, one can enable 'classical_nebular' to set all :math:`\beta_\textrm{Sobolev}` to 1. Both options are used for checking with other codes and should not be enabled in -normal operations. \ No newline at end of file +normal operations. + +NLTE Ionization +^^^^^^^^^^^^^^^ + +.. code-block:: yaml + + plasma: + nlte_ionization_species: [H I, H II, He I, He II] + +This option allows the user to specify which species should be included in the NLTE ionization treatment. Note that the +species must be present in the continuum interaction species as well. diff --git a/docs/physics/setup/plasma/index.rst b/docs/physics/setup/plasma/index.rst index 6d0c26ad7ff..06c4a2203c9 100644 --- a/docs/physics/setup/plasma/index.rst +++ b/docs/physics/setup/plasma/index.rst @@ -73,6 +73,9 @@ The next more complex class is `LTEPlasma` which will calculate the ionization b TARDIS also allows for NLTE treatments of specified species, as well as special NLTE treatments for Helium. +.. note:: + The NLTE treatment of specified species is currently incompatible with the NLTE treatment for helium and cannot be used simulataneously. + .. toctree:: :maxdepth: 2 diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index 848ba3cfe10..50be58a6d8f 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -217,7 +217,9 @@ def opacity_state_initialize(plasma, line_interaction_type): ].values phot_nus = phot_nus.values - ff_opacity_factor = plasma.ff_cooling_factor / np.sqrt(t_electrons) + ff_opacity_factor = ( + plasma.ff_cooling_factor / np.sqrt(t_electrons) + ).astype(np.float64) emissivities = plasma.fb_emission_cdf.loc[ plasma.level2continuum_idx.index ].values diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py index 9718efc266f..a2e8b5f028f 100644 --- a/tardis/plasma/properties/nlte_rate_equation_solver.py +++ b/tardis/plasma/properties/nlte_rate_equation_solver.py @@ -11,7 +11,7 @@ class NLTERateEquationSolver(ProcessingPlasmaProperty): - outputs = ("ion_number_density_nlte", "electron_densities_nlte") + outputs = ("ion_number_density", "electron_densities") def calculate( self, diff --git a/tardis/plasma/properties/transition_probabilities.py b/tardis/plasma/properties/transition_probabilities.py index 9def7e650df..e69d655ef6b 100644 --- a/tardis/plasma/properties/transition_probabilities.py +++ b/tardis/plasma/properties/transition_probabilities.py @@ -41,10 +41,16 @@ def normalize_trans_probs(p): all probabilites with the same source_level_idx sum to one. Indexed by source_level_idx, destination_level_idx. """ - p_summed = p.groupby(level=0).sum() + # Dtype conversion is needed for pandas to return nan instead of + # a ZeroDivisionError in cases where the sum is zero. + p = p.astype(np.float64) + p_summed = p.groupby(level=0).sum().astype(np.float64) index = p.index.get_level_values("source_level_idx") p_norm = p / p_summed.loc[index].values p_norm = p_norm.fillna(0.0) + # Convert back to original dtypes to avoid typing problems later on + # in the numba code. + p_norm = p_norm.convert_dtypes() return p_norm @@ -78,7 +84,9 @@ def series2matrix(series, idx2reduced_idx): idx2reduced_idx.loc[q_indices[1]].values, ) max_idx = idx2reduced_idx.max() + 1 - matrix = sp.coo_matrix((series, q_indices), shape=(max_idx, max_idx)) + matrix = sp.coo_matrix( + (series.astype(np.float64), q_indices), shape=(max_idx, max_idx) + ) return matrix @staticmethod diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 63c6ac2324e..82bbc54a263 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -6,6 +6,10 @@ import pandas as pd from tardis.io.atom_data import AtomData +from tardis.plasma.properties.level_population import LevelNumberDensity +from tardis.plasma.properties.nlte_rate_equation_solver import ( + NLTERateEquationSolver, +) from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper from tardis.util.base import species_string_to_tuple from tardis.plasma import BasePlasma @@ -265,6 +269,26 @@ def assemble_plasma(config, simulation_state, atom_data=None): delta_treatment=config.plasma.delta_treatment ) + if ( + config.plasma.helium_treatment == "recomb-nlte" + or config.plasma.helium_treatment == "numerical-nlte" + ) and ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + # Prevent the user from using helium NLTE treatment with + # NLTE ionization and excitation treatment. This is because + # the helium_nlte_properties could overwrite the NLTE ionization + # and excitation ion number and electron densities. + # helium_numerical_nlte_properties is also included here because + # it is currently in the same if else block, and thus may block + # the addition of the components from the else block. + raise PlasmaConfigError( + "Helium NLTE treatment is incompatible with the NLTE eonization and excitation treatment." + ) + + # TODO: Disentangle these if else block such that compatible components + # can be added independently. if config.plasma.helium_treatment == "recomb-nlte": plasma_modules += helium_nlte_properties elif config.plasma.helium_treatment == "numerical-nlte": @@ -277,7 +301,16 @@ def assemble_plasma(config, simulation_state, atom_data=None): heating_rate_data_file=config.plasma.heating_rate_data_file ) else: - plasma_modules += helium_lte_properties + # If nlte ionization species are present, we don't want to add the + # IonNumberDensity from helium_lte_properties, since we want + # to use the IonNumberDensity provided by the NLTE solver. + if ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + plasma_modules += [LevelNumberDensity] + else: + plasma_modules += helium_lte_properties if simulation_state._electron_densities is not None: electron_densities = pd.Series( @@ -287,6 +320,13 @@ def assemble_plasma(config, simulation_state, atom_data=None): property_kwargs[IonNumberDensityHeNLTE] = dict( electron_densities=electron_densities ) + elif ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + property_kwargs[NLTERateEquationSolver] = dict( + electron_densities=electron_densities + ) else: property_kwargs[IonNumberDensity] = dict( electron_densities=electron_densities diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index fe47b496e97..beeb345cf6d 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -260,7 +260,7 @@ def nlte_raw_plasma_w0( def test_critical_case_w1(nlte_raw_plasma_w1): """Check that the LTE and NLTE solution agree for w=1.0.""" - ion_number_density_nlte = nlte_raw_plasma_w1.ion_number_density_nlte.values + ion_number_density_nlte = nlte_raw_plasma_w1.ion_number_density.values ion_number_density_nlte[ion_number_density_nlte < 1e-10] = 0.0 ind = IonNumberDensity(nlte_raw_plasma_w1)