tool function for performance based split; new model training notebooks

2026-01-02 12:02:48 +01:00
parent f5796a5cdd
commit a3119e4fc3
3 changed files with 1187 additions and 0 deletions
@@ -0,0 +1,139 @@
+import pandas as pd
+import numpy as np
+
+
+def performance_based_split(
+    subject_ids,
+    performance_df,
+    split_ratio=0.33,
+    threshold=0.01,
+    max_iter=100,
+    random_seed=None
+):
+    """
+    Split subjects into two groups based on performance scores with balanced means.
+    
+    Parameters
+    ----------
+    subject_ids : array-like
+        List or array of subject IDs present in your dataset
+    performance_df : pd.DataFrame
+        DataFrame containing 'subjectID' and 'overall_score' columns
+    split_ratio : float, default=0.33
+        Proportion of subjects for the smaller group (0 < split_ratio < 1)
+    threshold : float, default=0.01
+        Target difference threshold between group means
+    max_iter : int, default=100
+        Maximum number of swap iterations
+    random_seed : int, optional
+        Random seed for reproducibility
+        
+    Returns
+    -------
+    group_small_ids : np.ndarray
+        Subject IDs for the smaller group
+    group_large_ids : np.ndarray
+        Subject IDs for the larger group
+    score_diff : float
+        Final absolute difference between group means
+        
+    Raises
+    ------
+    ValueError
+        If subjects are missing performance scores or no subjects match
+    """
+    if random_seed is not None:
+        np.random.seed(random_seed)
+    
+    # Filter performance data
+    perf_filtered = performance_df[
+        performance_df["subjectID"].isin(subject_ids)
+    ][["subjectID", "overall_score"]]
+    
+    # Merge to get only subjects present in both dataset and performance file
+    merged = (
+        pd.DataFrame({"subjectID": subject_ids})
+        .merge(perf_filtered, on="subjectID", how="inner")
+    )
+    
+    if len(merged) == 0:
+        raise ValueError("No subjects found in both dataset and performance file.")
+    
+    # Check for missing scores
+    if merged["overall_score"].isna().any():
+        raise ValueError("Missing score values for some subjects.")
+    
+    merged_sorted = merged.sort_values("overall_score", ascending=False).reset_index(drop=True)
+    
+    scores = merged_sorted["overall_score"].values
+    n_total = len(merged_sorted)
+    n_small = int(n_total * split_ratio)
+    n_large = n_total - n_small
+    
+    # Initial random split
+    idx = np.arange(n_total)
+    np.random.shuffle(idx)
+    
+    small_idx = idx[:n_small]
+    large_idx = idx[n_small:]
+    
+    def score_diff(small_idx, large_idx):
+        return abs(scores[small_idx].mean() - scores[large_idx].mean())
+    
+    diff = score_diff(small_idx, large_idx)
+    count = 0
+    
+    # Optimize via random swaps
+    while diff > threshold and count < max_iter:
+        si = np.random.choice(small_idx)
+        li = np.random.choice(large_idx)
+        
+        new_small_idx = small_idx.copy()
+        new_large_idx = large_idx.copy()
+        
+        new_small_idx[new_small_idx == si] = li
+        new_large_idx[new_large_idx == li] = si
+        
+        new_diff = score_diff(new_small_idx, new_large_idx)
+        
+        if new_diff < diff:
+            small_idx = new_small_idx
+            large_idx = new_large_idx
+            diff = new_diff
+        
+        count += 1
+    
+    # Extract subject IDs
+    group_small_ids = merged_sorted.loc[small_idx, "subjectID"].values
+    group_large_ids = merged_sorted.loc[large_idx, "subjectID"].values
+    
+    return group_small_ids, group_large_ids, diff
+
+
+# Example usage for 2-way split (train/test)
+# subjects = df["subjectID"].unique()
+# performance_df = pd.read_csv("performance.csv")
+# 
+# train_ids, test_ids, diff = performance_based_split(
+#     subject_ids=subjects,
+#     performance_df=performance_df,
+#     split_ratio=0.2,
+#     random_seed=42
+# )
+
+# Example usage for 3-way split (train/val/test)
+# Step 1: Split into train and temp
+# train_ids, temp_ids, diff1 = performance_based_split(
+#     subject_ids=subjects,
+#     performance_df=performance_df,
+#     split_ratio=0.6,  # 60% train, 40% temp
+#     random_seed=42
+# )
+# 
+# Step 2: Split temp into val and test
+# val_ids, test_ids, diff2 = performance_based_split(
+#     subject_ids=temp_ids,
+#     performance_df=performance_df,
+#     split_ratio=0.5,  # 50/50 split of remaining 40%
+#     random_seed=43
+# )