PR to separate main gamma ray loop (#2529)

* PR to separate main gamma ray loop * Added logging instead of printing * Added isotope abundance info before decay * Added a working example of the gamma-ray code * Added distribute packets which can replace decays_per_isotope and fractional energy. Much Cleaner * Fix most tests * Remove excess decay * Final cleanup --------- Co-authored-by: Andrew Fullard <[email protected]>
tardis-sn · Mar 12, 2024 · 20b2a64 · 20b2a64
1 parent 07a3083
commit 20b2a64
Show file tree

Hide file tree

Showing 8 changed files with 1,022 additions and 67 deletions.
diff --git a/docs/working_gamma_ray_test.ipynb b/docs/working_gamma_ray_test.ipynb
diff --git a/tardis/energy_input/GXPacket.py b/tardis/energy_input/GXPacket.py
@@ -19,6 +19,7 @@ class GXPacketStatus(IntEnum):
     PAIR_CREATION = 2
     IN_PROCESS = 3
     END = 4
+    ESCAPED = 5
 
 
 gxpacket_spec = [
@@ -92,7 +93,6 @@ def initialize_packet_properties(
     initial_radius,
     times,
     effective_times,
-    inventory,
     average_power_per_mass,
     uniform_packet_energies=True,
 ):

diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py
@@ -50,6 +50,8 @@ def gamma_packet_loop(
     energy_df_rows,
     energy_plot_df_rows,
     energy_out,
+    packets_out,
+    packets_info_array,
 ):
     """Propagates packets through the simulation
 
@@ -151,28 +153,31 @@ def gamma_packet_loop(
                     # electron count per isotope
                     photoabsorption_opacity = 0
                     # photoabsorption_opacity_calculation_kasen()
-                else:
+                elif photoabsorption_opacity_type == "tardis":
                     photoabsorption_opacity = (
                         photoabsorption_opacity_calculation(
                             comoving_energy,
                             mass_density_time[packet.shell, time_index],
                             iron_group_fraction_per_shell[packet.shell],
                         )
                     )
+                else:
+                    raise ValueError("Invalid photoabsorption opacity type!")
 
                 if pair_creation_opacity_type == "artis":
                     pair_creation_opacity = pair_creation_opacity_artis(
                         comoving_energy,
                         mass_density_time[packet.shell, time_index],
                         iron_group_fraction_per_shell[packet.shell],
                     )
-                else:
+                elif pair_creation_opacity_type == "tardis":
                     pair_creation_opacity = pair_creation_opacity_calculation(
                         comoving_energy,
                         mass_density_time[packet.shell, time_index],
                         iron_group_fraction_per_shell[packet.shell],
                     )
-
+                else:
+                    raise ValueError("Invalid pair creation opacity type!")
             else:
                 compton_opacity = 0.0
                 pair_creation_opacity = 0.0
@@ -235,7 +240,6 @@ def gamma_packet_loop(
                     )
 
             elif distance == distance_interaction:
-
                 packet.status = scatter_type(
                     compton_opacity,
                     photoabsorption_opacity,
@@ -277,14 +281,18 @@ def gamma_packet_loop(
 
                 if packet.shell > len(mass_density_time[:, 0]) - 1:
                     rest_energy = packet.nu_rf * H_CGS_KEV
+                    lum_rf = (packet.energy_rf * 1.6022e-9) / dt
                     bin_index = get_index(rest_energy, energy_bins)
                     bin_width = (
                         energy_bins[bin_index + 1] - energy_bins[bin_index]
                     )
                     energy_out[bin_index, time_index] += rest_energy / (
                         bin_width * dt
                     )
-                    packet.status = GXPacketStatus.END
+                    packets_out[bin_index] = np.array(
+                        [bin_index, i, rest_energy, packet.Z, packet.A]
+                    )
+                    packet.status = GXPacketStatus.ESCAPED
                     escaped_packets += 1
                     if scattered:
                         scattered_packets += 1
@@ -293,10 +301,31 @@ def gamma_packet_loop(
                     packet.energy_cmf = 0.0
                     packet.status = GXPacketStatus.END
 
+            packets_info_array[i] = np.array(
+                [
+                    i,
+                    packet.status,
+                    packet.nu_cmf,
+                    packet.nu_rf,
+                    packet.energy_cmf,
+                    lum_rf,
+                    packet.energy_rf,
+                    packet.shell,
+                ]
+            )
+
     print("Escaped packets:", escaped_packets)
     print("Scattered packets:", scattered_packets)
 
-    return energy_df_rows, energy_plot_df_rows, energy_out, deposition_estimator
+    return (
+        energy_df_rows,
+        energy_plot_df_rows,
+        energy_out,
+        deposition_estimator,
+        packets_out,
+        bin_width,
+        packets_info_array,
+    )
 
 
 @njit(**njit_dict_no_parallel)