From d0c567e8cf45c768f85a69dc0d1a6fec5e1045f9 Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Tue, 31 Oct 2023 21:52:37 -0400 Subject: [PATCH 01/11] Rename nlte ion and electron --- tardis/plasma/properties/nlte_rate_equation_solver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py index d76817b9a77..6781dd546c0 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, From 24cdbc377a503ea36dd9253aa02169d66b190dbd Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Tue, 31 Oct 2023 22:18:31 -0400 Subject: [PATCH 02/11] Enable NLTE --- tardis/plasma/standard_plasmas.py | 32 ++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 63c6ac2324e..fefce3e7a54 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -6,6 +6,7 @@ import pandas as pd from tardis.io.atom_data import AtomData +from tardis.plasma.properties.level_population import LevelNumberDensity 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 +266,26 @@ def assemble_plasma(config, simulation_state, atom_data=None): delta_treatment=config.plasma.delta_treatment ) + if ( + config.plama.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 +298,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( From 44fe04a813d39dfe027c45239b2957b8abb3e54e Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Tue, 31 Oct 2023 22:34:57 -0400 Subject: [PATCH 03/11] Fix typo --- tardis/plasma/standard_plasmas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index fefce3e7a54..b9cd5bcb243 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -267,7 +267,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): ) if ( - config.plama.helium_treatment == "recomb-nlte" + config.plasma.helium_treatment == "recomb-nlte" or config.plasma.helium_treatment == "numerical-nlte" ) and ( config.plasma.nlte_ionization_species From d62ccfd2555efc54394b3f75a02103f60451e820 Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Tue, 31 Oct 2023 22:43:28 -0400 Subject: [PATCH 04/11] Add docs --- docs/io/configuration/components/plasma.rst | 13 ++++++++++++- docs/physics/setup/plasma/index.rst | 3 +++ 2 files changed, 15 insertions(+), 1 deletion(-) 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 From cfc039d02d9f568f38278490b453b8ded6979b0f Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Wed, 1 Nov 2023 11:19:58 -0400 Subject: [PATCH 05/11] Fix transitionprob dtype --- tardis/plasma/properties/transition_probabilities.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tardis/plasma/properties/transition_probabilities.py b/tardis/plasma/properties/transition_probabilities.py index 9def7e650df..def893e552a 100644 --- a/tardis/plasma/properties/transition_probabilities.py +++ b/tardis/plasma/properties/transition_probabilities.py @@ -41,10 +41,14 @@ 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) + breakpoint() return p_norm From f40239e13e4814dd9ad94371f40d88d446012d73 Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Wed, 1 Nov 2023 11:28:32 -0400 Subject: [PATCH 06/11] dtype conversion --- tardis/plasma/properties/transition_probabilities.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tardis/plasma/properties/transition_probabilities.py b/tardis/plasma/properties/transition_probabilities.py index def893e552a..31de8ad708f 100644 --- a/tardis/plasma/properties/transition_probabilities.py +++ b/tardis/plasma/properties/transition_probabilities.py @@ -48,7 +48,9 @@ def normalize_trans_probs(p): index = p.index.get_level_values("source_level_idx") p_norm = p / p_summed.loc[index].values p_norm = p_norm.fillna(0.0) - breakpoint() + # Convert back to original dtypes to avoid typing problems later on + # in the number code. + p_norm = p_norm.convert_dtypes() return p_norm From 36189ee4ccdd7391f92e8c427a717ccad9ad2364 Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Wed, 1 Nov 2023 11:36:48 -0400 Subject: [PATCH 07/11] dtype for sparse matrix --- tardis/plasma/properties/transition_probabilities.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tardis/plasma/properties/transition_probabilities.py b/tardis/plasma/properties/transition_probabilities.py index 31de8ad708f..c36e18a0cf6 100644 --- a/tardis/plasma/properties/transition_probabilities.py +++ b/tardis/plasma/properties/transition_probabilities.py @@ -84,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 From 8eaf6bb394d54b8f17af9daebcf3861efb5750c6 Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Wed, 1 Nov 2023 11:49:52 -0400 Subject: [PATCH 08/11] Typing in numba --- tardis/montecarlo/montecarlo_numba/numba_interface.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 From a66bf8e91f8da483b3d011cd47739ebd68fc71f1 Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Wed, 1 Nov 2023 12:01:48 -0400 Subject: [PATCH 09/11] Fix test variable name --- tardis/plasma/tests/test_nlte_solver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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) From 36c0e4eb8e24b8b5303d5307c418ae7194d1a9da Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Wed, 1 Nov 2023 13:00:40 -0400 Subject: [PATCH 10/11] Add missing electron density case --- tardis/plasma/standard_plasmas.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index b9cd5bcb243..82bbc54a263 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -7,6 +7,9 @@ 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 @@ -317,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 From a289d014adf7d96c695fb50bd286b4bf9038359c Mon Sep 17 00:00:00 2001 From: Alexander Holas Date: Wed, 1 Nov 2023 13:01:37 -0400 Subject: [PATCH 11/11] Fix typo --- tardis/plasma/properties/transition_probabilities.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/plasma/properties/transition_probabilities.py b/tardis/plasma/properties/transition_probabilities.py index c36e18a0cf6..e69d655ef6b 100644 --- a/tardis/plasma/properties/transition_probabilities.py +++ b/tardis/plasma/properties/transition_probabilities.py @@ -49,7 +49,7 @@ def normalize_trans_probs(p): 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 number code. + # in the numba code. p_norm = p_norm.convert_dtypes() return p_norm