From 9ac481dc1bcf5bd6b7f908bfbcb913097f200530 Mon Sep 17 00:00:00 2001 From: Isaac Smith <71480393+smithis7@users.noreply.github.com> Date: Fri, 9 Jul 2021 13:07:08 -0400 Subject: [PATCH] Adding mu and r to virtual packet logging (#1696) * added r and mu to vpacket logging * fixing problem with merge conflict * fixing test failure * making documentation --- docs/io/output/vpacket_logging.rst | 4 +++ tardis/montecarlo/base.py | 8 +++-- tardis/montecarlo/montecarlo_numba/base.py | 32 +++++++++++++------ .../montecarlo_numba/numba_interface.py | 14 ++++++++ .../tests/test_numba_interface.py | 20 ++++++++++++ tardis/montecarlo/montecarlo_numba/vpacket.py | 2 ++ 6 files changed, 69 insertions(+), 11 deletions(-) diff --git a/docs/io/output/vpacket_logging.rst b/docs/io/output/vpacket_logging.rst index 75633bfef65..3c3a577f23b 100644 --- a/docs/io/output/vpacket_logging.rst +++ b/docs/io/output/vpacket_logging.rst @@ -25,6 +25,10 @@ After running the simulation, the following information can be retrieved: - List of virtual packet frequencies * - ``runner.virt_packet_energies`` - List of virtual packet energies + * - ``runner.virt_packet_initial_mus`` + - List of propagation directions that virtual packets are launched at + * - ``runner.virt_packet_initial_rs`` + - List of radii that virtual packets are launched at * - ``runner.virt_packet_last_interaction_type`` - Type of interaction that caused the virtual packet to be spawned * - ``runner.virt_packet_last_interaction_in_nu`` diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index cc0b7fa646f..a78d82c50b4 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -60,12 +60,14 @@ class MontecarloRunner(HDFWriterMixin): ] vpacket_hdf_properties = [ + "virt_packet_nus", + "virt_packet_energies", + "virt_packet_initial_rs", + "virt_packet_initial_mus", "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", - "virt_packet_nus", - "virt_packet_energies", ] hdf_name = "runner" @@ -129,6 +131,8 @@ def __init__( self.virt_packet_last_line_interaction_out_id = np.ones(2) * -1 self.virt_packet_nus = np.ones(2) * -1.0 self.virt_packet_energies = np.ones(2) * -1.0 + self.virt_packet_initial_rs = np.ones(2) * -1.0 + self.virt_packet_initial_mus = np.ones(2) * -1.0 # set up logger based on config mc_logger.DEBUG_MODE = debug_packets diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index 52221ec8246..ba107d6292a 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -62,6 +62,8 @@ def montecarlo_radial1d(model, plasma, runner): 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, @@ -89,6 +91,12 @@ def montecarlo_radial1d(model, plasma, runner): 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() @@ -162,6 +170,8 @@ def montecarlo_main_loop( # Arrays for vpacket logging 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 = [] @@ -204,6 +214,8 @@ def montecarlo_main_loop( vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx] vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx] + vpackets_initial_mu = vpacket_collection.initial_mus[: vpacket_collection.idx] + vpackets_initial_r = vpacket_collection.initial_rs[: vpacket_collection.idx] v_packets_idx = np.floor( (vpackets_nu - spectrum_frequency[0]) / delta_nu @@ -221,17 +233,17 @@ def montecarlo_main_loop( if montecarlo_configuration.VPACKET_LOGGING: for vpacket_collection in vpacket_collections: vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx] - vpackets_energy = vpacket_collection.energies[ - : vpacket_collection.idx - ] + vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx] + vpackets_initial_mu = vpacket_collection.initial_mus[: vpacket_collection.idx] + vpackets_initial_r = vpacket_collection.initial_rs[: vpacket_collection.idx] virt_packet_nus.append(np.ascontiguousarray(vpackets_nu)) virt_packet_energies.append(np.ascontiguousarray(vpackets_energy)) - virt_packet_last_interaction_in_nu.append( - np.ascontiguousarray( - vpacket_collection.last_interaction_in_nu[ - : vpacket_collection.idx - ] - ) + virt_packet_initial_mus.append(np.ascontiguousarray(vpackets_initial_mu)) + virt_packet_initial_rs.append(np.ascontiguousarray(vpackets_initial_r)) + virt_packet_last_interaction_in_nu.append(np.ascontiguousarray( + vpacket_collection.last_interaction_in_nu[ + : vpacket_collection.idx + ]) ) virt_packet_last_interaction_type.append( np.ascontiguousarray( @@ -266,6 +278,8 @@ def montecarlo_main_loop( last_line_interaction_out_ids, 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, diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index abffe5b8441..8a3bbf32622 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -196,6 +196,8 @@ def __init__( ("v_packet_spawn_end_frequency", float64), ("nus", float64[:]), ("energies", float64[:]), + ("initial_mus", float64[:]), + ("initial_rs", float64[:]), ("idx", int64), ("number_of_vpackets", int64), ("length", int64), @@ -222,6 +224,8 @@ def __init__( self.v_packet_spawn_end_frequency = v_packet_spawn_end_frequency self.nus = np.empty(temporary_v_packet_bins, dtype=np.float64) self.energies = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.initial_mus = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.initial_rs = np.empty(temporary_v_packet_bins, dtype=np.float64) self.number_of_vpackets = number_of_vpackets self.last_interaction_in_nu = np.zeros( temporary_v_packet_bins, dtype=np.float64 @@ -243,6 +247,8 @@ def set_properties( self, nu, energy, + initial_mu, + initial_r, last_interaction_in_nu, last_interaction_type, last_interaction_in_id, @@ -252,6 +258,8 @@ def set_properties( temp_length = self.length * 2 + self.number_of_vpackets temp_nus = np.empty(temp_length, dtype=np.float64) temp_energies = np.empty(temp_length, dtype=np.float64) + temp_initial_mus = np.empty(temp_length, dtype=np.float64) + temp_initial_rs = np.empty(temp_length, dtype=np.float64) temp_last_interaction_in_nu = np.empty( temp_length, dtype=np.float64 ) @@ -260,6 +268,8 @@ def set_properties( temp_last_interaction_out_id = np.empty(temp_length, dtype=np.int64) temp_nus[: self.length] = self.nus temp_energies[: self.length] = self.energies + temp_initial_mus[: self.length] = self.initial_mus + temp_initial_rs[: self.length] = self.initial_rs temp_last_interaction_in_nu[ : self.length ] = self.last_interaction_in_nu @@ -275,6 +285,8 @@ def set_properties( self.nus = temp_nus self.energies = temp_energies + self.initial_mus = temp_initial_mus + self.initial_rs = temp_initial_rs self.last_interaction_in_nu = temp_last_interaction_in_nu self.last_interaction_type = temp_last_interaction_type self.last_interaction_in_id = temp_last_interaction_in_id @@ -283,6 +295,8 @@ def set_properties( self.nus[self.idx] = nu self.energies[self.idx] = energy + self.initial_mus[self.idx] = initial_mu + self.initial_rs[self.idx] = initial_r self.last_interaction_in_nu[self.idx] = last_interaction_in_nu self.last_interaction_type[self.idx] = last_interaction_type self.last_interaction_in_id[self.idx] = last_interaction_in_id diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py index a4fd4a926dd..75b287e0faa 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py @@ -65,6 +65,8 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection): nus = [3.0e15, 0.0, 1e15, 1e5] energies = [0.4, 0.1, 0.6, 1e10] + initial_mus = [.1, 0, 1, .9] + initial_rs = [3e42, 4.5e45, 0, 9.0e40] last_interaction_in_nus = np.array( [3.0e15, 0.0, 1e15, 1e5], dtype=np.float64 ) @@ -75,6 +77,8 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection): for ( nu, energy, + initial_mu, + initial_r, last_interaction_in_nu, last_interaction_type, last_interaction_in_id, @@ -82,6 +86,8 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection): ) in zip( nus, energies, + initial_mus, + initial_rs, last_interaction_in_nus, last_interaction_types, last_interaction_in_ids, @@ -90,6 +96,8 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection): verysimple_3vpacket_collection.set_properties( nu, energy, + initial_mu, + initial_r, last_interaction_in_nu, last_interaction_type, last_interaction_in_id, @@ -108,6 +116,18 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection): ], energies, ) + npt.assert_array_equal( + verysimple_3vpacket_collection.initial_mus[ + : verysimple_3vpacket_collection.idx + ], + initial_mus, + ) + npt.assert_array_equal( + verysimple_3vpacket_collection.initial_rs[ + : verysimple_3vpacket_collection.idx + ], + initial_rs, + ) npt.assert_array_equal( verysimple_3vpacket_collection.last_interaction_in_nu[ : verysimple_3vpacket_collection.idx diff --git a/tardis/montecarlo/montecarlo_numba/vpacket.py b/tardis/montecarlo/montecarlo_numba/vpacket.py index e42de010110..ec8bca0254a 100644 --- a/tardis/montecarlo/montecarlo_numba/vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/vpacket.py @@ -273,6 +273,8 @@ def trace_vpacket_volley( vpacket_collection.set_properties( v_packet.nu, v_packet.energy, + v_packet_mu, + r_packet.r, r_packet.last_interaction_in_nu, r_packet.last_interaction_type, r_packet.last_line_interaction_in_id,