Fixed Iteration values getting converted to float from int,

Reordered the columns in DataFrame
tardis-sn · Aug 16, 2021 · 2770626 · 2770626
1 parent 53a194f
commit 2770626
Showing 1 changed file with 29 additions and 4 deletions.
diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py
@@ -340,9 +340,17 @@ def create_tracked_rpacket_df(tracked_rpacket, iteration):
     mu = tracked_rpacket.mu
     energy = tracked_rpacket.energy
     shell_id = tracked_rpacket.shell_id
-    iteration = iteration
 
-    vals = [index, seed, status, r, nu, mu, energy, shell_id]
+    vals = [
+        index,
+        seed,
+        status,
+        r,
+        nu,
+        mu,
+        energy,
+        shell_id,
+    ]
     columns_name = [
         "Packet Index",
         "Packet Seed",
@@ -354,11 +362,28 @@ def create_tracked_rpacket_df(tracked_rpacket, iteration):
         "shell_id",
     ]
 
-    rpacket_tracked_df = pd.DataFrame(zip(*vals), columns=columns_name)
+    rpacket_tracked_df = pd.DataFrame(
+        zip(*vals), columns=columns_name, dtype=object
+    )
     rpacket_tracked_df["Iteration"] = iteration
     return rpacket_tracked_df
 
 
 def track_rpacket_dataframe(tracked_rpacket_df, tracked_df):
-    tracked_rpacket_df = tracked_rpacket_df.append(tracked_df)
+    tracked_rpacket_df = tracked_rpacket_df.append(
+        tracked_df, ignore_index=True
+    )
+    tracked_rpacket_df = tracked_rpacket_df.convert_dtypes()
+    columns_reorder = [
+        "Iteration",
+        "Packet Index",
+        "Packet Seed",
+        "Packet Status",
+        "r",
+        "nu",
+        "mu",
+        "energy",
+        "shell_id",
+    ]
+    tracked_rpacket_df = tracked_rpacket_df[columns_reorder]
     return tracked_rpacket_df