Flexible packet initial radius (tardis-sn#1592)

* Packet collections are initialized with a radius

* Docstring for create_packets

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)

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,

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

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,

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],

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,15 +73,15 @@ 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
             temperature
         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