Gamma-ray packet source refactor (#2546)

* First steps

Radioactive packet source
GXPacket collection
Some initial methods

* Add more methods to the packet source

* Documentation, cleanup, positronium handling

* Black formatting, cleanup

* Fix kappa_calculation import loop

* Working packet source

Creates a Packet collection with all necessary packet properties

* Positron deposition and packet tracker arrays

* Fix parents dict

* Result now mostly consistent with master

* Fix kappa calculation test

* Adds missing docstring for calculate_energy_factors

* Fix util tests and black

* Testing skeleton for packet source

* black tests

* Address comments

* Beginnings of an alternate packet source for the dataframe

* Possible deposition fix, import fixes

* Added old functions

* Cleanup, comments, and more dataframe sampling setup

* Corrected packet time computation

* Modified sampling to use positron information

* Remove packet sampling call and replace with dataframe method

* Update method call to match current usage

* Comment responses

* More comments, improve packet dataframe index handling

---------

Co-authored-by: Knights-Templars <[email protected]>

tardis/energy_input/GXPacket.py

@@ -79,6 +79,36 @@ def get_location_r(self):
         )
 
 
+class GXPacketCollection:
+    """
+    Gamma-ray packet collection
+    """
+
+    def __init__(
+        self,
+        location,
+        direction,
+        energy_rf,
+        energy_cmf,
+        nu_rf,
+        nu_cmf,
+        status,
+        shell,
+        time_current,
+    ):
+        self.location = location
+        self.direction = direction
+        self.energy_rf = energy_rf
+        self.energy_cmf = energy_cmf
+        self.nu_rf = nu_rf
+        self.nu_cmf = nu_cmf
+        self.status = status
+        self.shell = shell
+        self.time_current = time_current
+        self.number_of_packets = len(self.energy_rf)
+        self.tau = -np.log(np.random.random(self.number_of_packets))
+
+
 # @njit(**njit_dict_no_parallel)
 def initialize_packet_properties(
     isotope_energy,

tardis/energy_input/__init__.py

@@ -1,5 +1,3 @@
 """
 Contains classes and functions to handle energy deposition and transport.
 """
-
-from tardis.energy_input.util import *

tardis/energy_input/gamma_packet_loop.py

@@ -7,6 +7,7 @@
     photoabsorption_opacity_calculation,
     pair_creation_opacity_calculation,
     photoabsorption_opacity_calculation_kasen,
+    kappa_calculation,
     pair_creation_opacity_artis,
     SIGMA_T,
 )
@@ -18,7 +19,6 @@
     doppler_factor_3d,
     C_CGS,
     H_CGS_KEV,
-    kappa_calculation,
     get_index,
 )
 from tardis.energy_input.GXPacket import GXPacketStatus
@@ -50,7 +50,6 @@ def gamma_packet_loop(
     energy_df_rows,
     energy_plot_df_rows,
     energy_out,
-    packets_out,
     packets_info_array,
 ):
     """Propagates packets through the simulation
@@ -289,9 +288,6 @@ def gamma_packet_loop(
                     energy_out[bin_index, time_index] += rest_energy / (
                         bin_width * dt
                     )
-                    packets_out[bin_index] = np.array(
-                        [bin_index, i, rest_energy, packet.Z, packet.A]
-                    )
                     packet.status = GXPacketStatus.ESCAPED
                     escaped_packets += 1
                     if scattered:
@@ -322,7 +318,6 @@ def gamma_packet_loop(
         energy_plot_df_rows,
         energy_out,
         deposition_estimator,
-        packets_out,
         bin_width,
         packets_info_array,
     )

tardis/energy_input/gamma_ray_estimators.py

@@ -6,13 +6,13 @@
     compton_opacity_calculation,
     SIGMA_T,
     photoabsorption_opacity_calculation,
+    kappa_calculation,
 )
 from tardis.energy_input.util import (
     angle_aberration_gamma,
     doppler_factor_3d,
     H_CGS_KEV,
     ELECTRON_MASS_ENERGY_KEV,
-    kappa_calculation,
 )
 
 

tardis/energy_input/gamma_ray_interactions.py

@@ -2,10 +2,12 @@
 from numba import njit
 
 from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
-from tardis.montecarlo.montecarlo_numba.opacities import compton_opacity_partial
+from tardis.montecarlo.montecarlo_numba.opacities import (
+    compton_opacity_partial,
+    kappa_calculation,
+)
 from tardis.energy_input.util import (
     get_random_unit_vector,
-    kappa_calculation,
     euler_rodrigues,
     compton_theta_distribution,
     get_perpendicular_vector,
@@ -157,7 +159,6 @@ def get_compton_fraction_urilight(energy):
 
     accept = False
     while not accept:
-
         z = np.random.random(3)
         alpha1 = np.log(1.0 / x0)
         alpha2 = (1.0 - x0**2.0) / 2.0