def __init__(self, model, plasma, runner, points=1000):
self.model = model
self.runner = runner
self.points = points
if plasma:
self.plasma = numba_plasma_initialize(plasma,
runner.line_interaction_type)
self.atomic_data = plasma.atomic_data
self.original_plasma = plasma
def generate_numba_objects(self):
'''instantiate the numba interface objects
needed for computing the formal integral'''
self.numba_model = NumbaModel(
self.runner.r_inner_i,
self.runner.r_outer_i,
self.model.time_explosion.to("s").value,
)
self.numba_plasma = numba_plasma_initialize(
self.original_plasma, self.runner.line_interaction_type)
self.numba_integrator = NumbaFormalIntegrator(self.numba_model,
self.numba_plasma,
self.points)
def test_numba_plasma_initialize(nb_simulation_verysimple, input_params):
line_interaction_type = input_params
plasma = nb_simulation_verysimple.plasma
actual = numba_interface.numba_plasma_initialize(plasma,
line_interaction_type)
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":
empty = np.zeros(1, dtype=np.int64)
npt.assert_allclose(actual.transition_probabilities,
np.zeros((1, 1), dtype=np.float64))
npt.assert_allclose(actual.line2macro_level_upper, empty)
npt.assert_allclose(actual.macro_block_references, empty)
npt.assert_allclose(actual.transition_type, empty)
npt.assert_allclose(actual.destination_level_id, empty)
npt.assert_allclose(actual.transition_line_id, empty)
else:
npt.assert_allclose(
actual.transition_probabilities,
plasma.transition_probabilities.values,
)
npt.assert_allclose(
actual.line2macro_level_upper,
plasma.atomic_data.lines_upper2macro_reference_idx,
)
npt.assert_allclose(
actual.macro_block_references,
plasma.atomic_data.macro_atom_references["block_references"].
values,
)
npt.assert_allclose(
actual.transition_type,
plasma.atomic_data.macro_atom_data["transition_type"].values,
)
npt.assert_allclose(
actual.destination_level_id,
plasma.atomic_data.macro_atom_data["destination_level_idx"].values,
)
npt.assert_allclose(
actual.transition_line_id,
plasma.atomic_data.macro_atom_data["lines_idx"].values,
)
def montecarlo_radial1d(model, plasma, runner):
packet_collection = PacketCollection(
runner.input_nu,
runner.input_mu,
runner.input_energy,
runner._output_nu,
runner._output_energy,
)
numba_model = NumbaModel(
runner.r_inner_cgs,
runner.r_outer_cgs,
model.time_explosion.to("s").value,
)
numba_plasma = numba_plasma_initialize(plasma,
runner.line_interaction_type)
estimators = Estimators(
runner.j_estimator,
runner.nu_bar_estimator,
runner.j_blue_estimator,
runner.Edotlu_estimator,
)
packet_seeds = montecarlo_configuration.packet_seeds
number_of_vpackets = montecarlo_configuration.number_of_vpackets
(
v_packets_energy_hist,
last_interaction_type,
last_interaction_in_nu,
last_line_interaction_in_id,
last_line_interaction_out_id,
virt_packet_nus,
virt_packet_energies,
virt_packet_last_interaction_in_nu,
virt_packet_last_interaction_type,
virt_packet_last_line_interaction_in_id,
virt_packet_last_line_interaction_out_id,
) = montecarlo_main_loop(
packet_collection,
numba_model,
numba_plasma,
estimators,
runner.spectrum_frequency.value,
number_of_vpackets,
packet_seeds,
)
runner._montecarlo_virtual_luminosity.value[:] = v_packets_energy_hist
runner.last_interaction_type = last_interaction_type
runner.last_interaction_in_nu = last_interaction_in_nu
runner.last_line_interaction_in_id = last_line_interaction_in_id
runner.last_line_interaction_out_id = last_line_interaction_out_id
if montecarlo_configuration.VPACKET_LOGGING and number_of_vpackets > 0:
runner.virt_packet_nus = np.concatenate(
np.array(virt_packet_nus)).ravel()
runner.virt_packet_energies = np.concatenate(
np.array(virt_packet_energies)).ravel()
runner.virt_packet_last_interaction_in_nu = np.concatenate(
np.array(virt_packet_last_interaction_in_nu)).ravel()
runner.virt_packet_last_interaction_type = np.concatenate(
np.array(virt_packet_last_interaction_type)).ravel()
runner.virt_packet_last_line_interaction_in_id = np.concatenate(
np.array(virt_packet_last_line_interaction_in_id)).ravel()
runner.virt_packet_last_line_interaction_out_id = np.concatenate(
np.array(virt_packet_last_line_interaction_out_id)).ravel()
def verysimple_numba_plasma(nb_simulation_verysimple):
return numba_plasma_initialize(
nb_simulation_verysimple.plasma, line_interaction_type="macroatom"
)
def montecarlo_radial1d(
model,
plasma,
iteration,
no_of_packets,
total_iterations,
show_progress_bars,
runner,
):
packet_collection = PacketCollection(
runner.input_r,
runner.input_nu,
runner.input_mu,
runner.input_energy,
runner._output_nu,
runner._output_energy,
)
numba_model = NumbaModel(
runner.r_inner_cgs,
runner.r_outer_cgs,
model.time_explosion.to("s").value,
)
numba_plasma = numba_plasma_initialize(plasma, runner.line_interaction_type)
estimators = Estimators(
runner.j_estimator,
runner.nu_bar_estimator,
runner.j_blue_estimator,
runner.Edotlu_estimator,
)
packet_seeds = montecarlo_configuration.packet_seeds
number_of_vpackets = montecarlo_configuration.number_of_vpackets
(
v_packets_energy_hist,
last_interaction_type,
last_interaction_in_nu,
last_line_interaction_in_id,
last_line_interaction_out_id,
virt_packet_nus,
virt_packet_energies,
virt_packet_initial_mus,
virt_packet_initial_rs,
virt_packet_last_interaction_in_nu,
virt_packet_last_interaction_type,
virt_packet_last_line_interaction_in_id,
virt_packet_last_line_interaction_out_id,
rpacket_trackers,
) = montecarlo_main_loop(
packet_collection,
numba_model,
numba_plasma,
estimators,
runner.spectrum_frequency.value,
number_of_vpackets,
packet_seeds,
montecarlo_configuration.VPACKET_LOGGING,
iteration=iteration,
show_progress_bars=show_progress_bars,
no_of_packets=no_of_packets,
total_iterations=total_iterations,
)
runner._montecarlo_virtual_luminosity.value[:] = v_packets_energy_hist
runner.last_interaction_type = last_interaction_type
runner.last_interaction_in_nu = last_interaction_in_nu
runner.last_line_interaction_in_id = last_line_interaction_in_id
runner.last_line_interaction_out_id = last_line_interaction_out_id
if montecarlo_configuration.VPACKET_LOGGING and number_of_vpackets > 0:
runner.virt_packet_nus = np.concatenate(
np.array(virt_packet_nus)
).ravel()
runner.virt_packet_energies = np.concatenate(
np.array(virt_packet_energies)
).ravel()
runner.virt_packet_initial_mus = np.concatenate(
np.array(virt_packet_initial_mus)
).ravel()
runner.virt_packet_initial_rs = np.concatenate(
np.array(virt_packet_initial_rs)
).ravel()
runner.virt_packet_last_interaction_in_nu = np.concatenate(
np.array(virt_packet_last_interaction_in_nu)
).ravel()
runner.virt_packet_last_interaction_type = np.concatenate(
np.array(virt_packet_last_interaction_type)
).ravel()
runner.virt_packet_last_line_interaction_in_id = np.concatenate(
np.array(virt_packet_last_line_interaction_in_id)
).ravel()
runner.virt_packet_last_line_interaction_out_id = np.concatenate(
np.array(virt_packet_last_line_interaction_out_id)
).ravel()
update_iterations_pbar(1)
# Condition for Checking if RPacket Tracking is enabled
if montecarlo_configuration.RPACKET_TRACKING:
runner.rpacket_tracker = rpacket_trackers
def test_montecarlo_main_loop(
config_verysimple,
atomic_dataset,
tardis_ref_path,
tmpdir,
set_seed_fixture,
random_call_fixture,
):
montecarlo_configuration.LEGACY_MODE_ENABLED = True
# Load C data from refdata
C_fname = os.path.join(tardis_ref_path, "montecarlo_1e5_compare_data.h5")
expected_nu = pd.read_hdf(C_fname,
key="/simulation/runner/output_nu").values
expected_energy = pd.read_hdf(
C_fname, key="/simulation/runner/output_energy").values
expected_nu_bar_estimator = pd.read_hdf(
C_fname, key="/simulation/runner/nu_bar_estimator").values
expected_j_estimator = pd.read_hdf(
C_fname, key="/simulation/runner/j_estimator").values
# Setup model config from verysimple
atomic_data = deepcopy(atomic_dataset)
config_verysimple.montecarlo.last_no_of_packets = 1e5
config_verysimple.montecarlo.no_of_virtual_packets = 0
config_verysimple.montecarlo.iterations = 1
config_verysimple.montecarlo.single_packet_seed = 0
del config_verysimple["config_dirname"]
sim = Simulation.from_config(config_verysimple, atom_data=atomic_data)
# Init model
numba_plasma = numba_plasma_initialize(sim.plasma,
line_interaction_type="macroatom")
runner = sim.runner
model = sim.model
runner._initialize_geometry_arrays(model)
runner._initialize_estimator_arrays(numba_plasma.tau_sobolev.shape)
runner._initialize_packets(model.t_inner.value, 100000, 0)
# Init parameters
montecarlo_configuration.v_packet_spawn_start_frequency = (
runner.virtual_spectrum_spawn_range.end.to(
u.Hz, equivalencies=u.spectral()).value)
montecarlo_configuration.v_packet_spawn_end_frequency = (
runner.virtual_spectrum_spawn_range.start.to(
u.Hz, equivalencies=u.spectral()).value)
montecarlo_configuration.temporary_v_packet_bins = 20000
montecarlo_configuration.full_relativity = runner.enable_full_relativity
montecarlo_configuration.single_packet_seed = 0
# Init packet collection from runner
packet_collection = PacketCollection(
runner.input_nu,
runner.input_mu,
runner.input_energy,
runner._output_nu,
runner._output_energy,
)
# Init model from runner
numba_model = NumbaModel(
runner.r_inner_cgs,
runner.r_outer_cgs,
model.time_explosion.to("s").value,
)
# Init estimators from runner
estimators = Estimators(
runner.j_estimator,
runner.nu_bar_estimator,
runner.j_blue_estimator,
runner.Edotlu_estimator,
)
# Empty vpacket collection
vpacket_collection = VPacketCollection(0, np.array([0, 0],
dtype=np.float64), 0,
np.inf, 0, 0)
# output arrays
output_nus = np.empty_like(packet_collection.packets_output_nu)
output_energies = np.empty_like(packet_collection.packets_output_nu)
# IMPORTANT: seeds RNG state within JIT
seed = 23111963
set_seed_fixture(seed)
for i in range(len(packet_collection.packets_input_nu)):
# Generate packet
packet = r_packet.RPacket(
numba_model.r_inner[0],
packet_collection.packets_input_mu[i],
packet_collection.packets_input_nu[i],
packet_collection.packets_input_energy[i],
seed,
i,
0,
)
# Loop packet
spl.single_packet_loop(packet, numba_model, numba_plasma, estimators,
vpacket_collection)
output_nus[i] = packet.nu
if packet.status == r_packet.PacketStatus.REABSORBED:
output_energies[i] = -packet.energy
elif packet.status == r_packet.PacketStatus.EMITTED:
output_energies[i] = packet.energy
# RNG to match C
random_call_fixture()
packet_collection.packets_output_energy[:] = output_energies[:]
packet_collection.packets_output_nu[:] = output_nus[:]
actual_energy = packet_collection.packets_output_energy
actual_nu = packet_collection.packets_output_nu
actual_nu_bar_estimator = estimators.nu_bar_estimator
actual_j_estimator = estimators.j_estimator
# Compare
npt.assert_allclose(actual_nu_bar_estimator,
expected_nu_bar_estimator,
rtol=1e-13)
npt.assert_allclose(actual_j_estimator, expected_j_estimator, rtol=1e-13)
npt.assert_allclose(actual_energy, expected_energy, rtol=1e-13)
npt.assert_allclose(actual_nu, expected_nu, rtol=1e-13)
