import logging
from typing import Self
from numba import cuda, set_num_threads
import tardis.transport.montecarlo.configuration.constants as constants
from tardis import constants as const
from tardis.io.configuration.config_reader import Configuration
from tardis.io.hdf_writer_mixin import HDFWriterMixin
from tardis.io.logger import montecarlo_tracking as mc_tracker
from tardis.transport.montecarlo.configuration import montecarlo_globals
from tardis.transport.montecarlo.configuration.base import (
MonteCarloConfiguration,
configuration_initialize,
)
from tardis.transport.montecarlo.estimators.mc_rad_field_solver import (
MCRadiationFieldPropertiesSolver,
)
from tardis.transport.montecarlo.modes.classic.packet_propagation import (
packet_propagation,
)
from tardis.transport.montecarlo.modes.montecarlo_transport import (
montecarlo_transport,
)
from tardis.transport.montecarlo.montecarlo_transport_state import (
MonteCarloTransportState,
)
from tardis.transport.montecarlo.packet_source.base import BasePacketSource
from tardis.transport.montecarlo.packets.trackers.tracker_full_util import (
generate_tracker_full_list,
tracker_full_df2tracker_last_interaction_df,
trackers_full_to_df,
)
from tardis.transport.montecarlo.packets.trackers.tracker_last_interaction_util import (
generate_tracker_last_interaction_list,
trackers_last_interaction_to_df,
)
from tardis.transport.montecarlo.progress_bars import (
refresh_packet_pbar,
reset_packet_pbar,
update_iterations_pbar,
)
logger = logging.getLogger(__name__)
# TODO: refactor this into more parts
[docs]
class MCTransportSolverClassic(HDFWriterMixin):
"""
This class modifies the MonteCarloTransportState to solve the radiative
transfer problem.
"""
hdf_properties = ["transport_state"]
hdf_name = "transport"
def __init__(
self,
radfield_prop_solver: MCRadiationFieldPropertiesSolver,
enable_full_relativity: bool,
line_interaction_type: str,
spectrum_method: str,
packet_source: BasePacketSource,
enable_rpacket_tracking: bool = False,
nthreads: int = 1,
debug_packets: bool = False,
logger_buffer: int = 1,
use_gpu: bool = False,
montecarlo_configuration: MonteCarloConfiguration | None = None,
) -> None:
self.radfield_prop_solver = radfield_prop_solver
# inject different packets
self.enable_full_relativity = enable_full_relativity
self.line_interaction_type = line_interaction_type
self.spectrum_method = spectrum_method
self.use_gpu = use_gpu
self.enable_rpacket_tracking = enable_rpacket_tracking
self.montecarlo_configuration = montecarlo_configuration
self.packet_source = packet_source
# Set number of threads
self.nthreads = nthreads
# set up logger based on config
mc_tracker.DEBUG_MODE = debug_packets
mc_tracker.BUFFER = logger_buffer
[docs]
def initialize_transport_state(
self,
simulation_state,
opacity_state,
macro_atom_state,
plasma,
no_of_packets,
iteration=0,
):
if not plasma.continuum_interaction_species.empty:
if plasma.gamma is not None:
n_levels_bf_species_by_n_cells_tuple = plasma.gamma.shape
else:
n_levels_bf_species_by_n_cells_tuple = plasma.phi_lucy.shape
else:
n_levels_bf_species_by_n_cells_tuple = (0, 0)
packet_collection = self.packet_source.create_packets(
no_of_packets, seed_offset=iteration
)
# Classic mode: continuum processes disabled
montecarlo_globals.CONTINUUM_PROCESSES_ENABLED = False
geometry_state = simulation_state.geometry.to_numba()
opacity_state_numba = opacity_state.to_numba(
macro_atom_state,
self.line_interaction_type,
)
transport_state = MonteCarloTransportState(
packet_collection,
geometry_state_numba=geometry_state,
opacity_state_numba=opacity_state_numba,
time_explosion=simulation_state.time_explosion,
n_levels_bf_species_by_n_cells_tuple=n_levels_bf_species_by_n_cells_tuple,
)
transport_state.enable_full_relativity = (
self.montecarlo_configuration.ENABLE_FULL_RELATIVITY
)
configuration_initialize(self.montecarlo_configuration, self)
return transport_state
[docs]
def run(
self,
transport_state: MonteCarloTransportState,
show_progress_bars: bool = True,
record_full_histories: bool = False,
) -> None:
"""
Run the montecarlo calculation.
Parameters
----------
transport_state : tardis.transport.montecarlo.transport_state.TransportState
Transport state containing all the data needed for the Monte Carlo simulation
show_progress_bars : bool
Show progress bars
record_full_histories : bool, optional
Record full real-packet histories in
``transport_state.tracker_full_df`` even without r-packet
tracking, e.g. for virtual-packet post-processing.
"""
self.run_classic(
transport_state, show_progress_bars, record_full_histories
)
[docs]
def run_classic(
self,
transport_state: MonteCarloTransportState,
show_progress_bars: bool = True,
record_full_histories: bool = False,
) -> None:
"""
Run the montecarlo calculation using classic mode (no continuum).
Parameters
----------
transport_state : tardis.transport.montecarlo.transport_state.TransportState
Transport state containing all the data needed for the Monte Carlo simulation
show_progress_bars : bool
Show progress bars
record_full_histories : bool, optional
Record full real-packet histories in
``transport_state.tracker_full_df`` even without r-packet
tracking, e.g. for virtual-packet post-processing.
"""
set_num_threads(self.nthreads)
self.transport_state = transport_state
number_of_rpackets = len(transport_state.packet_collection.initial_nus)
full_tracking = self.enable_rpacket_tracking or record_full_histories
if full_tracking:
trackers_list = generate_tracker_full_list(
number_of_rpackets,
self.montecarlo_configuration.INITIAL_TRACKING_ARRAY_LENGTH,
)
else:
# Initialize the last interaction tracker list directly
trackers_list = generate_tracker_last_interaction_list(
number_of_rpackets
)
# Reset packet progress bar for this iteration
if show_progress_bars:
reset_packet_pbar(number_of_rpackets)
estimators_bulk, estimators_line = montecarlo_transport(
transport_state.packet_collection,
transport_state.geometry_state_numba,
transport_state.time_explosion.cgs.value,
transport_state.opacity_state_numba,
self.montecarlo_configuration,
trackers_list,
show_progress_bars=show_progress_bars,
packet_propagation_function=packet_propagation,
)
# Attach estimators to transport state
transport_state.estimators_bulk = estimators_bulk
transport_state.estimators_line = estimators_line
update_iterations_pbar(1)
refresh_packet_pbar()
if full_tracking:
# Memory: the trackers are not needed once copied, so free them
# during the conversion rather than holding them alongside the
# full DataFrame.
self.transport_state.tracker_full_df = trackers_full_to_df(
trackers_list, release_trackers=True
)
self.transport_state.tracker_last_interaction_df = (
tracker_full_df2tracker_last_interaction_df(
self.transport_state.tracker_full_df
)
)
else:
self.transport_state.tracker_full_df = None
self.transport_state.tracker_last_interaction_df = (
trackers_last_interaction_to_df(trackers_list)
)
[docs]
@classmethod
def from_config(
cls,
config: Configuration,
packet_source: BasePacketSource,
) -> Self:
"""
Create a new MontecarloTransport instance from a Configuration object.
Parameters
----------
config : tardis.io.config_reader.Configuration
Returns
-------
MontecarloTransport
"""
if config.plasma.disable_electron_scattering:
logger.warning(
"Disabling electron scattering - this is not physical."
"Likely bug in formal integral - "
"will not give same results."
)
constants.SIGMA_THOMSON = 1e-200
else:
logger.debug("Electron scattering switched on")
constants.SIGMA_THOMSON = const.sigma_T.to("cm^2").value
running_mode = config.spectrum.integrated.compute.upper()
if running_mode == "GPU":
if cuda.is_available():
use_gpu = True
else:
raise ValueError(
"""The GPU option was selected for the formal_integral,
but no CUDA GPU is available."""
)
elif running_mode == "AUTOMATIC":
use_gpu = bool(cuda.is_available())
elif running_mode == "CPU":
use_gpu = False
else:
raise ValueError(
"""An invalid option for compute was passed. The three
valid values are 'GPU', 'CPU', and 'Automatic'."""
)
montecarlo_configuration = MonteCarloConfiguration()
montecarlo_configuration.DISABLE_LINE_SCATTERING = (
config.plasma.disable_line_scattering
)
montecarlo_configuration.DISABLE_ELECTRON_SCATTERING = (
config.plasma.disable_electron_scattering
)
montecarlo_configuration.INITIAL_TRACKING_ARRAY_LENGTH = (
config.montecarlo.tracking.initial_array_length
)
radfield_prop_solver = MCRadiationFieldPropertiesSolver(
config.plasma.w_epsilon
)
return cls(
radfield_prop_solver=radfield_prop_solver,
enable_full_relativity=config.montecarlo.enable_full_relativity,
line_interaction_type=config.plasma.line_interaction_type,
spectrum_method=config.spectrum.method,
packet_source=packet_source,
debug_packets=config.montecarlo.debug_packets,
logger_buffer=config.montecarlo.logger_buffer,
enable_rpacket_tracking=config.montecarlo.tracking.track_rpacket,
nthreads=config.montecarlo.nthreads,
use_gpu=use_gpu,
montecarlo_configuration=montecarlo_configuration,
)