diff --git a/.gitignore b/.gitignore index f8cc8cf6295..b28f0559fa1 100644 --- a/.gitignore +++ b/.gitignore @@ -71,3 +71,6 @@ pip-wheel-metadata/ # Mac OSX .DS_Store + +# Visual Studio Code +*.code-workspace 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/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/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..b1f33f7429b 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__) diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 01914fb8e6a..8e9420e2ea8 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, @@ -572,7 +574,7 @@ def parse_radiation_field_state( ] assert len(dilution_factor) == geometry.no_of_shells - return DiluteThermalRadiationFieldState(t_radiative, dilution_factor) + return DiluteBlackBodyRadiationFieldState(t_radiative, dilution_factor) def initialize_packet_source(config, geometry, packet_source): @@ -687,7 +689,7 @@ 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) def calculate_t_radiative_from_t_inner(geometry, packet_source): diff --git a/tardis/model/radiation_field_state.py b/tardis/model/radiation_field_state.py index f6b1356f76b..6395035c09c 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. @@ -30,3 +31,21 @@ def __init__( 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/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/estimators/radfield_mc_estimators.py b/tardis/montecarlo/estimators/radfield_mc_estimators.py new file mode 100644 index 00000000000..328a0c4d55b --- /dev/null +++ b/tardis/montecarlo/estimators/radfield_mc_estimators.py @@ -0,0 +1,239 @@ +from math import exp + +import numpy as np +from numba import float64, int64, njit +from numba.experimental import jitclass + +from tardis.montecarlo import montecarlo_configuration as nc +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_estimator_statistics(tau_sobolev_shape, gamma_shape): + """ + Initializes the estimators used in the Monte Carlo simulation. + + Parameters + ---------- + tau_sobolev_shape : tuple + Shape of the array with the Sobolev optical depth. + gamma_shape : tuple + Shape of the array with the photoionization rate coefficients. + + Returns + ------- + Estimators + The initialized estimators. + + Examples + -------- + >>> tau_sobolev_shape = (10, 20) + >>> gamma_shape = (5, 5) + >>> 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) + Edotlu_estimator = np.zeros(tau_sobolev_shape) + + photo_ion_estimator = np.zeros(gamma_shape, dtype=np.float64) + stim_recomb_estimator = np.zeros(gamma_shape, dtype=np.float64) + stim_recomb_cooling_estimator = np.zeros(gamma_shape, dtype=np.float64) + bf_heating_estimator = np.zeros(gamma_shape, dtype=np.float64) + + stim_recomb_cooling_estimator = np.zeros(gamma_shape, dtype=np.float64) + + photo_ion_estimator_statistics = np.zeros(gamma_shape, dtype=np.int64) + return RadiationFieldMCEstimators( + j_estimator, + nu_bar_estimator, + j_blue_estimator, + Edotlu_estimator, + photo_ion_estimator, + stim_recomb_estimator, + bf_heating_estimator, + stim_recomb_cooling_estimator, + photo_ion_estimator_statistics, + ) + + +base_estimators_spec = [ + ("j_estimator", float64[:]), + ("nu_bar_estimator", float64[:]), + ("j_blue_estimator", float64[:, :]), + ("Edotlu_estimator", float64[:, :]), +] + +continuum_estimators_spec = [ + ("photo_ion_estimator", float64[:, :]), + ("stim_recomb_estimator", float64[:, :]), + ("bf_heating_estimator", float64[:, :]), + ("stim_recomb_cooling_estimator", float64[:, :]), + ("photo_ion_estimator_statistics", int64[:, :]), +] + + +@jitclass(base_estimators_spec + continuum_estimators_spec) +class RadiationFieldMCEstimators: + def __init__( + self, + j_estimator, + nu_bar_estimator, + j_blue_estimator, + Edotlu_estimator, + photo_ion_estimator, + stim_recomb_estimator, + bf_heating_estimator, + stim_recomb_cooling_estimator, + photo_ion_estimator_statistics, + ): + self.j_estimator = j_estimator + self.nu_bar_estimator = nu_bar_estimator + self.j_blue_estimator = j_blue_estimator + self.Edotlu_estimator = Edotlu_estimator + self.photo_ion_estimator = photo_ion_estimator + self.stim_recomb_estimator = stim_recomb_estimator + self.bf_heating_estimator = bf_heating_estimator + self.stim_recomb_cooling_estimator = stim_recomb_cooling_estimator + 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 + self.Edotlu_estimator += other.Edotlu_estimator + self.photo_ion_estimator += other.photo_ion_estimator + self.stim_recomb_estimator += other.stim_recomb_estimator + self.bf_heating_estimator += other.bf_heating_estimator + self.stim_recomb_cooling_estimator += ( + other.stim_recomb_cooling_estimator + ) + self.photo_ion_estimator_statistics += ( + other.photo_ion_estimator_statistics + ) + + +@njit(**njit_dict_no_parallel) +def update_base_estimators( + r_packet, distance, estimator_state, comov_nu, comov_energy +): + """ + Updating the estimators + """ + estimator_state.j_estimator[r_packet.current_shell_id] += ( + comov_energy * distance + ) + estimator_state.nu_bar_estimator[r_packet.current_shell_id] += ( + comov_energy * distance * comov_nu + ) + + +@njit(**njit_dict_no_parallel) +def update_bound_free_estimators( + comov_nu, + comov_energy, + shell_id, + distance, + estimator_state, + t_electron, + x_sect_bfs, + current_continua, + bf_threshold_list_nu, +): + """ + Update the estimators for bound-free processes. + + Parameters + ---------- + comov_nu : float + comov_energy : float + shell_id : int + distance : float + numba_estimator : tardis.montecarlo.montecarlo_numba.numba_interface.Estimators + t_electron : float + Electron temperature in the current cell. + x_sect_bfs : numpy.ndarray, dtype float + Photoionization cross-sections of all bound-free continua for + which absorption is possible for frequency `comov_nu`. + current_continua : numpy.ndarray, dtype int + Continuum ids for which absorption is possible for frequency `comov_nu`. + bf_threshold_list_nu : numpy.ndarray, dtype float + Threshold frequencies for photoionization sorted by decreasing frequency. + """ + # TODO: Add full relativity mode + boltzmann_factor = exp(-(H * comov_nu) / (KB * t_electron)) + for i, current_continuum in enumerate(current_continua): + photo_ion_rate_estimator_increment = ( + comov_energy * distance * x_sect_bfs[i] / comov_nu + ) + estimator_state.photo_ion_estimator[ + current_continuum, shell_id + ] += photo_ion_rate_estimator_increment + estimator_state.stim_recomb_estimator[current_continuum, shell_id] += ( + photo_ion_rate_estimator_increment * boltzmann_factor + ) + estimator_state.photo_ion_estimator_statistics[ + current_continuum, shell_id + ] += 1 + + nu_th = bf_threshold_list_nu[current_continuum] + bf_heating_estimator_increment = ( + comov_energy * distance * x_sect_bfs[i] * (1 - nu_th / comov_nu) + ) + estimator_state.bf_heating_estimator[ + current_continuum, shell_id + ] += bf_heating_estimator_increment + estimator_state.stim_recomb_cooling_estimator[ + current_continuum, shell_id + ] += (bf_heating_estimator_increment * boltzmann_factor) + + +@njit(**njit_dict_no_parallel) +def update_line_estimators( + radfield_mc_estimators, + r_packet, + cur_line_id, + distance_trace, + time_explosion, +): + """ + Function to update the line estimators + + Parameters + ---------- + estimators : tardis.montecarlo.montecarlo_numba.numba_interface.Estimators + r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket + cur_line_id : int + distance_trace : float + time_explosion : float + """ + 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) + + 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/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..8d6221b8841 100644 --- a/tardis/plasma/properties/continuum_processes.py +++ b/tardis/plasma/properties/continuum_processes.py @@ -5,7 +5,11 @@ 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, @@ -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. @@ -1076,7 +1047,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/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..3e3de40b633 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -89,6 +89,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 +131,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( 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..df98c6ad01f 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -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/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, )