Create samples with packet energies

tardis/energy_input/gamma_ray_channel.py

@@ -167,3 +167,80 @@
     )
 
     return isotope_decay_df
+
+
+def distribute_packets(isotope_decay_df):
+    """
+    Function to distribute packets based on the energy of the decay.
+
+    Parameters
+    ----------
+    isotope_decay_df : pd.DataFrame
+        dataframe of isotopes for each shell with their decay mode, number of decays, radiation type,
+        radiation energy and radiation intensity.
+
+    Returns
+    -------
+    isotope_decay_df : pd.DataFrame
+        dataframe of isotopes for each shell with their decay mode, number of decays, radiation type,
+        radiation energy and radiation intensity, normalized energy and packets.
+    """
+
+    isotope_decay_df["normalized_energy"] = (
+        isotope_decay_df["decay_energy_erg"]
+        / isotope_decay_df["decay_energy_erg"].sum()
+    )
+
+    return isotope_decay_df
+
+
+def sample_single_packet(isotope_decay_df):
+    """
+    Function to sample packet energy from the energy distribution of the decay.
+
+    Parameters
+    ----------
+    isotope_decay_df : pd.DataFrame
+        dataframe of isotopes for each shell with their decay mode, number of decays, radiation type,
+        radiation energy and radiation intensity, normalized energy and packets.
+
+    Returns
+    -------
+    sampled_packet : pd.Series
+        A Series of sampled packets using the normalized energy in each channel.
+
+    """
+
+    sampled_packet = isotope_decay_df.sample(
+        weights="normalized_energy", random_state=np.random.RandomState(seed=29)
+    )  # random state for reproducibility
+
+    return sampled_packet
+
+
+def sample_packets(isotope_decay_df, number_of_packets):
+    """
+    Function to sample packets based on the energy of the decay.
+
+    Parameters
+    ----------
+    isotope_decay_df : pd.DataFrame
+        dataframe of isotopes for each shell with their decay mode, number of decays, radiation type,
+        radiation energy and radiation intensity, normalized energy and packets.
+    number_of_packets : int
+        number of packets to distribute.
+
+    Returns
+    -------
+    sampled_packets : numpy.ndarray
+        Creates an array of packets.
+    """
+
+    sampled_packets = np.array(
+        [
+            sample_single_packet(isotope_decay_df)
+            for i in range(number_of_packets)
+        ]
+    )
+
+    return sampled_packets