Merge branch 'fix_combine_dtype_windows' into add_matrix_testing

tobac-project · Oct 11, 2023 · d58111e · d58111e
2 parents 2f5c6da + b0f36c0
commit d58111e
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/tobac/utils/general.py b/tobac/utils/general.py
@@ -623,10 +623,16 @@ def combine_feature_dataframes(
     if old_feature_column_name is not None:
         combined_df[old_feature_column_name] = combined_df["feature"]
     # count_per_time = combined_feats.groupby('time')['index'].count()
-    combined_df["frame"] = combined_df["time"].rank(method="dense").astype(int) - 1
+    original_frame_dtype = combined_df["frame"].dtype
+    combined_df["frame"] = (
+        combined_df["time"].rank(method="dense").astype(original_frame_dtype) - 1
+    )
     combined_sorted = combined_df.sort_values(sort_features_by, ignore_index=True)
     if renumber_features:
-        combined_sorted["feature"] = np.arange(1, len(combined_sorted) + 1)
+        original_feature_dtype = combined_df["feature"].dtype
+        combined_sorted["feature"] = np.arange(
+            1, len(combined_sorted) + 1, dtype=original_feature_dtype
+        )
     combined_sorted = combined_sorted.reset_index(drop=True)
     return combined_sorted