From 1d0c04aab6188f74ba759c48f90bcab92478b9b5 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 15:59:34 -0500 Subject: [PATCH 01/15] MC configuration refactor begins - Converted MC configuration to a jitclass --- tardis/montecarlo/montecarlo_configuration.py | 99 ++++++++++--------- 1 file changed, 54 insertions(+), 45 deletions(-) diff --git a/tardis/montecarlo/montecarlo_configuration.py b/tardis/montecarlo/montecarlo_configuration.py index 6e813298f57..fa12591a591 100644 --- a/tardis/montecarlo/montecarlo_configuration.py +++ b/tardis/montecarlo/montecarlo_configuration.py @@ -1,74 +1,83 @@ from astropy import units as u -from tardis import constants as const -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) +from numba import float64, int64, boolean +from numba.experimental import jitclass + from tardis.montecarlo.montecarlo_numba.numba_interface import ( LineInteractionType, ) -ENABLE_FULL_RELATIVITY = False -TEMPORARY_V_PACKET_BINS = 0 -NUMBER_OF_VPACKETS = 0 -MONTECARLO_SEED = 0 -LINE_INTERACTION_TYPE = None -PACKET_SEEDS = [] -DISABLE_ELECTRON_SCATTERING = False -DISABLE_LINE_SCATTERING = False -SURVIVAL_PROBABILITY = 0.0 -VPACKET_TAU_RUSSIAN = 10.0 +numba_config_spec = [ + ("ENABLE_FULL_RELATIVITY", boolean), + ("TEMPORARY_V_PACKET_BINS", int64), + ("NUMBER_OF_VPACKETS", int64), + ("MONTECARLO_SEED", int64), + ("LINE_INTERACTION_TYPE", int64), + ("PACKET_SEEDS", int64[:]), + ("DISABLE_ELECTRON_SCATTERING", boolean), + ("DISABLE_LINE_SCATTERING", boolean), + ("SURVIVAL_PROBABILITY", float64), + ("VPACKET_TAU_RUSSIAN", float64), + ("INITIAL_TRACKING_ARRAY_LENGTH", int64), + ("LEGACY_MODE_ENABLED", boolean), + ("ENABLE_RPACKET_TRACKING", boolean), + ("CONTINUUM_PROCESSES_ENABLED", boolean), + ("VPACKET_SPAWN_START_FREQUENCY", float64), + ("VPACKET_SPAWN_END_FREQUENCY", float64), + ("ENABLE_VPACKET_TRACKING", boolean), +] -INITIAL_TRACKING_ARRAY_LENGTH = None -LEGACY_MODE_ENABLED = False +@jitclass(numba_config_spec) +class MonteCarloConfiguration: + def __init__(self): + self.ENABLE_FULL_RELATIVITY = False + self.TEMPORARY_V_PACKET_BINS = 0 + self.NUMBER_OF_VPACKETS = 0 + self.MONTECARLO_SEED = 0 + self.LINE_INTERACTION_TYPE = None + self.PACKET_SEEDS = [] + self.DISABLE_ELECTRON_SCATTERING = False + self.DISABLE_LINE_SCATTERING = False + self.SURVIVAL_PROBABILITY = 0.0 + self.VPACKET_TAU_RUSSIAN = 10.0 -ENABLE_RPACKET_TRACKING = False -CONTINUUM_PROCESSES_ENABLED = False + self.INITIAL_TRACKING_ARRAY_LENGTH = None + self.LEGACY_MODE_ENABLED = False + self.ENABLE_RPACKET_TRACKING = False + self.CONTINUUM_PROCESSES_ENABLED = False -VPACKET_SPAWN_START_FREQUENCY = 0 -VPACKET_SPAWN_END_FREQUENCY = 1e200 -ENABLE_VPACKET_TRACKING = False + self.VPACKET_SPAWN_START_FREQUENCY = 0 + self.VPACKET_SPAWN_END_FREQUENCY = 1e200 + self.ENABLE_VPACKET_TRACKING = False -def configuration_initialize(transport, number_of_vpackets): +def configuration_initialize(config, transport, number_of_vpackets): if transport.line_interaction_type == "macroatom": - montecarlo_configuration.LINE_INTERACTION_TYPE = ( - LineInteractionType.MACROATOM - ) + config.LINE_INTERACTION_TYPE = LineInteractionType.MACROATOM elif transport.line_interaction_type == "downbranch": - montecarlo_configuration.LINE_INTERACTION_TYPE = ( - LineInteractionType.DOWNBRANCH - ) + config.LINE_INTERACTION_TYPE = LineInteractionType.DOWNBRANCH elif transport.line_interaction_type == "scatter": - montecarlo_configuration.LINE_INTERACTION_TYPE = ( - LineInteractionType.SCATTER - ) + config.LINE_INTERACTION_TYPE = LineInteractionType.SCATTER else: raise ValueError( f'Line interaction type must be one of "macroatom",' f'"downbranch", or "scatter" but is ' f"{transport.line_interaction_type}" ) - montecarlo_configuration.NUMBER_OF_VPACKETS = number_of_vpackets - montecarlo_configuration.TEMPORARY_V_PACKET_BINS = number_of_vpackets - montecarlo_configuration.ENABLE_FULL_RELATIVITY = ( - transport.enable_full_relativity - ) - montecarlo_configuration.MONTECARLO_SEED = transport.packet_source.base_seed - montecarlo_configuration.VPACKET_SPAWN_START_FREQUENCY = ( + config.NUMBER_OF_VPACKETS = number_of_vpackets + config.TEMPORARY_V_PACKET_BINS = number_of_vpackets + config.ENABLE_FULL_RELATIVITY = transport.enable_full_relativity + config.MONTECARLO_SEED = transport.packet_source.base_seed + config.VPACKET_SPAWN_START_FREQUENCY = ( transport.virtual_spectrum_spawn_range.end.to( u.Hz, equivalencies=u.spectral() ).value ) - montecarlo_configuration.VPACKET_SPAWN_END_FREQUENCY = ( + config.VPACKET_SPAWN_END_FREQUENCY = ( transport.virtual_spectrum_spawn_range.start.to( u.Hz, equivalencies=u.spectral() ).value ) - montecarlo_configuration.ENABLE_VPACKET_TRACKING = ( - transport.enable_vpacket_tracking - ) - montecarlo_configuration.ENABLE_RPACKET_TRACKING = ( - transport.enable_rpacket_tracking - ) + config.ENABLE_VPACKET_TRACKING = transport.enable_vpacket_tracking + config.ENABLE_RPACKET_TRACKING = transport.enable_rpacket_tracking From c80c803752ac54a69261d1fd62b8956e4136b7c3 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 16:00:12 -0500 Subject: [PATCH 02/15] Initialize new MC config object in the simulation --- tardis/montecarlo/base.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index 08ff1612306..8f78c0ba077 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -7,12 +7,12 @@ from tardis import constants as const 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, + MonteCarloConfiguration, ) from tardis.montecarlo.montecarlo_numba import ( montecarlo_main_loop, @@ -68,6 +68,7 @@ def __init__( debug_packets=False, logger_buffer=1, use_gpu=False, + montecarlo_configuration=None, ): # inject different packets self.disable_electron_scattering = disable_electron_scattering @@ -86,6 +87,7 @@ def __init__( self.enable_vpacket_tracking = enable_virtual_packet_logging self.enable_rpacket_tracking = enable_rpacket_tracking + self.montecarlo_configuration = montecarlo_configuration self.packet_source = packet_source @@ -139,7 +141,9 @@ def initialize_transport_state( transport_state._integrator = FormalIntegrator( simulation_state, plasma, self ) - configuration_initialize(self, no_of_virtual_packets) + configuration_initialize( + self.montecarlo_configuration, self, no_of_virtual_packets + ) return transport_state @@ -172,7 +176,7 @@ def run( numba_model = NumbaModel(time_explosion.to("s").value) - number_of_vpackets = montecarlo_configuration.NUMBER_OF_VPACKETS + number_of_vpackets = self.montecarlo_configuration.NUMBER_OF_VPACKETS ( v_packets_energy_hist, @@ -184,6 +188,7 @@ def run( transport_state.geometry_state, numba_model, transport_state.opacity_state, + self.montecarlo_configuration, transport_state.radfield_mc_estimators, transport_state.spectrum_frequency.value, number_of_vpackets, @@ -208,7 +213,7 @@ def run( last_interaction_tracker.shell_ids ) - if montecarlo_configuration.ENABLE_VPACKET_TRACKING and ( + if self.montecarlo_configuration.ENABLE_VPACKET_TRACKING and ( number_of_vpackets > 0 ): transport_state.vpacket_tracker = vpacket_tracker @@ -216,7 +221,7 @@ def run( update_iterations_pbar(1) refresh_packet_pbar() # Condition for Checking if RPacket Tracking is enabled - if montecarlo_configuration.ENABLE_RPACKET_TRACKING: + if self.montecarlo_configuration.ENABLE_RPACKET_TRACKING: transport_state.rpacket_tracker = rpacket_trackers if self.transport_state.rpacket_tracker is not None: @@ -226,7 +231,7 @@ def run( ) ) transport_state.virt_logging = ( - montecarlo_configuration.ENABLE_VPACKET_TRACKING + self.montecarlo_configuration.ENABLE_VPACKET_TRACKING ) def legacy_return(self): @@ -300,6 +305,8 @@ def from_config( valid values are 'GPU', 'CPU', and 'Automatic'.""" ) + montecarlo_configuration = MonteCarloConfiguration + montecarlo_configuration.DISABLE_LINE_SCATTERING = ( config.plasma.disable_line_scattering ) @@ -329,4 +336,5 @@ def from_config( enable_rpacket_tracking=config.montecarlo.tracking.track_rpacket, nthreads=config.montecarlo.nthreads, use_gpu=use_gpu, + montecarlo_configuration=montecarlo_configuration, ) From 67ef7c59cd10a99fc582ad02f2124184d59de308 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 16:02:36 -0500 Subject: [PATCH 03/15] Refactor radfield estimators to solve circular import --- .../estimators/radfield_estimator_calcs.py | 120 +++++++++++++++ .../estimators/radfield_mc_estimators.py | 139 +++--------------- tardis/montecarlo/montecarlo_numba/base.py | 22 +-- 3 files changed, 143 insertions(+), 138 deletions(-) create mode 100644 tardis/montecarlo/estimators/radfield_estimator_calcs.py diff --git a/tardis/montecarlo/estimators/radfield_estimator_calcs.py b/tardis/montecarlo/estimators/radfield_estimator_calcs.py new file mode 100644 index 00000000000..f97cce03f66 --- /dev/null +++ b/tardis/montecarlo/estimators/radfield_estimator_calcs.py @@ -0,0 +1,120 @@ +from math import exp + +from numba import njit + +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, +) + + +@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, + enable_full_relativity, +): + """ + 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 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/radfield_mc_estimators.py b/tardis/montecarlo/estimators/radfield_mc_estimators.py index 328a0c4d55b..5e5b868a986 100644 --- a/tardis/montecarlo/estimators/radfield_mc_estimators.py +++ b/tardis/montecarlo/estimators/radfield_mc_estimators.py @@ -1,18 +1,7 @@ -from math import exp - import numpy as np -from numba import float64, int64, njit +from numba import float64, int64 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, -) +from numba.typed import List def initialize_estimator_statistics(tau_sobolev_shape, gamma_shape): @@ -131,109 +120,21 @@ def increment(self, other): 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 + def create_list(self, number): + estimator_list = List() + + for i in range(number): + estimator_list.append( + RadiationFieldMCEstimators( + np.copy(self.j_estimator), + np.copy(self.nu_bar_estimator), + np.copy(self.j_blue_estimator), + np.copy(self.Edotlu_estimator), + np.copy(self.photo_ion_estimator), + np.copy(self.stim_recomb_estimator), + np.copy(self.bf_heating_estimator), + np.copy(self.stim_recomb_cooling_estimator), + np.copy(self.photo_ion_estimator_statistics), + ) + ) + return estimator_list diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index 987a70fe7b1..dbd2d1ac872 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -3,11 +3,7 @@ from numba.np.ufunc.parallel import get_num_threads, get_thread_id from numba.typed import List -from tardis.montecarlo import montecarlo_configuration from tardis.montecarlo.montecarlo_numba import njit_dict -from tardis.montecarlo.estimators.radfield_mc_estimators import ( - RadiationFieldMCEstimators, -) from tardis.montecarlo.montecarlo_numba.numba_interface import ( NumbaModel, RPacketTracker, @@ -33,6 +29,7 @@ def montecarlo_main_loop( geometry_state, numba_model, opacity_state, + montecarlo_configuration, estimators, spectrum_frequency, number_of_vpackets, @@ -97,24 +94,10 @@ def montecarlo_main_loop( main_thread_id = get_thread_id() n_threads = get_num_threads() - estimator_list = List() # betting get thread_id goes from 0 to num threads # Note that get_thread_id() returns values from 0 to n_threads-1, # so we iterate from 0 to n_threads-1 to create the estimator_list - for i in range(n_threads): - estimator_list.append( - RadiationFieldMCEstimators( - np.copy(estimators.j_estimator), - np.copy(estimators.nu_bar_estimator), - np.copy(estimators.j_blue_estimator), - np.copy(estimators.Edotlu_estimator), - np.copy(estimators.photo_ion_estimator), - np.copy(estimators.stim_recomb_estimator), - np.copy(estimators.bf_heating_estimator), - np.copy(estimators.stim_recomb_cooling_estimator), - np.copy(estimators.photo_ion_estimator_statistics), - ) - ) + estimator_list = estimators.create_list(n_threads) for i in prange(no_of_packets): thread_id = get_thread_id() @@ -157,6 +140,7 @@ def montecarlo_main_loop( local_estimators, vpacket_collection, rpacket_tracker, + montecarlo_configuration, ) packet_collection.output_nus[i] = r_packet.nu From 9657857e6dd4380e9d96e89a1e0ad4a4699de737 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 17:03:48 -0500 Subject: [PATCH 04/15] Fix config object definition --- tardis/montecarlo/montecarlo_configuration.py | 64 +++++++++---------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/tardis/montecarlo/montecarlo_configuration.py b/tardis/montecarlo/montecarlo_configuration.py index fa12591a591..0bbcc88e9d5 100644 --- a/tardis/montecarlo/montecarlo_configuration.py +++ b/tardis/montecarlo/montecarlo_configuration.py @@ -1,6 +1,7 @@ from astropy import units as u from numba import float64, int64, boolean from numba.experimental import jitclass +import numpy as np from tardis.montecarlo.montecarlo_numba.numba_interface import ( LineInteractionType, @@ -28,14 +29,14 @@ @jitclass(numba_config_spec) -class MonteCarloConfiguration: +class MonteCarloConfiguration(object): def __init__(self): self.ENABLE_FULL_RELATIVITY = False self.TEMPORARY_V_PACKET_BINS = 0 self.NUMBER_OF_VPACKETS = 0 self.MONTECARLO_SEED = 0 - self.LINE_INTERACTION_TYPE = None - self.PACKET_SEEDS = [] + self.LINE_INTERACTION_TYPE = 0 + self.PACKET_SEEDS = np.empty(1, dtype=np.int64) self.DISABLE_ELECTRON_SCATTERING = False self.DISABLE_LINE_SCATTERING = False self.SURVIVAL_PROBABILITY = 0.0 @@ -51,33 +52,32 @@ def __init__(self): self.VPACKET_SPAWN_END_FREQUENCY = 1e200 self.ENABLE_VPACKET_TRACKING = False - -def configuration_initialize(config, transport, number_of_vpackets): - if transport.line_interaction_type == "macroatom": - config.LINE_INTERACTION_TYPE = LineInteractionType.MACROATOM - elif transport.line_interaction_type == "downbranch": - config.LINE_INTERACTION_TYPE = LineInteractionType.DOWNBRANCH - elif transport.line_interaction_type == "scatter": - config.LINE_INTERACTION_TYPE = LineInteractionType.SCATTER - else: - raise ValueError( - f'Line interaction type must be one of "macroatom",' - f'"downbranch", or "scatter" but is ' - f"{transport.line_interaction_type}" + def configuration_initialize(self, transport, number_of_vpackets): + if transport.line_interaction_type == "macroatom": + self.LINE_INTERACTION_TYPE = LineInteractionType.MACROATOM + elif transport.line_interaction_type == "downbranch": + self.LINE_INTERACTION_TYPE = LineInteractionType.DOWNBRANCH + elif transport.line_interaction_type == "scatter": + self.LINE_INTERACTION_TYPE = LineInteractionType.SCATTER + else: + raise ValueError( + f'Line interaction type must be one of "macroatom",' + f'"downbranch", or "scatter" but is ' + f"{transport.line_interaction_type}" + ) + self.NUMBER_OF_VPACKETS = number_of_vpackets + self.TEMPORARY_V_PACKET_BINS = number_of_vpackets + self.ENABLE_FULL_RELATIVITY = transport.enable_full_relativity + self.MONTECARLO_SEED = transport.packet_source.base_seed + self.VPACKET_SPAWN_START_FREQUENCY = ( + transport.virtual_spectrum_spawn_range.end.to( + u.Hz, equivalencies=u.spectral() + ).value + ) + self.VPACKET_SPAWN_END_FREQUENCY = ( + transport.virtual_spectrum_spawn_range.start.to( + u.Hz, equivalencies=u.spectral() + ).value ) - config.NUMBER_OF_VPACKETS = number_of_vpackets - config.TEMPORARY_V_PACKET_BINS = number_of_vpackets - config.ENABLE_FULL_RELATIVITY = transport.enable_full_relativity - config.MONTECARLO_SEED = transport.packet_source.base_seed - config.VPACKET_SPAWN_START_FREQUENCY = ( - transport.virtual_spectrum_spawn_range.end.to( - u.Hz, equivalencies=u.spectral() - ).value - ) - config.VPACKET_SPAWN_END_FREQUENCY = ( - transport.virtual_spectrum_spawn_range.start.to( - u.Hz, equivalencies=u.spectral() - ).value - ) - config.ENABLE_VPACKET_TRACKING = transport.enable_vpacket_tracking - config.ENABLE_RPACKET_TRACKING = transport.enable_rpacket_tracking + self.ENABLE_VPACKET_TRACKING = transport.enable_vpacket_tracking + self.ENABLE_RPACKET_TRACKING = transport.enable_rpacket_tracking From 5fa4f09d8138206f33cfe8d04ee54698209da532 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 17:04:56 -0500 Subject: [PATCH 05/15] Initial application of the configuration object (incomplete) --- tardis/io/model_reader.py | 1 + tardis/montecarlo/base.py | 12 ++- tardis/montecarlo/montecarlo_numba/base.py | 6 +- .../montecarlo_numba/formal_integral.py | 11 ++- .../montecarlo_numba/interaction.py | 99 +++++++++++++------ .../montecarlo_numba/numba_interface.py | 20 ++-- .../montecarlo_numba/packet_collections.py | 1 - .../montecarlo/montecarlo_numba/r_packet.py | 6 +- .../montecarlo_numba/single_packet_loop.py | 17 +++- .../montecarlo_numba/tests/conftest.py | 3 - .../montecarlo_numba/tests/test_base.py | 3 - .../montecarlo_numba/tests/test_packet.py | 6 +- .../montecarlo_numba/tests/test_r_packet.py | 3 - .../montecarlo_numba/tests/test_vpacket.py | 2 +- tardis/montecarlo/montecarlo_numba/vpacket.py | 39 +++++--- tardis/montecarlo/packet_source.py | 17 ++-- tardis/montecarlo/tests/test_montecarlo.py | 7 +- tardis/montecarlo/tests/test_packet_source.py | 10 +- tardis/simulation/base.py | 5 +- tardis/transport/frame_transformations.py | 9 +- .../transport/geometry/calculate_distances.py | 1 - tardis/transport/r_packet_transport.py | 14 +-- tardis/transport/tests/test_doppler_factor.py | 2 +- 23 files changed, 174 insertions(+), 120 deletions(-) diff --git a/tardis/io/model_reader.py b/tardis/io/model_reader.py index 67beb541ab2..826afc5b784 100644 --- a/tardis/io/model_reader.py +++ b/tardis/io/model_reader.py @@ -741,6 +741,7 @@ def transport_from_hdf(fname): nthreads=d["nthreads"], enable_virtual_packet_logging=d["virt_logging"], use_gpu=d["use_gpu"], + montecarlo_configuration=d["montecarlo_configuration"], ) new_transport.Edotlu_estimator = d["Edotlu_estimator"] diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index 8f78c0ba077..554961c3509 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -11,7 +11,6 @@ initialize_estimator_statistics, ) from tardis.montecarlo.montecarlo_configuration import ( - configuration_initialize, MonteCarloConfiguration, ) from tardis.montecarlo.montecarlo_numba import ( @@ -126,7 +125,10 @@ def initialize_transport_state( geometry_state = simulation_state.geometry.to_numba() opacity_state = opacity_state_initialize( - plasma, self.line_interaction_type + plasma, + self.line_interaction_type, + self.montecarlo_configuration.DISABLE_LINE_SCATTERING, + self.montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, ) transport_state = MonteCarloTransportState( packet_collection, @@ -141,8 +143,8 @@ def initialize_transport_state( transport_state._integrator = FormalIntegrator( simulation_state, plasma, self ) - configuration_initialize( - self.montecarlo_configuration, self, no_of_virtual_packets + self.montecarlo_configuration.configuration_initialize( + transport=self, number_of_vpackets=no_of_virtual_packets ) return transport_state @@ -305,7 +307,7 @@ def from_config( valid values are 'GPU', 'CPU', and 'Automatic'.""" ) - montecarlo_configuration = MonteCarloConfiguration + montecarlo_configuration = MonteCarloConfiguration() montecarlo_configuration.DISABLE_LINE_SCATTERING = ( config.plasma.disable_line_scattering diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index dbd2d1ac872..4a44e3bf33a 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -88,7 +88,11 @@ def montecarlo_main_loop( montecarlo_configuration.TEMPORARY_V_PACKET_BINS, ) ) - rpacket_trackers.append(RPacketTracker()) + rpacket_trackers.append( + RPacketTracker( + montecarlo_configuration.INITIAL_TRACKING_ARRAY_LENGTH + ) + ) # Get the ID of the main thread and the number of threads main_thread_id = get_thread_id() diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral.py b/tardis/montecarlo/montecarlo_numba/formal_integral.py index c49c706b373..162eeee97e3 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/formal_integral.py @@ -11,9 +11,6 @@ from tardis.montecarlo.montecarlo_numba.numba_config import SIGMA_THOMSON -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) from tardis.montecarlo.montecarlo_numba import njit_dict, njit_dict_no_parallel from tardis.montecarlo.montecarlo_numba.numba_interface import ( opacity_state_initialize, @@ -284,9 +281,13 @@ def __init__(self, simulation_state, plasma, transport, points=1000): self.simulation_state = simulation_state self.transport = transport self.points = points + self.montecarlo_configuration = self.transport.montecarlo_configuration if plasma: self.plasma = opacity_state_initialize( - plasma, transport.line_interaction_type + plasma, + transport.line_interaction_type, + self.montecarlo_configuration.DISABLE_LINE_SCATTERING, + self.montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, ) self.atomic_data = plasma.atomic_data self.original_plasma = plasma @@ -361,7 +362,7 @@ def raise_or_return(message): 'and line_interaction_type == "macroatom"' ) - if montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED: + if self.montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED: return raise_or_return( "The FormalIntegrator currently does not work for " "continuum interactions." diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py index 3da74637f02..8310b2a5094 100644 --- a/tardis/montecarlo/montecarlo_numba/interaction.py +++ b/tardis/montecarlo/montecarlo_numba/interaction.py @@ -2,9 +2,6 @@ from numba import njit from tardis import constants as const -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel from tardis.montecarlo.montecarlo_numba.macro_atom import ( MacroAtomTransitionType, @@ -150,6 +147,8 @@ def continuum_event( chi_ff, chi_bf_contributions, current_continua, + continuum_processes_enabled, + enable_full_relativity, ): """ continuum event handler - activate the macroatom and run the handler @@ -162,12 +161,12 @@ def continuum_event( continuum : tardis.montecarlo.montecarlo_numba.numba_interface.Continuum """ old_doppler_factor = get_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) r_packet.mu = get_random_mu() inverse_doppler_factor = get_inverse_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) comov_energy = r_packet.energy * old_doppler_factor comov_nu = ( @@ -186,13 +185,23 @@ def continuum_event( ) macro_atom_event( - destination_level_idx, r_packet, time_explosion, opacity_state + destination_level_idx, + r_packet, + time_explosion, + opacity_state, + continuum_processes_enabled, + enable_full_relativity, ) @njit(**njit_dict_no_parallel) def macro_atom_event( - destination_level_idx, r_packet, time_explosion, opacity_state + destination_level_idx, + r_packet, + time_explosion, + opacity_state, + continuum_processes_enabled, + enable_full_relativity, ): """ Macroatom event handler - run the macroatom and handle the result @@ -209,32 +218,45 @@ def macro_atom_event( ) if ( - montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + continuum_processes_enabled and transition_type == MacroAtomTransitionType.FF_EMISSION ): - free_free_emission(r_packet, time_explosion, opacity_state) + free_free_emission( + r_packet, time_explosion, opacity_state, enable_full_relativity + ) elif ( - montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + continuum_processes_enabled and transition_type == MacroAtomTransitionType.BF_EMISSION ): bound_free_emission( - r_packet, time_explosion, opacity_state, transition_id + r_packet, + time_explosion, + opacity_state, + transition_id, + enable_full_relativity, ) elif ( - montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + continuum_processes_enabled and transition_type == MacroAtomTransitionType.BF_COOLING ): bf_cooling(r_packet, time_explosion, opacity_state) elif ( - montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + continuum_processes_enabled and transition_type == MacroAtomTransitionType.ADIABATIC_COOLING ): adiabatic_cooling(r_packet) elif transition_type == MacroAtomTransitionType.BB_EMISSION: - line_emission(r_packet, transition_id, time_explosion, opacity_state) + line_emission( + r_packet, + transition_id, + time_explosion, + opacity_state, + continuum_processes_enabled, + enable_full_relativity, + ) else: raise Exception("No Interaction Found!") @@ -295,7 +317,9 @@ def get_current_line_id(nu, line_list): @njit(**njit_dict_no_parallel) -def free_free_emission(r_packet, time_explosion, opacity_state): +def free_free_emission( + r_packet, time_explosion, opacity_state, full_relativity +): """ Free-Free emission - set the frequency from electron-ion interaction @@ -313,14 +337,16 @@ def free_free_emission(r_packet, time_explosion, opacity_state): current_line_id = get_current_line_id(comov_nu, opacity_state.line_list_nu) r_packet.next_line_id = current_line_id - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if full_relativity: r_packet.mu = angle_aberration_CMF_to_LF( r_packet, time_explosion, r_packet.mu ) @njit(**njit_dict_no_parallel) -def bound_free_emission(r_packet, time_explosion, opacity_state, continuum_id): +def bound_free_emission( + r_packet, time_explosion, opacity_state, continuum_id, full_relativity +): """ Bound-Free emission - set the frequency from photo-ionization @@ -342,14 +368,14 @@ def bound_free_emission(r_packet, time_explosion, opacity_state, continuum_id): current_line_id = get_current_line_id(comov_nu, opacity_state.line_list_nu) r_packet.next_line_id = current_line_id - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if full_relativity: r_packet.mu = angle_aberration_CMF_to_LF( r_packet, time_explosion, r_packet.mu ) @njit(**njit_dict_no_parallel) -def thomson_scatter(r_packet, time_explosion): +def thomson_scatter(r_packet, time_explosion, enable_full_relativity): """ Thomson scattering — no longer line scattering \n1) get the doppler factor at that position with the old angle @@ -364,7 +390,7 @@ def thomson_scatter(r_packet, time_explosion): time since explosion in seconds """ old_doppler_factor = get_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) comov_nu = r_packet.nu * old_doppler_factor comov_energy = r_packet.energy * old_doppler_factor @@ -375,18 +401,23 @@ def thomson_scatter(r_packet, time_explosion): r_packet.nu = comov_nu * inverse_new_doppler_factor r_packet.energy = comov_energy * inverse_new_doppler_factor - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if enable_full_relativity: r_packet.mu = angle_aberration_CMF_to_LF( r_packet, time_explosion, r_packet.mu ) temp_doppler_factor = get_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) @njit(**njit_dict_no_parallel) def line_scatter( - r_packet, time_explosion, line_interaction_type, opacity_state + r_packet, + time_explosion, + line_interaction_type, + opacity_state, + continuum_processes_enabled, + enable_full_relativity, ): """ Line scatter function that handles the scattering itself, including new angle drawn, and calculating nu out using macro atom @@ -399,7 +430,7 @@ def line_scatter( opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState """ old_doppler_factor = get_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) r_packet.mu = get_random_mu() @@ -412,7 +443,12 @@ def line_scatter( if line_interaction_type == LineInteractionType.SCATTER: line_emission( - r_packet, r_packet.next_line_id, time_explosion, opacity_state + r_packet, + r_packet.next_line_id, + time_explosion, + opacity_state, + continuum_processes_enabled, + enable_full_relativity, ) else: # includes both macro atom and downbranch - encoded in the transition probabilities comov_nu = r_packet.nu * old_doppler_factor # Is this necessary? @@ -421,12 +457,19 @@ def line_scatter( r_packet.next_line_id ] macro_atom_event( - activation_level_id, r_packet, time_explosion, opacity_state + activation_level_id, + r_packet, + time_explosion, + opacity_state, + continuum_processes_enabled, + enable_full_relativity, ) @njit(**njit_dict_no_parallel) -def line_emission(r_packet, emission_line_id, time_explosion, opacity_state): +def line_emission( + r_packet, emission_line_id, time_explosion, opacity_state, full_relativity +): """ Sets the frequency of the RPacket properly given the emission channel @@ -451,7 +494,7 @@ def line_emission(r_packet, emission_line_id, time_explosion, opacity_state): r_packet.next_line_id = emission_line_id + 1 nu_line = opacity_state.line_list_nu[emission_line_id] - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if full_relativity: r_packet.mu = angle_aberration_CMF_to_LF( r_packet, time_explosion, r_packet.mu ) diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index 6d310565622..f6c9198444d 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -6,9 +6,6 @@ from tardis import constants as const -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) C_SPEED_OF_LIGHT = const.c.to("cm/s").value @@ -134,7 +131,12 @@ def __init__( self.k_packet_idx = k_packet_idx -def opacity_state_initialize(plasma, line_interaction_type): +def opacity_state_initialize( + plasma, + line_interaction_type, + disable_line_scattering, + continuum_processes_enabled, +): """ Initialize the OpacityState object and copy over the data over from TARDIS Plasma @@ -150,7 +152,7 @@ def opacity_state_initialize(plasma, line_interaction_type): tau_sobolev = np.ascontiguousarray( plasma.tau_sobolevs.values.copy(), dtype=np.float64 ) - if montecarlo_configuration.DISABLE_LINE_SCATTERING: + if disable_line_scattering: tau_sobolev *= 0 if line_interaction_type == "scatter": @@ -173,7 +175,7 @@ def opacity_state_initialize(plasma, line_interaction_type): ) # TODO: Fix setting of block references for non-continuum mode - if montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED: + if continuum_processes_enabled: macro_block_references = plasma.macro_block_references else: macro_block_references = plasma.atomic_data.macro_atom_references[ @@ -186,7 +188,7 @@ def opacity_state_initialize(plasma, line_interaction_type): "destination_level_idx" ].values transition_line_id = plasma.macro_atom_data["lines_idx"].values - if montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED: + if continuum_processes_enabled: bf_threshold_list_nu = plasma.nu_i.loc[ plasma.level2continuum_idx.index ].values @@ -312,8 +314,8 @@ class RPacketTracker(object): Internal counter for the interactions that a particular RPacket undergoes """ - def __init__(self): - self.length = montecarlo_configuration.INITIAL_TRACKING_ARRAY_LENGTH + def __init__(self, length): + self.length = length self.seed = np.int64(0) self.index = np.int64(0) self.status = np.empty(self.length, dtype=np.int64) diff --git a/tardis/montecarlo/montecarlo_numba/packet_collections.py b/tardis/montecarlo/montecarlo_numba/packet_collections.py index 0a30889962f..d373f172e3b 100644 --- a/tardis/montecarlo/montecarlo_numba/packet_collections.py +++ b/tardis/montecarlo/montecarlo_numba/packet_collections.py @@ -2,7 +2,6 @@ from numba import float64, int64, njit from numba.experimental import jitclass -from tardis.montecarlo import montecarlo_configuration from tardis.montecarlo.montecarlo_numba import ( njit_dict_no_parallel, ) diff --git a/tardis/montecarlo/montecarlo_numba/r_packet.py b/tardis/montecarlo/montecarlo_numba/r_packet.py index a6d927e5eee..842debe9320 100644 --- a/tardis/montecarlo/montecarlo_numba/r_packet.py +++ b/tardis/montecarlo/montecarlo_numba/r_packet.py @@ -64,10 +64,12 @@ def __init__(self, r, mu, nu, energy, seed, index=0): self.last_line_interaction_out_id = -1 self.last_line_interaction_shell_id = -1 - def initialize_line_id(self, opacity_state, numba_model): + def initialize_line_id( + self, opacity_state, numba_model, enable_full_relativity + ): inverse_line_list_nu = opacity_state.line_list_nu[::-1] doppler_factor = get_doppler_factor( - self.r, self.mu, numba_model.time_explosion + self.r, self.mu, numba_model.time_explosion, enable_full_relativity ) comov_nu = self.nu * doppler_factor next_line_id = len(opacity_state.line_list_nu) - np.searchsorted( diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py index e61ec6b5e08..bdc653964c4 100644 --- a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py +++ b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py @@ -1,10 +1,7 @@ from numba import njit from tardis import constants as const -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) -from tardis.montecarlo.estimators.radfield_mc_estimators import ( +from tardis.montecarlo.estimators.radfield_estimator_calcs import ( update_bound_free_estimators, ) from tardis.montecarlo.montecarlo_numba.interaction import ( @@ -43,6 +40,7 @@ def single_packet_loop( estimators, vpacket_collection, rpacket_tracker, + montecarlo_configuration, ): """ Parameters @@ -85,8 +83,12 @@ def single_packet_loop( # Compute continuum quantities # trace packet (takes opacities) doppler_factor = get_doppler_factor( - r_packet.r, r_packet.mu, numba_model.time_explosion + r_packet.r, + r_packet.mu, + numba_model.time_explosion, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) + comov_nu = r_packet.nu * doppler_factor chi_e = chi_electron_calculator( opacity_state, comov_nu, r_packet.current_shell_id @@ -139,6 +141,9 @@ def single_packet_loop( estimators, chi_continuum, escat_prob, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, + montecarlo_configuration.DISABLE_LINE_SCATTERING, ) # If continuum processes: update continuum estimators @@ -229,6 +234,7 @@ def set_packet_props_partial_relativity(r_packet, numba_model): r_packet.r, r_packet.mu, numba_model.time_explosion, + enable_full_relativity=False, ) r_packet.nu *= inverse_doppler_factor r_packet.energy *= inverse_doppler_factor @@ -249,6 +255,7 @@ def set_packet_props_full_relativity(r_packet, numba_model): r_packet.r, r_packet.mu, numba_model.time_explosion, + enable_full_relativity=True, ) r_packet.nu *= inverse_doppler_factor diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index 12c4133ccd2..a991da4fcde 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -3,9 +3,6 @@ import pytest import numpy as np from numba import njit -from tardis.montecarlo.estimators.radfield_mc_estimators import ( - RadiationFieldMCEstimators, -) from tardis.montecarlo.montecarlo_numba.packet_collections import ( VPacketCollection, ) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index 34d2d0926e2..7b80c8c764c 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -3,9 +3,6 @@ import numpy.testing as npt import pytest -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) from tardis.simulation import Simulation diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py index a3af0173047..1c25575c32d 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -12,7 +12,7 @@ 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.estimators.radfield_mc_estimators import ( +from tardis.montecarlo.estimators.radfield_estimator_calcs import ( update_line_estimators, ) @@ -114,7 +114,7 @@ def test_calculate_distance_line( time_explosion = model.time_explosion doppler_factor = frame_transformations.get_doppler_factor( - static_packet.r, static_packet.mu, time_explosion + static_packet.r, static_packet.mu, time_explosion, False ) comov_nu = static_packet.nu * doppler_factor @@ -282,7 +282,7 @@ def test_move_r_packet( numba_config.ENABLE_FULL_RELATIVITY = ENABLE_FULL_RELATIVITY r_packet_transport.move_r_packet.recompile() # This must be done as move_r_packet was jitted with ENABLE_FULL_RELATIVITY doppler_factor = frame_transformations.get_doppler_factor( - packet.r, packet.mu, model.time_explosion + packet.r, packet.mu, model.time_explosion, ENABLE_FULL_RELATIVITY ) r_packet_transport.move_r_packet( diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py index 1964dc48edc..15866a7c711 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_r_packet.py @@ -5,9 +5,6 @@ import pytest from tardis.base import run_tardis -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) from tardis.montecarlo.montecarlo_numba.r_packet import ( rpacket_trackers_to_dataframe, ) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py index 4581ded44a8..a032ce2ca73 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py @@ -30,7 +30,7 @@ def v_packet(): def v_packet_initialize_line_id(v_packet, opacity_state, numba_model): inverse_line_list_nu = opacity_state.line_list_nu[::-1] doppler_factor = get_doppler_factor( - v_packet.r, v_packet.mu, numba_model.time_explosion + v_packet.r, v_packet.mu, numba_model.time_explosion, False ) comov_nu = v_packet.nu * doppler_factor next_line_id = len(opacity_state.line_list_nu) - np.searchsorted( diff --git a/tardis/montecarlo/montecarlo_numba/vpacket.py b/tardis/montecarlo/montecarlo_numba/vpacket.py index 627cbd75e64..7c3aa15d878 100644 --- a/tardis/montecarlo/montecarlo_numba/vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/vpacket.py @@ -9,10 +9,6 @@ from numba.experimental import jitclass from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) - from tardis.montecarlo.montecarlo_numba.r_packet import ( PacketStatus, ) @@ -72,7 +68,12 @@ def __init__( @njit(**njit_dict_no_parallel) def trace_vpacket_within_shell( - v_packet, numba_radial_1d_geometry, numba_model, opacity_state + v_packet, + numba_radial_1d_geometry, + numba_model, + opacity_state, + enable_full_relativity, + continuum_processes_enabled, ): """ Trace VPacket within one shell (relatively simple operation) @@ -95,12 +96,15 @@ def trace_vpacket_within_shell( # Calculating doppler factor doppler_factor = get_doppler_factor( - v_packet.r, v_packet.mu, numba_model.time_explosion + v_packet.r, + v_packet.mu, + numba_model.time_explosion, + enable_full_relativity, ) comov_nu = v_packet.nu * doppler_factor - if montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED: + if continuum_processes_enabled: ( chi_bf_tot, chi_bf_contributions, @@ -115,7 +119,7 @@ def trace_vpacket_within_shell( else: chi_continuum = chi_e - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if enable_full_relativity: chi_continuum *= doppler_factor tau_continuum = chi_continuum * distance_boundary @@ -223,6 +227,7 @@ def trace_vpacket_volley( numba_radial_1d_geometry, numba_model, opacity_state, + enable_full_relativity, ): """ Shoot a volley of vpackets (the vpacket collection specifies how many) @@ -258,7 +263,7 @@ def trace_vpacket_volley( r_inner_over_r = numba_radial_1d_geometry.r_inner[0] / r_packet.r mu_min = -math.sqrt(1 - r_inner_over_r * r_inner_over_r) v_packet_on_inner_boundary = False - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if enable_full_relativity: mu_min = angle_aberration_LF_to_CMF( r_packet, numba_model.time_explosion, mu_min ) @@ -266,20 +271,23 @@ def trace_vpacket_volley( v_packet_on_inner_boundary = True mu_min = 0.0 - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if enable_full_relativity: inv_c = 1 / C_SPEED_OF_LIGHT inv_t = 1 / numba_model.time_explosion beta_inner = numba_radial_1d_geometry.r_inner[0] * inv_t * inv_c mu_bin = (1.0 - mu_min) / no_of_vpackets r_packet_doppler_factor = get_doppler_factor( - r_packet.r, r_packet.mu, numba_model.time_explosion + r_packet.r, + r_packet.mu, + numba_model.time_explosion, + enable_full_relativity, ) for i in range(no_of_vpackets): v_packet_mu = mu_min + i * mu_bin + np.random.random() * mu_bin if v_packet_on_inner_boundary: # The weights are described in K&S 2014 - if not montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if not enable_full_relativity: weight = 2 * v_packet_mu / no_of_vpackets else: weight = ( @@ -293,12 +301,15 @@ def trace_vpacket_volley( weight = (1 - mu_min) / (2 * no_of_vpackets) # C code: next line, angle_aberration_CMF_to_LF( & virt_packet, storage); - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if enable_full_relativity: v_packet_mu = angle_aberration_CMF_to_LF( r_packet, numba_model.time_explosion, v_packet_mu ) v_packet_doppler_factor = get_doppler_factor( - r_packet.r, v_packet_mu, numba_model.time_explosion + r_packet.r, + v_packet_mu, + numba_model.time_explosion, + enable_full_relativity, ) # transform between r_packet mu and v_packet_mu diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index 454d0f0aefd..de2449415b8 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -3,9 +3,6 @@ import numpy as np import numexpr as ne from tardis import constants as const -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) from tardis.montecarlo.montecarlo_numba.packet_collections import ( PacketCollection, ) @@ -30,12 +27,12 @@ class BasePacketSource(abc.ABC): # seed val to the maximum allowed by numpy. MAX_SEED_VAL = 2**32 - 1 - def __init__(self, base_seed=None, legacy_second_seed=None): + def __init__( + self, base_seed=None, legacy_mode_enabled=False, legacy_second_seed=None + ): self.base_seed = base_seed - if ( - montecarlo_configuration.LEGACY_MODE_ENABLED - and legacy_second_seed is not None - ): + self.legacy_mode_enabled = legacy_mode_enabled + if self.legacy_mode_enabled and legacy_second_seed is not None: np.random.seed(legacy_second_seed) else: np.random.seed(self.base_seed) @@ -213,7 +210,7 @@ def create_packet_nus(self, no_of_packets, l_samples=1000): l_coef = np.pi**4 / 90.0 # For testing purposes - if montecarlo_configuration.LEGACY_MODE_ENABLED: + if self.legacy_mode_enabled: xis = np.random.random((5, no_of_packets)) else: xis = self.rng.random((5, no_of_packets)) @@ -241,7 +238,7 @@ def create_packet_mus(self, no_of_packets): """ # For testing purposes - if montecarlo_configuration.LEGACY_MODE_ENABLED: + if self.legacy_mode_enabled: return np.sqrt(np.random.random(no_of_packets)) else: return np.sqrt(self.rng.random(no_of_packets)) diff --git a/tardis/montecarlo/tests/test_montecarlo.py b/tardis/montecarlo/tests/test_montecarlo.py index 8ab1356b33c..753878d5d00 100644 --- a/tardis/montecarlo/tests/test_montecarlo.py +++ b/tardis/montecarlo/tests/test_montecarlo.py @@ -4,7 +4,6 @@ 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.montecarlo.montecarlo_numba.utils as utils @@ -393,7 +392,9 @@ def test_compute_distance2line(packet_params, expected_params): time_explosion = 5.2e7 - doppler_factor = get_doppler_factor(packet.r, packet.mu, time_explosion) + doppler_factor = get_doppler_factor( + packet.r, packet.mu, time_explosion, False + ) comov_nu = packet.nu * doppler_factor d_line = 0 @@ -467,7 +468,7 @@ def test_move_packet(packet_params, expected_params, full_relativity): packet_params["j"], packet_params["nu_bar"], 0, 0 ) r_packet_transport.move_r_packet( - packet, distance, time_explosion, numba_estimator + packet, distance, time_explosion, numba_estimator, full_relativity ) assert_almost_equal(packet.mu, expected_params["mu"]) diff --git a/tardis/montecarlo/tests/test_packet_source.py b/tardis/montecarlo/tests/test_packet_source.py index 53b57b69d67..4a1f2bd4820 100644 --- a/tardis/montecarlo/tests/test_packet_source.py +++ b/tardis/montecarlo/tests/test_packet_source.py @@ -43,10 +43,10 @@ def blackbodysimplesource(self, request): tardis.montecarlo.packet_source.BlackBodySimpleSource """ cls = type(self) - montecarlo_configuration.LEGACY_MODE_ENABLED = True - cls.bb = BlackBodySimpleSource(base_seed=1963, legacy_second_seed=2508) + cls.bb = BlackBodySimpleSource( + base_seed=1963, legacy_mode_enabled=True, legacy_second_seed=2508 + ) yield cls.bb - montecarlo_configuration.LEGACY_MODE_ENABLED = False @pytest.fixture(scope="class") def blackbody_simplesource_relativistic(self, request): @@ -57,12 +57,10 @@ def blackbody_simplesource_relativistic(self, request): ------- tardis.montecarlo.packet_source.BlackBodySimpleSourceRelativistic """ - montecarlo_configuration.LEGACY_MODE_ENABLED = True bb_rel = BlackBodySimpleSourceRelativistic( - base_seed=1963, legacy_second_seed=2508 + base_seed=1963, legacy_mode_enabled=True, legacy_second_seed=2508 ) yield bb_rel - montecarlo_configuration.LEGACY_MODE_ENABLED = False def test_bb_packet_sampling( self, diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index aa569af2aac..dc404c1b354 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -15,9 +15,6 @@ from tardis.io.util import HDFWriterMixin from tardis.model import SimulationState from tardis.model.parse_input import initialize_packet_source -from tardis.montecarlo import ( - montecarlo_configuration as montecarlo_configuration, -) from tardis.montecarlo.base import MonteCarloTransportSolver from tardis.plasma.standard_plasmas import assemble_plasma from tardis.util.base import is_notebook @@ -199,7 +196,7 @@ def __init__( self._callbacks = OrderedDict() self._cb_next_id = 0 - montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED = ( + self.transport.montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED = ( not self.plasma.continuum_interaction_species.empty ) diff --git a/tardis/transport/frame_transformations.py b/tardis/transport/frame_transformations.py index 07bb9dc5b14..a67ad27be41 100644 --- a/tardis/transport/frame_transformations.py +++ b/tardis/transport/frame_transformations.py @@ -6,16 +6,15 @@ njit_dict_no_parallel, ) -from tardis.montecarlo import montecarlo_configuration as nc from tardis.montecarlo.montecarlo_numba.numba_config import C_SPEED_OF_LIGHT @njit(**njit_dict_no_parallel) -def get_doppler_factor(r, mu, time_explosion): +def get_doppler_factor(r, mu, time_explosion, enable_full_relativity): inv_c = 1 / C_SPEED_OF_LIGHT inv_t = 1 / time_explosion beta = r * inv_t * inv_c - if not nc.ENABLE_FULL_RELATIVITY: + if not enable_full_relativity: return get_doppler_factor_partial_relativity(mu, beta) else: return get_doppler_factor_full_relativity(mu, beta) @@ -32,7 +31,7 @@ def get_doppler_factor_full_relativity(mu, beta): @njit(**njit_dict_no_parallel) -def get_inverse_doppler_factor(r, mu, time_explosion): +def get_inverse_doppler_factor(r, mu, time_explosion, enable_full_relativity): """ Calculate doppler factor for frame transformation @@ -45,7 +44,7 @@ def get_inverse_doppler_factor(r, mu, time_explosion): inv_c = 1 / C_SPEED_OF_LIGHT inv_t = 1 / time_explosion beta = r * inv_t * inv_c - if not nc.ENABLE_FULL_RELATIVITY: + if not enable_full_relativity: return get_inverse_doppler_factor_partial_relativity(mu, beta) else: return get_inverse_doppler_factor_full_relativity(mu, beta) diff --git a/tardis/transport/geometry/calculate_distances.py b/tardis/transport/geometry/calculate_distances.py index 204b42da3f2..edef2e2479c 100644 --- a/tardis/transport/geometry/calculate_distances.py +++ b/tardis/transport/geometry/calculate_distances.py @@ -6,7 +6,6 @@ njit_dict_no_parallel, ) -import tardis.montecarlo.montecarlo_configuration as nc from tardis.montecarlo.montecarlo_numba.numba_config import ( C_SPEED_OF_LIGHT, MISS_DISTANCE, diff --git a/tardis/transport/r_packet_transport.py b/tardis/transport/r_packet_transport.py index 20e4bd252df..b3ad6eed7db 100644 --- a/tardis/transport/r_packet_transport.py +++ b/tardis/transport/r_packet_transport.py @@ -1,14 +1,13 @@ import numpy as np from numba import njit -from tardis.montecarlo import montecarlo_configuration from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel from tardis.transport.geometry.calculate_distances import ( calculate_distance_boundary, calculate_distance_electron, calculate_distance_line, ) -from tardis.montecarlo.estimators.radfield_mc_estimators import ( +from tardis.montecarlo.estimators.radfield_estimator_calcs import ( update_line_estimators, update_base_estimators, ) @@ -31,6 +30,9 @@ def trace_packet( estimators, chi_continuum, escat_prob, + continuum_processes_enabled, + full_relativity, + disable_line_scattering, ): """ Traces the RPacket through the ejecta and stops when an interaction happens (heart of the calculation) @@ -110,7 +112,7 @@ def trace_packet( r_packet.next_line_id = cur_line_id break elif distance == distance_continuum: - if not montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED: + if not continuum_processes_enabled: interaction_type = InteractionType.ESCATTERING else: zrand = np.random.random() @@ -131,12 +133,10 @@ def trace_packet( cur_line_id, distance_trace, numba_model.time_explosion, + full_relativity, ) - if ( - tau_trace_combined > tau_event - and not montecarlo_configuration.DISABLE_LINE_SCATTERING - ): + if tau_trace_combined > tau_event and not disable_line_scattering: interaction_type = InteractionType.LINE # Line r_packet.last_interaction_in_nu = r_packet.nu r_packet.last_line_interaction_in_id = cur_line_id diff --git a/tardis/transport/tests/test_doppler_factor.py b/tardis/transport/tests/test_doppler_factor.py index 4d1ff158970..cfe0749a1b1 100644 --- a/tardis/transport/tests/test_doppler_factor.py +++ b/tardis/transport/tests/test_doppler_factor.py @@ -128,7 +128,7 @@ def test_get_inverse_doppler_factor(mu, r, inv_t_exp, expected): # to other methods or introduction of some temporary variables obtained = frame_transformations.get_inverse_doppler_factor( - r, mu, time_explosion + r, mu, time_explosion, False ) # Perform required assertions From 050c42b12d48b943534eed4be59a46e121123b77 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 18:09:33 -0500 Subject: [PATCH 06/15] Continuing to fix config references --- tardis/montecarlo/base.py | 5 +- tardis/montecarlo/montecarlo_configuration.py | 59 +++++++++--------- .../montecarlo_numba/interaction.py | 60 ++++++------------- .../montecarlo_numba/single_packet_loop.py | 34 ++++++++--- .../tests/test_interaction.py | 9 ++- .../montecarlo_numba/tests/test_packet.py | 20 ++++++- .../montecarlo_numba/tests/test_vpacket.py | 1 + tardis/montecarlo/montecarlo_numba/vpacket.py | 1 + tardis/montecarlo/packet_source.py | 25 +++----- tardis/montecarlo/tests/test_montecarlo.py | 9 ++- .../transport/geometry/calculate_distances.py | 9 ++- tardis/transport/r_packet_transport.py | 22 ++++--- tardis/transport/tests/test_doppler_factor.py | 4 +- 13 files changed, 143 insertions(+), 115 deletions(-) diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index 554961c3509..dec2513ebed 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -12,6 +12,7 @@ ) from tardis.montecarlo.montecarlo_configuration import ( MonteCarloConfiguration, + configuration_initialize, ) from tardis.montecarlo.montecarlo_numba import ( montecarlo_main_loop, @@ -143,8 +144,8 @@ def initialize_transport_state( transport_state._integrator = FormalIntegrator( simulation_state, plasma, self ) - self.montecarlo_configuration.configuration_initialize( - transport=self, number_of_vpackets=no_of_virtual_packets + configuration_initialize( + self.montecarlo_configuration, self, no_of_virtual_packets ) return transport_state diff --git a/tardis/montecarlo/montecarlo_configuration.py b/tardis/montecarlo/montecarlo_configuration.py index 0bbcc88e9d5..38793d13d2c 100644 --- a/tardis/montecarlo/montecarlo_configuration.py +++ b/tardis/montecarlo/montecarlo_configuration.py @@ -42,7 +42,7 @@ def __init__(self): self.SURVIVAL_PROBABILITY = 0.0 self.VPACKET_TAU_RUSSIAN = 10.0 - self.INITIAL_TRACKING_ARRAY_LENGTH = None + self.INITIAL_TRACKING_ARRAY_LENGTH = 0 self.LEGACY_MODE_ENABLED = False self.ENABLE_RPACKET_TRACKING = False @@ -52,32 +52,33 @@ def __init__(self): self.VPACKET_SPAWN_END_FREQUENCY = 1e200 self.ENABLE_VPACKET_TRACKING = False - def configuration_initialize(self, transport, number_of_vpackets): - if transport.line_interaction_type == "macroatom": - self.LINE_INTERACTION_TYPE = LineInteractionType.MACROATOM - elif transport.line_interaction_type == "downbranch": - self.LINE_INTERACTION_TYPE = LineInteractionType.DOWNBRANCH - elif transport.line_interaction_type == "scatter": - self.LINE_INTERACTION_TYPE = LineInteractionType.SCATTER - else: - raise ValueError( - f'Line interaction type must be one of "macroatom",' - f'"downbranch", or "scatter" but is ' - f"{transport.line_interaction_type}" - ) - self.NUMBER_OF_VPACKETS = number_of_vpackets - self.TEMPORARY_V_PACKET_BINS = number_of_vpackets - self.ENABLE_FULL_RELATIVITY = transport.enable_full_relativity - self.MONTECARLO_SEED = transport.packet_source.base_seed - self.VPACKET_SPAWN_START_FREQUENCY = ( - transport.virtual_spectrum_spawn_range.end.to( - u.Hz, equivalencies=u.spectral() - ).value - ) - self.VPACKET_SPAWN_END_FREQUENCY = ( - transport.virtual_spectrum_spawn_range.start.to( - u.Hz, equivalencies=u.spectral() - ).value + +def configuration_initialize(config, transport, number_of_vpackets): + if transport.line_interaction_type == "macroatom": + config.LINE_INTERACTION_TYPE = LineInteractionType.MACROATOM + elif transport.line_interaction_type == "downbranch": + config.LINE_INTERACTION_TYPE = LineInteractionType.DOWNBRANCH + elif transport.line_interaction_type == "scatter": + config.LINE_INTERACTION_TYPE = LineInteractionType.SCATTER + else: + raise ValueError( + f'Line interaction type must be one of "macroatom",' + f'"downbranch", or "scatter" but is ' + f"{transport.line_interaction_type}" ) - self.ENABLE_VPACKET_TRACKING = transport.enable_vpacket_tracking - self.ENABLE_RPACKET_TRACKING = transport.enable_rpacket_tracking + config.NUMBER_OF_VPACKETS = number_of_vpackets + config.TEMPORARY_V_PACKET_BINS = number_of_vpackets + config.ENABLE_FULL_RELATIVITY = transport.enable_full_relativity + config.MONTECARLO_SEED = transport.packet_source.base_seed + config.VPACKET_SPAWN_START_FREQUENCY = ( + transport.virtual_spectrum_spawn_range.end.to( + u.Hz, equivalencies=u.spectral() + ).value + ) + config.VPACKET_SPAWN_END_FREQUENCY = ( + transport.virtual_spectrum_spawn_range.start.to( + u.Hz, equivalencies=u.spectral() + ).value + ) + config.ENABLE_VPACKET_TRACKING = transport.enable_vpacket_tracking + config.ENABLE_RPACKET_TRACKING = transport.enable_rpacket_tracking diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py index 8310b2a5094..9368b121a0e 100644 --- a/tardis/montecarlo/montecarlo_numba/interaction.py +++ b/tardis/montecarlo/montecarlo_numba/interaction.py @@ -54,12 +54,8 @@ def determine_bf_macro_activation_idx( # ionization energy is created nu_threshold = opacity_state.photo_ion_nu_threshold_mins[continuum_id] fraction_ionization = nu_threshold / nu - if ( - np.random.random() < fraction_ionization - ): # Create ionization energy (i-packet) - destination_level_idx = opacity_state.photo_ion_activation_idx[ - continuum_id - ] + if np.random.random() < fraction_ionization: # Create ionization energy (i-packet) + destination_level_idx = opacity_state.photo_ion_activation_idx[continuum_id] else: # Create thermal energy (k-packet) destination_level_idx = opacity_state.k_packet_idx return destination_level_idx @@ -169,9 +165,7 @@ def continuum_event( r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) comov_energy = r_packet.energy * old_doppler_factor - comov_nu = ( - r_packet.nu * old_doppler_factor - ) # make sure frequency should be updated + comov_nu = r_packet.nu * old_doppler_factor # make sure frequency should be updated r_packet.energy = comov_energy * inverse_doppler_factor r_packet.nu = comov_nu * inverse_doppler_factor @@ -254,7 +248,6 @@ def macro_atom_event( transition_id, time_explosion, opacity_state, - continuum_processes_enabled, enable_full_relativity, ) else: @@ -272,9 +265,7 @@ 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 - ] + fb_cooling_prob = opacity_state.p_fb_deactivation[:, r_packet.current_shell_id] p = fb_cooling_prob[0] i = 0 zrand = np.random.random() @@ -317,9 +308,7 @@ def get_current_line_id(nu, line_list): @njit(**njit_dict_no_parallel) -def free_free_emission( - r_packet, time_explosion, opacity_state, full_relativity -): +def free_free_emission(r_packet, time_explosion, opacity_state, enable_full_relativity): """ Free-Free emission - set the frequency from electron-ion interaction @@ -330,22 +319,20 @@ def free_free_emission( opacity_state : tardis.montecarlo.montecarlo_numba.numba_interface.OpacityState """ inverse_doppler_factor = get_inverse_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) comov_nu = sample_nu_free_free(opacity_state, r_packet.current_shell_id) r_packet.nu = comov_nu * inverse_doppler_factor current_line_id = get_current_line_id(comov_nu, opacity_state.line_list_nu) r_packet.next_line_id = current_line_id - if full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF( - r_packet, time_explosion, r_packet.mu - ) + if enable_full_relativity: + r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) @njit(**njit_dict_no_parallel) def bound_free_emission( - r_packet, time_explosion, opacity_state, continuum_id, full_relativity + r_packet, time_explosion, opacity_state, continuum_id, enable_full_relativity ): """ Bound-Free emission - set the frequency from photo-ionization @@ -358,7 +345,7 @@ def bound_free_emission( continuum_id : int """ inverse_doppler_factor = get_inverse_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) comov_nu = sample_nu_free_bound( @@ -368,10 +355,8 @@ def bound_free_emission( current_line_id = get_current_line_id(comov_nu, opacity_state.line_list_nu) r_packet.next_line_id = current_line_id - if full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF( - r_packet, time_explosion, r_packet.mu - ) + if enable_full_relativity: + r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) @njit(**njit_dict_no_parallel) @@ -402,9 +387,7 @@ def thomson_scatter(r_packet, time_explosion, enable_full_relativity): r_packet.nu = comov_nu * inverse_new_doppler_factor r_packet.energy = comov_energy * inverse_new_doppler_factor if enable_full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF( - r_packet, time_explosion, r_packet.mu - ) + r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) temp_doppler_factor = get_doppler_factor( r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) @@ -435,7 +418,7 @@ def line_scatter( r_packet.mu = get_random_mu() inverse_new_doppler_factor = get_inverse_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) comov_energy = r_packet.energy * old_doppler_factor @@ -447,7 +430,6 @@ def line_scatter( r_packet.next_line_id, time_explosion, opacity_state, - continuum_processes_enabled, enable_full_relativity, ) else: # includes both macro atom and downbranch - encoded in the transition probabilities @@ -468,7 +450,7 @@ def line_scatter( @njit(**njit_dict_no_parallel) def line_emission( - r_packet, emission_line_id, time_explosion, opacity_state, full_relativity + r_packet, emission_line_id, time_explosion, opacity_state, enable_full_relativity ): """ Sets the frequency of the RPacket properly given the emission channel @@ -486,15 +468,11 @@ def line_emission( if emission_line_id != r_packet.next_line_id: pass inverse_doppler_factor = get_inverse_doppler_factor( - r_packet.r, r_packet.mu, time_explosion - ) - r_packet.nu = ( - opacity_state.line_list_nu[emission_line_id] * inverse_doppler_factor + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) + r_packet.nu = opacity_state.line_list_nu[emission_line_id] * inverse_doppler_factor r_packet.next_line_id = emission_line_id + 1 nu_line = opacity_state.line_list_nu[emission_line_id] - if full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF( - r_packet, time_explosion, r_packet.mu - ) + if enable_full_relativity: + r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py index bdc653964c4..036d723783c 100644 --- a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py +++ b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py @@ -65,7 +65,11 @@ def single_packet_loop( set_packet_props_full_relativity(r_packet, numba_model) else: set_packet_props_partial_relativity(r_packet, numba_model) - r_packet.initialize_line_id(opacity_state, numba_model) + r_packet.initialize_line_id( + opacity_state, + numba_model, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, + ) trace_vpacket_volley( r_packet, @@ -73,6 +77,7 @@ def single_packet_loop( numba_radial_1d_geometry, numba_model, opacity_state, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) if montecarlo_configuration.ENABLE_RPACKET_TRACKING: @@ -150,7 +155,11 @@ def single_packet_loop( if interaction_type == InteractionType.BOUNDARY: move_r_packet( - r_packet, distance, numba_model.time_explosion, estimators + r_packet, + distance, + numba_model.time_explosion, + estimators, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) move_packet_across_shell_boundary( r_packet, delta_shell, len(numba_radial_1d_geometry.r_inner) @@ -159,13 +168,19 @@ def single_packet_loop( elif interaction_type == InteractionType.LINE: r_packet.last_interaction_type = 2 move_r_packet( - r_packet, distance, numba_model.time_explosion, estimators + r_packet, + distance, + numba_model.time_explosion, + estimators, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) line_scatter( r_packet, numba_model.time_explosion, line_interaction_type, opacity_state, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) trace_vpacket_volley( r_packet, @@ -173,14 +188,13 @@ def single_packet_loop( numba_radial_1d_geometry, numba_model, opacity_state, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) elif interaction_type == InteractionType.ESCATTERING: r_packet.last_interaction_type = 1 - move_r_packet( - r_packet, distance, numba_model.time_explosion, estimators - ) + move_r_packet(r_packet, distance, numba_model.time_explosion, estimators) thomson_scatter(r_packet, numba_model.time_explosion) trace_vpacket_volley( @@ -189,6 +203,7 @@ def single_packet_loop( numba_radial_1d_geometry, numba_model, opacity_state, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) elif ( montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED @@ -196,7 +211,11 @@ def single_packet_loop( ): r_packet.last_interaction_type = InteractionType.CONTINUUM_PROCESS move_r_packet( - r_packet, distance, numba_model.time_explosion, estimators + r_packet, + distance, + numba_model.time_explosion, + estimators, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) continuum_event( r_packet, @@ -214,6 +233,7 @@ def single_packet_loop( numba_radial_1d_geometry, numba_model, opacity_state, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) else: pass diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index 969292de543..f620d026800 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -41,7 +41,12 @@ def test_line_scatter( time_explosion = verysimple_numba_model.time_explosion interaction.line_scatter( - packet, time_explosion, line_interaction_type, verysimple_opacity_state + packet, + time_explosion, + line_interaction_type, + verysimple_opacity_state, + False, + False, ) assert np.abs(packet.mu - init_mu) > 1e-7 @@ -94,7 +99,7 @@ def test_line_emission( time_explosion = verysimple_numba_model.time_explosion interaction.line_emission( - packet, emission_line_id, time_explosion, verysimple_opacity_state + packet, emission_line_id, time_explosion, verysimple_opacity_state, False ) assert packet.next_line_id == emission_line_id + 1 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py index 1c25575c32d..7eee17fdaa8 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -122,7 +122,12 @@ def test_calculate_distance_line( obtained_tardis_error = None try: d_line = calculate_distances.calculate_distance_line( - static_packet, comov_nu, is_last_line, nu_line, time_explosion + static_packet, + comov_nu, + is_last_line, + nu_line, + time_explosion, + False, ) except utils.MonteCarloException: obtained_tardis_error = utils.MonteCarloException @@ -209,7 +214,12 @@ def test_update_line_estimators( expected_Edotlu, ): update_line_estimators( - estimators, static_packet, cur_line_id, distance_trace, time_explosion + estimators, + static_packet, + cur_line_id, + distance_trace, + time_explosion, + enable_full_relativity=False, ) assert_allclose(estimators.j_blue_estimator, expected_j_blue) @@ -286,7 +296,11 @@ def test_move_r_packet( ) r_packet_transport.move_r_packet( - packet, distance, model.time_explosion, estimators + packet, + distance, + model.time_explosion, + estimators, + ENABLE_FULL_RELATIVITY, ) assert_almost_equal(packet.mu, expected_params["mu"]) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py index a032ce2ca73..5504343718e 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py @@ -119,6 +119,7 @@ def test_trace_vpacket_volley( verysimple_numba_radial_1d_geometry, verysimple_numba_model, verysimple_opacity_state, + enable_full_relativity=False, ) diff --git a/tardis/montecarlo/montecarlo_numba/vpacket.py b/tardis/montecarlo/montecarlo_numba/vpacket.py index 7c3aa15d878..3ea43ef9f96 100644 --- a/tardis/montecarlo/montecarlo_numba/vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/vpacket.py @@ -146,6 +146,7 @@ def trace_vpacket_within_shell( is_last_line, nu_line, numba_model.time_explosion, + enable_full_relativity, ) if distance_boundary <= distance_trace_line: diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index de2449415b8..36f486970cf 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -81,20 +81,14 @@ def create_packets(self, no_of_packets, seed_offset=0, *args, **kwargs): # because we call random.sample, which references a different internal # state than in the numpy.random module. self._reseed(self.base_seed + seed_offset) - packet_seeds = self.rng.choice( - self.MAX_SEED_VAL, no_of_packets, replace=True - ) + packet_seeds = self.rng.choice(self.MAX_SEED_VAL, no_of_packets, replace=True) radii = self.create_packet_radii(no_of_packets, *args, **kwargs).value nus = self.create_packet_nus(no_of_packets, *args, **kwargs).value mus = self.create_packet_mus(no_of_packets, *args, **kwargs) - energies = self.create_packet_energies( - no_of_packets, *args, **kwargs - ).value + energies = self.create_packet_energies(no_of_packets, *args, **kwargs).value # Check if all arrays have the same length - assert ( - len(radii) == len(nus) == len(mus) == len(energies) == no_of_packets - ) + assert len(radii) == len(nus) == len(mus) == len(energies) == no_of_packets radiation_field_luminosity = self.calculate_radfield_luminosity().value return PacketCollection( radii, @@ -117,13 +111,9 @@ def calculate_radfield_luminosity(self): ------- astropy.units.Quantity """ - return ( - 4 - * np.pi - * const.sigma_sb - * self.radius**2 - * self.temperature**4 - ).to("erg/s") + return (4 * np.pi * const.sigma_sb * self.radius**2 * self.temperature**4).to( + "erg/s" + ) class BlackBodySimpleSource(BasePacketSource): @@ -271,8 +261,7 @@ def set_temperature_from_luminosity(self, luminosity: u.Quantity): """ self.temperature = ( - (luminosity / (4 * np.pi * self.radius**2 * const.sigma_sb)) - ** 0.25 + (luminosity / (4 * np.pi * self.radius**2 * const.sigma_sb)) ** 0.25 ).to("K") diff --git a/tardis/montecarlo/tests/test_montecarlo.py b/tardis/montecarlo/tests/test_montecarlo.py index 753878d5d00..e3178c4daa6 100644 --- a/tardis/montecarlo/tests/test_montecarlo.py +++ b/tardis/montecarlo/tests/test_montecarlo.py @@ -461,9 +461,10 @@ def test_move_packet(packet_params, expected_params, full_relativity): packet.energy = packet_params["energy"] packet.r = packet_params["r"] # model.full_relativity = full_relativity - mc.ENABLE_FULL_RELATIVITY = full_relativity - doppler_factor = get_doppler_factor(packet.r, packet.mu, time_explosion) + doppler_factor = get_doppler_factor( + packet.r, packet.mu, time_explosion, full_relativity + ) numba_estimator = RadiationFieldMCEstimators( packet_params["j"], packet_params["nu_bar"], 0, 0 ) @@ -633,7 +634,9 @@ def test_compute_distance2line_relativistic( distance = r_packet.calculate_distance_line( packet, comov_nu, nu_line, t_exp ) - r_packet_transport.move_r_packet(packet, distance, t_exp, numba_estimator) + r_packet_transport.move_r_packet( + packet, distance, t_exp, numba_estimator, bool(full_relativity) + ) doppler_factor = get_doppler_factor(r, mu, t_exp) comov_nu = packet.nu * doppler_factor diff --git a/tardis/transport/geometry/calculate_distances.py b/tardis/transport/geometry/calculate_distances.py index edef2e2479c..d3b3f34903f 100644 --- a/tardis/transport/geometry/calculate_distances.py +++ b/tardis/transport/geometry/calculate_distances.py @@ -63,7 +63,12 @@ def calculate_distance_boundary(r, mu, r_inner, r_outer): @njit(**njit_dict_no_parallel) def calculate_distance_line( - r_packet, comov_nu, is_last_line, nu_line, time_explosion + r_packet, + comov_nu, + is_last_line, + nu_line, + time_explosion, + enable_full_relativity, ): """ Calculate distance until RPacket is in resonance with the next line @@ -100,7 +105,7 @@ def calculate_distance_line( else: raise MonteCarloException("nu difference is less than 0.0") - if nc.ENABLE_FULL_RELATIVITY: + if enable_full_relativity: return calculate_distance_line_full_relativity( nu_line, nu, time_explosion, r_packet ) diff --git a/tardis/transport/r_packet_transport.py b/tardis/transport/r_packet_transport.py index b3ad6eed7db..46cd14197f4 100644 --- a/tardis/transport/r_packet_transport.py +++ b/tardis/transport/r_packet_transport.py @@ -31,7 +31,7 @@ def trace_packet( chi_continuum, escat_prob, continuum_processes_enabled, - full_relativity, + enable_full_relativity, disable_line_scattering, ): """ @@ -66,7 +66,10 @@ def trace_packet( # Calculating doppler factor doppler_factor = get_doppler_factor( - r_packet.r, r_packet.mu, numba_model.time_explosion + r_packet.r, + r_packet.mu, + numba_model.time_explosion, + enable_full_relativity, ) comov_nu = r_packet.nu * doppler_factor @@ -96,6 +99,7 @@ def trace_packet( is_last_line, nu_line, numba_model.time_explosion, + enable_full_relativity, ) # calculating the tau continuum of how far the trace has progressed @@ -133,7 +137,7 @@ def trace_packet( cur_line_id, distance_trace, numba_model.time_explosion, - full_relativity, + enable_full_relativity, ) if tau_trace_combined > tau_event and not disable_line_scattering: @@ -162,7 +166,7 @@ def trace_packet( cur_line_id += 1 if distance_continuum < distance_boundary: distance = distance_continuum - if not montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED: + if not continuum_processes_enabled: interaction_type = InteractionType.ESCATTERING else: zrand = np.random.random() @@ -178,7 +182,9 @@ def trace_packet( @njit(**njit_dict_no_parallel) -def move_r_packet(r_packet, distance, time_explosion, numba_estimator): +def move_r_packet( + r_packet, distance, time_explosion, numba_estimator, enable_full_relativity +): """ Move packet a distance and recalculate the new angle mu @@ -194,7 +200,9 @@ def move_r_packet(r_packet, distance, time_explosion, numba_estimator): distance in cm """ - doppler_factor = get_doppler_factor(r_packet.r, r_packet.mu, time_explosion) + doppler_factor = get_doppler_factor( + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity + ) r = r_packet.r if distance > 0.0: @@ -208,7 +216,7 @@ def move_r_packet(r_packet, distance, time_explosion, numba_estimator): comov_energy = r_packet.energy * doppler_factor # Account for length contraction - if montecarlo_configuration.ENABLE_FULL_RELATIVITY: + if enable_full_relativity: distance *= doppler_factor update_base_estimators( diff --git a/tardis/transport/tests/test_doppler_factor.py b/tardis/transport/tests/test_doppler_factor.py index cfe0749a1b1..b3219c6e71b 100644 --- a/tardis/transport/tests/test_doppler_factor.py +++ b/tardis/transport/tests/test_doppler_factor.py @@ -38,7 +38,9 @@ def test_get_doppler_factor(mu, r, inv_t_exp, expected): # Perform any other setups just before this, they can be additional calls # to other methods or introduction of some temporary variables - obtained = frame_transformations.get_doppler_factor(r, mu, time_explosion) + obtained = frame_transformations.get_doppler_factor( + r, mu, time_explosion, False + ) # Perform required assertions assert_almost_equal(obtained, expected) From 7c541f346f9d15a482b6a175670d6b7b386790e5 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 19:12:04 -0500 Subject: [PATCH 07/15] Further test fixes --- tardis/montecarlo/montecarlo_numba/base.py | 5 +- .../montecarlo_numba/interaction.py | 68 ++++++++++++++----- .../montecarlo_numba/packet_collections.py | 8 ++- .../montecarlo_numba/single_packet_loop.py | 31 ++++++++- .../montecarlo_numba/tests/test_base.py | 6 +- .../tests/test_interaction.py | 8 ++- .../montecarlo_numba/tests/test_packet.py | 3 + .../montecarlo_numba/tests/test_vpacket.py | 9 +++ tardis/montecarlo/montecarlo_numba/vpacket.py | 34 ++++++++-- tardis/transport/tests/test_doppler_factor.py | 2 +- 10 files changed, 141 insertions(+), 33 deletions(-) diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index 4a44e3bf33a..2d8e18b61ac 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -175,7 +175,10 @@ def montecarlo_main_loop( if enable_virtual_packet_logging: vpacket_tracker = consolidate_vpacket_tracker( - vpacket_collections, spectrum_frequency + vpacket_collections, + spectrum_frequency, + montecarlo_configuration.VPACKET_SPAWN_START_FREQUENCY, + montecarlo_configuration.VPACKET_SPAWN_END_FREQUENCY, ) else: vpacket_tracker = VPacketCollection( diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py index 9368b121a0e..df788469cdd 100644 --- a/tardis/montecarlo/montecarlo_numba/interaction.py +++ b/tardis/montecarlo/montecarlo_numba/interaction.py @@ -54,8 +54,12 @@ def determine_bf_macro_activation_idx( # ionization energy is created nu_threshold = opacity_state.photo_ion_nu_threshold_mins[continuum_id] fraction_ionization = nu_threshold / nu - if np.random.random() < fraction_ionization: # Create ionization energy (i-packet) - destination_level_idx = opacity_state.photo_ion_activation_idx[continuum_id] + if ( + np.random.random() < fraction_ionization + ): # Create ionization energy (i-packet) + destination_level_idx = opacity_state.photo_ion_activation_idx[ + continuum_id + ] else: # Create thermal energy (k-packet) destination_level_idx = opacity_state.k_packet_idx return destination_level_idx @@ -165,7 +169,9 @@ def continuum_event( r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) comov_energy = r_packet.energy * old_doppler_factor - comov_nu = r_packet.nu * old_doppler_factor # make sure frequency should be updated + comov_nu = ( + r_packet.nu * old_doppler_factor + ) # make sure frequency should be updated r_packet.energy = comov_energy * inverse_doppler_factor r_packet.nu = comov_nu * inverse_doppler_factor @@ -234,7 +240,9 @@ def macro_atom_event( continuum_processes_enabled and transition_type == MacroAtomTransitionType.BF_COOLING ): - bf_cooling(r_packet, time_explosion, opacity_state) + bf_cooling( + r_packet, time_explosion, opacity_state, enable_full_relativity + ) elif ( continuum_processes_enabled @@ -255,7 +263,7 @@ def macro_atom_event( @njit(**njit_dict_no_parallel) -def bf_cooling(r_packet, time_explosion, opacity_state): +def bf_cooling(r_packet, time_explosion, opacity_state, enable_full_relativity): """ Bound-Free Cooling - Determine and run bf emission from cooling @@ -265,7 +273,9 @@ 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] + fb_cooling_prob = opacity_state.p_fb_deactivation[ + :, r_packet.current_shell_id + ] p = fb_cooling_prob[0] i = 0 zrand = np.random.random() @@ -273,7 +283,13 @@ def bf_cooling(r_packet, time_explosion, opacity_state): i += 1 p += fb_cooling_prob[i] continuum_idx = i - bound_free_emission(r_packet, time_explosion, opacity_state, continuum_idx) + bound_free_emission( + r_packet, + time_explosion, + opacity_state, + continuum_idx, + enable_full_relativity, + ) @njit(**njit_dict_no_parallel) @@ -308,7 +324,9 @@ def get_current_line_id(nu, line_list): @njit(**njit_dict_no_parallel) -def free_free_emission(r_packet, time_explosion, opacity_state, enable_full_relativity): +def free_free_emission( + r_packet, time_explosion, opacity_state, enable_full_relativity +): """ Free-Free emission - set the frequency from electron-ion interaction @@ -327,12 +345,18 @@ def free_free_emission(r_packet, time_explosion, opacity_state, enable_full_rela r_packet.next_line_id = current_line_id if enable_full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) @njit(**njit_dict_no_parallel) def bound_free_emission( - r_packet, time_explosion, opacity_state, continuum_id, enable_full_relativity + r_packet, + time_explosion, + opacity_state, + continuum_id, + enable_full_relativity, ): """ Bound-Free emission - set the frequency from photo-ionization @@ -356,7 +380,9 @@ def bound_free_emission( r_packet.next_line_id = current_line_id if enable_full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) @njit(**njit_dict_no_parallel) @@ -381,13 +407,15 @@ def thomson_scatter(r_packet, time_explosion, enable_full_relativity): comov_energy = r_packet.energy * old_doppler_factor r_packet.mu = get_random_mu() inverse_new_doppler_factor = get_inverse_doppler_factor( - r_packet.r, r_packet.mu, time_explosion + r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) r_packet.nu = comov_nu * inverse_new_doppler_factor r_packet.energy = comov_energy * inverse_new_doppler_factor if enable_full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) temp_doppler_factor = get_doppler_factor( r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) @@ -450,7 +478,11 @@ def line_scatter( @njit(**njit_dict_no_parallel) def line_emission( - r_packet, emission_line_id, time_explosion, opacity_state, enable_full_relativity + r_packet, + emission_line_id, + time_explosion, + opacity_state, + enable_full_relativity, ): """ Sets the frequency of the RPacket properly given the emission channel @@ -470,9 +502,13 @@ def line_emission( inverse_doppler_factor = get_inverse_doppler_factor( r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) - r_packet.nu = opacity_state.line_list_nu[emission_line_id] * inverse_doppler_factor + r_packet.nu = ( + opacity_state.line_list_nu[emission_line_id] * inverse_doppler_factor + ) r_packet.next_line_id = emission_line_id + 1 nu_line = opacity_state.line_list_nu[emission_line_id] if enable_full_relativity: - r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion, r_packet.mu) + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) diff --git a/tardis/montecarlo/montecarlo_numba/packet_collections.py b/tardis/montecarlo/montecarlo_numba/packet_collections.py index d373f172e3b..2968d0ddcb9 100644 --- a/tardis/montecarlo/montecarlo_numba/packet_collections.py +++ b/tardis/montecarlo/montecarlo_numba/packet_collections.py @@ -277,7 +277,9 @@ def finalize_arrays(self): @njit(**njit_dict_no_parallel) -def consolidate_vpacket_tracker(vpacket_collections, spectrum_frequency): +def consolidate_vpacket_tracker( + vpacket_collections, spectrum_frequency, start_frequency, end_frequency +): """ Consolidate the vpacket trackers from multiple collections into a single vpacket tracker. @@ -301,8 +303,8 @@ def consolidate_vpacket_tracker(vpacket_collections, spectrum_frequency): vpacket_tracker = VPacketCollection( -1, spectrum_frequency, - montecarlo_configuration.VPACKET_SPAWN_START_FREQUENCY, - montecarlo_configuration.VPACKET_SPAWN_END_FREQUENCY, + start_frequency, + end_frequency, -1, vpacket_tracker_length, ) diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py index 036d723783c..3e3580f6d5c 100644 --- a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py +++ b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py @@ -78,6 +78,9 @@ def single_packet_loop( numba_model, opacity_state, montecarlo_configuration.ENABLE_FULL_RELATIVITY, + montecarlo_configuration.VPACKET_TAU_RUSSIAN, + montecarlo_configuration.SURVIVAL_PROBABILITY, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, ) if montecarlo_configuration.ENABLE_RPACKET_TRACKING: @@ -121,6 +124,9 @@ def single_packet_loop( estimators, chi_continuum, escat_prob, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, + montecarlo_configuration.DISABLE_LINE_SCATTERING, ) update_bound_free_estimators( comov_nu, @@ -189,13 +195,26 @@ def single_packet_loop( numba_model, opacity_state, montecarlo_configuration.ENABLE_FULL_RELATIVITY, + montecarlo_configuration.VPACKET_TAU_RUSSIAN, + montecarlo_configuration.SURVIVAL_PROBABILITY, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, ) elif interaction_type == InteractionType.ESCATTERING: r_packet.last_interaction_type = 1 - move_r_packet(r_packet, distance, numba_model.time_explosion, estimators) - thomson_scatter(r_packet, numba_model.time_explosion) + move_r_packet( + r_packet, + distance, + numba_model.time_explosion, + estimators, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, + ) + thomson_scatter( + r_packet, + numba_model.time_explosion, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, + ) trace_vpacket_volley( r_packet, @@ -204,6 +223,9 @@ def single_packet_loop( numba_model, opacity_state, montecarlo_configuration.ENABLE_FULL_RELATIVITY, + montecarlo_configuration.VPACKET_TAU_RUSSIAN, + montecarlo_configuration.SURVIVAL_PROBABILITY, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, ) elif ( montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED @@ -225,6 +247,8 @@ def single_packet_loop( chi_ff, chi_bf_contributions, current_continua, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, + montecarlo_configuration.ENABLE_FULL_RELATIVITY, ) trace_vpacket_volley( @@ -234,6 +258,9 @@ def single_packet_loop( numba_model, opacity_state, montecarlo_configuration.ENABLE_FULL_RELATIVITY, + montecarlo_configuration.VPACKET_TAU_RUSSIAN, + montecarlo_configuration.SURVIVAL_PROBABILITY, + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, ) else: pass diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index 7b80c8c764c..98c79bebc7a 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -15,7 +15,6 @@ def test_montecarlo_radial1d(): def montecarlo_main_loop_config( config_montecarlo_1e5_verysimple, ): - montecarlo_configuration.LEGACY_MODE_ENABLED = True # Setup model config from verysimple config_montecarlo_1e5_verysimple.montecarlo.last_no_of_packets = 1e5 @@ -89,7 +88,10 @@ def test_montecarlo_main_loop_vpacket_log( montecarlo_main_loop_simulation.run_convergence() montecarlo_main_loop_simulation.run_final() - assert montecarlo_configuration.ENABLE_VPACKET_TRACKING == True + assert ( + montecarlo_main_loop_simulation.transport.montecarlo_configuration.ENABLE_VPACKET_TRACKING + == True + ) expected_hdf_store = regression_data.sync_hdf_store( montecarlo_main_loop_simulation diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index f620d026800..47d67f945f8 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -13,7 +13,7 @@ def test_thomson_scatter(packet, verysimple_numba_model): init_energy = packet.energy time_explosion = verysimple_numba_model.time_explosion - interaction.thomson_scatter(packet, time_explosion) + interaction.thomson_scatter(packet, time_explosion, False) assert np.abs(packet.mu - init_mu) > 1e-7 assert np.abs(packet.nu - init_nu) > 1e-7 @@ -99,7 +99,11 @@ def test_line_emission( time_explosion = verysimple_numba_model.time_explosion interaction.line_emission( - packet, emission_line_id, time_explosion, verysimple_opacity_state, False + packet, + emission_line_id, + time_explosion, + verysimple_opacity_state, + False, ) assert packet.next_line_id == emission_line_id + 1 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py index 7eee17fdaa8..e05a5ef60fa 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -243,6 +243,9 @@ def test_trace_packet( verysimple_numba_model, verysimple_opacity_state, verysimple_estimators, + continuum_processes_enabled=False, + enable_full_relativity=False, + disable_line_scattering=False, ) assert delta_shell == 1 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py index 5504343718e..dd27baf0654 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py @@ -59,6 +59,8 @@ def test_trace_vpacket_within_shell( verysimple_numba_radial_1d_geometry, verysimple_numba_model, verysimple_opacity_state, + enable_full_relativity=False, + continuum_processes_enabled=False, ) npt.assert_almost_equal(tau_trace_combined, 8164850.891288479) @@ -86,6 +88,8 @@ def test_trace_vpacket( verysimple_numba_radial_1d_geometry, verysimple_numba_model, verysimple_opacity_state, + 10.0, + 0.0, ) npt.assert_almost_equal(tau_trace_combined, 8164850.891288479) @@ -120,6 +124,9 @@ def test_trace_vpacket_volley( verysimple_numba_model, verysimple_opacity_state, enable_full_relativity=False, + tau_russian=10.0, + survival_probability=0.0, + continuum_processes_enabled=False, ) @@ -147,4 +154,6 @@ def test_trace_bad_vpacket( verysimple_numba_radial_1d_geometry, verysimple_numba_model, verysimple_opacity_state, + 10.0, + 0.0, ) diff --git a/tardis/montecarlo/montecarlo_numba/vpacket.py b/tardis/montecarlo/montecarlo_numba/vpacket.py index 3ea43ef9f96..184dfc18e21 100644 --- a/tardis/montecarlo/montecarlo_numba/vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/vpacket.py @@ -164,7 +164,14 @@ def trace_vpacket_within_shell( @njit(**njit_dict_no_parallel) def trace_vpacket( - v_packet, numba_radial_1d_geometry, numba_model, opacity_state + v_packet, + numba_radial_1d_geometry, + numba_model, + opacity_state, + tau_russian, + survival_probability, + enable_full_relativity, + continuum_processes_enabled, ): """ Trace single vpacket. @@ -186,7 +193,12 @@ def trace_vpacket( distance_boundary, delta_shell, ) = trace_vpacket_within_shell( - v_packet, numba_radial_1d_geometry, numba_model, opacity_state + v_packet, + numba_radial_1d_geometry, + numba_model, + opacity_state, + enable_full_relativity, + continuum_processes_enabled, ) tau_trace_combined += tau_trace_combined_shell @@ -194,15 +206,15 @@ def trace_vpacket( v_packet, delta_shell, len(numba_radial_1d_geometry.r_inner) ) - if tau_trace_combined > montecarlo_configuration.VPACKET_TAU_RUSSIAN: + if tau_trace_combined > tau_russian: event_random = np.random.random() - if event_random > montecarlo_configuration.SURVIVAL_PROBABILITY: + if event_random > survival_probability: v_packet.energy = 0.0 v_packet.status = PacketStatus.EMITTED else: v_packet.energy = ( v_packet.energy - / montecarlo_configuration.SURVIVAL_PROBABILITY + / survival_probability * math.exp(-tau_trace_combined) ) tau_trace_combined = 0.0 @@ -229,6 +241,9 @@ def trace_vpacket_volley( numba_model, opacity_state, enable_full_relativity, + tau_russian, + survival_probability, + continuum_processes_enabled, ): """ Shoot a volley of vpackets (the vpacket collection specifies how many) @@ -335,7 +350,14 @@ def trace_vpacket_volley( ) tau_vpacket = trace_vpacket( - v_packet, numba_radial_1d_geometry, numba_model, opacity_state + v_packet, + numba_radial_1d_geometry, + numba_model, + opacity_state, + tau_russian, + survival_probability, + enable_full_relativity, + continuum_processes_enabled, ) v_packet.energy *= math.exp(-tau_vpacket) diff --git a/tardis/transport/tests/test_doppler_factor.py b/tardis/transport/tests/test_doppler_factor.py index b3219c6e71b..e9c6442acf3 100644 --- a/tardis/transport/tests/test_doppler_factor.py +++ b/tardis/transport/tests/test_doppler_factor.py @@ -130,7 +130,7 @@ def test_get_inverse_doppler_factor(mu, r, inv_t_exp, expected): # to other methods or introduction of some temporary variables obtained = frame_transformations.get_inverse_doppler_factor( - r, mu, time_explosion, False + r, mu, time_explosion, enable_full_relativity=False ) # Perform required assertions From 3bd32141159496092edcf43e0ef7e48a497bb7f3 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 19:42:33 -0500 Subject: [PATCH 08/15] Further testing fixes --- tardis/model/base.py | 12 ++- tardis/model/parse_input.py | 22 ++++-- .../montecarlo_numba/formal_integral.py | 73 ++++++------------- .../montecarlo_numba/tests/conftest.py | 13 ++-- .../montecarlo_numba/tests/test_base.py | 10 +-- .../tests/test_interaction.py | 4 +- .../montecarlo_numba/tests/test_macro_atom.py | 2 +- .../tests/test_numba_interface.py | 25 +++---- .../montecarlo_numba/tests/test_vpacket.py | 4 + tardis/simulation/base.py | 6 +- 10 files changed, 81 insertions(+), 90 deletions(-) diff --git a/tardis/model/base.py b/tardis/model/base.py index b53f202fe52..cf4e2da7548 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -235,7 +235,7 @@ def no_of_raw_shells(self): return self.geometry.no_of_shells @classmethod - def from_config(cls, config, atom_data): + def from_config(cls, config, atom_data, legacy_mode_enabled=False): """ Create a new SimulationState instance from a Configuration object. @@ -269,7 +269,9 @@ def from_config(cls, config, atom_data): atom_data.atom_data.mass.copy(), ) - packet_source = parse_packet_source(config, geometry) + packet_source = parse_packet_source( + config, geometry, legacy_mode_enabled + ) radiation_field_state = parse_radiation_field_state( config, t_radiative, @@ -288,7 +290,7 @@ def from_config(cls, config, atom_data): ) @classmethod - def from_csvy(cls, config, atom_data=None): + def from_csvy(cls, config, atom_data=None, legacy_mode_enabled=False): """ Create a new SimulationState instance from a Configuration object. @@ -366,7 +368,9 @@ def from_csvy(cls, config, atom_data=None): geometry, ) - packet_source = parse_packet_source(config, geometry) + packet_source = parse_packet_source( + config, geometry, legacy_mode_enabled + ) radiation_field_state = parse_csvy_radiation_field_state( config, csvy_model_config, csvy_model_data, geometry, packet_source diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 9dd4aa8baf1..4b47a77d858 100644 --- a/tardis/model/parse_input.py +++ b/tardis/model/parse_input.py @@ -586,7 +586,9 @@ def parse_radiation_field_state( ) -def initialize_packet_source(config, geometry, packet_source): +def initialize_packet_source( + config, geometry, packet_source, legacy_mode_enabled +): """ Initialize the packet source based on config and geometry @@ -613,9 +615,13 @@ def initialize_packet_source(config, geometry, packet_source): packet_source = BlackBodySimpleSourceRelativistic( base_seed=config.montecarlo.seed, time_explosion=config.supernova.time_explosion, + legacy_mode_enabled=legacy_mode_enabled, ) else: - packet_source = BlackBodySimpleSource(base_seed=config.montecarlo.seed) + packet_source = BlackBodySimpleSource( + base_seed=config.montecarlo.seed, + legacy_mode_enabled=legacy_mode_enabled, + ) luminosity_requested = config.supernova.luminosity_requested if config.plasma.initial_t_inner > 0.0 * u.K: @@ -635,7 +641,7 @@ def initialize_packet_source(config, geometry, packet_source): return packet_source -def parse_packet_source(config, geometry): +def parse_packet_source(config, geometry, legacy_mode_enabled): """ Parse the packet source based on the given configuration and geometry. @@ -655,11 +661,17 @@ def parse_packet_source(config, geometry): packet_source = BlackBodySimpleSourceRelativistic( base_seed=config.montecarlo.seed, time_explosion=config.supernova.time_explosion, + legacy_mode_enabled=legacy_mode_enabled, ) else: - packet_source = BlackBodySimpleSource(base_seed=config.montecarlo.seed) + packet_source = BlackBodySimpleSource( + base_seed=config.montecarlo.seed, + legacy_mode_enabled=legacy_mode_enabled, + ) - return initialize_packet_source(config, geometry, packet_source) + return initialize_packet_source( + config, geometry, packet_source, legacy_mode_enabled + ) def parse_csvy_radiation_field_state( diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral.py b/tardis/montecarlo/montecarlo_numba/formal_integral.py index 162eeee97e3..3daab4b123e 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/formal_integral.py @@ -108,9 +108,7 @@ def numba_formal_integral( p = pp[p_idx] # initialize z intersections for p values - size_z = populate_z( - geometry, model, p, z, shell_id - ) # check returns + size_z = populate_z(geometry, model, p, z, shell_id) # check returns # initialize I_nu if p <= R_ph: I_nu[p_idx] = intensity_black_body(nu * z[0], iT) @@ -145,9 +143,7 @@ def numba_formal_integral( for _ in range(max(nu_end_idx - pline, 0)): # calculate e-scattering optical depth to next resonance point zend = ( - model.time_explosion - / C_INV - * (1.0 - line_list_nu[pline] / nu) + model.time_explosion / C_INV * (1.0 - line_list_nu[pline] / nu) ) # check if first == 1: @@ -187,12 +183,8 @@ def numba_formal_integral( # calculate e-scattering optical depth to grid cell boundary Jkkp = 0.5 * (Jred_lu[pJred_lu] + Jblue_lu[pJblue_lu]) - zend = ( - model.time_explosion / C_INV * (1.0 - nu_end / nu) - ) # check - escat_contrib += ( - (zend - zstart) * escat_op * (Jkkp - I_nu[p_idx]) - ) + zend = model.time_explosion / C_INV * (1.0 - nu_end / nu) # check + escat_contrib += (zend - zstart) * escat_op * (Jkkp - I_nu[p_idx]) zstart = zend # advance pointers @@ -281,7 +273,8 @@ def __init__(self, simulation_state, plasma, transport, points=1000): self.simulation_state = simulation_state self.transport = transport self.points = points - self.montecarlo_configuration = self.transport.montecarlo_configuration + if transport: + self.montecarlo_configuration = self.transport.montecarlo_configuration if plasma: self.plasma = opacity_state_initialize( plasma, @@ -310,7 +303,10 @@ def generate_numba_objects(self): self.simulation_state.time_explosion.cgs.value, ) self.opacity_state = opacity_state_initialize( - self.original_plasma, self.transport.line_interaction_type + self.original_plasma, + self.transport.line_interaction_type, + self.montecarlo_configuration.DISABLE_LINE_SCATTERING, + self.montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED, ) if self.transport.use_gpu: self.integrator = CudaFormalIntegrator( @@ -434,9 +430,7 @@ def make_source_function(self): if transport.line_interaction_type == "macroatom": internal_jump_mask = (macro_data.transition_type >= 0).values ma_int_data = macro_data[internal_jump_mask] - internal = self.original_plasma.transition_probabilities[ - internal_jump_mask - ] + internal = self.original_plasma.transition_probabilities[internal_jump_mask] source_level_idx = ma_int_data.source_level_idx.values destination_level_idx = ma_int_data.destination_level_idx.values @@ -475,17 +469,13 @@ def make_source_function(self): * Jbluelu_norm_factor ) - upper_level_index = self.atomic_data.lines.index.droplevel( - "level_number_lower" - ) + upper_level_index = self.atomic_data.lines.index.droplevel("level_number_lower") e_dot_lu = pd.DataFrame(Edotlu.values, index=upper_level_index) e_dot_u = e_dot_lu.groupby(level=[0, 1, 2]).sum() e_dot_u_src_idx = macro_ref.loc[e_dot_u.index].references_idx.values if transport.line_interaction_type == "macroatom": - C_frame = pd.DataFrame( - columns=np.arange(no_shells), index=macro_ref.index - ) + C_frame = pd.DataFrame(columns=np.arange(no_shells), index=macro_ref.index) q_indices = (source_level_idx, destination_level_idx) for shell in range(no_shells): Q = sp.coo_matrix( @@ -502,8 +492,7 @@ def make_source_function(self): "source_level_number", ] # To make the q_ul e_dot_u product work, could be cleaner transitions = self.original_plasma.atomic_data.macro_atom_data[ - self.original_plasma.atomic_data.macro_atom_data.transition_type - == -1 + self.original_plasma.atomic_data.macro_atom_data.transition_type == -1 ].copy() transitions_index = transitions.set_index( ["atomic_number", "ion_number", "source_level_number"] @@ -515,9 +504,9 @@ def make_source_function(self): t = simulation_state.time_explosion.value t = simulation_state.time_explosion.value lines = self.atomic_data.lines.set_index("line_id") - wave = lines.wavelength_cm.loc[ - transitions.transition_line_id - ].values.reshape(-1, 1) + wave = lines.wavelength_cm.loc[transitions.transition_line_id].values.reshape( + -1, 1 + ) if transport.line_interaction_type == "macroatom": e_dot_u = C_frame.loc[e_dot_u.index] att_S_ul = wave * (q_ul * e_dot_u) * t / (4 * np.pi) @@ -540,24 +529,16 @@ def make_source_function(self): att_S_ul, Jredlu, Jbluelu, e_dot_u ) else: - transport.r_inner_i = ( - montecarlo_transport_state.geometry_state.r_inner - ) - transport.r_outer_i = ( - montecarlo_transport_state.geometry_state.r_outer - ) - transport.tau_sobolevs_integ = ( - self.original_plasma.tau_sobolevs.values - ) + transport.r_inner_i = montecarlo_transport_state.geometry_state.r_inner + transport.r_outer_i = montecarlo_transport_state.geometry_state.r_outer + transport.tau_sobolevs_integ = self.original_plasma.tau_sobolevs.values transport.electron_densities_integ = ( self.original_plasma.electron_densities.values ) return att_S_ul, Jredlu, Jbluelu, e_dot_u - def interpolate_integrator_quantities( - self, att_S_ul, Jredlu, Jbluelu, e_dot_u - ): + def interpolate_integrator_quantities(self, att_S_ul, Jredlu, Jbluelu, e_dot_u): transport = self.transport mct_state = transport.transport_state plasma = self.original_plasma @@ -596,12 +577,8 @@ def interpolate_integrator_quantities( Jredlu = pd.DataFrame( interp1d(r_middle, Jredlu, fill_value="extrapolate")(r_middle_integ) ) - Jbluelu = interp1d(r_middle, Jbluelu, fill_value="extrapolate")( - r_middle_integ - ) - e_dot_u = interp1d(r_middle, e_dot_u, fill_value="extrapolate")( - r_middle_integ - ) + Jbluelu = interp1d(r_middle, Jbluelu, fill_value="extrapolate")(r_middle_integ) + e_dot_u = interp1d(r_middle, e_dot_u, fill_value="extrapolate")(r_middle_integ) # Set negative values from the extrapolation to zero att_S_ul = att_S_ul.clip(0.0) @@ -647,9 +624,7 @@ def formal_integral(self, nu, N): ] ) error = np.max(np.abs((L_test - L) / L)) - assert ( - error < 1e-7 - ), f"Incorrect I_nu_p values, max relative difference:{error}" + assert error < 1e-7, f"Incorrect I_nu_p values, max relative difference:{error}" return np.array(L, np.float64) diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index a991da4fcde..7241109d259 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -31,7 +31,10 @@ def nb_simulation_verysimple(config_verysimple, atomic_dataset): @pytest.fixture(scope="package") def verysimple_opacity_state(nb_simulation_verysimple): return opacity_state_initialize( - nb_simulation_verysimple.plasma, line_interaction_type="macroatom" + nb_simulation_verysimple.plasma, + line_interaction_type="macroatom", + disable_line_scattering=False, + continuum_processes_enabled=False, ) @@ -67,9 +70,7 @@ def verysimple_estimators(nb_simulation_verysimple): @pytest.fixture(scope="package") def verysimple_vpacket_collection(nb_simulation_verysimple): - spectrum_frequency = ( - nb_simulation_verysimple.transport.spectrum_frequency.value - ) + spectrum_frequency = nb_simulation_verysimple.transport.spectrum_frequency.value return VPacketCollection( source_rpacket_index=0, spectrum_frequency=spectrum_frequency, @@ -82,9 +83,7 @@ def verysimple_vpacket_collection(nb_simulation_verysimple): @pytest.fixture(scope="package") def verysimple_3vpacket_collection(nb_simulation_verysimple): - spectrum_frequency = ( - nb_simulation_verysimple.transport.spectrum_frequency.value - ) + spectrum_frequency = nb_simulation_verysimple.transport.spectrum_frequency.value return VPacketCollection( source_rpacket_index=0, spectrum_frequency=spectrum_frequency, diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index 98c79bebc7a..8bdd44ab16b 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -36,6 +36,7 @@ def test_montecarlo_main_loop( montecarlo_main_loop_config, atom_data=atomic_dataset, virtual_packet_logging=False, + legacy_mode_enabled=True, ) montecarlo_main_loop_simulation.run_convergence() montecarlo_main_loop_simulation.run_final() @@ -84,14 +85,14 @@ def test_montecarlo_main_loop_vpacket_log( montecarlo_main_loop_config, atom_data=atomic_dataset, virtual_packet_logging=True, + legacy_mode_enabled=True, ) montecarlo_main_loop_simulation.run_convergence() montecarlo_main_loop_simulation.run_final() - assert ( - montecarlo_main_loop_simulation.transport.montecarlo_configuration.ENABLE_VPACKET_TRACKING - == True - ) + transport = montecarlo_main_loop_simulation.transport + + assert transport.montecarlo_configuration.ENABLE_VPACKET_TRACKING is True expected_hdf_store = regression_data.sync_hdf_store( montecarlo_main_loop_simulation @@ -112,7 +113,6 @@ def test_montecarlo_main_loop_vpacket_log( "/simulation/transport/virt_packet_energies" ] - transport = montecarlo_main_loop_simulation.transport transport_state = transport.transport_state actual_energy = transport_state.packet_collection.output_energies diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index 47d67f945f8..2e71bcee07e 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -37,7 +37,7 @@ def test_line_scatter( init_mu = packet.mu init_nu = packet.nu init_energy = packet.energy - packet.initialize_line_id(verysimple_opacity_state, verysimple_numba_model) + packet.initialize_line_id(verysimple_opacity_state, verysimple_numba_model, False) time_explosion = verysimple_numba_model.time_explosion interaction.line_scatter( @@ -94,7 +94,7 @@ def test_line_emission( emission_line_id = test_packet["emission_line_id"] packet.mu = test_packet["mu"] packet.energy = test_packet["energy"] - packet.initialize_line_id(verysimple_opacity_state, verysimple_numba_model) + packet.initialize_line_id(verysimple_opacity_state, verysimple_numba_model, False) time_explosion = verysimple_numba_model.time_explosion diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py index e2ab2a016d0..36b0b4f72bd 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py @@ -17,7 +17,7 @@ def test_macro_atom( ): set_seed_fixture(seed) static_packet.initialize_line_id( - verysimple_opacity_state, verysimple_numba_model + verysimple_opacity_state, verysimple_numba_model, False ) activation_level_id = verysimple_opacity_state.line2macro_level_upper[ static_packet.next_line_id diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py index ef6492efdc1..fd8c86ce0a1 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py @@ -9,12 +9,13 @@ def test_opacity_state_initialize(nb_simulation_verysimple, input_params): line_interaction_type = input_params plasma = nb_simulation_verysimple.plasma actual = numba_interface.opacity_state_initialize( - plasma, line_interaction_type + plasma, + line_interaction_type, + disable_line_scattering=False, + continuum_processes_enabled=False, ) - npt.assert_allclose( - actual.electron_density, plasma.electron_densities.values - ) + npt.assert_allclose(actual.electron_density, plasma.electron_densities.values) npt.assert_allclose(actual.line_list_nu, plasma.atomic_data.lines.nu.values) npt.assert_allclose(actual.tau_sobolev, plasma.tau_sobolevs.values) if line_interaction_type == "scatter": @@ -66,9 +67,7 @@ def test_VPacketCollection_add_packet(verysimple_3vpacket_collection): energies = [0.4, 0.1, 0.6, 1e10] initial_mus = [0.1, 0, 1, 0.9] initial_rs = [3e42, 4.5e45, 0, 9.0e40] - last_interaction_in_nus = np.array( - [3.0e15, 0.0, 1e15, 1e5], dtype=np.float64 - ) + last_interaction_in_nus = np.array([3.0e15, 0.0, 1e15, 1e5], dtype=np.float64) last_interaction_types = np.array([1, 1, 3, 2], dtype=np.int64) last_interaction_in_ids = np.array([100, 0, 1, 1000], dtype=np.int64) last_interaction_out_ids = np.array([1201, 123, 545, 1232], dtype=np.int64) @@ -108,15 +107,11 @@ def test_VPacketCollection_add_packet(verysimple_3vpacket_collection): ) npt.assert_array_equal( - verysimple_3vpacket_collection.nus[ - : verysimple_3vpacket_collection.idx - ], + verysimple_3vpacket_collection.nus[: verysimple_3vpacket_collection.idx], nus, ) npt.assert_array_equal( - verysimple_3vpacket_collection.energies[ - : verysimple_3vpacket_collection.idx - ], + verysimple_3vpacket_collection.energies[: verysimple_3vpacket_collection.idx], energies, ) npt.assert_array_equal( @@ -126,9 +121,7 @@ def test_VPacketCollection_add_packet(verysimple_3vpacket_collection): initial_mus, ) npt.assert_array_equal( - verysimple_3vpacket_collection.initial_rs[ - : verysimple_3vpacket_collection.idx - ], + verysimple_3vpacket_collection.initial_rs[: verysimple_3vpacket_collection.idx], initial_rs, ) npt.assert_array_equal( diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py index dd27baf0654..0968e5c0cea 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py @@ -90,6 +90,8 @@ def test_trace_vpacket( verysimple_opacity_state, 10.0, 0.0, + False, + False, ) npt.assert_almost_equal(tau_trace_combined, 8164850.891288479) @@ -156,4 +158,6 @@ def test_trace_bad_vpacket( verysimple_opacity_state, 10.0, 0.0, + False, + False, ) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index dc404c1b354..2277adb4580 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -639,6 +639,7 @@ def from_config( virtual_packet_logging=False, show_convergence_plots=False, show_progress_bars=True, + legacy_mode_enabled=False, **kwargs, ): """ @@ -700,7 +701,10 @@ def from_config( ) if packet_source is not None: simulation_state.packet_source = initialize_packet_source( - config, simulation_state.geometry, packet_source + config, + simulation_state.geometry, + packet_source, + legacy_mode_enabled, ) if "plasma" in kwargs: plasma = kwargs["plasma"] From 66f6cf7e59cf546f3252d937a737c0da2d9a390c Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 19:55:48 -0500 Subject: [PATCH 09/15] All tests pass --- tardis/simulation/base.py | 64 +++++++++++---------------------------- 1 file changed, 17 insertions(+), 47 deletions(-) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 2277adb4580..4104734c89b 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -37,9 +37,7 @@ class PlasmaStateStorerMixin: def __init__(self, iterations, no_of_shells): self.iterations_w = np.zeros((iterations, no_of_shells)) self.iterations_t_rad = np.zeros((iterations, no_of_shells)) * u.K - self.iterations_electron_densities = np.zeros( - (iterations, no_of_shells) - ) + self.iterations_electron_densities = np.zeros((iterations, no_of_shells)) self.iterations_t_inner = np.zeros(iterations) * u.K def store_plasma_state(self, i, w, t_rad, electron_densities, t_inner): @@ -74,15 +72,11 @@ def reshape_plasma_state_store(self, executed_iterations): iteration index, i.e. number of iterations executed minus one! """ self.iterations_w = self.iterations_w[: executed_iterations + 1, :] - self.iterations_t_rad = self.iterations_t_rad[ - : executed_iterations + 1, : - ] + self.iterations_t_rad = self.iterations_t_rad[: executed_iterations + 1, :] self.iterations_electron_densities = self.iterations_electron_densities[ : executed_iterations + 1, : ] - self.iterations_t_inner = self.iterations_t_inner[ - : executed_iterations + 1 - ] + self.iterations_t_inner = self.iterations_t_inner[: executed_iterations + 1] class Simulation(PlasmaStateStorerMixin, HDFWriterMixin): @@ -133,9 +127,7 @@ def __init__( convergence_plots_kwargs, show_progress_bars, ): - super(Simulation, self).__init__( - iterations, simulation_state.no_of_shells - ) + super(Simulation, self).__init__(iterations, simulation_state.no_of_shells) self.converged = False self.iterations = iterations @@ -209,9 +201,7 @@ def estimate_t_inner( ) ) - luminosity_ratios = ( - (emitted_luminosity / luminosity_requested).to(1).value - ) + luminosity_ratios = (emitted_luminosity / luminosity_requested).to(1).value return input_t_inner * luminosity_ratios**t_inner_update_exponent @@ -227,9 +217,7 @@ def _get_convergence_status( # FIXME: Move the convergence checking in its own class. no_of_shells = self.simulation_state.no_of_shells - convergence_t_rad = ( - abs(t_rad - estimated_t_rad) / estimated_t_rad - ).value + convergence_t_rad = (abs(t_rad - estimated_t_rad) / estimated_t_rad).value convergence_w = abs(w - estimated_w) / estimated_w convergence_t_inner = ( abs(t_inner - estimated_t_inner) / estimated_t_inner @@ -242,14 +230,10 @@ def _get_convergence_status( / no_of_shells ) - t_rad_converged = ( - fraction_t_rad_converged > self.convergence_strategy.fraction - ) + t_rad_converged = fraction_t_rad_converged > self.convergence_strategy.fraction fraction_w_converged = ( - np.count_nonzero( - convergence_w < self.convergence_strategy.w.threshold - ) + np.count_nonzero(convergence_w < self.convergence_strategy.w.threshold) / no_of_shells ) @@ -412,9 +396,7 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): show_progress_bars=self.show_progress_bars, ) - output_energy = ( - self.transport.transport_state.packet_collection.output_energies - ) + output_energy = self.transport.transport_state.packet_collection.output_energies if np.sum(output_energy < 0) == len(output_energy): logger.critical("No r-packet escaped through the outer boundary.") @@ -548,15 +530,11 @@ def log_plasma_state( if logger.level <= logging.INFO: if not logger.filters: display( - plasma_state_log.iloc[::log_sampling].style.format( - "{:.3g}" - ) + plasma_state_log.iloc[::log_sampling].style.format("{:.3g}") ) elif logger.filters[0].log_level == 20: display( - plasma_state_log.iloc[::log_sampling].style.format( - "{:.3g}" - ) + plasma_state_log.iloc[::log_sampling].style.format("{:.3g}") ) else: output_df = "" @@ -667,14 +645,10 @@ def from_config( if Path(config.atom_data).is_absolute(): atom_data_fname = Path(config.atom_data) else: - atom_data_fname = ( - Path(config.config_dirname) / config.atom_data - ) + atom_data_fname = Path(config.config_dirname) / config.atom_data else: - raise ValueError( - "No atom_data option found in the configuration." - ) + raise ValueError("No atom_data option found in the configuration.") logger.info(f"\n\tReading Atomic Data from {atom_data_fname}") @@ -693,11 +667,11 @@ def from_config( else: if hasattr(config, "csvy_model"): simulation_state = SimulationState.from_csvy( - config, atom_data=atom_data + config, atom_data=atom_data, legacy_mode_enabled=legacy_mode_enabled ) else: simulation_state = SimulationState.from_config( - config, atom_data=atom_data + config, atom_data=atom_data, legacy_mode_enabled=legacy_mode_enabled ) if packet_source is not None: simulation_state.packet_source = initialize_packet_source( @@ -735,18 +709,14 @@ def from_config( "export_convergence_plots", ] convergence_plots_kwargs = {} - for item in set(convergence_plots_config_options).intersection( - kwargs.keys() - ): + for item in set(convergence_plots_config_options).intersection(kwargs.keys()): convergence_plots_kwargs[item] = kwargs[item] luminosity_nu_start = config.supernova.luminosity_wavelength_end.to( u.Hz, u.spectral() ) - if u.isclose( - config.supernova.luminosity_wavelength_start, 0 * u.angstrom - ): + if u.isclose(config.supernova.luminosity_wavelength_start, 0 * u.angstrom): luminosity_nu_end = np.inf * u.Hz else: luminosity_nu_end = ( From f1cd7ddf52a952d8d222e8440630c2b4368dedd2 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 28 Feb 2024 20:08:56 -0500 Subject: [PATCH 10/15] All tests passing, black applied --- .../tests/test_continuum_property_solver.py | 1 - .../montecarlo_numba/formal_integral.py | 69 ++++++++++++++----- .../montecarlo_numba/tests/conftest.py | 8 ++- .../montecarlo_numba/tests/test_continuum.py | 1 - .../tests/test_interaction.py | 8 ++- .../tests/test_numba_interface.py | 20 ++++-- tardis/simulation/base.py | 68 +++++++++++++----- 7 files changed, 128 insertions(+), 47 deletions(-) diff --git a/tardis/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/montecarlo/estimators/tests/test_continuum_property_solver.py index 7ef65daa481..8b0ea4d77bd 100644 --- a/tardis/montecarlo/estimators/tests/test_continuum_property_solver.py +++ b/tardis/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -11,7 +11,6 @@ from tardis.simulation import Simulation -@pytest.mark.skip() def test_continuum_estimators( continuum_config, nlte_atomic_dataset, diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral.py b/tardis/montecarlo/montecarlo_numba/formal_integral.py index 3daab4b123e..063a9266052 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/formal_integral.py @@ -108,7 +108,9 @@ def numba_formal_integral( p = pp[p_idx] # initialize z intersections for p values - size_z = populate_z(geometry, model, p, z, shell_id) # check returns + size_z = populate_z( + geometry, model, p, z, shell_id + ) # check returns # initialize I_nu if p <= R_ph: I_nu[p_idx] = intensity_black_body(nu * z[0], iT) @@ -143,7 +145,9 @@ def numba_formal_integral( for _ in range(max(nu_end_idx - pline, 0)): # calculate e-scattering optical depth to next resonance point zend = ( - model.time_explosion / C_INV * (1.0 - line_list_nu[pline] / nu) + model.time_explosion + / C_INV + * (1.0 - line_list_nu[pline] / nu) ) # check if first == 1: @@ -183,8 +187,12 @@ def numba_formal_integral( # calculate e-scattering optical depth to grid cell boundary Jkkp = 0.5 * (Jred_lu[pJred_lu] + Jblue_lu[pJblue_lu]) - zend = model.time_explosion / C_INV * (1.0 - nu_end / nu) # check - escat_contrib += (zend - zstart) * escat_op * (Jkkp - I_nu[p_idx]) + zend = ( + model.time_explosion / C_INV * (1.0 - nu_end / nu) + ) # check + escat_contrib += ( + (zend - zstart) * escat_op * (Jkkp - I_nu[p_idx]) + ) zstart = zend # advance pointers @@ -274,7 +282,9 @@ def __init__(self, simulation_state, plasma, transport, points=1000): self.transport = transport self.points = points if transport: - self.montecarlo_configuration = self.transport.montecarlo_configuration + self.montecarlo_configuration = ( + self.transport.montecarlo_configuration + ) if plasma: self.plasma = opacity_state_initialize( plasma, @@ -430,7 +440,9 @@ def make_source_function(self): if transport.line_interaction_type == "macroatom": internal_jump_mask = (macro_data.transition_type >= 0).values ma_int_data = macro_data[internal_jump_mask] - internal = self.original_plasma.transition_probabilities[internal_jump_mask] + internal = self.original_plasma.transition_probabilities[ + internal_jump_mask + ] source_level_idx = ma_int_data.source_level_idx.values destination_level_idx = ma_int_data.destination_level_idx.values @@ -469,13 +481,17 @@ def make_source_function(self): * Jbluelu_norm_factor ) - upper_level_index = self.atomic_data.lines.index.droplevel("level_number_lower") + upper_level_index = self.atomic_data.lines.index.droplevel( + "level_number_lower" + ) e_dot_lu = pd.DataFrame(Edotlu.values, index=upper_level_index) e_dot_u = e_dot_lu.groupby(level=[0, 1, 2]).sum() e_dot_u_src_idx = macro_ref.loc[e_dot_u.index].references_idx.values if transport.line_interaction_type == "macroatom": - C_frame = pd.DataFrame(columns=np.arange(no_shells), index=macro_ref.index) + C_frame = pd.DataFrame( + columns=np.arange(no_shells), index=macro_ref.index + ) q_indices = (source_level_idx, destination_level_idx) for shell in range(no_shells): Q = sp.coo_matrix( @@ -492,7 +508,8 @@ def make_source_function(self): "source_level_number", ] # To make the q_ul e_dot_u product work, could be cleaner transitions = self.original_plasma.atomic_data.macro_atom_data[ - self.original_plasma.atomic_data.macro_atom_data.transition_type == -1 + self.original_plasma.atomic_data.macro_atom_data.transition_type + == -1 ].copy() transitions_index = transitions.set_index( ["atomic_number", "ion_number", "source_level_number"] @@ -504,9 +521,9 @@ def make_source_function(self): t = simulation_state.time_explosion.value t = simulation_state.time_explosion.value lines = self.atomic_data.lines.set_index("line_id") - wave = lines.wavelength_cm.loc[transitions.transition_line_id].values.reshape( - -1, 1 - ) + wave = lines.wavelength_cm.loc[ + transitions.transition_line_id + ].values.reshape(-1, 1) if transport.line_interaction_type == "macroatom": e_dot_u = C_frame.loc[e_dot_u.index] att_S_ul = wave * (q_ul * e_dot_u) * t / (4 * np.pi) @@ -529,16 +546,24 @@ def make_source_function(self): att_S_ul, Jredlu, Jbluelu, e_dot_u ) else: - transport.r_inner_i = montecarlo_transport_state.geometry_state.r_inner - transport.r_outer_i = montecarlo_transport_state.geometry_state.r_outer - transport.tau_sobolevs_integ = self.original_plasma.tau_sobolevs.values + transport.r_inner_i = ( + montecarlo_transport_state.geometry_state.r_inner + ) + transport.r_outer_i = ( + montecarlo_transport_state.geometry_state.r_outer + ) + transport.tau_sobolevs_integ = ( + self.original_plasma.tau_sobolevs.values + ) transport.electron_densities_integ = ( self.original_plasma.electron_densities.values ) return att_S_ul, Jredlu, Jbluelu, e_dot_u - def interpolate_integrator_quantities(self, att_S_ul, Jredlu, Jbluelu, e_dot_u): + def interpolate_integrator_quantities( + self, att_S_ul, Jredlu, Jbluelu, e_dot_u + ): transport = self.transport mct_state = transport.transport_state plasma = self.original_plasma @@ -577,8 +602,12 @@ def interpolate_integrator_quantities(self, att_S_ul, Jredlu, Jbluelu, e_dot_u): Jredlu = pd.DataFrame( interp1d(r_middle, Jredlu, fill_value="extrapolate")(r_middle_integ) ) - Jbluelu = interp1d(r_middle, Jbluelu, fill_value="extrapolate")(r_middle_integ) - e_dot_u = interp1d(r_middle, e_dot_u, fill_value="extrapolate")(r_middle_integ) + Jbluelu = interp1d(r_middle, Jbluelu, fill_value="extrapolate")( + r_middle_integ + ) + e_dot_u = interp1d(r_middle, e_dot_u, fill_value="extrapolate")( + r_middle_integ + ) # Set negative values from the extrapolation to zero att_S_ul = att_S_ul.clip(0.0) @@ -624,7 +653,9 @@ def formal_integral(self, nu, N): ] ) error = np.max(np.abs((L_test - L) / L)) - assert error < 1e-7, f"Incorrect I_nu_p values, max relative difference:{error}" + assert ( + error < 1e-7 + ), f"Incorrect I_nu_p values, max relative difference:{error}" return np.array(L, np.float64) diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index 7241109d259..5f7a5aece92 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -70,7 +70,9 @@ def verysimple_estimators(nb_simulation_verysimple): @pytest.fixture(scope="package") def verysimple_vpacket_collection(nb_simulation_verysimple): - spectrum_frequency = nb_simulation_verysimple.transport.spectrum_frequency.value + spectrum_frequency = ( + nb_simulation_verysimple.transport.spectrum_frequency.value + ) return VPacketCollection( source_rpacket_index=0, spectrum_frequency=spectrum_frequency, @@ -83,7 +85,9 @@ def verysimple_vpacket_collection(nb_simulation_verysimple): @pytest.fixture(scope="package") def verysimple_3vpacket_collection(nb_simulation_verysimple): - spectrum_frequency = nb_simulation_verysimple.transport.spectrum_frequency.value + spectrum_frequency = ( + nb_simulation_verysimple.transport.spectrum_frequency.value + ) return VPacketCollection( source_rpacket_index=0, spectrum_frequency=spectrum_frequency, diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_continuum.py b/tardis/montecarlo/montecarlo_numba/tests/test_continuum.py index 5d0327081cf..d11b394b2cf 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_continuum.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_continuum.py @@ -6,7 +6,6 @@ from tardis.simulation import Simulation -@pytest.mark.skip() def test_montecarlo_continuum( continuum_config, regression_data, diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index 2e71bcee07e..3b4e7a9512d 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -37,7 +37,9 @@ def test_line_scatter( init_mu = packet.mu init_nu = packet.nu init_energy = packet.energy - packet.initialize_line_id(verysimple_opacity_state, verysimple_numba_model, False) + packet.initialize_line_id( + verysimple_opacity_state, verysimple_numba_model, False + ) time_explosion = verysimple_numba_model.time_explosion interaction.line_scatter( @@ -94,7 +96,9 @@ def test_line_emission( emission_line_id = test_packet["emission_line_id"] packet.mu = test_packet["mu"] packet.energy = test_packet["energy"] - packet.initialize_line_id(verysimple_opacity_state, verysimple_numba_model, False) + packet.initialize_line_id( + verysimple_opacity_state, verysimple_numba_model, False + ) time_explosion = verysimple_numba_model.time_explosion diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py index fd8c86ce0a1..2e74bef78ed 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py @@ -15,7 +15,9 @@ def test_opacity_state_initialize(nb_simulation_verysimple, input_params): continuum_processes_enabled=False, ) - npt.assert_allclose(actual.electron_density, plasma.electron_densities.values) + npt.assert_allclose( + actual.electron_density, plasma.electron_densities.values + ) npt.assert_allclose(actual.line_list_nu, plasma.atomic_data.lines.nu.values) npt.assert_allclose(actual.tau_sobolev, plasma.tau_sobolevs.values) if line_interaction_type == "scatter": @@ -67,7 +69,9 @@ def test_VPacketCollection_add_packet(verysimple_3vpacket_collection): energies = [0.4, 0.1, 0.6, 1e10] initial_mus = [0.1, 0, 1, 0.9] initial_rs = [3e42, 4.5e45, 0, 9.0e40] - last_interaction_in_nus = np.array([3.0e15, 0.0, 1e15, 1e5], dtype=np.float64) + last_interaction_in_nus = np.array( + [3.0e15, 0.0, 1e15, 1e5], dtype=np.float64 + ) last_interaction_types = np.array([1, 1, 3, 2], dtype=np.int64) last_interaction_in_ids = np.array([100, 0, 1, 1000], dtype=np.int64) last_interaction_out_ids = np.array([1201, 123, 545, 1232], dtype=np.int64) @@ -107,11 +111,15 @@ def test_VPacketCollection_add_packet(verysimple_3vpacket_collection): ) npt.assert_array_equal( - verysimple_3vpacket_collection.nus[: verysimple_3vpacket_collection.idx], + verysimple_3vpacket_collection.nus[ + : verysimple_3vpacket_collection.idx + ], nus, ) npt.assert_array_equal( - verysimple_3vpacket_collection.energies[: verysimple_3vpacket_collection.idx], + verysimple_3vpacket_collection.energies[ + : verysimple_3vpacket_collection.idx + ], energies, ) npt.assert_array_equal( @@ -121,7 +129,9 @@ def test_VPacketCollection_add_packet(verysimple_3vpacket_collection): initial_mus, ) npt.assert_array_equal( - verysimple_3vpacket_collection.initial_rs[: verysimple_3vpacket_collection.idx], + verysimple_3vpacket_collection.initial_rs[ + : verysimple_3vpacket_collection.idx + ], initial_rs, ) npt.assert_array_equal( diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 4104734c89b..e748f8a8f75 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -37,7 +37,9 @@ class PlasmaStateStorerMixin: def __init__(self, iterations, no_of_shells): self.iterations_w = np.zeros((iterations, no_of_shells)) self.iterations_t_rad = np.zeros((iterations, no_of_shells)) * u.K - self.iterations_electron_densities = np.zeros((iterations, no_of_shells)) + self.iterations_electron_densities = np.zeros( + (iterations, no_of_shells) + ) self.iterations_t_inner = np.zeros(iterations) * u.K def store_plasma_state(self, i, w, t_rad, electron_densities, t_inner): @@ -72,11 +74,15 @@ def reshape_plasma_state_store(self, executed_iterations): iteration index, i.e. number of iterations executed minus one! """ self.iterations_w = self.iterations_w[: executed_iterations + 1, :] - self.iterations_t_rad = self.iterations_t_rad[: executed_iterations + 1, :] + self.iterations_t_rad = self.iterations_t_rad[ + : executed_iterations + 1, : + ] self.iterations_electron_densities = self.iterations_electron_densities[ : executed_iterations + 1, : ] - self.iterations_t_inner = self.iterations_t_inner[: executed_iterations + 1] + self.iterations_t_inner = self.iterations_t_inner[ + : executed_iterations + 1 + ] class Simulation(PlasmaStateStorerMixin, HDFWriterMixin): @@ -127,7 +133,9 @@ def __init__( convergence_plots_kwargs, show_progress_bars, ): - super(Simulation, self).__init__(iterations, simulation_state.no_of_shells) + super(Simulation, self).__init__( + iterations, simulation_state.no_of_shells + ) self.converged = False self.iterations = iterations @@ -201,7 +209,9 @@ def estimate_t_inner( ) ) - luminosity_ratios = (emitted_luminosity / luminosity_requested).to(1).value + luminosity_ratios = ( + (emitted_luminosity / luminosity_requested).to(1).value + ) return input_t_inner * luminosity_ratios**t_inner_update_exponent @@ -217,7 +227,9 @@ def _get_convergence_status( # FIXME: Move the convergence checking in its own class. no_of_shells = self.simulation_state.no_of_shells - convergence_t_rad = (abs(t_rad - estimated_t_rad) / estimated_t_rad).value + convergence_t_rad = ( + abs(t_rad - estimated_t_rad) / estimated_t_rad + ).value convergence_w = abs(w - estimated_w) / estimated_w convergence_t_inner = ( abs(t_inner - estimated_t_inner) / estimated_t_inner @@ -230,10 +242,14 @@ def _get_convergence_status( / no_of_shells ) - t_rad_converged = fraction_t_rad_converged > self.convergence_strategy.fraction + t_rad_converged = ( + fraction_t_rad_converged > self.convergence_strategy.fraction + ) fraction_w_converged = ( - np.count_nonzero(convergence_w < self.convergence_strategy.w.threshold) + np.count_nonzero( + convergence_w < self.convergence_strategy.w.threshold + ) / no_of_shells ) @@ -396,7 +412,9 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): show_progress_bars=self.show_progress_bars, ) - output_energy = self.transport.transport_state.packet_collection.output_energies + output_energy = ( + self.transport.transport_state.packet_collection.output_energies + ) if np.sum(output_energy < 0) == len(output_energy): logger.critical("No r-packet escaped through the outer boundary.") @@ -530,11 +548,15 @@ def log_plasma_state( if logger.level <= logging.INFO: if not logger.filters: display( - plasma_state_log.iloc[::log_sampling].style.format("{:.3g}") + plasma_state_log.iloc[::log_sampling].style.format( + "{:.3g}" + ) ) elif logger.filters[0].log_level == 20: display( - plasma_state_log.iloc[::log_sampling].style.format("{:.3g}") + plasma_state_log.iloc[::log_sampling].style.format( + "{:.3g}" + ) ) else: output_df = "" @@ -645,10 +667,14 @@ def from_config( if Path(config.atom_data).is_absolute(): atom_data_fname = Path(config.atom_data) else: - atom_data_fname = Path(config.config_dirname) / config.atom_data + atom_data_fname = ( + Path(config.config_dirname) / config.atom_data + ) else: - raise ValueError("No atom_data option found in the configuration.") + raise ValueError( + "No atom_data option found in the configuration." + ) logger.info(f"\n\tReading Atomic Data from {atom_data_fname}") @@ -667,11 +693,15 @@ def from_config( else: if hasattr(config, "csvy_model"): simulation_state = SimulationState.from_csvy( - config, atom_data=atom_data, legacy_mode_enabled=legacy_mode_enabled + config, + atom_data=atom_data, + legacy_mode_enabled=legacy_mode_enabled, ) else: simulation_state = SimulationState.from_config( - config, atom_data=atom_data, legacy_mode_enabled=legacy_mode_enabled + config, + atom_data=atom_data, + legacy_mode_enabled=legacy_mode_enabled, ) if packet_source is not None: simulation_state.packet_source = initialize_packet_source( @@ -709,14 +739,18 @@ def from_config( "export_convergence_plots", ] convergence_plots_kwargs = {} - for item in set(convergence_plots_config_options).intersection(kwargs.keys()): + for item in set(convergence_plots_config_options).intersection( + kwargs.keys() + ): convergence_plots_kwargs[item] = kwargs[item] luminosity_nu_start = config.supernova.luminosity_wavelength_end.to( u.Hz, u.spectral() ) - if u.isclose(config.supernova.luminosity_wavelength_start, 0 * u.angstrom): + if u.isclose( + config.supernova.luminosity_wavelength_start, 0 * u.angstrom + ): luminosity_nu_end = np.inf * u.Hz else: luminosity_nu_end = ( From 18d7d2df7ae45eaadaa8d52be6e9e536d0cbea5f Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 10 Apr 2024 14:54:11 -0400 Subject: [PATCH 11/15] Apply black --- tardis/montecarlo/packet_source.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index 36f486970cf..de2449415b8 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -81,14 +81,20 @@ def create_packets(self, no_of_packets, seed_offset=0, *args, **kwargs): # because we call random.sample, which references a different internal # state than in the numpy.random module. self._reseed(self.base_seed + seed_offset) - packet_seeds = self.rng.choice(self.MAX_SEED_VAL, no_of_packets, replace=True) + packet_seeds = self.rng.choice( + self.MAX_SEED_VAL, no_of_packets, replace=True + ) radii = self.create_packet_radii(no_of_packets, *args, **kwargs).value nus = self.create_packet_nus(no_of_packets, *args, **kwargs).value mus = self.create_packet_mus(no_of_packets, *args, **kwargs) - energies = self.create_packet_energies(no_of_packets, *args, **kwargs).value + energies = self.create_packet_energies( + no_of_packets, *args, **kwargs + ).value # Check if all arrays have the same length - assert len(radii) == len(nus) == len(mus) == len(energies) == no_of_packets + assert ( + len(radii) == len(nus) == len(mus) == len(energies) == no_of_packets + ) radiation_field_luminosity = self.calculate_radfield_luminosity().value return PacketCollection( radii, @@ -111,9 +117,13 @@ def calculate_radfield_luminosity(self): ------- astropy.units.Quantity """ - return (4 * np.pi * const.sigma_sb * self.radius**2 * self.temperature**4).to( - "erg/s" - ) + return ( + 4 + * np.pi + * const.sigma_sb + * self.radius**2 + * self.temperature**4 + ).to("erg/s") class BlackBodySimpleSource(BasePacketSource): @@ -261,7 +271,8 @@ def set_temperature_from_luminosity(self, luminosity: u.Quantity): """ self.temperature = ( - (luminosity / (4 * np.pi * self.radius**2 * const.sigma_sb)) ** 0.25 + (luminosity / (4 * np.pi * self.radius**2 * const.sigma_sb)) + ** 0.25 ).to("K") From 9e0eb2bc642cd999eee52e0d5071de5f892b1144 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 17 Apr 2024 12:44:32 -0400 Subject: [PATCH 12/15] Resolves comments --- tardis/montecarlo/estimators/radfield_mc_estimators.py | 2 +- .../montecarlo_numba/tests/test_interaction.py | 9 +++++---- .../montecarlo/montecarlo_numba/tests/test_macro_atom.py | 4 +++- tardis/montecarlo/montecarlo_numba/tests/test_packet.py | 2 +- tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py | 8 ++++---- 5 files changed, 14 insertions(+), 11 deletions(-) diff --git a/tardis/montecarlo/estimators/radfield_mc_estimators.py b/tardis/montecarlo/estimators/radfield_mc_estimators.py index 5e5b868a986..de91643e5a5 100644 --- a/tardis/montecarlo/estimators/radfield_mc_estimators.py +++ b/tardis/montecarlo/estimators/radfield_mc_estimators.py @@ -120,7 +120,7 @@ def increment(self, other): other.photo_ion_estimator_statistics ) - def create_list(self, number): + def create_estimator_list(self, number): estimator_list = List() for i in range(number): diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index 3b4e7a9512d..a8559aa723f 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -47,8 +47,7 @@ def test_line_scatter( time_explosion, line_interaction_type, verysimple_opacity_state, - False, - False, + enable_full_relativity=False, ) assert np.abs(packet.mu - init_mu) > 1e-7 @@ -97,7 +96,9 @@ def test_line_emission( packet.mu = test_packet["mu"] packet.energy = test_packet["energy"] packet.initialize_line_id( - verysimple_opacity_state, verysimple_numba_model, False + verysimple_opacity_state, + verysimple_numba_model, + enable_full_relativity=False, ) time_explosion = verysimple_numba_model.time_explosion @@ -107,7 +108,7 @@ def test_line_emission( emission_line_id, time_explosion, verysimple_opacity_state, - False, + enable_full_relativity=False, ) assert packet.next_line_id == emission_line_id + 1 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py index 36b0b4f72bd..e46a6aae13c 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py @@ -17,7 +17,9 @@ def test_macro_atom( ): set_seed_fixture(seed) static_packet.initialize_line_id( - verysimple_opacity_state, verysimple_numba_model, False + verysimple_opacity_state, + verysimple_numba_model, + enable_full_relativity=False, ) activation_level_id = verysimple_opacity_state.line2macro_level_upper[ static_packet.next_line_id diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py index e05a5ef60fa..6822ca13173 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -127,7 +127,7 @@ def test_calculate_distance_line( is_last_line, nu_line, time_explosion, - False, + enable_full_relativity=False, ) except utils.MonteCarloException: obtained_tardis_error = utils.MonteCarloException diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py index 0968e5c0cea..294c6c4233e 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py @@ -90,8 +90,8 @@ def test_trace_vpacket( verysimple_opacity_state, 10.0, 0.0, - False, - False, + enable_full_relativity=False, + continuum_processes_enabled=False, ) npt.assert_almost_equal(tau_trace_combined, 8164850.891288479) @@ -158,6 +158,6 @@ def test_trace_bad_vpacket( verysimple_opacity_state, 10.0, 0.0, - False, - False, + enable_full_relativity=False, + continuum_processes_enabled=False, ) From b17adede7e1b4f749166f66ae2d6ff5ecb408a6b Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 17 Apr 2024 13:12:09 -0400 Subject: [PATCH 13/15] Fix reference to function --- tardis/montecarlo/montecarlo_numba/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index 2d8e18b61ac..14f671b3288 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -101,7 +101,7 @@ def montecarlo_main_loop( # betting get thread_id goes from 0 to num threads # Note that get_thread_id() returns values from 0 to n_threads-1, # so we iterate from 0 to n_threads-1 to create the estimator_list - estimator_list = estimators.create_list(n_threads) + estimator_list = estimators.create_estimator_list(n_threads) for i in prange(no_of_packets): thread_id = get_thread_id() From 24ed0583f47b2dc588ef8635b4eaaa2c07751a83 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 17 Apr 2024 15:03:05 -0400 Subject: [PATCH 14/15] Fix some parts numba was upset by --- .../montecarlo_numba/tests/test_interaction.py | 9 ++++++--- .../montecarlo/montecarlo_numba/tests/test_macro_atom.py | 3 ++- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index a8559aa723f..35ee1f8896d 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -37,8 +37,9 @@ def test_line_scatter( init_mu = packet.mu init_nu = packet.nu init_energy = packet.energy + full_relativity=False packet.initialize_line_id( - verysimple_opacity_state, verysimple_numba_model, False + verysimple_opacity_state, verysimple_numba_model, full_relativity ) time_explosion = verysimple_numba_model.time_explosion @@ -47,6 +48,7 @@ def test_line_scatter( time_explosion, line_interaction_type, verysimple_opacity_state, + continuum_processes_enabled=False, enable_full_relativity=False, ) @@ -95,10 +97,11 @@ def test_line_emission( emission_line_id = test_packet["emission_line_id"] packet.mu = test_packet["mu"] packet.energy = test_packet["energy"] + full_relativity = False packet.initialize_line_id( verysimple_opacity_state, verysimple_numba_model, - enable_full_relativity=False, + full_relativity, ) time_explosion = verysimple_numba_model.time_explosion @@ -108,7 +111,7 @@ def test_line_emission( emission_line_id, time_explosion, verysimple_opacity_state, - enable_full_relativity=False, + full_relativity ) assert packet.next_line_id == emission_line_id + 1 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py index e46a6aae13c..fef47651bd6 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py @@ -16,10 +16,11 @@ def test_macro_atom( expected, ): set_seed_fixture(seed) + full_relativity=False static_packet.initialize_line_id( verysimple_opacity_state, verysimple_numba_model, - enable_full_relativity=False, + full_relativity, ) activation_level_id = verysimple_opacity_state.line2macro_level_upper[ static_packet.next_line_id From d2f1c58fab6793c10d310436730f8b6d9c28137e Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 17 Apr 2024 16:33:55 -0400 Subject: [PATCH 15/15] Fix formatting --- tardis/montecarlo/montecarlo_numba/tests/test_interaction.py | 4 ++-- tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index 35ee1f8896d..2b31b9e2eb9 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -37,7 +37,7 @@ def test_line_scatter( init_mu = packet.mu init_nu = packet.nu init_energy = packet.energy - full_relativity=False + full_relativity = False packet.initialize_line_id( verysimple_opacity_state, verysimple_numba_model, full_relativity ) @@ -111,7 +111,7 @@ def test_line_emission( emission_line_id, time_explosion, verysimple_opacity_state, - full_relativity + full_relativity, ) assert packet.next_line_id == emission_line_id + 1 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py index fef47651bd6..c63a45c8007 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py @@ -16,7 +16,7 @@ def test_macro_atom( expected, ): set_seed_fixture(seed) - full_relativity=False + full_relativity = False static_packet.initialize_line_id( verysimple_opacity_state, verysimple_numba_model,