diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index 7d5c65185b4..cc0b7fa646f 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -166,7 +166,7 @@ def _initialize_geometry_arrays(self, model): self.r_outer_cgs = model.r_outer.to("cm").value self.v_inner_cgs = model.v_inner.to("cm/s").value - def _initialize_packets(self, T, no_of_packets, iteration): + def _initialize_packets(self, T, no_of_packets, iteration, radius): # the iteration is added each time to preserve randomness # across different simulations with the same temperature, # for example. We seed the random module instead of the numpy module @@ -175,10 +175,11 @@ def _initialize_packets(self, T, no_of_packets, iteration): seed = self.seed + iteration rng = np.random.default_rng(seed=seed) seeds = rng.choice(MAX_SEED_VAL, no_of_packets, replace=True) - nus, mus, energies = self.packet_source.create_packets( - T, no_of_packets, rng + radii, nus, mus, energies = self.packet_source.create_packets( + T, no_of_packets, rng, radius ) mc_config_module.packet_seeds = seeds + self.input_r = radii self.input_nu = nus self.input_mu = mus self.input_energy = energies @@ -291,7 +292,9 @@ def run( self._initialize_geometry_arrays(model) - self._initialize_packets(model.t_inner.value, no_of_packets, iteration) + self._initialize_packets( + model.t_inner.value, no_of_packets, iteration, model.r_inner[0] + ) configuration_initialize(self, no_of_virtual_packets) montecarlo_radial1d(model, plasma, self) diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index 0343f33a267..020de1f04b4 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -29,6 +29,7 @@ def montecarlo_radial1d(model, plasma, runner): packet_collection = PacketCollection( + runner.input_r, runner.input_nu, runner.input_mu, runner.input_energy, diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index c45a679626b..abffe5b8441 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -67,7 +67,7 @@ def __init__( ): """ Plasma for the Numba code - + Parameters ---------- electron_density : numpy.ndarray @@ -161,6 +161,7 @@ def numba_plasma_initialize(plasma, line_interaction_type): packet_collection_spec = [ + ("packets_input_radius", float64[:]), ("packets_input_nu", float64[:]), ("packets_input_mu", float64[:]), ("packets_input_energy", float64[:]), @@ -173,12 +174,14 @@ def numba_plasma_initialize(plasma, line_interaction_type): class PacketCollection(object): def __init__( self, + packets_input_radius, packets_input_nu, packets_input_mu, packets_input_energy, packets_output_nu, packets_output_energy, ): + self.packets_input_radius = packets_input_radius self.packets_input_nu = packets_input_nu self.packets_input_mu = packets_input_mu self.packets_input_energy = packets_input_energy @@ -220,10 +223,18 @@ def __init__( self.nus = np.empty(temporary_v_packet_bins, dtype=np.float64) self.energies = 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) - self.last_interaction_type = -1 * np.ones(temporary_v_packet_bins, dtype=np.int64) - self.last_interaction_in_id = -1 * np.ones(temporary_v_packet_bins, dtype=np.int64) - self.last_interaction_out_id = -1 * np.ones(temporary_v_packet_bins, dtype=np.int64) + self.last_interaction_in_nu = np.zeros( + temporary_v_packet_bins, dtype=np.float64 + ) + self.last_interaction_type = -1 * np.ones( + temporary_v_packet_bins, dtype=np.int64 + ) + self.last_interaction_in_id = -1 * np.ones( + temporary_v_packet_bins, dtype=np.int64 + ) + self.last_interaction_out_id = -1 * np.ones( + temporary_v_packet_bins, dtype=np.int64 + ) self.idx = 0 self.rpacket_index = rpacket_index self.length = temporary_v_packet_bins @@ -241,21 +252,31 @@ 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_last_interaction_in_nu = np.empty(temp_length, dtype=np.float64) + temp_last_interaction_in_nu = np.empty( + temp_length, dtype=np.float64 + ) temp_last_interaction_type = np.empty(temp_length, dtype=np.int64) temp_last_interaction_in_id = np.empty(temp_length, dtype=np.int64) 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_last_interaction_in_nu[: self.length] = self.last_interaction_in_nu - temp_last_interaction_type[: self.length] = self.last_interaction_type - temp_last_interaction_in_id[: self.length] = self.last_interaction_in_id - temp_last_interaction_out_id[: self.length] = self.last_interaction_out_id + temp_last_interaction_in_nu[ + : self.length + ] = self.last_interaction_in_nu + temp_last_interaction_type[ + : self.length + ] = self.last_interaction_type + temp_last_interaction_in_id[ + : self.length + ] = self.last_interaction_in_id + temp_last_interaction_out_id[ + : self.length + ] = self.last_interaction_out_id self.nus = temp_nus self.energies = temp_energies self.last_interaction_in_nu = temp_last_interaction_in_nu - self.last_interaction_type = temp_last_interaction_type + self.last_interaction_type = temp_last_interaction_type self.last_interaction_in_id = temp_last_interaction_in_id self.last_interaction_out_id = temp_last_interaction_out_id self.length = temp_length diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index 20223587ed4..9f08c1d802e 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -28,6 +28,7 @@ def nb_simulation_verysimple(config_verysimple, atomic_dataset): def verysimple_collection(nb_simulation_verysimple): runner = nb_simulation_verysimple.runner return PacketCollection( + runner.input_r, runner.input_nu, runner.input_mu, runner.input_energy, @@ -99,6 +100,7 @@ def verysimple_3vpacket_collection(nb_simulation_verysimple): def verysimple_packet_collection(nb_simulation_verysimple): runner = nb_simulation_verysimple.runner return PacketCollection( + runner.input_r, runner.input_nu, runner.input_mu, runner.input_energy, diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index 33b4bb1d78c..10a31f0dfb6 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -74,7 +74,9 @@ def test_montecarlo_main_loop( runner._initialize_geometry_arrays(model) runner._initialize_estimator_arrays(numba_plasma.tau_sobolev.shape) - runner._initialize_packets(model.t_inner.value, 100000, 0) + runner._initialize_packets( + model.t_inner.value, 100000, 0, model.r_inner[0].value + ) # Init parameters montecarlo_configuration.v_packet_spawn_start_frequency = ( @@ -93,6 +95,7 @@ def test_montecarlo_main_loop( # Init packet collection from runner packet_collection = PacketCollection( + runner.input_r, runner.input_nu, runner.input_mu, runner.input_energy, @@ -130,7 +133,7 @@ def test_montecarlo_main_loop( 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_radius[i], packet_collection.packets_input_mu[i], packet_collection.packets_input_nu[i], packet_collection.packets_input_energy[i], diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index 32d74f7c35b..6efe4443749 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -38,9 +38,9 @@ def create_zero_limb_darkening_packet_mus(no_of_packets, rng): @staticmethod def create_uniform_packet_energies(no_of_packets, rng): """ - Uniformly distribute energy in arbitrary units where the ensemble of - packets has energy of 1. - + Uniformly distribute energy in arbitrary units where the ensemble of + packets has energy of 1. + Parameters ---------- no_of_packets : int @@ -73,7 +73,7 @@ def create_blackbody_packet_nus(T, no_of_packets, rng, l_samples=1000): .. math:: x = -\\ln{(\\xi_1\\xi_2\\xi_3\\xi_4)}/l_{\\rm min}\\; . where :math:`x=h\\nu/kT` - + Parameters ---------- T : float @@ -81,7 +81,7 @@ def create_blackbody_packet_nus(T, no_of_packets, rng, l_samples=1000): no_of_packets : int l_samples : int number of l_samples needed in the algorithm - + Returns ------- : numpy.ndarray @@ -110,9 +110,33 @@ class BlackBodySimpleSource(BasePacketSource): part. """ - def create_packets(self, T, no_of_packets, rng): + def create_packets(self, T, no_of_packets, rng, radius): + """Generate black-body packet properties as arrays + + Parameters + ---------- + T : float64 + Temperature + no_of_packets : int + Number of packets + rng : numpy random number generator + radius : float64 + Initial packet radius + + Returns + ------- + array + Packet radii + array + Packet frequencies + array + Packet directions + array + Packet energies + """ + radii = np.ones(no_of_packets) * radius nus = self.create_blackbody_packet_nus(T, no_of_packets, rng) mus = self.create_zero_limb_darkening_packet_mus(no_of_packets, rng) energies = self.create_uniform_packet_energies(no_of_packets, rng) - return nus, mus, energies + return radii, nus, mus, energies