Source code for tardis.montecarlo.montecarlo_numba.macro_atom

import numpy as np
from enum import IntEnum

from numba import njit
from tardis.montecarlo.montecarlo_numba import njit_dict ,njit_dict_no_parallel


[docs]class MacroAtomError(ValueError): pass
[docs]class MacroAtomTransitionType(IntEnum): INTERNAL_UP = 1 INTERNAL_DOWN = 0 BB_EMISSION = -1 BF_EMISSION = -2 FF_EMISSION = -3 ADIABATIC_COOLING = -4
[docs]@njit(**njit_dict_no_parallel) def macro_atom(r_packet, numba_plasma): """ Parameters ---------- r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket numba_plasma : tardis.montecarlo.numba_interface.numba_plasma.NumbaPlasma Returns ------- """ activation_level_id = numba_plasma.line2macro_level_upper[ r_packet.next_line_id ] current_transition_type = 0 while current_transition_type >= 0: probability = 0.0 probability_event = np.random.random() block_start = numba_plasma.macro_block_references[activation_level_id] block_end = numba_plasma.macro_block_references[activation_level_id + 1] # looping through the transition probabilities for transition_id in range(block_start, block_end): transition_probability = numba_plasma.transition_probabilities[ transition_id, r_packet.current_shell_id ] probability += transition_probability if probability > probability_event: activation_level_id = numba_plasma.destination_level_id[ transition_id ] current_transition_type = numba_plasma.transition_type[ transition_id ] break else: raise MacroAtomError( "MacroAtom ran out of the block. This should not happen as " "the sum of probabilities is normalized to 1 and " "the probability_event should be less than 1" ) if current_transition_type == MacroAtomTransitionType.BB_EMISSION: return numba_plasma.transition_line_id[transition_id] else: raise MacroAtomError("MacroAtom currently only allows BB transitions")