Parallelize data processing tasks with ThreadPoolExecutor

stats/correlation.py

@@ -4,6 +4,7 @@
 
 import argparse
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
 
 import numpy as np
 import pandas as pd
@@ -69,8 +70,7 @@ def calculate_correlation(df: pd.DataFrame, method: str, target_id: str | None =
     else:
         groups = []
 
-    out = np.zeros((len(unique_ids), n_cols), dtype=np.float64)
-    for gi, sub in enumerate(groups):
+    def _process_group(sub):
         mask = ~np.isnan(sub).any(axis=1)
         sub = sub[mask]
         if len(sub) > 1:
@@ -80,7 +80,11 @@ def calculate_correlation(df: pd.DataFrame, method: str, target_id: str | None =
                 c = np.abs(np.corrcoef(sub, rowvar=False))
             np.nan_to_num(c, copy=False, nan=0.0)
             np.fill_diagonal(c, 0.0)
-            out[gi] = c.max(axis=0)
+            return c.max(axis=0)
+        return np.zeros(n_cols, dtype=np.float64)
+
+    with ThreadPoolExecutor() as executor:
+        out = np.array(list(executor.map(_process_group, groups)))
 
     result = pd.DataFrame(out, index=unique_ids, columns=available_cols)
     result.index.name = 'Identifier'

stats/entropy.py

@@ -4,6 +4,7 @@
 
 import argparse
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
 
 import numpy as np
 import pandas as pd
@@ -63,10 +64,14 @@ def calculate_byte_entropy(df: pd.DataFrame) -> pd.DataFrame:
     else:
         groups = []
 
-    out = np.zeros((len(unique_ids), n_cols), dtype=np.float64)
-    for gi, sub in enumerate(groups):
+    def _process_group(sub):
+        res = np.zeros(n_cols, dtype=np.float64)
         for ci in range(n_cols):
-            out[gi, ci] = _entropy_col(sub[:, ci])
+            res[ci] = _entropy_col(sub[:, ci])
+        return res
+
+    with ThreadPoolExecutor() as executor:
+        out = np.array(list(executor.map(_process_group, groups)))
 
     result = pd.DataFrame(out, index=unique_ids, columns=available_cols)
     result.index.name = 'Identifier'