diff --git a/.gitignore b/.gitignore index f8cc8cf6295..7b0c1a23a4a 100644 --- a/.gitignore +++ b/.gitignore @@ -71,3 +71,9 @@ pip-wheel-metadata/ # Mac OSX .DS_Store + +# Visual Studio Code +*.code-workspace + +# Random files +.hypothesis/unicode_data/11.0.0/charmap.json.gz diff --git a/.mailmap b/.mailmap index 0525b2ec384..09eccfe3b4f 100644 --- a/.mailmap +++ b/.mailmap @@ -66,6 +66,9 @@ Dhruv Sondhi DhruvSondhi Epson Heringer Epson Heringer Heringer-Epson +Erin Visser +Erin Visser erinvisser + Ezequiel Pássaro Frederik Beaujean diff --git a/.zenodo.json b/.zenodo.json index 8720d4b208a..7bee94b7df7 100644 --- a/.zenodo.json +++ b/.zenodo.json @@ -113,28 +113,28 @@ "orcid": "0000-0002-8310-0829" }, { - "name": "Sondhi, Dhruv" + "name": "Shields, Joshua" }, { - "name": "O'Brien, Jack" + "name": "Barbosa, Talytha" }, { - "name": "Barbosa, Talytha" + "name": "Sondhi, Dhruv" }, { - "name": "Yu, Jenny" + "name": "O'Brien, Jack" }, { - "name": "Shields, Joshua" + "name": "Yu, Jenny" }, { "name": "Patel, Maryam" }, { - "name": "Varanasi, Kaushik" + "name": "Rathi, Shikha" }, { - "name": "Rathi, Shikha" + "name": "Varanasi, Kaushik" }, { "name": "Chitchyan, Sona" @@ -154,38 +154,38 @@ "name": "Reinecke, Martin" }, { - "name": "Eweis, Youssef" + "name": "Holas, Alexander" }, { - "name": "Holas, Alexander" + "name": "Eweis, Youssef" }, { "name": "Bylund, Tomas" }, { - "name": "Black, William" + "name": "Bentil, Laud" }, { - "name": "Bentil, Laud" + "name": "Black, William" }, { "name": "Eguren, Jordi", "orcid": "0000-0002-2328-8030" }, { - "name": "Bartnik, Matthew" + "name": "Kumar, Ansh" }, { "name": "Alam, Arib" }, { - "name": "Kumar, Ansh" + "name": "Bartnik, Matthew" }, { - "name": "Varma Buddaraju, Rohith" + "name": "Magee, Mark" }, { - "name": "Magee, Mark" + "name": "Varma Buddaraju, Rohith" }, { "name": "Kambham, Satwik" @@ -193,9 +193,6 @@ { "name": "Livneh, Ran" }, - { - "name": "Rajagopalan, Srinath" - }, { "name": "Daksh, Ayushi" }, @@ -204,37 +201,46 @@ "orcid": "0000-0001-8302-1584" }, { - "name": "Dutta, Anirban" + "name": "Rajagopalan, Srinath" }, { - "name": "Floers, Andreas" + "name": "Dutta, Anirban" }, { "name": "Jain, Rinkle" }, { - "name": "Reichenbach, John" + "name": "Actions, GitHub" }, { - "name": "Actions, GitHub" + "name": "Floers, Andreas" + }, + { + "name": "Reichenbach, John" }, { "name": "Bhakar, Jayant" }, { - "name": "Brar, Antreev" + "name": "Singh, Sourav" }, { "name": "Chaumal, Aarya" }, { - "name": "Singh, Sourav" + "name": "Brar, Antreev" }, { - "name": "Selsing, Jonatan" + "name": "Lu, Jing" }, { - "name": "Kowalski, Nathan" + "name": "Matsumura, Yuki" + }, + { + "name": "Talegaonkar, Chinmay" + }, + { + "name": "Patidar, Abhishek" }, { "name": "Kumar, Aman" @@ -243,22 +249,19 @@ "name": "Gupta, Harshul" }, { - "name": "Talegaonkar, Chinmay" + "name": "Kowalski, Nathan" }, { - "name": "Matsumura, Yuki" + "name": "Selsing, Jonatan" }, { "name": "Sofiatti, Caroline" }, { - "name": "Patidar, Abhishek" + "name": "Visser, Erin" }, { - "name": "Venkat, Shashank" - }, - { - "name": "Buchner, Johannes" + "name": "Prasad, Shilpi" }, { "name": "Yap, Kevin" @@ -279,10 +282,10 @@ "name": "Sarafina, Nance" }, { - "name": "Volodin, Dmitry" + "name": "Patra, Nilesh" }, { - "name": "Patra, Nilesh" + "name": "Singh Rathore, Parikshit" }, { "name": "Patel, Pratik" @@ -291,10 +294,10 @@ "name": "Sharma, Sampark" }, { - "name": "Lu, Jing" + "name": "Venkat, Shashank" }, { - "name": "Prasad, Shilpi" + "name": "Buchner, Johannes" }, { "name": "Gupta, Suyash" @@ -309,7 +312,7 @@ "name": "Aggarwal, Yash" }, { - "name": "Singh Rathore, Parikshit" + "name": "Volodin, Dmitry" }, { "name": "Dasgupta, Debajyoti" @@ -324,10 +327,10 @@ "name": "Kolliboyina, Chaitanya" }, { - "name": "Kumar, Atul" + "name": "Kharkar, Atharwa" }, { - "name": "Kharkar, Atharwa" + "name": "Kumar, Atul" } ] } \ No newline at end of file diff --git a/astropy_helpers b/astropy_helpers new file mode 160000 index 00000000000..9f82aac6c21 --- /dev/null +++ b/astropy_helpers @@ -0,0 +1 @@ +Subproject commit 9f82aac6c2141b425e2d639560f7260189d90b54 diff --git a/docs/installation.rst b/docs/installation.rst index 8c2c54a956b..25c077d9618 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -30,7 +30,7 @@ this method by following the steps described below. :: - $ wget -q https://github.com/tardis-sn/tardis/releases/latest/download/conda-{platform}-64.lock + $ wget -q https://raw.githubusercontent.com/tardis-sn/tardis/master/conda-{platform}-64.lock 2. Create and activate the ``tardis`` environment. @@ -39,29 +39,7 @@ this method by following the steps described below. $ conda create --name tardis --file conda-{platform}-64.lock $ conda activate tardis -3. a. Non-developers can install the latest release from ``conda-forge`` with the ``--no-deps`` flag, - - .. warning:: - - Currently the conda forge installation doesn't work. It's recommended to install from the specific releases using pip- - - `$ pip install git+https://github.com/tardis-sn/tardis.git@{tag}` - - For example- - - `pip install git+https://github.com/tardis-sn/tardis.git@release-2023.09.17` - - :: - - $ conda install tardis-sn --channel conda-forge --no-deps - - or trying the most recent, unreleased changes from upstream. - - :: - - $ pip install git+https://github.com/tardis-sn/tardis.git@master - - b. Instead, developers should `fork the repository `_, configure +3. a. Developers should `fork the repository `_, configure GitHub to `work with SSH keys `_, set up the `upstream remote `_, and install the package in development mode. @@ -78,7 +56,24 @@ this method by following the steps described below. .. note:: The complete developer guidelines can be found :ref:`here `. + + b. Non-developers can install from specific releases using pip- + + :: + + $ pip install git+https://github.com/tardis-sn/tardis.git@{tag} + + For example- + + :: + + $ pip install git+https://github.com/tardis-sn/tardis.git@release-latest + + or trying the most recent, unreleased changes from upstream. + :: + + $ pip install git+https://github.com/tardis-sn/tardis.git@master 4. Once finished working, you can deactivate your environment. @@ -111,6 +106,7 @@ To update the environment after a new release, download the latest lockfile and :: + $ wget -q https://github.com/tardis-sn/tardis/releases/latest/download/conda-{platform}-64.lock $ conda update --name tardis --file conda-{platform}-64.lock .. note:: diff --git a/docs/io/configuration/components/plasma.rst b/docs/io/configuration/components/plasma.rst index 4a27598d2fe..7ff017c9776 100644 --- a/docs/io/configuration/components/plasma.rst +++ b/docs/io/configuration/components/plasma.rst @@ -72,6 +72,22 @@ NLTE Ionization plasma: nlte_ionization_species: [H I, H II, He I, He II] + nlte_solver: root 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. +Here, ``nlte_solver`` can be set to ``root`` or ``lu``. ``root`` is the default and uses a root solver to calculate the +NLTE populations. ``lu`` uses an iterative LU decomposition scheme to calculate the NLTE populations. + +.. note :: + + ``lu`` iterates over the solutions up to a set tolerance. This tolerance is currently hard-coded to 1e-3. This + can be changed in the code by changing the ``NLTE_POPULATION_SOLVER_TOLERANCE`` constant in ``tardis/plasma/properties/nlte_rate_equation_solver.py``. + Furthermore, the maximum number of iterations is set to 1000. This can be changed in the code by changing the ``NLTE_POPULATION_SOLVER_MAX_ITERATIONS`` + constant in ``tardis/plasma/properties/nlte_rate_equation_solver.py``. + +.. warning :: + + ``lu`` is generally faster than ``root`` but does not solve explicitly for the electron density. Therefore, it is + not recommended to use ``lu`` for simulations where the electron density is important (e.g. for simulations where + NLTE excitation is important). diff --git a/docs/physics/montecarlo/initialization.ipynb b/docs/physics/montecarlo/initialization.ipynb index 2ab262e6da1..99028dbac34 100644 --- a/docs/physics/montecarlo/initialization.ipynb +++ b/docs/physics/montecarlo/initialization.ipynb @@ -237,7 +237,7 @@ "def planck_function(nu):\n", " return (\n", " 8\n", - " * np.pi**2\n", + " * np.pi\n", " * r_boundary_inner**2\n", " * h\n", " * nu**3\n", diff --git a/docs/physics/update_and_conv/update_and_conv.ipynb b/docs/physics/update_and_conv/update_and_conv.ipynb index b0c8da0d24c..cb8f874d945 100644 --- a/docs/physics/update_and_conv/update_and_conv.ipynb +++ b/docs/physics/update_and_conv/update_and_conv.ipynb @@ -311,7 +311,7 @@ }, "outputs": [], "source": [ - "j_estimator = transport.transport_state.estimators.j_estimator * (u.erg * u.cm) \n", + "j_estimator = transport.transport_state.j_estimator * (u.erg * u.cm) \n", "j_estimator" ] }, @@ -324,7 +324,7 @@ }, "outputs": [], "source": [ - "nu_bar_estimator = transport.transport_state.estimators.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n", + "nu_bar_estimator = transport.transport_state.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n", "nu_bar_estimator" ] }, diff --git a/tardis/energy_input/tests/test_energy_source.py b/tardis/energy_input/tests/test_energy_source.py index c6eb1e26226..106124f1a68 100644 --- a/tardis/energy_input/tests/test_energy_source.py +++ b/tardis/energy_input/tests/test_energy_source.py @@ -1,6 +1,6 @@ -import pytest import numpy as np import numpy.testing as npt +import pytest from tardis.energy_input.samplers import ( create_energy_cdf, diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index d5cc7dc79ce..47af303d0c9 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -1,16 +1,17 @@ -# atomic model - - import logging +from collections import OrderedDict + import numpy as np import pandas as pd - -from scipy import interpolate -from collections import OrderedDict from astropy import units as u -from tardis import constants as const from astropy.units import Quantity +from scipy import interpolate + +from tardis import constants as const from tardis.io.atom_data.util import resolve_atom_data_fname +from tardis.plasma.properties.continuum_processes import ( + get_ground_state_multi_index, +) class AtomDataNotPreparedError(Exception): @@ -24,7 +25,7 @@ class AtomDataMissingError(Exception): logger = logging.getLogger(__name__) -class AtomData(object): +class AtomData: """ Class for storing atomic data @@ -167,7 +168,6 @@ def from_hdf(cls, fname=None): Path to the HDFStore file or name of known atom data file (default: None) """ - dataframes = dict() nonavailable = list() @@ -208,7 +208,7 @@ def from_hdf(cls, fname=None): "Atomic Data Collisions Not a Valid Chanti or CMFGEN Carsus Data File" ) except KeyError as e: - logger.warn( + logger.warning( "Atomic Data is not a Valid Carsus Atomic Data File" ) raise @@ -253,7 +253,7 @@ def from_hdf(cls, fname=None): ) atom_data.version = None - # ToDo: strore data sources as attributes in carsus + # TODO: strore data sources as attributes in carsus logger.info( f"Reading Atom Data with: UUID = {atom_data.uuid1} MD5 = {atom_data.md5} " @@ -373,8 +373,9 @@ def _check_related(self): def prepare_atom_data( self, selected_atomic_numbers, - line_interaction_type="scatter", - nlte_species=[], + line_interaction_type, + nlte_species, + continuum_interaction_species, ): """ Prepares the atom data to set the lines, levels and if requested macro @@ -437,6 +438,84 @@ def prepare_atom_data( .values ) + self.prepare_macro_atom_data( + line_interaction_type, + tmp_lines_lower2level_idx, + tmp_lines_upper2level_idx, + ) + if len(continuum_interaction_species) > 0: + self.prepare_continuum_interaction_data( + continuum_interaction_species + ) + + self.nlte_data = NLTEData(self, nlte_species) + + def prepare_continuum_interaction_data(self, continuum_interaction_species): + """ + Prepares the atom data for the continuum interaction + + Parameters + ---------- + continuum_interaction : ContinuumInteraction + The continuum interaction object + """ + # photoionization_data = atomic_data.photoionization_data.set_index( + # ["atomic_number", "ion_number", "level_number"] + # ) + mask_selected_species = self.photoionization_data.index.droplevel( + "level_number" + ).isin(continuum_interaction_species) + self.photoionization_data = self.photoionization_data[ + mask_selected_species + ] + self.photo_ion_block_references = np.pad( + self.photoionization_data.nu.groupby(level=[0, 1, 2]) + .count() + .values.cumsum(), + [1, 0], + ) + self.photo_ion_unique_index = self.photoionization_data.index.unique() + self.nu_ion_threshold = ( + self.photoionization_data.groupby(level=[0, 1, 2]).first().nu + ) + self.energy_photo_ion_lower_level = self.levels.loc[ + self.photo_ion_unique_index + ].energy + + source_idx = self.macro_atom_references.loc[ + self.photo_ion_unique_index + ].references_idx + destination_idx = self.macro_atom_references.loc[ + get_ground_state_multi_index(self.photo_ion_unique_index) + ].references_idx + self.photo_ion_levels_idx = pd.DataFrame( + { + "source_level_idx": source_idx.values, + "destination_level_idx": destination_idx.values, + }, + index=self.photo_ion_unique_index, + ) + + self.level2continuum_edge_idx = pd.Series( + np.arange(len(self.nu_ion_threshold)), + self.nu_ion_threshold.sort_values(ascending=False).index, + name="continuum_idx", + ) + + level_idxs2continuum_idx = self.photo_ion_levels_idx.copy() + level_idxs2continuum_idx[ + "continuum_idx" + ] = self.level2continuum_edge_idx + self.level_idxs2continuum_idx = level_idxs2continuum_idx.set_index( + ["source_level_idx", "destination_level_idx"] + ) + + def prepare_macro_atom_data( + self, + line_interaction_type, + tmp_lines_lower2level_idx, + tmp_lines_upper2level_idx, + ): if ( self.macro_atom_data_all is not None and not line_interaction_type == "scatter" @@ -505,6 +584,13 @@ def prepare_atom_data( ) if line_interaction_type == "macroatom": + self.lines_lower2macro_reference_idx = ( + self.macro_atom_references.loc[ + tmp_lines_lower2level_idx, "references_idx" + ] + .astype(np.int64) + .values + ) # Sets all tmp_macro_destination_level_idx = pd.MultiIndex.from_arrays( [ @@ -548,8 +634,6 @@ def prepare_atom_data( self.selected_atomic_numbers, level=0 ) - self.nlte_data = NLTEData(self, nlte_species) - def _check_selected_atomic_numbers(self): selected_atomic_numbers = self.selected_atomic_numbers available_atomic_numbers = np.unique( @@ -571,7 +655,7 @@ def __repr__(self): return f"" -class NLTEData(object): +class NLTEData: def __init__(self, atom_data, nlte_species): self.atom_data = atom_data self.lines = atom_data.lines.reset_index() diff --git a/tardis/io/configuration/schemas/plasma.yml b/tardis/io/configuration/schemas/plasma.yml index 9d39845db9d..47189725bc0 100644 --- a/tardis/io/configuration/schemas/plasma.yml +++ b/tardis/io/configuration/schemas/plasma.yml @@ -119,6 +119,13 @@ properties: type: array default: [] description: List of species treated with nlte excitation. In the format ["H I", "He II"] etc. + nlte_solver: + type: string + default: root + enum: + - root + - lu + description: Selects NLTE population equation solver approach. required: - ionization - excitation diff --git a/tardis/io/configuration/tests/data/tardis_configv1_nlte.yml b/tardis/io/configuration/tests/data/tardis_configv1_nlte.yml index ee3c640e4f5..9e232f3c140 100644 --- a/tardis/io/configuration/tests/data/tardis_configv1_nlte.yml +++ b/tardis/io/configuration/tests/data/tardis_configv1_nlte.yml @@ -35,6 +35,7 @@ plasma: initial_t_inner: 12000 K link_t_rad_t_electron: 1.0 nlte_ionization_species: [H I, He II, Ti II] + nlte_solver: root continuum_interaction: species: diff --git a/tardis/io/configuration/tests/test_config_reader.py b/tardis/io/configuration/tests/test_config_reader.py index 1fb373cf6b2..f525557ca63 100644 --- a/tardis/io/configuration/tests/test_config_reader.py +++ b/tardis/io/configuration/tests/test_config_reader.py @@ -169,9 +169,9 @@ def test_plasma_section_config(key, tardis_config_verysimple): ) -def test_plasma_nlte_section_config( +def test_plasma_nlte_section_root_config( tardis_config_verysimple_nlte, - nlte_raw_simulation_state, + nlte_raw_model_root, nlte_atom_data, ): """ @@ -182,7 +182,7 @@ def test_plasma_nlte_section_config( Parameter --------- - `tardis_config_verysimple_nlte` : YAML File + `tardis_config_verysimple_nlte_root` : YAML File `nlte_raw_model` : A simple model `nlte_atom_data` : An example atomic dataset @@ -198,11 +198,44 @@ def test_plasma_nlte_section_config( tardis_config_verysimple_nlte["plasma"]["nlte_ionization_species"] = ["H I"] config = Configuration.from_config_dict(tardis_config_verysimple_nlte) with pytest.raises(PlasmaConfigError) as ve: - assemble_plasma(config, nlte_raw_simulation_state, nlte_atom_data) + assemble_plasma(config, nlte_raw_model_root, nlte_atom_data) -def test_plasma_nlte_exc_section_config( - tardis_config_verysimple_nlte, nlte_raw_simulation_state, nlte_atom_data +def test_plasma_nlte_section_lu_config( + tardis_config_verysimple_nlte, + nlte_raw_model_lu, + nlte_atom_data, +): + """ + 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_nlte_root` : 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_ionization_species"] = ["H I"] + tardis_config_verysimple_nlte["plasma"]["nlte_solver"] = "lu" + config = Configuration.from_config_dict(tardis_config_verysimple_nlte) + with pytest.raises(PlasmaConfigError) as ve: + assemble_plasma(config, nlte_raw_model_lu, nlte_atom_data) + + +def test_plasma_nlte_root_exc_section_config( + tardis_config_verysimple_nlte, nlte_raw_model_root, nlte_atom_data ): """ Configuration Validation Test for Plasma Section of the Tardis Config YAML File. @@ -212,7 +245,7 @@ def test_plasma_nlte_exc_section_config( Parameter --------- - `tardis_config_verysimple_nlte` : YAML File + `tardis_config_verysimple_nlte_root` : YAML File `nlte_raw_model` : A simple model `nlte_atom_data` : An example atomic dataset @@ -226,11 +259,10 @@ def test_plasma_nlte_exc_section_config( "He I", ] tardis_config_verysimple_nlte["plasma"]["nlte_excitation_species"] = ["H I"] + tardis_config_verysimple_nlte["plasma"]["nlte_solver"] = "root" config = Configuration.from_config_dict(tardis_config_verysimple_nlte) with pytest.raises(PlasmaConfigError): - plasma = assemble_plasma( - config, nlte_raw_simulation_state, nlte_atom_data - ) + plasma = assemble_plasma(config, nlte_raw_model_root, nlte_atom_data) def test_spectrum_section_config(tardis_config_verysimple): diff --git a/tardis/io/tests/test_atomic.py b/tardis/io/tests/test_atomic.py index 537f3a033ef..5fdf4c9ba41 100644 --- a/tardis/io/tests/test_atomic.py +++ b/tardis/io/tests/test_atomic.py @@ -55,7 +55,12 @@ def test_atom_data_lines(lines): def test_atomic_reprepare(kurucz_atomic_data): - kurucz_atomic_data.prepare_atom_data([14, 20]) + kurucz_atomic_data.prepare_atom_data( + [14, 20], + line_interaction_type="scatter", + nlte_species=[], + continuum_interaction_species=[], + ) lines = kurucz_atomic_data.lines.reset_index() assert lines["atomic_number"].isin([14, 20]).all() assert len(lines.loc[lines["atomic_number"] == 14]) > 0 diff --git a/tardis/model/base.py b/tardis/model/base.py index 1820670b7ba..d197d0f8139 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -23,7 +23,9 @@ parse_packet_source, ) from tardis.montecarlo.packet_source import BlackBodySimpleSource -from tardis.model.radiation_field_state import DiluteThermalRadiationFieldState +from tardis.model.radiation_field_state import ( + DiluteBlackBodyRadiationFieldState, +) from tardis.util.base import is_valid_nuclide_or_elem logger = logging.getLogger(__name__) @@ -133,7 +135,9 @@ def dilution_factor(self): @dilution_factor.setter def dilution_factor(self, new_dilution_factor): if len(new_dilution_factor) == self.no_of_shells: - self.radiation_field_state.dilution_factor = new_dilution_factor + self.radiation_field_state.dilution_factor[ + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index + ] = new_dilution_factor else: raise ValueError( "Trying to set dilution_factor for unmatching number" @@ -149,7 +153,9 @@ def t_radiative(self): @t_radiative.setter def t_radiative(self, new_t_radiative): if len(new_t_radiative) == self.no_of_shells: - self.radiation_field_state.t_radiative = new_t_radiative + self.radiation_field_state.t_radiative[ + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index + ] = new_t_radiative else: raise ValueError( "Trying to set t_radiative for different number of shells." @@ -222,7 +228,7 @@ def volume(self): @property def no_of_shells(self): - return self.geometry.no_of_shells + return self.geometry.no_of_shells_active @property def no_of_raw_shells(self): @@ -260,15 +266,6 @@ def from_config(cls, config, atom_data): density, nuclide_mass_fraction, atom_data.atom_data.mass.copy() ) - nuclide_mass_fraction = parse_abundance_config( - config, geometry, time_explosion - ) - - # using atom_data.mass.copy() to ensure that the original atom_data is not modified - composition = Composition( - density, nuclide_mass_fraction, atom_data.atom_data.mass.copy() - ) - packet_source = parse_packet_source(config, geometry) radiation_field_state = parse_radiation_field_state( config, diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 01914fb8e6a..745def52e1b 100644 --- a/tardis/model/parse_input.py +++ b/tardis/model/parse_input.py @@ -17,7 +17,9 @@ from tardis.model.geometry.radial1d import HomologousRadial1DGeometry from tardis.model.matter.composition import Composition from tardis.model.matter.decay import IsotopicMassFraction -from tardis.model.radiation_field_state import DiluteThermalRadiationFieldState +from tardis.model.radiation_field_state import ( + DiluteBlackBodyRadiationFieldState, +) from tardis.montecarlo.packet_source import ( BlackBodySimpleSource, BlackBodySimpleSourceRelativistic, @@ -566,13 +568,12 @@ def parse_radiation_field_state( if dilution_factor is None: dilution_factor = calculate_geometric_dilution_factor(geometry) - elif len(dilution_factor) != geometry.no_of_shells: - dilution_factor = dilution_factor[ - geometry.v_inner_boundary_index : geometry.v_outer_boundary_index - ] - assert len(dilution_factor) == geometry.no_of_shells - return DiluteThermalRadiationFieldState(t_radiative, dilution_factor) + assert len(dilution_factor) == geometry.no_of_shells + + return DiluteBlackBodyRadiationFieldState( + t_radiative, dilution_factor, geometry + ) def initialize_packet_source(config, geometry, packet_source): @@ -608,13 +609,13 @@ def initialize_packet_source(config, geometry, packet_source): luminosity_requested = config.supernova.luminosity_requested if config.plasma.initial_t_inner > 0.0 * u.K: - packet_source.radius = geometry.r_inner[0] + packet_source.radius = geometry.r_inner_active[0] packet_source.temperature = config.plasma.initial_t_inner elif (config.plasma.initial_t_inner < 0.0 * u.K) and ( luminosity_requested is not None ): - packet_source.radius = geometry.r_inner[0] + packet_source.radius = geometry.r_inner_active[0] packet_source.set_temperature_from_luminosity(luminosity_requested) else: raise ValueError( @@ -672,9 +673,6 @@ def parse_csvy_radiation_field_state( t_radiative = ( np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad ) - t_radiative = ( - np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad - ) else: t_radiative = calculate_t_radiative_from_t_inner( geometry, packet_source @@ -687,7 +685,9 @@ def parse_csvy_radiation_field_state( else: dilution_factor = calculate_geometric_dilution_factor(geometry) - return DiluteThermalRadiationFieldState(t_radiative, dilution_factor) + return DiluteBlackBodyRadiationFieldState( + t_radiative, dilution_factor, geometry + ) def calculate_t_radiative_from_t_inner(geometry, packet_source): @@ -718,6 +718,12 @@ def calculate_geometric_dilution_factor(geometry): return 0.5 * ( 1 - np.sqrt( - 1 - (geometry.r_inner[0] ** 2 / geometry.r_middle**2).to(1).value + 1 + - ( + geometry.r_inner[geometry.v_inner_boundary_index] ** 2 + / geometry.r_middle**2 + ) + .to(1) + .value ) ) diff --git a/tardis/model/radiation_field_state.py b/tardis/model/radiation_field_state.py index f6b1356f76b..dff0bd65bbc 100644 --- a/tardis/model/radiation_field_state.py +++ b/tardis/model/radiation_field_state.py @@ -1,11 +1,12 @@ -from tardis.montecarlo.montecarlo_numba.numba_interface import OpacityState - - import numpy as np from astropy import units as u +from tardis.util.base import intensity_black_body + +from typing import Union + -class DiluteThermalRadiationFieldState: +class DiluteBlackBodyRadiationFieldState: """ Represents the state of a dilute thermal radiation field. @@ -16,17 +17,54 @@ class DiluteThermalRadiationFieldState: Radiative temperature in each shell dilution_factor : numpy.ndarray Dilution Factors in each shell + geometry: tardis.model.Radial1DModel + The geometry of the model that uses to constrains the active shells """ def __init__( self, t_radiative: u.Quantity, dilution_factor: np.ndarray, + geometry=None, ): # ensuring that the radiation_field has both # dilution_factor and t_radiative equal length assert len(t_radiative) == len(dilution_factor) - assert np.all(t_radiative > 0 * u.K) - assert np.all(dilution_factor > 0) + if ( + geometry is not None + ): # check the active shells only (this is used when setting up the radiation_field_state) + assert np.all( + t_radiative[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + > 0 * u.K + ) + assert np.all( + dilution_factor[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + > 0 + ) + else: + assert np.all(t_radiative > 0 * u.K) + assert np.all(dilution_factor > 0) self.t_radiative = t_radiative self.dilution_factor = dilution_factor + + def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): + """ + Calculate the intensity of the radiation field at a given frequency. + + Parameters + ---------- + nu : u.Quantity + Frequency at which the intensity is to be calculated + + Returns + ------- + intensity : u.Quantity + Intensity of the radiation field at the given frequency + """ + return self.dilution_factor * intensity_black_body( + nu[np.newaxis].T, self.t_radiative + ) diff --git a/tardis/model/tests/test_csvy_model.py b/tardis/model/tests/test_csvy_model.py index 19e60716bcd..c82c23b8e26 100644 --- a/tardis/model/tests/test_csvy_model.py +++ b/tardis/model/tests/test_csvy_model.py @@ -86,6 +86,32 @@ def test_compare_models(model_config_fnames, atomic_dataset): ) +def test_dimensionality_after_update_v_inner_boundary( + example_csvy_file_dir, atomic_dataset +): + """Test that the dimensionality of SimulationState parameters after updating v_inner_boundary + in the context of csvy models specifically""" + csvy_config_file = example_csvy_file_dir / "radiative_csvy.yml" + config = Configuration.from_yaml(csvy_config_file) + csvy_model = SimulationState.from_csvy(config, atom_data=atomic_dataset) + + new_config = config + new_config.model.v_inner_boundary = csvy_model.velocity[1] + new_csvy_model = SimulationState.from_csvy( + new_config, atom_data=atomic_dataset + ) + + assert new_csvy_model.no_of_raw_shells == csvy_model.no_of_raw_shells + assert new_csvy_model.no_of_shells == csvy_model.no_of_shells - 1 + assert new_csvy_model.velocity.shape[0] == csvy_model.velocity.shape[0] - 1 + assert new_csvy_model.density.shape[0] == csvy_model.density.shape[0] - 1 + assert new_csvy_model.volume.shape[0] == csvy_model.volume.shape[0] - 1 + assert ( + new_csvy_model.t_radiative.shape[0] + == csvy_model.t_radiative.shape[0] - 1 + ) + + @pytest.fixture(scope="module") def csvy_model_test_abundances(example_csvy_file_dir, atomic_dataset): """Returns SimulationState to use to test abundances dataframes""" diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index 2faf4e8d000..08ff1612306 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -8,6 +8,9 @@ from tardis.io.logger import montecarlo_tracking as mc_tracker from tardis.io.util import HDFWriterMixin from tardis.montecarlo import montecarlo_configuration +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + initialize_estimator_statistics, +) from tardis.montecarlo.montecarlo_configuration import ( configuration_initialize, ) @@ -15,7 +18,6 @@ montecarlo_main_loop, numba_config, ) -from tardis.montecarlo.montecarlo_numba.estimators import initialize_estimators from tardis.montecarlo.montecarlo_numba.formal_integral import FormalIntegrator from tardis.montecarlo.montecarlo_numba.numba_interface import ( NumbaModel, @@ -97,8 +99,8 @@ def __init__( mc_tracker.DEBUG_MODE = debug_packets mc_tracker.BUFFER = logger_buffer - if self.spectrum_method == "integrated": - self.optional_hdf_properties.append("spectrum_integrated") + # if self.spectrum_method == "integrated": + # self.optional_hdf_properties.append("spectrum_integrated") def initialize_transport_state( self, @@ -116,7 +118,7 @@ def initialize_transport_state( packet_collection = self.packet_source.create_packets( no_of_packets, seed_offset=iteration ) - estimators = initialize_estimators( + estimators = initialize_estimator_statistics( plasma.tau_sobolevs.shape, gamma_shape ) @@ -182,7 +184,7 @@ def run( transport_state.geometry_state, numba_model, transport_state.opacity_state, - transport_state.estimators, + transport_state.radfield_mc_estimators, transport_state.spectrum_frequency.value, number_of_vpackets, iteration=iteration, @@ -231,8 +233,8 @@ def legacy_return(self): return ( self.transport_state.packet_collection.output_nus, self.transport_state.packet_collection.output_energies, - self.transport_state.estimators.j_estimator, - self.transport_state.estimators.nu_bar_estimator, + self.transport_state.j_estimator, + self.transport_state.nu_bar_estimator, self.transport_state.last_line_interaction_in_id, self.transport_state.last_line_interaction_out_id, self.transport_state.last_interaction_type, diff --git a/tardis/montecarlo/conftest.py b/tardis/montecarlo/conftest.py new file mode 100644 index 00000000000..8be326055f5 --- /dev/null +++ b/tardis/montecarlo/conftest.py @@ -0,0 +1,16 @@ +import pytest + +from tardis.io.configuration.config_reader import Configuration + + +@pytest.fixture(scope="function") +def continuum_config( + tardis_config_verysimple_nlte, +): + continuum_config = Configuration.from_config_dict( + tardis_config_verysimple_nlte + ) + continuum_config.plasma.continuum_interaction.species = ["H I", "Ti II"] + continuum_config.plasma.nlte_ionization_species = [] + + return continuum_config diff --git a/tardis/montecarlo/estimators/__init__.py b/tardis/montecarlo/estimators/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/montecarlo/estimators/continuum_radfield_properties.py b/tardis/montecarlo/estimators/continuum_radfield_properties.py new file mode 100644 index 00000000000..7ea3ee02ff1 --- /dev/null +++ b/tardis/montecarlo/estimators/continuum_radfield_properties.py @@ -0,0 +1,225 @@ +from dataclasses import dataclass + +import numpy as np +import pandas as pd +from astropy import units as u + +import tardis.constants as const +from tardis.io.atom_data import AtomData +from tardis.model.radiation_field_state import ( + DiluteBlackBodyRadiationFieldState, +) +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + RadiationFieldMCEstimators, +) +from tardis.montecarlo.estimators.util import ( + bound_free_estimator_array2frame, + integrate_array_by_blocks, +) +from tardis.plasma.properties.continuum_processes import PhotoIonBoltzmannFactor + +H = const.h.cgs.value + + +class MCContinuumPropertiesSolver: + def __init__( + self, + atom_data: AtomData, + ): + self.atom_data = atom_data + + def solve( + self, + radfield_mc_estimators: RadiationFieldMCEstimators, + time_simulation, + volume, + ): + """ + Solve for the continuum properties. + + Parameters + ---------- + radfield_mc_estimators : RadiationFieldMCEstimators + The Monte Carlo estimators for the radiation field. + time_simulation : float + The simulation time. + volume : float + The volume of the cells. + + Returns + ------- + ContinuumProperties + The calculated continuum properties. + """ + photo_ion_norm_factor = (time_simulation * volume * H) ** -1 + + photo_ionization_rate_coefficient = bound_free_estimator_array2frame( + radfield_mc_estimators.photo_ion_estimator, + self.atom_data.level2continuum_edge_idx, + ) + photo_ionization_rate_coefficient *= photo_ion_norm_factor + + stimulated_recomb_rate_factor = bound_free_estimator_array2frame( + radfield_mc_estimators.stim_recomb_estimator, + self.atom_data.level2continuum_edge_idx, + ) + stimulated_recomb_rate_factor *= photo_ion_norm_factor + + return ContinuumProperties( + stimulated_recomb_rate_factor, photo_ionization_rate_coefficient + ) + + +class DiluteBlackBodyContinuumPropertiesSolver: + def __init__(self, atom_data: AtomData) -> None: + self.atom_data = atom_data + + def solve( + self, + dilute_blackbody_radiationfield_state: DiluteBlackBodyRadiationFieldState, + t_electrons: u.Quantity, + ): + """ + Solve for the continuum properties. + + Parameters + ---------- + dilute_blackbody_radiationfield_state : DiluteBlackBodyRadiationFieldState + The state of the dilute blackbody radiation field. + t_electrons : u.Quantity + The temperature of the electrons. + + Returns + ------- + ContinuumProperties + The calculated continuum properties. + + """ + photo_ion_boltzmann_factor = PhotoIonBoltzmannFactor.calculate( + self.atom_data.photoionization_data, t_electrons + ) + mean_intensity_photo_ion_df = ( + self.calculate_mean_intensity_photo_ion_table( + dilute_blackbody_radiationfield_state + ) + ) + + photo_ion_rate_coeff = self.calculate_photo_ionization_rate_coefficient( + mean_intensity_photo_ion_df + ) + stimulated_recomb_rate_coeff = ( + self.calculate_stimulated_recomb_rate_factor( + mean_intensity_photo_ion_df, + photo_ion_boltzmann_factor, + ) + ) + + return ContinuumProperties( + stimulated_recomb_rate_coeff, photo_ion_rate_coeff + ) + + def calculate_photo_ionization_rate_coefficient( + self, + mean_intensity_photo_ion_df: pd.DataFrame, + ): + """ + Calculate the photoionization rate coefficient. + + Parameters + ---------- + mean_intensity_photo_ion_df : pd.DataFrame + The mean intensity of the radiation field at the frequencies of the photoionization table. + + Returns + ------- + pd.DataFrame + The calculated photoionization rate coefficient. + + Notes + ----- + Equation 16 in Lucy 2003. + """ + gamma = mean_intensity_photo_ion_df.multiply( + 4.0 + * np.pi + * self.atom_data.photoionization_data.x_sect + / (self.atom_data.photoionization_data.nu * H), + axis=0, + ) + gamma = integrate_array_by_blocks( + gamma.values, + self.atom_data.photoionization_data.nu.values, + self.atom_data.photo_ion_block_references, + ) + gamma = pd.DataFrame(gamma, index=self.atom_data.photo_ion_unique_index) + return gamma + + def calculate_stimulated_recomb_rate_factor( + self, + mean_intensity_photo_ion_df: pd.DataFrame, + photo_ion_boltzmann_factor: np.ndarray, + ): + """ + Calculate the stimulated recombination rate factor (the phi_ik are multiplied in the plasma). + + Parameters + ---------- + mean_intensity_photo_ion_df : pd.DataFrame + The mean intensity of the radiation field at the frequencies of the photoionization table. + photo_ion_boltzmann_factor: np.ndarray + The Boltzmann factor for the photoionization table. + + Returns + ------- + pd.DataFrame + The calculated stimulated recombination rate factor. + + Notes + ----- + Equation 15 in Lucy 2003. + """ + alpha_stim_factor = ( + mean_intensity_photo_ion_df * photo_ion_boltzmann_factor + ) + alpha_stim_factor = alpha_stim_factor.multiply( + 4.0 + * np.pi + * self.atom_data.photoionization_data.x_sect + / (self.atom_data.photoionization_data.nu * H), + axis=0, + ) + alpha_stim_factor = integrate_array_by_blocks( + alpha_stim_factor.values, + self.atom_data.photoionization_data.nu.values, + self.atom_data.photo_ion_block_references, + ) + + alpha_stim_factor = pd.DataFrame( + alpha_stim_factor, index=self.atom_data.photo_ion_unique_index + ) + + return alpha_stim_factor + + def calculate_mean_intensity_photo_ion_table( + self, + dilute_blackbody_radiationfield_state: DiluteBlackBodyRadiationFieldState, + ): + mean_intensity = ( + dilute_blackbody_radiationfield_state.calculate_mean_intensity( + self.atom_data.photoionization_data.nu.values + ) + ) + mean_intensity_df = pd.DataFrame( + mean_intensity, + index=self.atom_data.photoionization_data.index, + columns=np.arange( + len(dilute_blackbody_radiationfield_state.t_radiative) + ), + ) + return mean_intensity_df + + +@dataclass +class ContinuumProperties: + stimulated_recomb_rate_factor: pd.DataFrame + photo_ionization_rate_coefficient: pd.DataFrame diff --git a/tardis/montecarlo/estimators/dilute_blackbody_properties.py b/tardis/montecarlo/estimators/dilute_blackbody_properties.py new file mode 100644 index 00000000000..8e7c98c7ed7 --- /dev/null +++ b/tardis/montecarlo/estimators/dilute_blackbody_properties.py @@ -0,0 +1,57 @@ +import numpy as np +from astropy import units as u +from scipy.special import zeta + +from tardis import constants as const +from tardis.model.radiation_field_state import ( + DiluteBlackBodyRadiationFieldState, +) + +DILUTION_FACTOR_ESTIMATOR_CONSTANT = ( + (const.c**2 / (2 * const.h)) + * (15 / np.pi**4) + * (const.h / const.k_B) ** 4 + / (4 * np.pi) +).cgs.value +T_RADIATIVE_ESTIMATOR_CONSTANT = ( + (np.pi**4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B) +).cgs.value + + +class MCDiluteBlackBodyRadFieldSolver: + def __init__(self) -> None: + pass + + def solve(self, radfield_mc_estimators, time_of_simulation, volume): + """ + Calculate an updated radiation field from the :math: + `\\bar{nu}_\\textrm{estimator}` and :math:`\\J_\\textrm{estimator}` + calculated in the montecarlo simulation. + The details of the calculation can be found in the documentation. + + Parameters + ---------- + nubar_estimator : np.ndarray (float) + j_estimator : np.ndarray (float) + + Returns + ------- + t_radiative : astropy.units.Quantity (float) + dilution_factor : numpy.ndarray (float) + """ + temperature_radiative = ( + T_RADIATIVE_ESTIMATOR_CONSTANT + * radfield_mc_estimators.nu_bar_estimator + / radfield_mc_estimators.j_estimator + ) * u.K + dilution_factor = radfield_mc_estimators.j_estimator / ( + 4 + * const.sigma_sb.cgs.value + * temperature_radiative.value**4 + * time_of_simulation.value + * volume + ) + + return DiluteBlackBodyRadiationFieldState( + temperature_radiative, dilution_factor + ) diff --git a/tardis/montecarlo/montecarlo_numba/estimators.py b/tardis/montecarlo/estimators/radfield_mc_estimators.py similarity index 89% rename from tardis/montecarlo/montecarlo_numba/estimators.py rename to tardis/montecarlo/estimators/radfield_mc_estimators.py index 5e3ddf6514f..328a0c4d55b 100644 --- a/tardis/montecarlo/montecarlo_numba/estimators.py +++ b/tardis/montecarlo/estimators/radfield_mc_estimators.py @@ -1,24 +1,21 @@ from math import exp import numpy as np - -from numba import njit, float64, int64 +from numba import float64, int64, njit from numba.experimental import jitclass from tardis.montecarlo import montecarlo_configuration as nc -from tardis.montecarlo.montecarlo_numba.numba_config import H, KB - from tardis.montecarlo.montecarlo_numba import ( njit_dict_no_parallel, ) - +from tardis.montecarlo.montecarlo_numba.numba_config import KB, H from tardis.transport.frame_transformations import ( calc_packet_energy, calc_packet_energy_full_relativity, ) -def initialize_estimators(tau_sobolev_shape, gamma_shape): +def initialize_estimator_statistics(tau_sobolev_shape, gamma_shape): """ Initializes the estimators used in the Monte Carlo simulation. @@ -41,7 +38,6 @@ def initialize_estimators(tau_sobolev_shape, gamma_shape): >>> initialize_estimators(tau_sobolev_shape, gamma_shape) """ - j_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) nu_bar_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) j_blue_estimator = np.zeros(tau_sobolev_shape) @@ -55,7 +51,7 @@ def initialize_estimators(tau_sobolev_shape, gamma_shape): stim_recomb_cooling_estimator = np.zeros(gamma_shape, dtype=np.float64) photo_ion_estimator_statistics = np.zeros(gamma_shape, dtype=np.int64) - return Estimators( + return RadiationFieldMCEstimators( j_estimator, nu_bar_estimator, j_blue_estimator, @@ -85,7 +81,7 @@ def initialize_estimators(tau_sobolev_shape, gamma_shape): @jitclass(base_estimators_spec + continuum_estimators_spec) -class Estimators(object): +class RadiationFieldMCEstimators: def __init__( self, j_estimator, @@ -109,6 +105,18 @@ def __init__( self.photo_ion_estimator_statistics = photo_ion_estimator_statistics def increment(self, other): + """ + Increments each estimator with the corresponding estimator from another instance of the class. + + Parameters + ---------- + other : RadiationFieldMCEstimators + Another instance of the RadiationFieldMCEstimators class. + + Returns + ------- + None + """ self.j_estimator += other.j_estimator self.nu_bar_estimator += other.nu_bar_estimator self.j_blue_estimator += other.j_blue_estimator @@ -201,7 +209,11 @@ def update_bound_free_estimators( @njit(**njit_dict_no_parallel) def update_line_estimators( - estimators, r_packet, cur_line_id, distance_trace, time_explosion + radfield_mc_estimators, + r_packet, + cur_line_id, + distance_trace, + time_explosion, ): """ Function to update the line estimators @@ -214,23 +226,14 @@ def update_line_estimators( distance_trace : float time_explosion : float """ - - """ Actual calculation - simplified below - r_interaction = math.sqrt(r_packet.r**2 + distance_trace**2 + - 2 * r_packet.r * distance_trace * r_packet.mu) - mu_interaction = (r_packet.mu * r_packet.r + distance_trace) / r_interaction - doppler_factor = 1.0 - mu_interaction * r_interaction / - ( time_explosion * C) - """ - if not nc.ENABLE_FULL_RELATIVITY: energy = calc_packet_energy(r_packet, distance_trace, time_explosion) else: energy = calc_packet_energy_full_relativity(r_packet) - estimators.j_blue_estimator[cur_line_id, r_packet.current_shell_id] += ( - energy / r_packet.nu - ) - estimators.Edotlu_estimator[ + radfield_mc_estimators.j_blue_estimator[ + cur_line_id, r_packet.current_shell_id + ] += (energy / r_packet.nu) + radfield_mc_estimators.Edotlu_estimator[ cur_line_id, r_packet.current_shell_id ] += energy diff --git a/tardis/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/montecarlo/estimators/tests/test_continuum_property_solver.py new file mode 100644 index 00000000000..7ef65daa481 --- /dev/null +++ b/tardis/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -0,0 +1,86 @@ +from copy import deepcopy + +import numpy.testing as npt +import pandas.testing as pdt +import pytest + +from tardis.montecarlo.estimators.continuum_radfield_properties import ( + DiluteBlackBodyContinuumPropertiesSolver, + MCContinuumPropertiesSolver, +) +from tardis.simulation import Simulation + + +@pytest.mark.skip() +def test_continuum_estimators( + continuum_config, + nlte_atomic_dataset, +): + nlte_atomic_dataset = deepcopy(nlte_atomic_dataset) + continuum_simulation = Simulation.from_config( + continuum_config, + atom_data=nlte_atomic_dataset, + virtual_packet_logging=False, + ) + # continuum_simulation.run_convergence() + continuum_properties_solver_dilute_bb = ( + DiluteBlackBodyContinuumPropertiesSolver( + continuum_simulation.plasma.atomic_data + ) + ) + + continuum_properties_dilute_bb = ( + continuum_properties_solver_dilute_bb.solve( + continuum_simulation.simulation_state.radiation_field_state, + continuum_simulation.plasma.t_electrons, + ) + ) + + continuum_plasma = continuum_simulation.plasma + + pdt.assert_frame_equal( + continuum_properties_dilute_bb.photo_ionization_rate_coefficient, + continuum_simulation.plasma.gamma, + ) + stimulated_recomb_rate_coeff = ( + continuum_properties_dilute_bb.stimulated_recomb_rate_factor + * continuum_plasma.phi_ik.loc[ + continuum_properties_dilute_bb.stimulated_recomb_rate_factor.index + ] + ) + pdt.assert_frame_equal( + stimulated_recomb_rate_coeff, continuum_plasma.alpha_stim + ) + + continuum_simulation.run_final() + + continuum_properties_solver_mc = MCContinuumPropertiesSolver( + continuum_simulation.plasma.atomic_data + ) + transport_state = continuum_simulation.transport.transport_state + continuum_properties_mc = continuum_properties_solver_mc.solve( + transport_state.radfield_mc_estimators, + transport_state.time_of_simulation, + transport_state.geometry_state.volume, + ) + + continuum_plasma.update( + gamma_estimator=transport_state.radfield_mc_estimators.photo_ion_estimator, + alpha_stim_estimator=transport_state.radfield_mc_estimators.stim_recomb_estimator, + bf_heating_coeff_estimator=transport_state.radfield_mc_estimators.bf_heating_estimator, + stim_recomb_cooling_coeff_estimator=transport_state.radfield_mc_estimators.stim_recomb_cooling_estimator, + ) + + pdt.assert_frame_equal( + continuum_properties_mc.photo_ionization_rate_coefficient, + continuum_simulation.plasma.gamma, + ) + stimulated_recomb_rate_coeff = ( + continuum_properties_mc.stimulated_recomb_rate_factor + * continuum_plasma.phi_ik.loc[ + continuum_properties_dilute_bb.stimulated_recomb_rate_factor.index + ] + ) + pdt.assert_frame_equal( + stimulated_recomb_rate_coeff, continuum_plasma.alpha_stim + ) diff --git a/tardis/montecarlo/estimators/util.py b/tardis/montecarlo/estimators/util.py new file mode 100644 index 00000000000..f7b69acb992 --- /dev/null +++ b/tardis/montecarlo/estimators/util.py @@ -0,0 +1,67 @@ +import numpy as np +import pandas as pd +from numba import njit, prange + +from tardis.montecarlo.montecarlo_numba import njit_dict + + +def bound_free_estimator_array2frame( + bound_free_estimator_array, level2continuum_idx +): + """ + Transform a bound-free estimator array to a DataFrame. + + This function transforms a bound-free estimator array with entries + sorted by frequency to a multi-indexed DataFrame sorted by level. + + Parameters + ---------- + bf_estimator_array : numpy.ndarray, dtype float + Array of bound-free estimators (e.g., for the stimulated recombination rate) + with entries sorted by the threshold frequency of the bound-free continuum. + level2continuum_idx : pandas.Series, dtype int + Maps a level MultiIndex (atomic_number, ion_number, level_number) to + the continuum_idx of the corresponding bound-free continuum (which are + sorted by decreasing frequency). + + Returns + ------- + pandas.DataFrame, dtype float + Bound-free estimators indexed by (atomic_number, ion_number, level_number). + """ + bf_estimator_frame = pd.DataFrame( + bound_free_estimator_array, index=level2continuum_idx.index + ).sort_index() + bf_estimator_frame.columns.name = "Shell No." + return bf_estimator_frame + + +@njit(**njit_dict) +def integrate_array_by_blocks(f, x, block_references): + """ + Integrate a function over blocks. + + This function integrates a function `f` defined at locations `x` + over blocks given in `block_references`. + + Parameters + ---------- + f : numpy.ndarray, dtype float + 2D input array to integrate. + x : numpy.ndarray, dtype float + 1D array with the sample points corresponding to the `f` values. + block_references : numpy.ndarray, dtype int + 1D array with the start indices of the blocks to be integrated. + + Returns + ------- + numpy.ndarray, dtype float + 2D array with integrated values. + """ + integrated = np.zeros((len(block_references) - 1, f.shape[1])) + for i in prange(f.shape[1]): # columns + for j in prange(len(integrated)): # rows + start = block_references[j] + stop = block_references[j + 1] + integrated[j, i] = np.trapz(f[start:stop, i], x[start:stop]) + return integrated diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index 455423fbe88..987a70fe7b1 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -5,7 +5,9 @@ from tardis.montecarlo import montecarlo_configuration from tardis.montecarlo.montecarlo_numba import njit_dict -from tardis.montecarlo.montecarlo_numba.estimators import Estimators +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + RadiationFieldMCEstimators, +) from tardis.montecarlo.montecarlo_numba.numba_interface import ( NumbaModel, RPacketTracker, @@ -101,7 +103,7 @@ def montecarlo_main_loop( # so we iterate from 0 to n_threads-1 to create the estimator_list for i in range(n_threads): estimator_list.append( - Estimators( + RadiationFieldMCEstimators( np.copy(estimators.j_estimator), np.copy(estimators.nu_bar_estimator), np.copy(estimators.j_blue_estimator), diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral.py b/tardis/montecarlo/montecarlo_numba/formal_integral.py index 27aa54d283e..c49c706b373 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/formal_integral.py @@ -448,7 +448,7 @@ def make_source_function(self): Edotlu = ( Edotlu_norm_factor * exptau - * montecarlo_transport_state.estimators.Edotlu_estimator + * montecarlo_transport_state.radfield_mc_estimators.Edotlu_estimator ) # The following may be achieved by calling the appropriate plasma @@ -470,7 +470,7 @@ def make_source_function(self): # Jbluelu should already by in the correct order, i.e. by wavelength of # the transition l->u Jbluelu = ( - transport.transport_state.estimators.j_blue_estimator + transport.transport_state.radfield_mc_estimators.j_blue_estimator * Jbluelu_norm_factor ) diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py index c8968c7ce75..3da74637f02 100644 --- a/tardis/montecarlo/montecarlo_numba/interaction.py +++ b/tardis/montecarlo/montecarlo_numba/interaction.py @@ -1,28 +1,27 @@ -from numba import njit import numpy as np -from tardis.montecarlo.montecarlo_numba import njit_dict, njit_dict_no_parallel -from tardis.montecarlo.montecarlo_numba.numba_interface import ( - LineInteractionType, -) +from numba import njit +from tardis import constants as const from tardis.montecarlo import ( montecarlo_configuration as montecarlo_configuration, ) -from tardis.transport.frame_transformations import ( - get_doppler_factor, - get_inverse_doppler_factor, - angle_aberration_CMF_to_LF, +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.montecarlo.montecarlo_numba.macro_atom import ( + MacroAtomTransitionType, + macro_atom, +) +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + LineInteractionType, ) from tardis.montecarlo.montecarlo_numba.r_packet import ( - InteractionType, PacketStatus, ) from tardis.montecarlo.montecarlo_numba.utils import get_random_mu -from tardis.montecarlo.montecarlo_numba.macro_atom import ( - macro_atom, - MacroAtomTransitionType, +from tardis.transport.frame_transformations import ( + angle_aberration_CMF_to_LF, + get_doppler_factor, + get_inverse_doppler_factor, ) -from tardis import constants as const K_B = const.k_B.cgs.value H = const.h.cgs.value @@ -129,7 +128,6 @@ def sample_nu_free_bound(opacity_state, shell, continuum_id): nu_fb_sampler : float Frequency of the free-bounds emission process """ - start = opacity_state.photo_ion_block_references[continuum_id] end = opacity_state.photo_ion_block_references[continuum_id + 1] phot_nus_block = opacity_state.phot_nus[start:end] @@ -206,7 +204,6 @@ def macro_atom_event( time_explosion : float opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState """ - transition_id, transition_type = macro_atom( destination_level_idx, r_packet.current_shell_id, opacity_state ) @@ -253,7 +250,6 @@ def bf_cooling(r_packet, time_explosion, opacity_state): time_explosion : float opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState """ - fb_cooling_prob = opacity_state.p_fb_deactivation[ :, r_packet.current_shell_id ] @@ -276,7 +272,6 @@ def adiabatic_cooling(r_packet): ---------- r_packet: tardis.montecarlo.montecarlo_numba.r_packet.RPacket """ - r_packet.status = PacketStatus.ADIABATIC_COOLING @@ -290,7 +285,6 @@ def get_current_line_id(nu, line_list): nu : float line_list : np.ndarray """ - # Note: Since this reverses the array, # it may be faster to just write our own reverse-binary-search @@ -311,7 +305,6 @@ def free_free_emission(r_packet, time_explosion, opacity_state): time_explosion : float opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState """ - inverse_doppler_factor = get_inverse_doppler_factor( r_packet.r, r_packet.mu, time_explosion ) @@ -338,7 +331,6 @@ def bound_free_emission(r_packet, time_explosion, opacity_state, continuum_id): opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState continuum_id : int """ - inverse_doppler_factor = get_inverse_doppler_factor( r_packet.r, r_packet.mu, time_explosion ) @@ -371,7 +363,6 @@ def thomson_scatter(r_packet, time_explosion): time_explosion : float time since explosion in seconds """ - old_doppler_factor = get_doppler_factor( r_packet.r, r_packet.mu, time_explosion ) @@ -407,7 +398,6 @@ def line_scatter( line_interaction_type : enum opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState """ - old_doppler_factor = get_doppler_factor( r_packet.r, r_packet.mu, time_explosion ) @@ -447,7 +437,6 @@ def line_emission(r_packet, emission_line_id, time_explosion, opacity_state): time_explosion : float opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState """ - r_packet.last_line_interaction_out_id = emission_line_id r_packet.last_line_interaction_shell_id = r_packet.current_shell_id diff --git a/tardis/montecarlo/montecarlo_numba/opacities.py b/tardis/montecarlo/montecarlo_numba/opacities.py index e043b280c06..dd10595281f 100644 --- a/tardis/montecarlo/montecarlo_numba/opacities.py +++ b/tardis/montecarlo/montecarlo_numba/opacities.py @@ -1,17 +1,15 @@ -from numba import njit import numpy as np import radioactivedecay as rd +from numba import njit +from tardis import constants as const +from tardis.energy_input.util import kappa_calculation from tardis.montecarlo.montecarlo_numba import ( njit_dict_no_parallel, ) - from tardis.montecarlo.montecarlo_numba.numba_config import ( SIGMA_THOMSON, ) -from tardis.energy_input.util import kappa_calculation - -from tardis import constants as const H = const.h.cgs.value M_E = const.m_e.cgs.value @@ -119,7 +117,6 @@ def chi_bf_interpolator(opacity_state, nu, shell): Photoionization cross-sections of all bound-free continua for which absorption is possible for frequency `nu`. """ - current_continua = get_current_bound_free_continua(opacity_state, nu) chi_bfs = np.zeros(len(current_continua)) x_sect_bfs = np.zeros(len(current_continua)) @@ -201,7 +198,6 @@ def chi_continuum_calculator(opacity_state, nu, shell): Returns ------- - chi_bf_tot : float Total bound-free opacity at frequency `nu`. chi_bf_contributions : numpy.ndarray, dtype float @@ -267,9 +263,9 @@ def compton_opacity_calculation(energy, electron_density): (Rybicki & Lightman, 1979) $ - \\rho / 2 p_m \\times 3/4 \\sigma_T ((1 + \kappa) / \kappa^3 - ((2\kappa(1 + \kappa)) / (1 + 2\kappa) - \ln(1 + 2\kappa) + 1/(2\kappa) \ln(1 + 2\kappa) - - (1 + 3\kappa) / (1 + 2\kappa)^2) + \\rho / 2 p_m \\times 3/4 \\sigma_T ((1 + \\kappa) / \\kappa^3 + ((2\\kappa(1 + \\kappa)) / (1 + 2\\kappa) - \\ln(1 + 2\\kappa) + 1/(2\\kappa) \\ln(1 + 2\\kappa) + - (1 + 3\\kappa) / (1 + 2\\kappa)^2) $ Parameters diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py index 00efd249b78..e61ec6b5e08 100644 --- a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py +++ b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py @@ -1,37 +1,34 @@ from numba import njit -from tardis.montecarlo.montecarlo_numba.r_packet import ( - PacketStatus, -) -from tardis.transport.r_packet_transport import ( - move_r_packet, - move_packet_across_shell_boundary, - trace_packet, +from tardis import constants as const +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, ) - -from tardis.transport.frame_transformations import ( - get_inverse_doppler_factor, - get_doppler_factor, +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + update_bound_free_estimators, ) from tardis.montecarlo.montecarlo_numba.interaction import ( - InteractionType, - thomson_scatter, - line_scatter, continuum_event, + line_scatter, + thomson_scatter, ) -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) - -from tardis.montecarlo.montecarlo_numba.vpacket import trace_vpacket_volley - -from tardis import constants as const from tardis.montecarlo.montecarlo_numba.opacities import ( chi_continuum_calculator, chi_electron_calculator, ) -from tardis.montecarlo.montecarlo_numba.estimators import ( - update_bound_free_estimators, +from tardis.montecarlo.montecarlo_numba.r_packet import ( + InteractionType, + PacketStatus, +) +from tardis.montecarlo.montecarlo_numba.vpacket import trace_vpacket_volley +from tardis.transport.frame_transformations import ( + get_doppler_factor, + get_inverse_doppler_factor, +) +from tardis.transport.r_packet_transport import ( + move_packet_across_shell_boundary, + move_r_packet, + trace_packet, ) C_SPEED_OF_LIGHT = const.c.to("cm/s").value diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index b5a68c0876f..12c4133ccd2 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -3,14 +3,18 @@ import pytest import numpy as np from numba import njit -from tardis.montecarlo.montecarlo_numba.estimators import Estimators +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + RadiationFieldMCEstimators, +) from tardis.montecarlo.montecarlo_numba.packet_collections import ( VPacketCollection, ) from tardis.simulation import Simulation from tardis.montecarlo.montecarlo_numba import RPacket -from tardis.montecarlo.montecarlo_numba.estimators import Estimators +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + RadiationFieldMCEstimators, +) from tardis.montecarlo.montecarlo_numba.numba_interface import ( @@ -51,7 +55,7 @@ def verysimple_numba_model(nb_simulation_verysimple): def verysimple_estimators(nb_simulation_verysimple): transport = nb_simulation_verysimple.transport - return Estimators( + return RadiationFieldMCEstimators( transport.j_estimator, transport.nu_bar_estimator, transport.j_blue_estimator, diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index 5741412b7cb..34d2d0926e2 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -62,8 +62,10 @@ def test_montecarlo_main_loop( transport_state = montecarlo_main_loop_simulation.transport.transport_state actual_energy = transport_state.packet_collection.output_energies actual_nu = transport_state.packet_collection.output_nus - actual_nu_bar_estimator = transport_state.estimators.nu_bar_estimator - actual_j_estimator = transport_state.estimators.j_estimator + actual_nu_bar_estimator = ( + transport_state.radfield_mc_estimators.nu_bar_estimator + ) + actual_j_estimator = transport_state.radfield_mc_estimators.j_estimator # Compare npt.assert_allclose( @@ -116,8 +118,8 @@ def test_montecarlo_main_loop_vpacket_log( actual_energy = transport_state.packet_collection.output_energies actual_nu = transport_state.packet_collection.output_nus - actual_nu_bar_estimator = transport_state.estimators.nu_bar_estimator - actual_j_estimator = transport_state.estimators.j_estimator + actual_nu_bar_estimator = transport_state.nu_bar_estimator + actual_j_estimator = transport_state.j_estimator actual_vpacket_log_nus = transport_state.vpacket_tracker.nus actual_vpacket_log_energies = transport_state.vpacket_tracker.energies diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_continuum.py b/tardis/montecarlo/montecarlo_numba/tests/test_continuum.py new file mode 100644 index 00000000000..5d0327081cf --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_continuum.py @@ -0,0 +1,40 @@ +from copy import deepcopy + +import numpy.testing as npt +import pytest + +from tardis.simulation import Simulation + + +@pytest.mark.skip() +def test_montecarlo_continuum( + continuum_config, + regression_data, + nlte_atomic_dataset, +): + nlte_atomic_dataset = deepcopy(nlte_atomic_dataset) + continuum_simulation = Simulation.from_config( + continuum_config, + atom_data=nlte_atomic_dataset, + virtual_packet_logging=False, + ) + # continuum_simulation.run_convergence() + + continuum_simulation.run_final() + + expected_hdf_store = regression_data.sync_hdf_store(continuum_simulation) + + expected_nu = expected_hdf_store[ + "/simulation/transport/transport_state/output_nu" + ] + expected_energy = expected_hdf_store[ + "/simulation/transport/transport_state/output_energy" + ] + expected_hdf_store.close() + + transport_state = continuum_simulation.transport.transport_state + actual_energy = transport_state.packet_collection.output_energies + actual_nu = transport_state.packet_collection.output_nus + + npt.assert_allclose(actual_energy, expected_energy, rtol=1e-13) + npt.assert_allclose(actual_nu, expected_nu, rtol=1e-13) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py index 9a5e8b148e1..a3af0173047 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -2,7 +2,7 @@ import pytest import tardis.montecarlo.montecarlo_configuration as numba_config -import tardis.montecarlo.montecarlo_numba.estimators +import tardis.montecarlo.estimators.radfield_mc_estimators import tardis.montecarlo.montecarlo_numba.numba_interface as numba_interface import tardis.montecarlo.montecarlo_numba.opacities as opacities import tardis.montecarlo.montecarlo_numba.r_packet as r_packet @@ -12,7 +12,9 @@ import tardis.transport.r_packet_transport as r_packet_transport from tardis import constants as const from tardis.model.geometry.radial1d import NumbaRadial1DGeometry -from tardis.montecarlo.montecarlo_numba.estimators import update_line_estimators +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + update_line_estimators, +) C_SPEED_OF_LIGHT = const.c.to("cm/s").value SIGMA_THOMSON = const.sigma_T.to("cm^2").value @@ -42,7 +44,7 @@ def model(): @pytest.fixture(scope="function") def estimators(): - return tardis.montecarlo.montecarlo_numba.estimators.Estimators( + return tardis.montecarlo.estimators.radfield_mc_estimators.RadiationFieldMCEstimators( j_estimator=np.array([0.0, 0.0], dtype=np.float64), nu_bar_estimator=np.array([0.0, 0.0], dtype=np.float64), j_blue_estimator=np.array( diff --git a/tardis/montecarlo/montecarlo_transport_state.py b/tardis/montecarlo/montecarlo_transport_state.py index e6565d09aaa..441f26a3f79 100644 --- a/tardis/montecarlo/montecarlo_transport_state.py +++ b/tardis/montecarlo/montecarlo_transport_state.py @@ -2,22 +2,12 @@ import numpy as np from astropy import units as u -from scipy.special import zeta -from tardis import constants as const from tardis.io.util import HDFWriterMixin from tardis.montecarlo.spectrum import TARDISSpectrum - -DILUTION_FACTOR_ESTIMATOR_CONSTANT = ( - (const.c**2 / (2 * const.h)) - * (15 / np.pi**4) - * (const.h / const.k_B) ** 4 - / (4 * np.pi) -).cgs.value - -T_RADIATIVE_ESTIMATOR_CONSTANT = ( - (np.pi**4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B) -).cgs.value +from tardis.montecarlo.estimators.dilute_blackbody_properties import ( + MCDiluteBlackBodyRadFieldSolver, +) class MonteCarloTransportState(HDFWriterMixin): @@ -65,7 +55,7 @@ class MonteCarloTransportState(HDFWriterMixin): def __init__( self, packet_collection, - estimators, + radfield_mc_estimators, spectrum_frequency, geometry_state, opacity_state, @@ -73,7 +63,7 @@ def __init__( vpacket_tracker=None, ): self.packet_collection = packet_collection - self.estimators = estimators + self.radfield_mc_estimators = radfield_mc_estimators self.spectrum_frequency = spectrum_frequency self._montecarlo_virtual_luminosity = u.Quantity( np.zeros_like(self.spectrum_frequency.value), "erg / s" @@ -104,20 +94,17 @@ def calculate_radiationfield_properties(self): t_radiative : astropy.units.Quantity (float) dilution_factor : numpy.ndarray (float) """ - estimated_t_radiative = ( - T_RADIATIVE_ESTIMATOR_CONSTANT - * self.estimators.nu_bar_estimator - / self.estimators.j_estimator - ) * u.K - dilution_factor = self.estimators.j_estimator / ( - 4 - * const.sigma_sb.cgs.value - * estimated_t_radiative.value**4 - * (self.packet_collection.time_of_simulation) - * self.geometry_state.volume + dilute_bb_solver = MCDiluteBlackBodyRadFieldSolver() + dilute_bb_radfield = dilute_bb_solver.solve( + self.radfield_mc_estimators, + self.time_of_simulation, + self.geometry_state.volume, ) - return estimated_t_radiative, dilution_factor + return ( + dilute_bb_radfield.t_radiative, + dilute_bb_radfield.dilution_factor, + ) @property def output_nu(self): @@ -129,11 +116,11 @@ def output_energy(self): @property def nu_bar_estimator(self): - return self.estimators.nu_bar_estimator + return self.radfield_mc_estimators.nu_bar_estimator @property def j_estimator(self): - return self.estimators.j_estimator + return self.radfield_mc_estimators.j_estimator @property def time_of_simulation(self): diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index 3350f9ab7a1..86d3becadbf 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -100,7 +100,12 @@ def create_packets(self, no_of_packets, seed_offset=0, *args, **kwargs): self.calculate_radfield_luminosity().to(u.erg / u.s).value ) return PacketCollection( - radii, nus, mus, energies, packet_seeds, radiation_field_luminosity + radii, + nus, + mus, + energies, + packet_seeds, + radiation_field_luminosity, ) def calculate_radfield_luminosity(self): diff --git a/tardis/montecarlo/tests/test_montecarlo.py b/tardis/montecarlo/tests/test_montecarlo.py index 5aa5f4eedc5..8ab1356b33c 100644 --- a/tardis/montecarlo/tests/test_montecarlo.py +++ b/tardis/montecarlo/tests/test_montecarlo.py @@ -1,32 +1,33 @@ import os -import pytest + import numpy as np import pandas as pd +import pytest + +import tardis.montecarlo.montecarlo_configuration as mc import tardis.montecarlo.montecarlo_numba.formal_integral as formal_integral import tardis.montecarlo.montecarlo_numba.r_packet as r_packet -import tardis.transport.r_packet_transport as r_packet_transport import tardis.montecarlo.montecarlo_numba.utils as utils -import tardis.montecarlo.montecarlo_configuration as mc +import tardis.transport.r_packet_transport as r_packet_transport from tardis import constants as const -from tardis.montecarlo.montecarlo_numba.estimators import Estimators +from tardis.montecarlo.estimators.radfield_mc_estimators import ( + RadiationFieldMCEstimators, +) from tardis.montecarlo.montecarlo_numba.numba_interface import RPacketTracker - - from tardis.transport.frame_transformations import ( - get_doppler_factor, - angle_aberration_LF_to_CMF, angle_aberration_CMF_to_LF, + angle_aberration_LF_to_CMF, + get_doppler_factor, ) - C_SPEED_OF_LIGHT = const.c.to("cm/s").value pytestmark = pytest.mark.skip(reason="Port from C to numba") from numpy.testing import ( - assert_almost_equal, assert_allclose, + assert_almost_equal, ) from tardis import __path__ as path @@ -462,7 +463,7 @@ def test_move_packet(packet_params, expected_params, full_relativity): mc.ENABLE_FULL_RELATIVITY = full_relativity doppler_factor = get_doppler_factor(packet.r, packet.mu, time_explosion) - numba_estimator = Estimators( + numba_estimator = RadiationFieldMCEstimators( packet_params["j"], packet_params["nu_bar"], 0, 0 ) r_packet_transport.move_r_packet( @@ -618,7 +619,7 @@ def test_compute_distance2line_relativistic( ): packet = r_packet.RPacket(r=r, nu=nu, mu=mu, energy=0.9) # packet.nu_line = nu_line - numba_estimator = Estimators( + numba_estimator = RadiationFieldMCEstimators( transport.j_estimator, transport.nu_bar_estimator, transport.j_blue_estimator, diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index 6ade1ae1377..f206264d3e4 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -1,28 +1,27 @@ import logging +from collections import Counter as counter import numpy as np import pandas as pd +import radioactivedecay as rd from numba import njit -from scipy.special import exp1 from scipy.interpolate import PchipInterpolator -from collections import Counter as counter -import radioactivedecay as rd -from tardis import constants as const +from scipy.special import exp1 +from tardis import constants as const +from tardis.plasma.exceptions import IncompleteAtomicData from tardis.plasma.properties.base import ( - ProcessingPlasmaProperty, - HiddenPlasmaProperty, BaseAtomicDataProperty, + HiddenPlasmaProperty, + ProcessingPlasmaProperty, ) -from tardis.plasma.exceptions import IncompleteAtomicData from tardis.plasma.properties.continuum_processes import ( - get_ground_state_multi_index, - K_B, - BETA_COLL, - H, A0, + BETA_COLL, + K_B, M_E, - C, + H, + get_ground_state_multi_index, ) logger = logging.getLogger(__name__) @@ -459,8 +458,8 @@ class LevelIdxs2LineIdx(HiddenPlasmaProperty): def calculate(self, atomic_data): index = pd.MultiIndex.from_arrays( [ - atomic_data.lines_upper2level_idx, - atomic_data.lines_lower2level_idx, + atomic_data.lines_upper2macro_reference_idx, + atomic_data.lines_lower2macro_reference_idx, ], names=["source_level_idx", "destination_level_idx"], ) @@ -470,7 +469,7 @@ def calculate(self, atomic_data): # Check for duplicate indices if level_idxs2line_idx.index.duplicated().any(): - logger.warn( + logger.warning( "Duplicate indices in level_idxs2line_idx. " "Dropping duplicates. " "This is an issue with the atomic data & carsus. " @@ -621,7 +620,7 @@ def _filter_atomic_property(self, ionization_data, selected_atoms): ionization_data = ionization_data[mask] counts = ionization_data.groupby(level="atomic_number").count() - if np.alltrue(counts.index == counts): + if np.all(counts.index == counts): return ionization_data else: raise IncompleteAtomicData( @@ -652,7 +651,7 @@ def _filter_atomic_property(self, zeta_data, selected_atoms): zeta_data_check = counter(zeta_data.atomic_number.values) keys = np.array(list(zeta_data_check.keys())) values = np.array(zeta_data_check.values()) - if np.alltrue(keys + 1 == values) and keys: + if np.all(keys + 1 == values) and keys: return zeta_data else: # raise IncompleteAtomicData('zeta data') @@ -666,7 +665,7 @@ def _filter_atomic_property(self, zeta_data, selected_atoms): if (atom, ion) not in zeta_data.index: missing_ions.append((atom, ion)) updated_index.append([atom, ion]) - logger.warn( + logger.warning( f"Zeta_data missing - replaced with 1s. Missing ions: {missing_ions}" ) updated_index = np.array(updated_index) diff --git a/tardis/plasma/properties/continuum_processes.py b/tardis/plasma/properties/continuum_processes.py index 3dbc138e992..25e5790ca9f 100644 --- a/tardis/plasma/properties/continuum_processes.py +++ b/tardis/plasma/properties/continuum_processes.py @@ -2,14 +2,18 @@ import numpy as np import pandas as pd +from numba import njit, prange -from numba import prange, njit from tardis import constants as const - +from tardis.montecarlo.estimators.util import ( + bound_free_estimator_array2frame, + integrate_array_by_blocks, +) +from tardis.montecarlo.montecarlo_numba import njit_dict from tardis.plasma.exceptions import PlasmaException from tardis.plasma.properties.base import ( - ProcessingPlasmaProperty, Input, + ProcessingPlasmaProperty, TransitionProbabilitiesProperty, ) from tardis.plasma.properties.j_blues import JBluesDiluteBlackBody @@ -65,39 +69,6 @@ logger = logging.getLogger(__name__) -njit_dict = {"fastmath": False, "parallel": False} - - -@njit(**njit_dict) -def integrate_array_by_blocks(f, x, block_references): - """ - Integrate a function over blocks. - - This function integrates a function `f` defined at locations `x` - over blocks given in `block_references`. - - Parameters - ---------- - f : numpy.ndarray, dtype float - 2D input array to integrate. - x : numpy.ndarray, dtype float - 1D array with the sample points corresponding to the `f` values. - block_references : numpy.ndarray, dtype int - 1D array with the start indices of the blocks to be integrated. - - Returns - ------- - numpy.ndarray, dtype float - 2D array with integrated values. - """ - integrated = np.zeros((len(block_references) - 1, f.shape[1])) - for i in prange(f.shape[1]): # columns - for j in prange(len(integrated)): # rows - start = block_references[j] - stop = block_references[j + 1] - integrated[j, i] = np.trapz(f[start:stop, i], x[start:stop]) - return integrated - # It is currently not possible to use scipy.integrate.cumulative_trapezoid in # numba. So here is my own implementation. @@ -240,36 +211,7 @@ def cooling_rate_series2dataframe(cooling_rate_series, destination_level_idx): return cooling_rate_frame -def bf_estimator_array2frame(bf_estimator_array, level2continuum_idx): - """ - Transform a bound-free estimator array to a DataFrame. - - This function transforms a bound-free estimator array with entries - sorted by frequency to a multi-indexed DataFrame sorted by level. - - Parameters - ---------- - bf_estimator_array : numpy.ndarray, dtype float - Array of bound-free estimators (e.g., for the stimulated recombination rate) - with entries sorted by the threshold frequency of the bound-free continuum. - level2continuum_idx : pandas.Series, dtype int - Maps a level MultiIndex (atomic_number, ion_number, level_number) to - the continuum_idx of the corresponding bound-free continuum (which are - sorted by decreasing frequency). - - Returns - ------- - pandas.DataFrame, dtype float - Bound-free estimators indexed by (atomic_number, ion_number, level_number). - """ - bf_estimator_frame = pd.DataFrame( - bf_estimator_array, index=level2continuum_idx.index - ).sort_index() - bf_estimator_frame.columns.name = "Shell No." - return bf_estimator_frame - - -class IndexSetterMixin(object): +class IndexSetterMixin: @staticmethod def set_index(p, photo_ion_idx, transition_type=0, reverse=True): idx = photo_ion_idx.loc[p.index] @@ -428,7 +370,7 @@ def calculate( w, ) else: - gamma_estimator = bf_estimator_array2frame( + gamma_estimator = bound_free_estimator_array2frame( gamma_estimator, level2continuum_idx ) gamma = gamma_estimator * photo_ion_norm_factor.value @@ -493,7 +435,7 @@ def calculate( boltzmann_factor_photo_ion, ) else: - alpha_stim_estimator = bf_estimator_array2frame( + alpha_stim_estimator = bound_free_estimator_array2frame( alpha_stim_estimator, level2continuum_idx ) alpha_stim = alpha_stim_estimator * photo_ion_norm_factor @@ -953,7 +895,6 @@ class TwoPhotonFrequencySampler(ProcessingPlasmaProperty): outputs = ("nu_two_photon_sampler",) def calculate(self, two_photon_emission_cdf): - nus = two_photon_emission_cdf["nu"].values em = two_photon_emission_cdf["cdf"].values @@ -1076,7 +1017,8 @@ class PhotoIonBoltzmannFactor(ProcessingPlasmaProperty): outputs = ("boltzmann_factor_photo_ion",) - def calculate(self, photo_ion_cross_sections, t_electrons): + @staticmethod + def calculate(photo_ion_cross_sections, t_electrons): nu = photo_ion_cross_sections["nu"].values boltzmann_factor = np.exp(-nu[np.newaxis].T / t_electrons * (H / K_B)) diff --git a/tardis/plasma/properties/ion_population.py b/tardis/plasma/properties/ion_population.py index c70bd06965a..794ecde510d 100644 --- a/tardis/plasma/properties/ion_population.py +++ b/tardis/plasma/properties/ion_population.py @@ -1,19 +1,19 @@ import logging +import sys import warnings -import sys import numpy as np import pandas as pd - from scipy import interpolate +from tardis.plasma.exceptions import PlasmaIonizationError from tardis.plasma.properties.base import ProcessingPlasmaProperty from tardis.plasma.properties.continuum_processes import get_ion_multi_index -from tardis.plasma.exceptions import PlasmaIonizationError - logger = logging.getLogger(__name__) +ION_ZERO_THRESHOLD = 1e-20 + __all__ = [ "PhiSahaNebular", "PhiSahaLTE", @@ -279,7 +279,10 @@ class IonNumberDensity(ProcessingPlasmaProperty): ) def __init__( - self, plasma_parent, ion_zero_threshold=1e-20, electron_densities=None + self, + plasma_parent, + ion_zero_threshold=ION_ZERO_THRESHOLD, + electron_densities=None, ): super(IonNumberDensity, self).__init__(plasma_parent) self.ion_zero_threshold = ion_zero_threshold @@ -359,7 +362,7 @@ def calculate(self, phi, partition_function, number_density): ) n_electron_iterations += 1 if n_electron_iterations > 100: - logger.warn( + logger.warning( f"n_electron iterations above 100 ({n_electron_iterations}) -" f" something is probably wrong" ) @@ -480,7 +483,7 @@ def calculate( ) n_electron_iterations += 1 if n_electron_iterations > 100: - logger.warn( + logger.warning( f"n_electron iterations above 100 ({n_electron_iterations}) -" f" something is probably wrong" ) diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py index 487b93ad5a3..c15a46bf196 100644 --- a/tardis/plasma/properties/nlte_rate_equation_solver.py +++ b/tardis/plasma/properties/nlte_rate_equation_solver.py @@ -1,16 +1,27 @@ -import pandas as pd +import logging + import numpy as np +import pandas as pd from scipy.optimize import root from tardis.plasma.properties.base import ProcessingPlasmaProperty -from tardis.plasma.properties.nlte_excitation_data import NLTEExcitationData __all__ = [ - "NLTERateEquationSolver", + "NLTEPopulationSolverRoot", + "NLTEPopulationSolverLU", ] +logger = logging.getLogger(__name__) -class NLTERateEquationSolver(ProcessingPlasmaProperty): +NLTE_POPULATION_SOLVER_MAX_ITERATIONS = 100 +NLTE_POPULATION_SOLVER_TOLERANCE = 1e-3 +NLTE_POPULATION_NEGATIVE_RELATIVE_POPULATION_TOLERANCE = ( + -1e-10 +) # Minimum relative negative population allowed before solver fails +NLTE_POPULATION_SOLVER_CHARGE_CONSERVATION_TOLERANCE = 1e-6 # Arbitrary tolerance for charge conservation, should be changed to a more reasonable value + + +class NLTEPopulationSolverRoot(ProcessingPlasmaProperty): outputs = ("ion_number_density", "electron_densities") def calculate( @@ -63,19 +74,18 @@ def calculate( Returns ------- - ion_number_densities_nlte : pandas.DataFrame + ion_number_densities : pandas.DataFrame Number density with NLTE ionization treatment. - electron_densities_nlte : Series + electron_densities : Series Electron density with NLTE ionization treatment. """ - # nlte_data = NLTEExcitationData(atomic_data.lines, nlte_excitation_species) - will be used in a future PR ( total_photo_ion_coefficients, total_rad_recomb_coefficients, total_coll_ion_coefficients, total_coll_recomb_coefficients, - ) = self.prepare_ion_recomb_coefficients_nlte_ion( + ) = prepare_ion_recomb_coefficients_nlte_ion( gamma, alpha_sp, alpha_stim, @@ -94,17 +104,16 @@ def calculate( ) # dropping the n_e index, as rate_matrix_index's first index is (atomic_numbers, "n_e") 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) + ion_number_density = pd.DataFrame(0.0, index=index, columns=phi.columns) + electron_densities = pd.Series(0.0, index=phi.columns) + ion_numbers = index.get_level_values("ion_number") for shell in phi.columns: - solution_vector = self.prepare_solution_vector( + solution_vector = calculate_balance_vector( number_density[shell], rate_matrix_index, ) - first_guess = self.prepare_first_guess( + first_guess = calculate_first_guess( rate_matrix_index, atomic_numbers, number_density[shell], @@ -115,7 +124,7 @@ def calculate( np.greater_equal(first_guess, 0.0).all() ).all(), "First guess for NLTE solver has negative values, something went wrong." solution = root( - self.population_objective_function, + population_objective_function, first_guess, args=( atomic_numbers, @@ -132,248 +141,36 @@ def calculate( assert ( solution.success ), "No solution for NLTE population equation found or solver takes too long to converge" - ion_number_density_nlte[shell] = solution.x[:-1] - electron_densities_nlte[shell] = 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( - atomic_numbers, - phi_shell, - electron_density, - rate_matrix_index, - total_photo_ion_coefficients, - total_rad_recomb_coefficients, - total_coll_ion_coefficients, - total_coll_recomb_coefficients, - ): - """ - - Parameters - ---------- - phi_shell : pandas.DataFrame - Saha Factors in the current shell - electron_density : float - Guess for electron density in the current shell - rate_matrix_index : pandas.MultiIndex - Index used for constructing the rate matrix - total_photo_ion_coefficients : pandas.DataFrame - Photo ionization coefficients - total_rad_recomb_coefficients : pandas.DataFrame - Radiative recombination coefficients (should get multiplied by electron density) - total_coll_ion_coefficients : pandas.DataFrame - Collisional ionization coefficients (should get multiplied by electron density) - total_coll_recomb_coefficients : pandas.DataFrame - Collisional recombination coefficients (should get multiplied by electron density^2) - - Returns - ------- - pandas.DataFrame - Rate matrix used for NLTE solver. - """ - rate_matrix = pd.DataFrame( - 0.0, columns=rate_matrix_index, index=rate_matrix_index - ) - total_rad_recomb_coefficients = ( - total_rad_recomb_coefficients * electron_density - ) - total_coll_ion_coefficients = ( - total_coll_ion_coefficients * electron_density - ) - total_coll_recomb_coefficients = ( - total_coll_recomb_coefficients * electron_density**2 - ) - for atomic_number in atomic_numbers: - ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values( - "ion_number" - ) - phi_block = phi_shell.loc[atomic_number] - rate_matrix_block = NLTERateEquationSolver.lte_rate_matrix_block( - phi_block, electron_density + ( + ion_number_density[shell], + electron_densities[shell], + ) = check_negative_population( + pd.Series(solution.x[:-1], index=index), + solution.x[-1], + number_density[shell], + ion_numbers, ) - - nlte_ion_numbers = ion_numbers[ - rate_matrix.loc[atomic_number].index.get_level_values( - "level_number" - ) - == "nlte_ion" - ] - # >>> lte_ion_numbers is for future use in NLTE excitation treatment - lte_ion_numbers = ion_numbers[ - rate_matrix.loc[atomic_number].index.get_level_values( - "level_number" - ) - == "lte_ion" - ] - # <<< - for ion_number in nlte_ion_numbers: - rate_matrix_block = NLTERateEquationSolver.set_nlte_ion_rate( - rate_matrix_block, - atomic_number, - ion_number, - total_rad_recomb_coefficients.loc[(atomic_number,)], - total_photo_ion_coefficients.loc[(atomic_number,)], - total_coll_ion_coefficients.loc[(atomic_number,)], - total_coll_recomb_coefficients.loc[(atomic_number,)], + # Check that the electron density is still in line with charge conservation + # after removing negative populations + assert ( + np.abs( + np.sum(ion_numbers * ion_number_density[shell]) + - electron_densities[shell] ) - rate_matrix.loc[ - (atomic_number, slice(None)), (atomic_number) - ] = rate_matrix_block - - charge_conservation_row = ( - NLTERateEquationSolver.prepare_charge_conservation_row( - atomic_numbers - ) - ) - rate_matrix.loc[("n_e", slice(None))] = charge_conservation_row - return rate_matrix - - @staticmethod - def set_nlte_ion_rate( - rate_matrix_block, - atomic_number, - ion_number, - total_rad_recomb_coefficients, - total_photo_ion_coefficients, - total_coll_ion_coefficients, - total_coll_recomb_coefficients, - ): - """Calculates the row for the species treated in NLTE ionization - - Parameters - ---------- - rate_matrix_block : numpy.array - The diagonal block corresponding to current atomic number. - atomic_number : int - Current atomic number - ion_number : int - Current ion number - total_rad_recomb_coefficients : pandas.DataFrame - Rad. recomb. coefficients for current atomic number - total_photo_ion_coefficients : pandas.DataFrame - Photo ionization coefficients for current atomic number - total_coll_ion_coefficients : pandas.DataFrame - Collisional ionization coefficients for current atomic number - total_coll_recomb_coefficients : pandas.DataFrame - Collisional recombination coefficients for current atomic number - - Returns - ------- - numpy.array - Rate matrix block with a changed row for NLTE ionization treatment - """ - ion_coefficients = ( - total_photo_ion_coefficients + total_coll_ion_coefficients - ) - recomb_coefficients = ( - total_rad_recomb_coefficients + total_coll_recomb_coefficients - ) - if atomic_number != ion_number: - ion_coeff_matrix_ion_row = NLTERateEquationSolver.ion_matrix( - ion_coefficients, atomic_number, ion_number - ) - recomb_coeff_matrix_ion_row = NLTERateEquationSolver.recomb_matrix( - recomb_coefficients, atomic_number, ion_number - ) - rate_matrix_block[ion_number, :] = ( - ion_coeff_matrix_ion_row + recomb_coeff_matrix_ion_row - ) - return rate_matrix_block - - @staticmethod - def lte_rate_matrix_block(phi_block, electron_density): - """Creates the generic LTE block for rate matrix. - - Parameters - ---------- - phi_block : pandas.DataFrame - Saha Factors for current atomic number - electron_density : float - Current guess for electron density - - Returns - ------- - numpy.array - LTE block for rate matrix - """ - lte_rate_vector_block = -1.0 * np.hstack([*phi_block.values, -1.0]) - lte_rate_matrix_block = np.diag(lte_rate_vector_block) - n_e_initial = np.ones(len(phi_block)) * electron_density - n_e_matrix = np.diag(n_e_initial, 1) - lte_rate_matrix_block += n_e_matrix - lte_rate_matrix_block[-1, :] = 1.0 - return lte_rate_matrix_block - - @staticmethod - def prepare_phi(phi): - """ - Makes sure that phi does not have any 0 entries. - """ - phi[phi == 0.0] = 1.0e-10 * phi[phi > 0.0].min().min() - return phi - - @staticmethod - def recomb_matrix(recomb_coefficients, atomic_number, ion_number): - """Constructs a recombination rate matrix from the recombination rates. - - Parameters - ---------- - recomb_coefficients : pandas.DataFrame - Recombination coefficients. - atomic_number : int64 - Current atomic number. Used for the dimension of a square matrix. - ion_number : int64 - Current ion number. Used for returning the correct row. + / electron_densities[shell] + < NLTE_POPULATION_SOLVER_CHARGE_CONSERVATION_TOLERANCE + ), "Charge conservation not fulfilled after correcting for negative populations, solver failed." - Returns - ------- - numpy.ndarray - """ - offdiag = np.zeros(atomic_number) - index = recomb_coefficients.index - for i in index: - offdiag[i] = recomb_coefficients.loc[i] - diag = np.hstack([np.zeros(1), -offdiag]) - return (np.diag(diag) + np.diag(offdiag, k=1))[ion_number, :] + # TODO: change the jacobian and rate matrix to use shell id and get coefficients from the attribute of the class. - @staticmethod - def ion_matrix(ion_coefficients, atomic_number, ion_number): - """Constructs an ionization rate matrix from the ionization rates. + return ion_number_density, electron_densities - Parameters - ---------- - ion_coefficients : pandas.DataFrame - Recombination coefficients. - atomic_number : int64 - Current atomic number. Used for the dimension of a square matrix. - ion_number : int64 - Current ion number. Used for returning the correct row. - Returns - ------- - numpy.ndarray - """ - offdiag = np.zeros(atomic_number) - index = ion_coefficients.index - for i in index: - offdiag[i] = ion_coefficients.loc[i] - diag = np.hstack([-offdiag, np.zeros(1)]) - return (np.diag(diag) + np.diag(offdiag, k=-1))[ion_number, :] - - @staticmethod - def prepare_charge_conservation_row(atomic_numbers): - """Prepares the last row of the rate_matrix. This row corresponds to the charge - density equation.""" - charge_conservation_row = [] - for atomic_number in atomic_numbers: - charge_conservation_row.append(np.arange(0, atomic_number + 1)) - charge_conservation_row = np.hstack([*charge_conservation_row, -1]) - # TODO needs to be modified for use in nlte_excitation - return charge_conservation_row +class NLTEPopulationSolverLU(ProcessingPlasmaProperty): + outputs = ("ion_number_density", "electron_densities") - @staticmethod - def prepare_ion_recomb_coefficients_nlte_ion( + def calculate( + self, gamma, alpha_sp, alpha_stim, @@ -382,10 +179,12 @@ def prepare_ion_recomb_coefficients_nlte_ion( partition_function, levels, level_boltzmann_factor, + phi, + rate_matrix_index, + number_density, + nlte_excitation_species, ): - """ - Prepares the ionization and recombination coefficients by grouping them for - ion numbers. + """Calculates ion number densities and electron densities using NLTE ionization using an iterative LU decomposition approach. Parameters ---------- @@ -407,526 +206,968 @@ def prepare_ion_recomb_coefficients_nlte_ion( Index of filtered atomic data. level_boltzmann_factor : pandas.DataFrame General Boltzmann factor. + phi : pandas.DataFrame + Saha Factors. + 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". + number_density : pandas.DataFrame + Number density in each shell for each species. + nlte_excitation_species : list + List of species treated in NLTE excitation. + Returns ------- - total_photo_ion_coefficients - Photoionization coefficients grouped by atomic number and ion number. - total_rad_recomb_coefficients - Radiative recombination coefficients grouped by atomic number and ion number. - total_coll_ion_coefficients - Collisional ionization coefficients grouped by atomic number and ion number. - total_coll_recomb_coefficients - Collisional recombination coefficients grouped by atomic number and ion number. + ion_number_densities : pandas.DataFrame + Number density with NLTE ionization treatment. + electron_densities : Series + Electron density with NLTE ionization treatment. """ - indexer = pd.Series( - np.arange(partition_function.shape[0]), - index=partition_function.index, - ) - _ion2level_idx = indexer.loc[levels.droplevel("level_number")].values - partition_function_broadcast = partition_function.values[_ion2level_idx] - level_population_fraction = pd.DataFrame( - level_boltzmann_factor.values / partition_function_broadcast, - index=levels, - ) - total_photo_ion_coefficients = ( - (level_population_fraction.loc[gamma.index] * gamma) - .groupby(level=("atomic_number", "ion_number")) - .sum() - ) - - total_rad_recomb_coefficients = ( - (alpha_sp + alpha_stim) - .groupby(level=["atomic_number", "ion_number"]) - .sum() - ) - total_coll_ion_coefficients = ( - ( - level_population_fraction.loc[coll_ion_coeff.index] - * coll_ion_coeff - ) - .groupby(level=("atomic_number", "ion_number")) - .sum() - ) - total_coll_recomb_coefficients = ( - (coll_recomb_coeff) - .groupby(level=("atomic_number", "ion_number")) - .sum() - ) - return ( + ( total_photo_ion_coefficients, total_rad_recomb_coefficients, total_coll_ion_coefficients, total_coll_recomb_coefficients, + ) = prepare_ion_recomb_coefficients_nlte_ion( + gamma, + alpha_sp, + alpha_stim, + coll_ion_coeff, + coll_recomb_coeff, + partition_function, + levels, + level_boltzmann_factor, ) + # TODO: Don't create the rate_matrix_index with n_e in the first place + rate_matrix_index = rate_matrix_index.drop("n_e") + + initial_electron_densities = number_density.sum(axis=0) + atomic_numbers = rate_matrix_index.get_level_values( + "atomic_number" + ).unique() + + index = rate_matrix_index.droplevel("level_number") + electron_densities = initial_electron_densities.copy() + ion_number_density = pd.DataFrame(0.0, index=index, columns=phi.columns) + ion_numbers = index.get_level_values("ion_number") + + # Ordering of loops is important to allow for + # easier parallelization in the future + logger.info("Starting NLTE ionization solver") + for shell in phi.columns: + iteration = 0 + balance_vector = calculate_balance_vector( + number_density[shell], + rate_matrix_index, + include_electron_density=False, + ) + while iteration < NLTE_POPULATION_SOLVER_MAX_ITERATIONS: + rate_matrix = calculate_rate_matrix( + atomic_numbers, + phi[shell], + electron_densities[shell], + rate_matrix_index, + total_photo_ion_coefficients[shell], + total_rad_recomb_coefficients[shell], + total_coll_ion_coefficients[shell], + total_coll_recomb_coefficients[shell], + set_charge_conservation=False, + ) + # TODO: Solve for each element individually + # and handle errors in the solver + ion_solution = np.linalg.solve(rate_matrix, balance_vector) + + ion_solution, electron_solution = check_negative_population( + pd.Series(ion_solution, index=index), + None, # Electron density will be recalculated from ion number densities + number_density[shell], + ion_numbers, + ) + delta_ion, delta_electron = self.calculate_lu_solver_delta( + ion_solution, + electron_solution, + ion_number_density[shell], + electron_densities[shell], + ) + ion_number_density[shell] = ion_solution + electron_densities[shell] = electron_solution + if ( + np.all(np.abs(delta_ion) < NLTE_POPULATION_SOLVER_TOLERANCE) + and np.abs(delta_electron) + < NLTE_POPULATION_SOLVER_TOLERANCE + ): + logger.debug( + "NLTE ionization solver converged after {} iterations for shell {}".format( + iteration, shell + ) + ) + break + + iteration += 1 + if iteration == NLTE_POPULATION_SOLVER_MAX_ITERATIONS: + logger.warning( + f"NLTE ionization solver did not converge for shell {shell} " + ) + break + + logger.info("NLTE ionization solver finished") + + return ion_number_density, electron_densities + @staticmethod - def jacobian_matrix( - atomic_numbers, - populations, - rate_matrix, - rate_matrix_index, - total_rad_recomb_coefficients, - total_coll_ion_coefficients, - total_coll_recomb_coefficients, + def calculate_lu_solver_delta( + ion_solution, electron_solution, ion_number_density, electron_densities ): - """Creates the jacobian matrix used for NLTE ionization solver + """Calculates relative change between new solution and old value. Parameters ---------- - populations : numpy.array - Ion populations, electron density - rate_matrix : pandas.DataFrame - Rate matrix used for NLTE solver. - 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_rad_recomb_coefficients : pandas.DataFrame - Radiative recombination coefficients grouped by atomic number and ion number. - total_coll_ion_coefficients : pandas.DataFrame - Collisional ionization coefficients(should get multiplied by electron density). - total_coll_recomb_coefficients : pandas.DataFrame - Collisional recombination coefficients(should get multiplied by electron density). + ion_solution : numpy.array + Solution vector for the NLTE ionization solver. + electron_solution : float + Solution for the electron density. Returns ------- numpy.array - Jacobian matrix used for NLTE ionization solver + Change in ion number densities. + float + Change in electron density. """ - # TODO: for future use, can be vectorized. - index = 0 - jacobian_matrix = rate_matrix.copy().values - jacobian_matrix[:-1, -1] = populations[1:] - for atomic_number in atomic_numbers: - for i in range(index, index + atomic_number): - if rate_matrix_index[i][2] == "nlte_ion": - jacobian_matrix[ - i, -1 - ] = NLTERateEquationSolver.deriv_matrix_block( - atomic_number, - total_rad_recomb_coefficients.loc[(atomic_number,)], - total_coll_ion_coefficients.loc[(atomic_number,)], - total_coll_recomb_coefficients.loc[(atomic_number,)], - populations[index : index + atomic_number + 1], - populations[-1], - )[ - i - index - ] - index += atomic_number + 1 - jacobian_matrix[index - 1, -1] = 0 # number conservation row - return jacobian_matrix + assert np.all( + ion_solution >= 0.0 + ), "Negative ion number density found, this should not happen." + assert ( + electron_solution >= 0.0 + ), "Negative electron density found, this should not happen." + delta_ion = (ion_number_density - ion_solution) / ion_solution + delta_electron = ( + electron_densities - electron_solution + ) / electron_solution + return delta_ion, delta_electron - @staticmethod - def deriv_matrix_block( - atomic_number, + +def check_negative_population( + ion_number_density, electron_densities, number_density, ion_numbers +): + """ + Checks if negative populations are present in the solution. If the relative + negative population is smaller than the tolerance, the negative population + is set to zero. If the relative negative population is larger than the + tolerance, the solver failed. + + Parameters + ---------- + ion_number_density : pandas.Series + Number density with NLTE ionization treatment. + electron_densities : float + Electron density with NLTE ionization treatment. + number_density : pandas.Series + Number density of all present species. + ion_numbers : numpy.array + Ion numbers of all present species. + + Returns + ------- + ion_number_density : pandas.Series + Number density with NLTE ionization treatment. + electron_densities : float + Electron density with NLTE ionization treatment. + """ + # This can probably be done in a more concise way, but since the number + # densities can be zero, we need to check that we don't divide by zero. + for atom_number in number_density.index: + if number_density.loc[atom_number] == 0.0: + # Not sure what to do here, but if the number density is zero, the + # ion number density should also be zero. + # There could be an edge case where the number density is zero, but + # e.g. two ion number densities are -1e-30 and 1e-30, which would + # result in a zero number density. The above value would be acceptable + # within numerical tolerances, but I don't know how to distinguish + # between this case and the case where the ion number density is + # to large to be numerical. + continue + else: + for ion_number in ion_number_density.loc[atom_number].index: + if (ion_number_density.loc[atom_number, ion_number] < 0.0) and ( + ion_number_density.loc[atom_number, ion_number] + / number_density.loc[atom_number] + > NLTE_POPULATION_NEGATIVE_RELATIVE_POPULATION_TOLERANCE + ): + ion_number_density.loc[atom_number, ion_number] = 0.0 + + assert np.greater_equal( + ion_number_density, 0.0 + ).all(), "Negative ion number density found, solver failed." + + if electron_densities is None: + # For the root solver, we don't want to recalculate the electron density + electron_densities = np.sum(ion_numbers * ion_number_density) + assert ( + electron_densities >= 0.0 + ), "Negative electron density found, solver failed." + return ion_number_density, electron_densities + + +def prepare_ion_recomb_coefficients_nlte_ion( + gamma, + alpha_sp, + alpha_stim, + coll_ion_coeff, + coll_recomb_coeff, + partition_function, + levels, + level_boltzmann_factor, +): + """ + Prepares the ionization and recombination coefficients by grouping them for + ion numbers. + + Parameters + ---------- + gamma : pandas.DataFrame + The rate coefficient for radiative ionization. + alpha_sp : pandas.DataFrame + The rate coefficient for spontaneous recombination. + alpha_stim : pandas.DataFrame + The rate coefficient for stimulated recombination. + coll_ion_coeff : pandas.DataFrame + The rate coefficient for collisional ionization in the Seaton + approximation. + coll_recomb_coeff : pandas.DataFrame + The rate coefficient for collisional recombination. + partition_function : pandas.DataFrame + General partition function. Indexed by atomic number, ion number. + levels : MultiIndex + (atomic_number, ion_number, level_number) + Index of filtered atomic data. + level_boltzmann_factor : pandas.DataFrame + General Boltzmann factor. + + Returns + ------- + total_photo_ion_coefficients + Photoionization coefficients grouped by atomic number and ion number. + total_rad_recomb_coefficients + Radiative recombination coefficients grouped by atomic number and ion number. + total_coll_ion_coefficients + Collisional ionization coefficients grouped by atomic number and ion number. + total_coll_recomb_coefficients + Collisional recombination coefficients grouped by atomic number and ion number. + """ + indexer = pd.Series( + np.arange(partition_function.shape[0]), + index=partition_function.index, + ) + _ion2level_idx = indexer.loc[levels.droplevel("level_number")].values + partition_function_broadcast = partition_function.values[_ion2level_idx] + level_population_fraction = pd.DataFrame( + level_boltzmann_factor.values / partition_function_broadcast, + index=levels, + ) + total_photo_ion_coefficients = ( + (level_population_fraction.loc[gamma.index] * gamma) + .groupby(level=("atomic_number", "ion_number")) + .sum() + ) + + total_rad_recomb_coefficients = ( + (alpha_sp + alpha_stim) + .groupby(level=["atomic_number", "ion_number"]) + .sum() + ) + total_coll_ion_coefficients = ( + (level_population_fraction.loc[coll_ion_coeff.index] * coll_ion_coeff) + .groupby(level=("atomic_number", "ion_number")) + .sum() + ) + total_coll_recomb_coefficients = ( + (coll_recomb_coeff).groupby(level=("atomic_number", "ion_number")).sum() + ) + return ( + total_photo_ion_coefficients, total_rad_recomb_coefficients, total_coll_ion_coefficients, total_coll_recomb_coefficients, - current_ion_number_densities, - current_electron_density, - ): - """Calculates the dot product of the derivative of rate matrix and ion number densities+electron density column. + ) - Parameters - ---------- - atomic_number : int64 - Current atomic number - total_rad_recomb_coefficients : pandas.DataFrame - Radiative recombination coefficients grouped by atomic number and ion number. - total_coll_ion_coefficients : pandas.DataFrame - Collisional ionization coefficients. - total_coll_recomb_coefficients : pandas.DataFrame - Collisional recombination coefficients. - current_ion_number_densities : numpy.array - Current ion number densities for the current atomic number. - current_electron_density : float64 - Current electron density - Returns - ------- - numpy.array - Returns the part of the last column of the jacobian matrix, corresponding to atomic number. - """ - ion_numbers = np.arange(0, atomic_number) - radiative_rate_coeff_matrix = NLTERateEquationSolver.recomb_matrix( - total_rad_recomb_coefficients, atomic_number, ion_numbers - ) - coll_recomb_matrix = ( - NLTERateEquationSolver.recomb_matrix( - total_coll_recomb_coefficients, atomic_number, ion_numbers - ) - * current_electron_density - * 2 - ) - coll_ion_coeff_matrix = NLTERateEquationSolver.ion_matrix( - total_coll_ion_coefficients, atomic_number, ion_numbers - ) - deriv_matrix = ( - radiative_rate_coeff_matrix - + coll_ion_coeff_matrix - + coll_recomb_matrix - ) - return np.dot(deriv_matrix, current_ion_number_densities) +def calculate_balance_vector( + number_density, + rate_matrix_index, + include_electron_density=True, +): + """Constructs the balance vector for the NLTE ionization solver set of equations by combining + all solution verctor blocks. - def prepare_first_guess( - self, - rate_matrix_index, - atomic_numbers, - number_density, - electron_density, - ): - """Constructs a first guess for ion number densities and electron density, where all species are singly ionized. + Parameters + ---------- + number_density : pandas.DataFrame + Number densities of all present species. + rate_matrix_index : pandas.MultiIndex + (atomic_number, ion_number, treatment type) - Parameters - ---------- - rate_matrix_index : pandas.MultiIndex - (atomic_number, ion_number, treatment type) - atomic_numbers : numpy.array - All atomic numbers present in the plasma. - number_density : pandas.DataFrame - Number density of present species. - electron_density : float - Current value of electron density. + Returns + ------- + numpy.array + Solution vector for the NLTE ionization solver. + """ + atomic_numbers = number_density.index + balance_array = [] + for atomic_number in atomic_numbers: + needed_number_of_elements = ( + rate_matrix_index.get_level_values("atomic_number") == atomic_number + ).sum() + balance_vector_block = np.zeros(needed_number_of_elements) + balance_vector_block[-1] = number_density.loc[atomic_number] + balance_array.append(balance_vector_block) + if include_electron_density: + balance_array.append(np.array([0.0])) + return np.hstack(balance_array) - Returns - ------- - numpy.array - Guess for ion number densities and electron density. - """ - first_guess = pd.Series(0.0, index=rate_matrix_index) - for atomic_number in atomic_numbers: - first_guess.loc[(atomic_number, 1)].iloc[0] = number_density.loc[ - atomic_number - ] - # TODO: After the first iteration, the new guess can be the old solution. - first_guess = first_guess.values - first_guess[-1] = electron_density - return first_guess - def population_objective_function( - self, - populations, +def calculate_first_guess( + rate_matrix_index, + atomic_numbers, + number_density, + electron_density, +): + """Constructs a first guess for ion number densities and electron density, where all species are singly ionized. + + Parameters + ---------- + rate_matrix_index : pandas.MultiIndex + (atomic_number, ion_number, treatment type) + atomic_numbers : numpy.array + All atomic numbers present in the plasma. + number_density : pandas.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. + """ + first_guess = pd.Series(0.0, index=rate_matrix_index) + for atomic_number in atomic_numbers: + first_guess.loc[(atomic_number, 1)].iloc[0] = number_density.loc[ + atomic_number + ] + # TODO: After the first iteration, the new guess can be the old solution. + first_guess = first_guess.values + first_guess[-1] = electron_density + return first_guess + + +def population_objective_function( + 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, + set_charge_conservation=True, +): + """Main set of equations for the NLTE ionization solver. + + To solve the statistical equilibrium equations, we need to find the root + of the objective function A*x - B, where x are 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 : pandas.DataFrame + Saha Factors of the current shell. + solution_vector : numpy.array + Solution vector for the set of equations. + rate_matrix_index : pandas.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 : pandas.DataFrame + Photo ion. coefficients for current atomic number + total_rad_recomb_coefficients : pandas.DataFrame + Radiative recombination coefficients for current atomic number + total_coll_ion_coefficients : pandas.DataFrame + Collisional ionization coefficients for current atomic number + total_coll_recomb_coefficients : pandas.DataFrame + Coll. recomb. coefficients for current atomic number + set_charge_conservation : bool + If True, sets the last row of the rate matrix to the charge conservation equation. + + Returns + ------- + (numpy.array, numpy.array) + Returns the objective function and jacobian of the rate matrix in a tuple. + """ + electron_density = populations[-1] + rate_matrix = calculate_rate_matrix( atomic_numbers, phi, - solution_vector, + electron_density, 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. + set_charge_conservation=set_charge_conservation, + ) + jacobian_matrix = calculate_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, + ) - To solve the statistical equilibrium equations, we need to find the root - of the objective function A*x - B, where x are 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 : pandas.DataFrame - Saha Factors of the current shell. - solution_vector : numpy.array - Solution vector for the set of equations. - rate_matrix_index : pandas.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 : pandas.DataFrame - Photo ion. coefficients for current atomic number - total_rad_recomb_coefficients : pandas.DataFrame - Radiative recombination coefficients for current atomic number - total_coll_ion_coefficients : pandas.DataFrame - Collisional ionization coefficients for current atomic number - total_coll_recomb_coefficients : pandas.DataFrame - Coll. recomb. coefficients for current atomic number - Returns - ------- - (numpy.array, numpy.array) - Returns the objective function and jacobian of the rate matrix in a tuple. - """ - 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, +def calculate_rate_matrix( + atomic_numbers, + phi_shell, + electron_density, + rate_matrix_index, + total_photo_ion_coefficients, + total_rad_recomb_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, + set_charge_conservation=True, +): + """ + + Parameters + ---------- + phi_shell : pandas.DataFrame + Saha Factors in the current shell + electron_density : float + Guess for electron density in the current shell + rate_matrix_index : pandas.MultiIndex + Index used for constructing the rate matrix + total_photo_ion_coefficients : pandas.DataFrame + Photo ionization coefficients + total_rad_recomb_coefficients : pandas.DataFrame + Radiative recombination coefficients (should get multiplied by electron density) + total_coll_ion_coefficients : pandas.DataFrame + Collisional ionization coefficients (should get multiplied by electron density) + total_coll_recomb_coefficients : pandas.DataFrame + Collisional recombination coefficients (should get multiplied by electron density^2) + set_charge_conservation : bool + If True, sets the last row of the rate matrix to the charge conservation equation. + + Returns + ------- + pandas.DataFrame + Rate matrix used for NLTE solver. + """ + rate_matrix = pd.DataFrame( + 0.0, columns=rate_matrix_index, index=rate_matrix_index + ) + total_rad_recomb_coefficients = ( + total_rad_recomb_coefficients * electron_density + ) + total_coll_ion_coefficients = total_coll_ion_coefficients * electron_density + total_coll_recomb_coefficients = ( + total_coll_recomb_coefficients * electron_density**2 + ) + for atomic_number in atomic_numbers: + ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values( + "ion_number" ) - 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, + phi_block = phi_shell.loc[atomic_number] + rate_matrix_block = lte_rate_matrix_block(phi_block, electron_density) + + nlte_ion_numbers = ion_numbers[ + rate_matrix.loc[atomic_number].index.get_level_values( + "level_number" + ) + == "nlte_ion" + ] + # >>> lte_ion_numbers is for future use in NLTE excitation treatment + if False: + lte_ion_numbers = ion_numbers[ + rate_matrix.loc[atomic_number].index.get_level_values( + "level_number" + ) + == "lte_ion" + ] + # <<< + for ion_number in nlte_ion_numbers: + rate_matrix_block = set_nlte_ion_rate( + rate_matrix_block, + atomic_number, + ion_number, + total_rad_recomb_coefficients.loc[(atomic_number,)], + total_photo_ion_coefficients.loc[(atomic_number,)], + total_coll_ion_coefficients.loc[(atomic_number,)], + total_coll_recomb_coefficients.loc[(atomic_number,)], + ) + rate_matrix.loc[ + (atomic_number, slice(None)), (atomic_number) + ] = rate_matrix_block + + charge_conservation_row = calculate_charge_conservation_row(atomic_numbers) + if set_charge_conservation: + rate_matrix.loc[("n_e", slice(None))] = charge_conservation_row + return rate_matrix + + +def calculate_jacobian_matrix( + atomic_numbers, + populations, + rate_matrix, + rate_matrix_index, + total_rad_recomb_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, +): + """Creates the jacobian matrix used for NLTE ionization solver + + Parameters + ---------- + populations : numpy.array + Ion populations, electron density + rate_matrix : pandas.DataFrame + Rate matrix used for NLTE solver. + 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_rad_recomb_coefficients : pandas.DataFrame + Radiative recombination coefficients grouped by atomic number and ion number. + total_coll_ion_coefficients : pandas.DataFrame + Collisional ionization coefficients(should get multiplied by electron density). + total_coll_recomb_coefficients : pandas.DataFrame + Collisional recombination coefficients(should get multiplied by electron density). + + Returns + ------- + numpy.array + Jacobian matrix used for NLTE ionization solver + """ + # TODO: for future use, can be vectorized. + index = 0 + jacobian_matrix = rate_matrix.copy().values + jacobian_matrix[:-1, -1] = populations[1:] + for atomic_number in atomic_numbers: + for i in range(index, index + atomic_number): + if rate_matrix_index[i][2] == "nlte_ion": + jacobian_matrix[i, -1] = deriv_matrix_block( + atomic_number, + total_rad_recomb_coefficients.loc[(atomic_number,)], + total_coll_ion_coefficients.loc[(atomic_number,)], + total_coll_recomb_coefficients.loc[(atomic_number,)], + populations[index : index + atomic_number + 1], + populations[-1], + )[i - index] + index += atomic_number + 1 + jacobian_matrix[index - 1, -1] = 0 # number conservation row + return jacobian_matrix + + +def lte_rate_matrix_block(phi_block, electron_density): + """Creates the generic LTE block for rate matrix. + + Parameters + ---------- + phi_block : pandas.DataFrame + Saha Factors for current atomic number + electron_density : float + Current guess for electron density + + Returns + ------- + numpy.array + LTE block for rate matrix + """ + lte_rate_vector_block = -1.0 * np.hstack([*phi_block.values, -1.0]) + lte_rate_matrix_block = np.diag(lte_rate_vector_block) + n_e_initial = np.ones(len(phi_block)) * electron_density + n_e_matrix = np.diag(n_e_initial, 1) + lte_rate_matrix_block += n_e_matrix + lte_rate_matrix_block[-1, :] = 1.0 + return lte_rate_matrix_block + + +def set_nlte_ion_rate( + rate_matrix_block, + atomic_number, + ion_number, + total_rad_recomb_coefficients, + total_photo_ion_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, +): + """Calculates the row for the species treated in NLTE ionization + + Parameters + ---------- + rate_matrix_block : numpy.array + The diagonal block corresponding to current atomic number. + atomic_number : int + Current atomic number + ion_number : int + Current ion number + total_rad_recomb_coefficients : pandas.DataFrame + Rad. recomb. coefficients for current atomic number + total_photo_ion_coefficients : pandas.DataFrame + Photo ionization coefficients for current atomic number + total_coll_ion_coefficients : pandas.DataFrame + Collisional ionization coefficients for current atomic number + total_coll_recomb_coefficients : pandas.DataFrame + Collisional recombination coefficients for current atomic number + + Returns + ------- + numpy.array + Rate matrix block with a changed row for NLTE ionization treatment + """ + ion_coefficients = ( + total_photo_ion_coefficients + total_coll_ion_coefficients + ) + recomb_coefficients = ( + total_rad_recomb_coefficients + total_coll_recomb_coefficients + ) + if atomic_number != ion_number: + ion_coeff_matrix_ion_row = ion_matrix( + ion_coefficients, atomic_number, ion_number ) - return ( - np.dot(rate_matrix.values, populations) - solution_vector, - jacobian_matrix, + recomb_coeff_matrix_ion_row = recomb_matrix( + recomb_coefficients, atomic_number, ion_number ) + rate_matrix_block[ion_number, :] = ( + ion_coeff_matrix_ion_row + recomb_coeff_matrix_ion_row + ) + return rate_matrix_block - def solution_vector_block(self, number_density, needed_number_of_elements): - """Block of the solution vector for the current atomic number. - Block for the solution vector has the form (0, 0, ..., 0, number_density). - Length is equal to the number of present ions and number of - levels, if treated with NLTE excitation. +def calculate_charge_conservation_row(atomic_numbers): + """calculate the last row of the rate_matrix. This row corresponds to the charge + density equation. + """ + charge_conservation_row = [] + for atomic_number in atomic_numbers: + charge_conservation_row.append(np.arange(0, atomic_number + 1)) + charge_conservation_row = np.hstack([*charge_conservation_row, -1]) + # TODO needs to be modified for use in nlte_excitation + return charge_conservation_row - Parameters - ---------- - number_density : float - Number density of the current atomic number. - needed_number_of_elements : int - Number of needed elements in the solution vector for current atomic number. - Returns - ------- - numpy.array - Block of the solution vector corresponding to the current atomic number. - """ - solution_vector = np.zeros(needed_number_of_elements) - solution_vector[-1] = number_density - return solution_vector +def deriv_matrix_block( + atomic_number, + total_rad_recomb_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, + current_ion_number_densities, + current_electron_density, +): + """Calculates the dot product of the derivative of rate matrix and ion number densities+electron density column. - def prepare_solution_vector(self, number_density, rate_matrix_index): - """Constructs the solution vector for the NLTE ionization solver set of equations by combining - all solution verctor blocks. + Parameters + ---------- + atomic_number : int64 + Current atomic number + total_rad_recomb_coefficients : pandas.DataFrame + Radiative recombination coefficients grouped by atomic number and ion number. + total_coll_ion_coefficients : pandas.DataFrame + Collisional ionization coefficients. + total_coll_recomb_coefficients : pandas.DataFrame + Collisional recombination coefficients. + current_ion_number_densities : numpy.array + Current ion number densities for the current atomic number. + current_electron_density : float64 + Current electron density - Parameters - ---------- - number_density : pandas.DataFrame - Number densities of all present species. - rate_matrix_index : pandas.MultiIndex - (atomic_number, ion_number, treatment type) + Returns + ------- + numpy.array + Returns the part of the last column of the jacobian matrix, corresponding to atomic number. + """ + ion_numbers = np.arange(0, atomic_number) + radiative_rate_coeff_matrix = recomb_matrix( + total_rad_recomb_coefficients, atomic_number, ion_numbers + ) + coll_recomb_matrix = ( + recomb_matrix( + total_coll_recomb_coefficients, atomic_number, ion_numbers + ) + * current_electron_density + * 2 + ) + coll_ion_coeff_matrix = ion_matrix( + total_coll_ion_coefficients, atomic_number, ion_numbers + ) + deriv_matrix = ( + radiative_rate_coeff_matrix + coll_ion_coeff_matrix + coll_recomb_matrix + ) + return np.dot(deriv_matrix, current_ion_number_densities) - Returns - ------- - numpy.array - Solution vector for the NLTE ionization solver. - """ - atomic_numbers = number_density.index - solution_array = [] - for atomic_number in atomic_numbers: - needed_number_of_elements = ( - rate_matrix_index.get_level_values("atomic_number") - == atomic_number - ).sum() - solution_array.append( - self.solution_vector_block( - number_density.loc[atomic_number], - needed_number_of_elements, - ) - ) - solution_vector = np.hstack(solution_array + [0]) - return solution_vector - @staticmethod - def prepare_bound_bound_rate_matrix( - number_of_levels, +def ion_matrix(ion_coefficients, atomic_number, ion_number): + """Constructs an ionization rate matrix from the ionization rates. + + Parameters + ---------- + ion_coefficients : pandas.DataFrame + Recombination coefficients. + atomic_number : int64 + Current atomic number. Used for the dimension of a square matrix. + ion_number : int64 + Current ion number. Used for returning the correct row. + + Returns + ------- + numpy.ndarray + """ + offdiag = np.zeros(atomic_number) + index = ion_coefficients.index + for i in index: + offdiag[i] = ion_coefficients.loc[i] + diag = np.hstack([-offdiag, np.zeros(1)]) + return (np.diag(diag) + np.diag(offdiag, k=-1))[ion_number, :] + + +def recomb_matrix(recomb_coefficients, atomic_number, ion_number): + """Constructs a recombination rate matrix from the recombination rates. + + Parameters + ---------- + recomb_coefficients : pandas.DataFrame + Recombination coefficients. + atomic_number : int64 + Current atomic number. Used for the dimension of a square matrix. + ion_number : int64 + Current ion number. Used for returning the correct row. + + Returns + ------- + numpy.ndarray + """ + offdiag = np.zeros(atomic_number) + index = recomb_coefficients.index + for i in index: + offdiag[i] = recomb_coefficients.loc[i] + diag = np.hstack([np.zeros(1), -offdiag]) + return (np.diag(diag) + np.diag(offdiag, k=1))[ion_number, :] + + +#################### +# Unused functions # +# vvvvvvvvvvvvvvvv # +#################### +def prepare_phi(phi): + """ + Makes sure that phi does not have any 0 entries. + """ + phi[phi == 0.0] = 1.0e-10 * phi[phi > 0.0].min().min() + return phi + + +def prepare_bound_bound_rate_matrix( + number_of_levels, + lines_index, + r_ul_index, + r_ul_matrix, + r_lu_index, + r_lu_matrix, + beta_sobolev, +): + """Calculates a matrix with bound-bound rates for NLTE excitation treatment. + + Parameters + ---------- + number_of_levels : int + Number of levels for the specified species. + lines_index : numpy.array + Index of lines in nlte_data. + r_ul_index : numpy.array + Index used for r_ul matrix + r_ul_matrix : numpy.array + Matrix with the rates(upper to lower transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) + (number_of_levels, number_of_levels, number_of_shells) + r_lu_index : numpy.array + Index used for r_lu matrix + r_lu_matrix : numpy.array + r_lu_matrix : numpy.array + Matrix with the rates(lower to upper transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) + (number_of_levels, number_of_levels, number_of_shells) + beta_sobolev : pandas.DataFrame + Beta Sobolev factors. + + Returns + ------- + numpy.array (number of levels, number of levels) + Matrix with excitation-deexcitation rates(should be added to NLTE rate matrix for excitation treatment). + NOTE: only works with ONE ion number treated in NLTE excitation AT ONCE. + """ + number_of_shells = beta_sobolev.shape[1] + try: + beta_sobolev_filtered = beta_sobolev.iloc[lines_index] + except AttributeError: + beta_sobolev_filtered = beta_sobolev + r_ul_matrix_reshaped = r_ul_matrix.reshape( + (number_of_levels**2, number_of_shells) + ) + r_lu_matrix_reshaped = r_lu_matrix.reshape( + (number_of_levels**2, number_of_shells) + ) + r_ul_matrix_reshaped[r_ul_index] *= beta_sobolev_filtered + r_lu_matrix_reshaped[r_lu_index] *= beta_sobolev_filtered + rates_matrix_bound_bound = r_ul_matrix + r_lu_matrix + for i in range(number_of_levels): + rates_matrix_bound_bound[i, i] = -rates_matrix_bound_bound[:, i].sum( + axis=0 + ) + return rates_matrix_bound_bound + + +def prepare_r_uls_r_lus( + number_of_levels, + number_of_shells, + j_blues, + excitation_species, + nlte_data, +): + """Calculates part of rate matrices for bound bound interactions. + + Parameters + ---------- + number_of_levels : int + Number of levels for the NLTE excitation species. + number_of_shells : int + Number of shells. + j_blues : pandas.DataFrame, dtype float + Mean intensities in the blue wings of the line transitions. + excitation_species : tuple + Species treated in NLTE excitation. + nlte_data : NLTEExcitationData + Data relevant to NLTE excitation species. + + Returns + ------- + lines_index : numpy.array + Index of lines in nlte_data. + number_of_levels : int + Number of levels for the specified species. + r_ul_index : numpy.array + Index used for r_ul matrix + r_ul_matrix_reshaped : numpy.array + Matrix with the rates(upper to lower transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) + r_lu_index : numpy.array + Index used for r_lu matrix + r_lu_matrix_reshaped : numpy.array + Matrix with the rates(lower to upper transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) + """ + # number_of_levels = atomic_data_levels.energy.loc[ + # excitation_species + # ].count() do this in the solver + lnl = nlte_data.lines_level_number_lower[excitation_species] + lnu = nlte_data.lines_level_number_upper[excitation_species] + (lines_index,) = nlte_data.lines_idx[excitation_species] + + try: + j_blues_filtered = j_blues.iloc[lines_index] + except AttributeError: + j_blues_filtered = j_blues + A_uls = nlte_data.A_uls[excitation_species] + B_uls = nlte_data.B_uls[excitation_species] + B_lus = nlte_data.B_lus[excitation_species] + r_lu_index = lnu * number_of_levels + lnl + r_ul_index = lnl * number_of_levels + lnu + r_ul_matrix = np.zeros( + (number_of_levels, number_of_levels, number_of_shells), + dtype=np.float64, + ) + r_ul_matrix_reshaped = r_ul_matrix.reshape( + (number_of_levels**2, number_of_shells) + ) + r_ul_matrix_reshaped[r_ul_index] = ( + A_uls[np.newaxis].T + B_uls[np.newaxis].T * j_blues_filtered + ) + r_lu_matrix = np.zeros_like(r_ul_matrix) + r_lu_matrix_reshaped = r_lu_matrix.reshape( + (number_of_levels**2, number_of_shells) + ) + r_lu_matrix_reshaped[r_lu_index] = B_lus[np.newaxis].T * j_blues_filtered + return ( lines_index, r_ul_index, r_ul_matrix, r_lu_index, r_lu_matrix, - beta_sobolev, - ): - """Calculates a matrix with bound-bound rates for NLTE excitation treatment. + ) - Parameters - ---------- - number_of_levels : int - Number of levels for the specified species. - lines_index : numpy.array - Index of lines in nlte_data. - r_ul_index : numpy.array - Index used for r_ul matrix - r_ul_matrix : numpy.array - Matrix with the rates(upper to lower transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) - (number_of_levels, number_of_levels, number_of_shells) - r_lu_index : numpy.array - Index used for r_lu matrix - r_lu_matrix : numpy.array - r_lu_matrix : numpy.array - Matrix with the rates(lower to upper transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) - (number_of_levels, number_of_levels, number_of_shells) - beta_sobolev : pandas.DataFrame - Beta Sobolev factors. - Returns - ------- - numpy.array (number of levels, number of levels) - Matrix with excitation-deexcitation rates(should be added to NLTE rate matrix for excitation treatment). - NOTE: only works with ONE ion number treated in NLTE excitation AT ONCE. - """ - number_of_shells = beta_sobolev.shape[1] - try: - beta_sobolev_filtered = beta_sobolev.iloc[lines_index] - except AttributeError: - beta_sobolev_filtered = beta_sobolev - r_ul_matrix_reshaped = r_ul_matrix.reshape( - (number_of_levels**2, number_of_shells) - ) - r_lu_matrix_reshaped = r_lu_matrix.reshape( - (number_of_levels**2, number_of_shells) - ) - r_ul_matrix_reshaped[r_ul_index] *= beta_sobolev_filtered - r_lu_matrix_reshaped[r_lu_index] *= beta_sobolev_filtered - rates_matrix_bound_bound = r_ul_matrix + r_lu_matrix - for i in range(number_of_levels): - rates_matrix_bound_bound[i, i] = -rates_matrix_bound_bound[ - :, i - ].sum(axis=0) - return rates_matrix_bound_bound - - def prepare_r_uls_r_lus( - number_of_levels, - number_of_shells, - j_blues, - excitation_species, - nlte_data, - ): - """Calculates part of rate matrices for bound bound interactions. - Parameters - ---------- - number_of_levels : int - Number of levels for the NLTE excitation species. - number_of_shells : int - Number of shells. - j_blues : pandas.DataFrame, dtype float - Mean intensities in the blue wings of the line transitions. - excitation_species : tuple - Species treated in NLTE excitation. - nlte_data : NLTEExcitationData - Data relevant to NLTE excitation species. +def create_coll_exc_deexc_matrix( + coll_exc_coefficient, + coll_deexc_coefficient, + number_of_levels, +): + """Generates a coefficient matrix from collisional excitation/deexcitation coefficients. - Returns - ------- - lines_index : numpy.array - Index of lines in nlte_data. - number_of_levels : int - Number of levels for the specified species. - r_ul_index : numpy.array - Index used for r_ul matrix - r_ul_matrix_reshaped : numpy.array - Matrix with the rates(upper to lower transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) - r_lu_index : numpy.array - Index used for r_lu matrix - r_lu_matrix_reshaped : numpy.array - Matrix with the rates(lower to upper transition) of bound-bound interaction(DOES NOT INCLUDE THE BETA SOBOLEVS) - """ - # number_of_levels = atomic_data_levels.energy.loc[ - # excitation_species - # ].count() do this in the solver - lnl = nlte_data.lines_level_number_lower[excitation_species] - lnu = nlte_data.lines_level_number_upper[excitation_species] - (lines_index,) = nlte_data.lines_idx[excitation_species] - - try: - j_blues_filtered = j_blues.iloc[lines_index] - except AttributeError: - j_blues_filtered = j_blues - A_uls = nlte_data.A_uls[excitation_species] - B_uls = nlte_data.B_uls[excitation_species] - B_lus = nlte_data.B_lus[excitation_species] - r_lu_index = lnu * number_of_levels + lnl - r_ul_index = lnl * number_of_levels + lnu - r_ul_matrix = np.zeros( - (number_of_levels, number_of_levels, number_of_shells), - dtype=np.float64, - ) - r_ul_matrix_reshaped = r_ul_matrix.reshape( - (number_of_levels**2, number_of_shells) - ) - r_ul_matrix_reshaped[r_ul_index] = ( - A_uls[np.newaxis].T + B_uls[np.newaxis].T * j_blues_filtered - ) - r_lu_matrix = np.zeros_like(r_ul_matrix) - r_lu_matrix_reshaped = r_lu_matrix.reshape( - (number_of_levels**2, number_of_shells) - ) - r_lu_matrix_reshaped[r_lu_index] = ( - B_lus[np.newaxis].T * j_blues_filtered - ) - return ( - lines_index, - r_ul_index, - r_ul_matrix, - r_lu_index, - r_lu_matrix, - ) - # TODO: beta sobolev needs to be recalculated for each iteration, because it depends on number density + Needs to be multiplied by electron density when added to the overall rate_matrix. - @staticmethod - def create_coll_exc_deexc_matrix( - coll_exc_coefficient, - coll_deexc_coefficient, - number_of_levels, - ): - """Generates a coefficient matrix from collisional excitation/deexcitation coefficients. + Parameters + ---------- + coll_exc_coefficient : pandas.Series + Series of collisional excitation coefficients for current (atomic number, ion_number) + in the current shell. + coll_deexc_coefficient : pandas.Series + Series of collisional deexcitation coefficients for (atomic number, ion_number) + in the current shell. + number_of_levels : int + Number of levels for the current atomic number, ion number. - Needs to be multiplied by electron density when added to the overall rate_matrix. - Parameters - ---------- - coll_exc_coefficient : pandas.Series - Series of collisional excitation coefficients for current (atomic number, ion_number) - in the current shell. - coll_deexc_coefficient : pandas.Series - Series of collisional deexcitation coefficients for (atomic number, ion_number) - in the current shell. - number_of_levels : int - Number of levels for the current atomic number, ion number. - - Returns - ------- - coeff_matrix : np.array (number of levels, number of levels) - Square matrix constructed by collisional exc./deexc. coefficients. - """ - diagonal_exc = np.zeros(number_of_levels) - diagonal_deexc = np.zeros(number_of_levels) - col_exc_coefficient_sum_lower = coll_exc_coefficient.groupby( - "level_number_lower" - ).sum() - col_deexc_coefficient_sum_upper = coll_deexc_coefficient.groupby( - "level_number_upper" - ).sum() + Returns + ------- + coeff_matrix : np.array (number of levels, number of levels) + Square matrix constructed by collisional exc./deexc. coefficients. + """ + diagonal_exc = np.zeros(number_of_levels) + diagonal_deexc = np.zeros(number_of_levels) + col_exc_coefficient_sum_lower = coll_exc_coefficient.groupby( + "level_number_lower" + ).sum() + col_deexc_coefficient_sum_upper = coll_deexc_coefficient.groupby( + "level_number_upper" + ).sum() - diagonal_exc[col_exc_coefficient_sum_lower.index] = ( - -1 * col_exc_coefficient_sum_lower.values + diagonal_exc[col_exc_coefficient_sum_lower.index] = ( + -1 * col_exc_coefficient_sum_lower.values + ) + diagonal_deexc[col_deexc_coefficient_sum_upper.index] = ( + -1 * col_deexc_coefficient_sum_upper.values + ) + exc_matrix = np.zeros((number_of_levels, number_of_levels)) + deexc_matrix = np.zeros((number_of_levels, number_of_levels)) + exc_matrix[ + ( + coll_exc_coefficient.index.get_level_values("level_number_upper"), + coll_exc_coefficient.index.get_level_values("level_number_lower"), ) - diagonal_deexc[col_deexc_coefficient_sum_upper.index] = ( - -1 * col_deexc_coefficient_sum_upper.values + ] = coll_exc_coefficient.values + deexc_matrix[ + ( + coll_exc_coefficient.index.get_level_values("level_number_lower"), + coll_exc_coefficient.index.get_level_values("level_number_upper"), ) - exc_matrix = np.zeros((number_of_levels, number_of_levels)) - deexc_matrix = np.zeros((number_of_levels, number_of_levels)) - exc_matrix[ - ( - coll_exc_coefficient.index.get_level_values( - "level_number_upper" - ), - coll_exc_coefficient.index.get_level_values( - "level_number_lower" - ), - ) - ] = coll_exc_coefficient.values - deexc_matrix[ - ( - coll_exc_coefficient.index.get_level_values( - "level_number_lower" - ), - coll_exc_coefficient.index.get_level_values( - "level_number_upper" - ), - ) - ] = coll_deexc_coefficient.values - np.fill_diagonal(exc_matrix, diagonal_exc) - np.fill_diagonal(deexc_matrix, diagonal_deexc) - coeff_matrix = exc_matrix + deexc_matrix - return coeff_matrix + ] = coll_deexc_coefficient.values + np.fill_diagonal(exc_matrix, diagonal_exc) + np.fill_diagonal(deexc_matrix, diagonal_deexc) + coeff_matrix = exc_matrix + deexc_matrix + return coeff_matrix diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index ab4d1f1cac4..deb2405347b 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -59,8 +59,11 @@ class PlasmaPropertyCollection(list): BetaSobolev, ] ) -nlte_solver_properties = PlasmaPropertyCollection( - [NLTEIndexHelper, NLTERateEquationSolver] +nlte_root_solver_properties = PlasmaPropertyCollection( + [NLTEIndexHelper, NLTEPopulationSolverRoot] +) +nlte_lu_solver_properties = PlasmaPropertyCollection( + [NLTEIndexHelper, NLTEPopulationSolverLU] ) helium_nlte_properties = PlasmaPropertyCollection( [ diff --git a/tardis/plasma/properties/transition_probabilities.py b/tardis/plasma/properties/transition_probabilities.py index e69d655ef6b..6f1af8f8308 100644 --- a/tardis/plasma/properties/transition_probabilities.py +++ b/tardis/plasma/properties/transition_probabilities.py @@ -41,16 +41,12 @@ def normalize_trans_probs(p): all probabilites with the same source_level_idx sum to one. Indexed by source_level_idx, destination_level_idx. """ - # 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) + p_summed = p.groupby(level=0).sum() + p_summed[p_summed == 0] = 1 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() + assert np.all(np.isfinite(p_norm)) return p_norm diff --git a/tardis/plasma/properties/util/integrate_array_by_blocks.py b/tardis/plasma/properties/util/integrate_array_by_blocks.py new file mode 100644 index 00000000000..e64509769e0 --- /dev/null +++ b/tardis/plasma/properties/util/integrate_array_by_blocks.py @@ -0,0 +1,35 @@ +import numpy as np +from numba import njit, prange + +from tardis.montecarlo.montecarlo_numba import njit_dict + + +@njit(**njit_dict) +def integrate_array_by_blocks(f, x, block_references): + """ + Integrate a function over blocks. + + This function integrates a function `f` defined at locations `x` + over blocks given in `block_references`. + + Parameters + ---------- + f : numpy.ndarray, dtype float + 2D input array to integrate. + x : numpy.ndarray, dtype float + 1D array with the sample points corresponding to the `f` values. + block_references : numpy.ndarray, dtype int + 1D array with the start indices of the blocks to be integrated. + + Returns + ------- + numpy.ndarray, dtype float + 2D array with integrated values. + """ + integrated = np.zeros((len(block_references) - 1, f.shape[1])) + for i in prange(f.shape[1]): # columns + for j in prange(len(integrated)): # rows + start = block_references[j] + stop = block_references[j + 1] + integrated[j, i] = np.trapz(f[start:stop, i], x[start:stop]) + return integrated diff --git a/tardis/plasma/properties/util/macro_atom.py b/tardis/plasma/properties/util/macro_atom.py index b89971776ff..3f20036fa7c 100644 --- a/tardis/plasma/properties/util/macro_atom.py +++ b/tardis/plasma/properties/util/macro_atom.py @@ -1,7 +1,8 @@ -from numba import njit -from tardis.montecarlo.montecarlo_numba import njit_dict import numpy as np +from numba import njit + from tardis import constants as const +from tardis.montecarlo.montecarlo_numba import njit_dict h_cgs = const.h.cgs.value c = const.c.to("cm/s").value @@ -28,7 +29,6 @@ def calculate_transition_probabilities( transition_type, lines_idx, and block_references must be int-type arrays beta_sobolev, j_blues,stimulated_emission_factor, and transition_probabilities must be 2D array """ - norm_factor = np.zeros(transition_probabilities.shape[1]) for i in range(transition_probabilities.shape[0]): @@ -57,5 +57,5 @@ def calculate_transition_probabilities( else: norm_factor[k] = 1.0 for j in range(block_references[i], block_references[i + 1]): - for k in range(0, transition_probabilities.shape[1]): + for k in range(transition_probabilities.shape[1]): transition_probabilities[j, k] *= norm_factor[k] diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 3bb9dbd3f6c..42d81af1e1b 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -8,7 +8,8 @@ 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, + NLTEPopulationSolverLU, + NLTEPopulationSolverRoot, ) from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper from tardis.util.base import species_string_to_tuple @@ -34,7 +35,8 @@ adiabatic_cooling_properties, two_photon_properties, isotope_properties, - nlte_solver_properties, + nlte_lu_solver_properties, + nlte_root_solver_properties, ) from tardis.plasma.exceptions import PlasmaConfigError @@ -89,6 +91,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): atom_data.prepare_atom_data( simulation_state.abundance.index, line_interaction_type=config.plasma.line_interaction_type, + continuum_interaction_species=continuum_interaction_species, nlte_species=nlte_species, ) @@ -130,7 +133,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): plasma_modules = basic_inputs + basic_properties property_kwargs = {} - if config.plasma.continuum_interaction.species: + if len(config.plasma.continuum_interaction.species) > 0: line_interaction_type = config.plasma.line_interaction_type if line_interaction_type != "macroatom": raise PlasmaConfigError( @@ -187,7 +190,19 @@ def assemble_plasma(config, simulation_state, atom_data=None): "nlte_ionization_species": config.plasma.nlte_ionization_species, "nlte_excitation_species": config.plasma.nlte_excitation_species, } - plasma_modules += nlte_solver_properties + if config.plasma.nlte_solver == "lu": + plasma_modules += nlte_lu_solver_properties + logger.warning( + "LU solver will be inaccurate for NLTE excitation, proceed with caution." + ) + elif config.plasma.nlte_solver == "root": + plasma_modules += nlte_root_solver_properties + else: + raise PlasmaConfigError( + "NLTE solver type unknown - {}".format( + config.plasma.nlte_solver + ) + ) kwargs.update( gamma_estimator=None, @@ -299,8 +314,15 @@ def assemble_plasma(config, simulation_state, atom_data=None): elif ( config.plasma.nlte_ionization_species or config.plasma.nlte_excitation_species - ): - property_kwargs[NLTERateEquationSolver] = dict( + ) and config.plasma.nlte_solver == "root": + property_kwargs[NLTEPopulationSolverRoot] = dict( + electron_densities=electron_densities + ) + elif ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ) and config.plasma.nlte_solver == "lu": + property_kwargs[NLTEPopulationSolverLU] = dict( electron_densities=electron_densities ) else: diff --git a/tardis/plasma/tests/test_nlte_excitation.py b/tardis/plasma/tests/test_nlte_excitation.py index 19f5c4bf7f5..fcefce08c73 100644 --- a/tardis/plasma/tests/test_nlte_excitation.py +++ b/tardis/plasma/tests/test_nlte_excitation.py @@ -2,10 +2,13 @@ import numpy.testing as npt import pandas as pd import pytest +from numpy.testing import assert_allclose from tardis.plasma.properties.nlte_excitation_data import NLTEExcitationData from tardis.plasma.properties.nlte_rate_equation_solver import ( - NLTERateEquationSolver, + prepare_r_uls_r_lus, + prepare_bound_bound_rate_matrix, + create_coll_exc_deexc_matrix, ) @@ -46,7 +49,7 @@ def test_prepare_bound_bound_rate_matrix(nlte_atomic_dataset, regression_data): r_ul_matrix, r_lu_index, r_lu_matrix, - ) = NLTERateEquationSolver.prepare_r_uls_r_lus( + ) = prepare_r_uls_r_lus( simple_number_of_levels, simple_number_of_shells, simple_j_blues, @@ -61,7 +64,7 @@ def test_prepare_bound_bound_rate_matrix(nlte_atomic_dataset, regression_data): index=copy_atomic_dataset.lines.index, columns=["0"], ) - actual_rate_matrix = NLTERateEquationSolver.prepare_bound_bound_rate_matrix( + actual_rate_matrix = prepare_bound_bound_rate_matrix( simple_number_of_levels, lines_index, r_ul_index, @@ -112,7 +115,7 @@ def test_coll_exc_deexc_matrix( ) exc_coeff = pd.Series(coll_exc_coeff_values, index=index) deexc_coeff = pd.Series(coll_deexc_coeff_values, index=index) - obtained_coeff_matrix = NLTERateEquationSolver.create_coll_exc_deexc_matrix( + obtained_coeff_matrix = create_coll_exc_deexc_matrix( exc_coeff, deexc_coeff, number_of_levels ) expected_obtained_coeff_matrix = regression_data.sync_ndarray( diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index e70152b9062..48045d3cf09 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -1,11 +1,19 @@ import numpy as np +import numpy.testing as npt import pandas as pd import pytest import numpy.testing as npt -from tardis.plasma.properties import NLTERateEquationSolver +from tardis.plasma.properties import ( + NLTEPopulationSolverLU, + NLTEPopulationSolverRoot, +) from tardis.plasma.properties.ion_population import IonNumberDensity +from tardis.plasma.properties.nlte_rate_equation_solver import ( + calculate_jacobian_matrix, + calculate_rate_matrix, +) from tardis.plasma.standard_plasmas import assemble_plasma @@ -146,7 +154,7 @@ def test_rate_matrix( Using a simple case of nlte_ion for HI and HeII, checks if the calculate_rate_matrix generates the correct data. """ atomic_numbers = [1, 2] - actual_rate_matrix = NLTERateEquationSolver.calculate_rate_matrix( + actual_rate_matrix = calculate_rate_matrix( atomic_numbers, simple_phi, simple_electron_density, @@ -186,7 +194,7 @@ def test_jacobian_matrix( 0.2878574499274399, simple_electron_density, ] - simple_rate_matrix = NLTERateEquationSolver.calculate_rate_matrix( + simple_rate_matrix = calculate_rate_matrix( atomic_numbers, simple_phi, simple_electron_density, @@ -197,7 +205,7 @@ def test_jacobian_matrix( simple_total_col_recomb_coefficients, ) - actual_jacobian_matrix = NLTERateEquationSolver.jacobian_matrix( + actual_jacobian_matrix = calculate_jacobian_matrix( atomic_numbers, initial_guess, simple_rate_matrix, @@ -215,46 +223,103 @@ def test_jacobian_matrix( @pytest.fixture -def nlte_raw_plasma_w1( - tardis_model_config_nlte, nlte_raw_simulation_state, nlte_atom_data +def nlte_raw_plasma_dilution_factor_1_root( + tardis_model_config_nlte_root, nlte_raw_model_root, nlte_atom_data ): """ Plasma assembled with dilution factors set to 1.0. """ - new_w = np.ones_like(nlte_raw_simulation_state.dilution_factor) - nlte_raw_simulation_state.dilution_factor = new_w + new_dilution_factor = np.ones_like(nlte_raw_model_root.dilution_factor) + nlte_raw_model_root.dilution_factor = new_dilution_factor plasma = assemble_plasma( - tardis_model_config_nlte, nlte_raw_simulation_state, nlte_atom_data + tardis_model_config_nlte_root, nlte_raw_model_root, nlte_atom_data ) return plasma @pytest.fixture -def nlte_raw_plasma_w0( - tardis_model_config_nlte, nlte_raw_simulation_state, nlte_atom_data +def nlte_raw_plasma_dilution_factor_1_lu( + tardis_model_config_nlte_lu, nlte_raw_model_lu, nlte_atom_data +): + """ + Plasma assembled with dilution factors set to 1.0. + """ + new_dilution_factor = np.ones_like(nlte_raw_model_lu.dilution_factor) + nlte_raw_model_lu.dilution_factor = new_dilution_factor + plasma = assemble_plasma( + tardis_model_config_nlte_lu, nlte_raw_model_lu, nlte_atom_data + ) + return plasma + + +@pytest.fixture +def nlte_raw_plasma_dilution_factor_0_root( + tardis_model_config_nlte_root, nlte_raw_model_root, nlte_atom_data ): """ Plasma assembled with dilution factors set to 0.0. """ - nlte_raw_simulation_state.dilution_factor = np.zeros_like( - nlte_raw_simulation_state.dilution_factor + new_dilution_factor = np.zeros_like(nlte_raw_model_root.dilution_factor) + nlte_raw_model_root.dilution_factor = new_dilution_factor + plasma = assemble_plasma( + tardis_model_config_nlte_root, nlte_raw_model_root, nlte_atom_data ) + return plasma + + +@pytest.fixture +def nlte_raw_plasma_dilution_factor_0_lu( + tardis_model_config_nlte_lu, nlte_raw_model_lu, nlte_atom_data +): + """ + Plasma assembled with dilution factors set to 0.0. + """ + new_dilution_factor = np.zeros_like(nlte_raw_model_lu.dilution_factor) + nlte_raw_model_lu.dilution_factor = new_dilution_factor plasma = assemble_plasma( - tardis_model_config_nlte, nlte_raw_simulation_state, nlte_atom_data + tardis_model_config_nlte_lu, nlte_raw_model_lu, nlte_atom_data ) return plasma -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.values - ion_number_density_nlte[ion_number_density_nlte < 1e-10] = 0.0 +def test_critical_case_dilution_factor_1_root( + nlte_raw_plasma_dilution_factor_1_root, +): + """Check that the LTE and NLTE solution agree for dilution_factor=1.0.""" + ion_number_density_nlte = ( + nlte_raw_plasma_dilution_factor_1_root.ion_number_density.values + ) + # ion_number_density_nlte[np.abs(ion_number_density_nlte) < 1e-10] = 0.0 - ind = IonNumberDensity(nlte_raw_plasma_w1) + ind = IonNumberDensity(nlte_raw_plasma_dilution_factor_1_root) ion_number_density_lte = ind.calculate( - nlte_raw_plasma_w1.thermal_phi_lte, - nlte_raw_plasma_w1.partition_function, - nlte_raw_plasma_w1.number_density, + nlte_raw_plasma_dilution_factor_1_root.thermal_phi_lte, + nlte_raw_plasma_dilution_factor_1_root.partition_function, + nlte_raw_plasma_dilution_factor_1_root.number_density, + )[0] + + npt.assert_allclose( + ion_number_density_lte, + ion_number_density_nlte, + atol=1e-10, # seems fair for the test + rtol=1e-2, + ) + + +def test_critical_case_dilution_factor_1_lu( + nlte_raw_plasma_dilution_factor_1_lu, +): + """Check that the LTE and NLTE solution agree for dilution_factor=1.0.""" + ion_number_density_nlte = ( + nlte_raw_plasma_dilution_factor_1_lu.ion_number_density.values + ) + # ion_number_density_nlte[np.abs(ion_number_density_nlte) < 1e-10] = 0.0 + + ind = IonNumberDensity(nlte_raw_plasma_dilution_factor_1_lu) + ion_number_density_lte = ind.calculate( + nlte_raw_plasma_dilution_factor_1_lu.thermal_phi_lte, + nlte_raw_plasma_dilution_factor_1_lu.partition_function, + nlte_raw_plasma_dilution_factor_1_lu.number_density, )[0] ion_number_density_lte = ion_number_density_lte.values @@ -264,35 +329,76 @@ def test_critical_case_w1(nlte_raw_plasma_w1): npt.assert_allclose( ion_number_density_lte, ion_number_density_nlte, + atol=1e-10, # seems fair for the test rtol=1e-2, ) -def test_critical_case_w0(nlte_raw_plasma_w0): - """Check that the LTE and NLTE solution agree for w=0.0.""" - nlte_solver = NLTERateEquationSolver(nlte_raw_plasma_w0) +def test_critical_case_dilution_factor_0_root( + nlte_raw_plasma_dilution_factor_0_root, +): + """Check that the LTE and NLTE solution agree for dilution_factor=0.0.""" + nlte_solver = NLTEPopulationSolverRoot( + nlte_raw_plasma_dilution_factor_0_root + ) + ion_number_density_nlte = nlte_solver.calculate( + nlte_raw_plasma_dilution_factor_0_root.gamma, + 0.0, # to test collisions only, we set the radiative recombination rate to 0 + nlte_raw_plasma_dilution_factor_0_root.alpha_stim, + nlte_raw_plasma_dilution_factor_0_root.coll_ion_coeff, + nlte_raw_plasma_dilution_factor_0_root.coll_recomb_coeff, + nlte_raw_plasma_dilution_factor_0_root.partition_function, + nlte_raw_plasma_dilution_factor_0_root.levels, + nlte_raw_plasma_dilution_factor_0_root.level_boltzmann_factor, + nlte_raw_plasma_dilution_factor_0_root.phi, + nlte_raw_plasma_dilution_factor_0_root.rate_matrix_index, + nlte_raw_plasma_dilution_factor_0_root.number_density, + nlte_raw_plasma_dilution_factor_0_root.nlte_excitation_species, + )[0] + ion_number_density_nlte = ion_number_density_nlte.values + # ion_number_density_nlte[np.abs(ion_number_density_nlte) < 1e-10] = 0.0 + + ind = IonNumberDensity(nlte_raw_plasma_dilution_factor_0_root) + ion_number_density_lte = ind.calculate( + nlte_raw_plasma_dilution_factor_0_root.thermal_phi_lte, + nlte_raw_plasma_dilution_factor_0_root.partition_function, + nlte_raw_plasma_dilution_factor_0_root.number_density, + )[0] + + ion_number_density_lte = ion_number_density_lte.values + npt.assert_allclose( + ion_number_density_lte, ion_number_density_nlte, rtol=1e-2, atol=1e-10 + ) + + +@pytest.mark.xfail +def test_critical_case_dilution_factor_0_lu( + nlte_raw_plasma_dilution_factor_0_lu, +): + """Check that the LTE and NLTE solution agree for dilution_factor=0.0.""" + nlte_solver = NLTEPopulationSolverLU(nlte_raw_plasma_dilution_factor_0_lu) ion_number_density_nlte = nlte_solver.calculate( - nlte_raw_plasma_w0.gamma, + nlte_raw_plasma_dilution_factor_0_lu.gamma, 0.0, # to test collisions only, we set the radiative recombination rate to 0 - nlte_raw_plasma_w0.alpha_stim, - nlte_raw_plasma_w0.coll_ion_coeff, - nlte_raw_plasma_w0.coll_recomb_coeff, - nlte_raw_plasma_w0.partition_function, - nlte_raw_plasma_w0.levels, - nlte_raw_plasma_w0.level_boltzmann_factor, - nlte_raw_plasma_w0.phi, - nlte_raw_plasma_w0.rate_matrix_index, - nlte_raw_plasma_w0.number_density, - nlte_raw_plasma_w0.nlte_excitation_species, + nlte_raw_plasma_dilution_factor_0_lu.alpha_stim, + nlte_raw_plasma_dilution_factor_0_lu.coll_ion_coeff, + nlte_raw_plasma_dilution_factor_0_lu.coll_recomb_coeff, + nlte_raw_plasma_dilution_factor_0_lu.partition_function, + nlte_raw_plasma_dilution_factor_0_lu.levels, + nlte_raw_plasma_dilution_factor_0_lu.level_boltzmann_factor, + nlte_raw_plasma_dilution_factor_0_lu.phi, + nlte_raw_plasma_dilution_factor_0_lu.rate_matrix_index, + nlte_raw_plasma_dilution_factor_0_lu.number_density, + nlte_raw_plasma_dilution_factor_0_lu.nlte_excitation_species, )[0] ion_number_density_nlte = ion_number_density_nlte.values - ion_number_density_nlte[ion_number_density_nlte < 1e-10] = 0.0 + # ion_number_density_nlte[np.abs(ion_number_density_nlte) < 1e-10] = 0.0 - ind = IonNumberDensity(nlte_raw_plasma_w0) + ind = IonNumberDensity(nlte_raw_plasma_dilution_factor_0_lu) ion_number_density_lte = ind.calculate( - nlte_raw_plasma_w0.thermal_phi_lte, - nlte_raw_plasma_w0.partition_function, - nlte_raw_plasma_w0.number_density, + nlte_raw_plasma_dilution_factor_0_lu.thermal_phi_lte, + nlte_raw_plasma_dilution_factor_0_lu.partition_function, + nlte_raw_plasma_dilution_factor_0_lu.number_density, )[0] ion_number_density_lte = ion_number_density_lte.values @@ -303,4 +409,5 @@ def test_critical_case_w0(nlte_raw_plasma_w0): ion_number_density_lte, ion_number_density_nlte, rtol=1e-2, + atol=1e-10, ) diff --git a/tardis/plasma/tests/test_plasma_contiuum.py b/tardis/plasma/tests/test_plasma_continuum.py similarity index 100% rename from tardis/plasma/tests/test_plasma_contiuum.py rename to tardis/plasma/tests/test_plasma_continuum.py diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 43ab0be8345..72abba91498 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -282,7 +282,7 @@ def advance_state(self): """ ( estimated_t_rad, - estimated_w, + estimated_dilution_factor, ) = self.transport.transport_state.calculate_radiationfield_properties() estimated_t_inner = self.estimate_t_inner( self.simulation_state.t_inner, @@ -295,7 +295,7 @@ def advance_state(self): self.simulation_state.dilution_factor, self.simulation_state.t_inner, estimated_t_rad, - estimated_w, + estimated_dilution_factor, estimated_t_inner, ) @@ -308,7 +308,7 @@ def advance_state(self): ) next_dilution_factor = self.damped_converge( self.simulation_state.dilution_factor, - estimated_w, + estimated_dilution_factor, self.convergence_strategy.w.damping_constant, ) if ( @@ -369,7 +369,7 @@ def advance_state(self): # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. - estimators = self.transport.transport_state.estimators + estimators = self.transport.transport_state.radfield_mc_estimators if "j_blue_estimator" in self.plasma.outputs_dict: update_properties.update( t_inner=next_t_inner, diff --git a/tardis/simulation/tests/test_simulation.py b/tardis/simulation/tests/test_simulation.py index c18c7b1ebb3..2fd80c8a403 100644 --- a/tardis/simulation/tests/test_simulation.py +++ b/tardis/simulation/tests/test_simulation.py @@ -90,7 +90,8 @@ def simulation_one_loop( def test_plasma_estimates(simulation_one_loop, refdata, name): if name in ["nu_bar_estimator", "j_estimator"]: actual = getattr( - simulation_one_loop.transport.transport_state.estimators, name + simulation_one_loop.transport.transport_state.radfield_mc_estimators, + name, ) elif name in ["t_radiative", "dilution_factor"]: actual = getattr(simulation_one_loop.simulation_state, name) diff --git a/tardis/tests/fixtures/atom_data.py b/tardis/tests/fixtures/atom_data.py index d912d0b6493..d2b4c7b084a 100644 --- a/tardis/tests/fixtures/atom_data.py +++ b/tardis/tests/fixtures/atom_data.py @@ -76,14 +76,32 @@ def nlte_atom_data(nlte_atomic_dataset): @pytest.fixture # (scope="session") -def tardis_model_config_nlte(example_configuration_dir): - return Configuration.from_yaml( +def tardis_model_config_nlte_root(example_configuration_dir): + config = Configuration.from_yaml( example_configuration_dir / "tardis_configv1_nlte.yml" ) + config.plasma.nlte_solver = "root" + return config @pytest.fixture # (scope="session") -def nlte_raw_simulation_state(tardis_model_config_nlte, nlte_atomic_dataset): +def tardis_model_config_nlte_lu(example_configuration_dir): + config = Configuration.from_yaml( + example_configuration_dir / "tardis_configv1_nlte.yml" + ) + config.plasma.nlte_solver = "lu" + return config + + +@pytest.fixture # (scope="session") +def nlte_raw_model_root(tardis_model_config_nlte_root, nlte_atom_data): + return SimulationState.from_config( + tardis_model_config_nlte_root, nlte_atom_data + ) + + +@pytest.fixture # (scope="session") +def nlte_raw_model_lu(tardis_model_config_nlte_lu, nlte_atom_data): return SimulationState.from_config( - tardis_model_config_nlte, atom_data=nlte_atomic_dataset + tardis_model_config_nlte_lu, nlte_atom_data ) diff --git a/tardis/tests/test_tardis_full.py b/tardis/tests/test_tardis_full.py index 5ababcbb6bc..9909ea8a358 100644 --- a/tardis/tests/test_tardis_full.py +++ b/tardis/tests/test_tardis_full.py @@ -78,7 +78,7 @@ def test_j_blue_estimators(self, transport, refdata): j_blue_estimator = refdata("j_blue_estimator").values npt.assert_allclose( - transport.transport_state.estimators.j_blue_estimator, + transport.transport_state.radfield_mc_estimators.j_blue_estimator, j_blue_estimator, ) diff --git a/tardis/tests/test_tardis_full_formal_integral.py b/tardis/tests/test_tardis_full_formal_integral.py index 318ce1dcaa9..e45ca01bb91 100644 --- a/tardis/tests/test_tardis_full_formal_integral.py +++ b/tardis/tests/test_tardis_full_formal_integral.py @@ -79,7 +79,7 @@ def test_j_blue_estimators(self, transport, refdata): j_blue_estimator = refdata("j_blue_estimator").values npt.assert_allclose( - transport.transport_state.estimators.j_blue_estimator, + transport.transport_state.radfield_mc_estimators.j_blue_estimator, j_blue_estimator, ) diff --git a/tardis/transport/r_packet_transport.py b/tardis/transport/r_packet_transport.py index 2520733c14d..20e4bd252df 100644 --- a/tardis/transport/r_packet_transport.py +++ b/tardis/transport/r_packet_transport.py @@ -8,7 +8,7 @@ calculate_distance_electron, calculate_distance_line, ) -from tardis.montecarlo.montecarlo_numba.estimators import ( +from tardis.montecarlo.estimators.radfield_mc_estimators import ( update_line_estimators, update_base_estimators, ) diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py index 8490a58b990..578c64fe66a 100644 --- a/tardis/visualization/tools/sdec_plot.py +++ b/tardis/visualization/tools/sdec_plot.py @@ -165,7 +165,7 @@ def from_simulation(cls, sim, packets_mode): "line_id" ) transport_state = sim.transport.transport_state - r_inner = sim.simulation_state.geometry.r_inner + r_inner = sim.simulation_state.geometry.r_inner_active t_inner = sim.simulation_state.packet_source.temperature time_of_simulation = ( transport_state.packet_collection.time_of_simulation * u.s