diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index ebe1ba8abbb..3f97ab125d5 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -59,8 +59,7 @@ def montecarlo_radial1d(model, plasma, runner): packet_seeds = montecarlo_configuration.packet_seeds number_of_vpackets = montecarlo_configuration.number_of_vpackets - ContinuumObject = create_continuum_class(plasma.chi_continuum_calculator) - + ContinuumObject = create_continuum_class(plasma) ( v_packets_energy_hist, last_interaction_type, @@ -205,15 +204,11 @@ def montecarlo_main_loop( vpacket_collection = vpacket_collections[i] loop = single_packet_loop( - r_packet, numba_model, numba_plasma, estimators, vpacket_collection, - continuum, - - r_packet, numba_model, numba_plasma, estimators, vpacket_collection, continuum ) # if loop and 'stop' in loop: diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py index 40288fb1751..8e7c5ffdf9a 100644 --- a/tardis/montecarlo/montecarlo_numba/interaction.py +++ b/tardis/montecarlo/montecarlo_numba/interaction.py @@ -1,4 +1,5 @@ from numba import njit +import numpy as np from tardis.montecarlo.montecarlo_numba import njit_dict, njit_dict_no_parallel from tardis.montecarlo.montecarlo_numba.numba_interface import ( LineInteractionType, @@ -16,8 +17,13 @@ InteractionType, ) from tardis.montecarlo.montecarlo_numba.utils import get_random_mu -from tardis.montecarlo.montecarlo_numba.macro_atom import macro_atom +from tardis.montecarlo.montecarlo_numba.macro_atom import ( + macro_atom, MacroAtomTransitionType +) + + +@njit(**njit_dict_no_parallel) def scatter(r_packet, time_explosion): old_doppler_factor = get_doppler_factor( @@ -34,7 +40,9 @@ def scatter(r_packet, time_explosion): return comov_nu, inverse_new_doppler_factor -def continuum_event(r_packet, time_explosion, plasma, chi_continuum_calculator): +# Maybe make the continuum selection a method of the continuum? +@njit(**njit_dict_no_parallel) +def continuum_event(r_packet, time_explosion, continuum, numba_plasma): comov_nu, inverse_new_doppler_factor = scatter(r_packet, time_explosion) @@ -43,28 +51,83 @@ def continuum_event(r_packet, time_explosion, plasma, chi_continuum_calculator): zrand = np.random.random() - ( - chi_bf, - chi_bf_contributions, - current_continua, - chi_ff, - ) = chi_continuum_calculator(comov_nu, r_packet.current_shell_id) + # Does this need to be re-calculated? + # continuum.calculate(comov_nu, r_packet.current_shell_id) + + # Need to determine if a collisional process or not. If not: + + # This somehow determines the transition type and continuum id needed + # but does not sample new frequencies + # This reruns continuum.calculate + destination_level_idx = continuum.determine_macro_activation_idx( + comov_nu, r_packet.current_shell_id) + + transition_id, transition_type = macro_atom( + destination_level_idx, + r_packet.current_shell_id, + numba_plasma + ) + + # Then use this to get a transition id from the macroatom + if transition_type == MacroAtomTransitionType.FF_EMISSION: + free_free_emission(r_packet, time_explosion, numba_plasma, continuum) + elif transition_type == MacroAtomTransitionType.BF_EMISSION: + bound_free_emission(r_packet, + time_explosion, + numba_plasma, + continuum, + transition_id) + +@njit(**njit_dict_no_parallel) +def get_current_line_id(nu, numba_plasma): + ''' + Get the next line id corresponding to a packet at frequency nu + ''' + + reverse_line_list = numba_plasma.line_list_nu[::-1] + number_of_lines = len(numba_plasma.line_list_nu) + line_id = number_of_lines - np.searchsorted(reverse_line_list, nu) + return line_id + + +@njit(**njit_dict_no_parallel) +def free_free_emission(r_packet, time_explosion, numba_plasma, continuum): + + inverse_doppler_factor = get_inverse_doppler_factor( + r_packet.r, r_packet.mu, time_explosion + ) + # Need to get the sampler into numba somehow + # maybe I'll update the numba_plasma? + comov_nu = continuum.sample_nu_free_free(r_packet.current_shell_id) + r_packet.nu = comov_nu * inverse_doppler_factor + current_line_id = get_current_line_id(r_packet.nu, numba_plasma) + r_packet.next_line_id = current_line_id + + if montecarlo_configuration.full_relativity: + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) - cur_electron_density = numba_plasma.electron_density[ - r_packet.current_shell_id - ] +@njit(**njit_dict_no_parallel) +def bound_free_emission(r_packet, time_explosion, numba_plasma, continuum, continuum_id): - chi_e = cur_electron_density * SIGMA_THOMSON + inverse_doppler_factor = get_inverse_doppler_factor( + r_packet.r, r_packet.mu, time_explosion + ) - chi_nu = chi_bf + chi_ff + chi_e + comov_nu = continuum.sample_nu_free_bound( + r_packet.current_shell_id, + continuum_id) + + r_packet.nu = comov_nu * inverse_doppler_factor + current_line_id = get_current_line_id(r_packet.nu, numba_plasma) + r_packet.next_line_id = current_line_id + if montecarlo_configuration.full_relativity: + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) - if zrand < SIGMA_THOMSON / chi_nu: - thomson_scatter(r_packet, time_explosion) - elif zrand < (SIGMA_THOMSON + chi_bf) / chi_nu: - bound_free_absorption(r_packet, time_explosion, plasma) - else: - free_free_absorption(r_packet, time_explosion) @njit(**njit_dict_no_parallel) diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index d6762d310b6..627d1fee660 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -318,8 +318,13 @@ def set_properties( -def create_continuum_class(chi_continuum_calculator): +def create_continuum_class(plasma): """called before mainloop""" + + chi_continuum_calculator = plasma.chi_continuum_calculator + nu_fb_sampler = plasma.nu_fb_sampler + nu_ff_sampler = plasma.nu_ff_sampler + get_macro_activation_idx = plasma.determine_continuum_macro_activation_idx continuum_spec = [ ("chi_bf_tot", float64), ("chi_bf_contributions", float64[:]), @@ -348,6 +353,24 @@ def calculate(self, nu, shell): self.chi_ff, ) = chi_continuum_calculator(nu, shell) + def sample_nu_free_bound(self, shell, continuum_id): + + return nu_fb_sampler(shell, continuum_id) + + def sample_nu_free_free(self, shell): + + return nu_ff_sampler(shell) + + def determine_macro_activation_idx(self, nu, shell): + + self.calculate(nu, shell) # maybe don't do this here + idx = get_macro_activation_idx( + nu, self.chi_bf_tot, self.chi_ff, + self.chi_bf_contributions, self.current_continua + ) + return idx + + @njit(**njit_dict_no_parallel) def continuum_constructor(): return Continuum() diff --git a/tardis/montecarlo/montecarlo_numba/r_packet.py b/tardis/montecarlo/montecarlo_numba/r_packet.py index 8ff94d19c0c..eec2f19b341 100644 --- a/tardis/montecarlo/montecarlo_numba/r_packet.py +++ b/tardis/montecarlo/montecarlo_numba/r_packet.py @@ -161,17 +161,6 @@ def trace_packet( chi_continuum = chi_e + chi_bf + chi_ff distance_continuum = tau_event / chi_continuum - ( - chi_bf, - chi_bf_contributions, - current_continua, - x_sect_bfs, - chi_ff, - ) = chi_continuum_calculator(comov_nu, r_packet.current_shell_id) - - chi_continuum = chi_e + chi_bf + chi_ff - distance_continuum = tau_event / chi_continuum - cur_line_id = start_line_id # initializing varibale for Numba # - do not remove last_line_id = len(numba_plasma.line_list_nu) - 1 @@ -210,6 +199,7 @@ def trace_packet( distance = min(distance_trace, distance_boundary, distance_continuum) if distance_trace != 0: + if distance == distance_boundary: interaction_type = InteractionType.BOUNDARY # BOUNDARY r_packet.next_line_id = cur_line_id diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py index f308a1d91fe..a269fab899b 100644 --- a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py +++ b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py @@ -17,6 +17,7 @@ InteractionType, thomson_scatter, line_scatter, + continuum_event, ) from tardis.montecarlo.montecarlo_numba.numba_interface import ( LineInteractionType, @@ -38,7 +39,7 @@ @njit def single_packet_loop( r_packet, numba_model, numba_plasma, estimators, vpacket_collection, - chi_bf_interpolator + continuum ): """ Parameters @@ -78,7 +79,7 @@ def single_packet_loop( while r_packet.status == PacketStatus.IN_PROCESS: distance, interaction_type, delta_shell = trace_packet( r_packet, numba_model, numba_plasma, - estimators, chi_bf_interpolator + estimators, continuum ) if interaction_type == InteractionType.BOUNDARY: @@ -116,6 +117,19 @@ def single_packet_loop( trace_vpacket_volley( r_packet, vpacket_collection, numba_model, numba_plasma ) + elif interaction_type == InteractionType.CONTINUUM_PROCESS: + r_packet.last_interaction_type = InteractionType.CONTINUUM_PROCESS + move_r_packet( + r_packet, distance, numba_model.time_explosion, estimators + ) + continuum_event(r_packet, numba_model.time_explosion, + continuum, numba_plasma) + + trace_vpacket_volley( + r_packet, vpacket_collection, numba_model, numba_plasma + ) + + if mc_tracker.DEBUG_MODE: r_packet_track_nu.append(r_packet.nu) r_packet_track_mu.append(r_packet.mu) diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index 53cf34788d7..480ca530b72 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -24,6 +24,17 @@ def nb_simulation_verysimple(config_verysimple, atomic_dataset): return sim +@pytest.fixture(scope="package") +def verysimple_continuum(nu_simulation_verysimple): + plasma = nu_simulation_verysimple + numba_plasma = numba_interface.numba_plasma_initialize(plasma, "macroatom") + + ContinuumClass = create_continuum_class(plasma) + + continuum = ContinuumClass() + return continuum + + @pytest.fixture() def verysimple_collection(nb_simulation_verysimple): runner = nb_simulation_verysimple.runner diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py index 75b287e0faa..0311b677d13 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py @@ -4,6 +4,47 @@ import numpy as np +@pytest.mark.parametrize( + ["current_shell_id", "nu"], + [(0, 0.6), (1, 0.4)] + ) +def test_coninuum_opacities( + verysimple_continuum, current_shell_id, nu): + + continuum = verysimple_continuum + continuum.calculate(nu, current_shell_id) + print(current_shell_id, nu) + print(continuum.chi_bf_tot) + print(continuum.chi_bf_contributions) + print(continuum.current_continuua) + print(continuum.x_sect_bfs) + print(continuum.chi_ff) + + +@pytest.mark.parametrize( + "current_shell_id", + [0, 1, 2] + ) +def test_continuum_free_free_sampler( + verysimple_continuum, current_shell_id): + + continuum = verysimple_continuum + nu = continuum.sample_nu_free_free(current_shell_id) + print(nu) + +@pytest.mark.parametrize( + ["current_shell_id", "continuum_id"], + [(0, 0), (1, 0)] + ) +def test_coninuum_opacities( + verysimple_continuum, current_shell_id, continuum_id): + + continuum = verysimple_continuum + nu = continuum.sample_nu_free_bound(current_shell_id, continuum_id) + print(nu) + + + @pytest.mark.parametrize("input_params", ["scatter", "macroatom", "downbranch"]) def test_numba_plasma_initialize(nb_simulation_verysimple, input_params): line_interaction_type = input_params diff --git a/tardis/plasma/properties/continuum_processes.py b/tardis/plasma/properties/continuum_processes.py index 51daf9bc44c..584c169e212 100644 --- a/tardis/plasma/properties/continuum_processes.py +++ b/tardis/plasma/properties/continuum_processes.py @@ -1029,7 +1029,7 @@ def nu_fb(shell, continuum_id): zrand = np.random.random() idx = np.searchsorted(em[start:end], zrand, side="right") - return phot_nus[idx] - (em[idx] - zrand) / (em[idx] - em[idx-1]) * (phot_nus[idx] - phot_nus[idx-1]) + return phot_nus[idx] - (em[idx] - zrand) / (em[idx] - em[idx - 1]) * (phot_nus[idx] - phot_nus[idx - 1]) return nu_fb diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index dba6bbcebd7..c6d281b69fd 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -93,7 +93,7 @@ class PlasmaPropertyCollection(list): ) continuum_interaction_properties = PlasmaPropertyCollection( [ - MacroAtomData, # TODO: Remove this once everything works + MacroAtomData, # TODO: Remove this once everything works PhotoIonizationData, SpontRecombRateCoeff, PhotoIonRateCoeff,