restructure feature preprocessing

config/config.json

@@ -2,6 +2,8 @@
     "seed": 0,
     "ecfp_radius": 3,
     "ecfp_fpsize": 8192,
+    "feature_minvar": 0.01,
+    "feature_maxcorr": 0.95,
     "model_path": "checkpoints/rf_alltasks.joblib",
     "data_folder": "data/",
     "log_folder": "logs/",

src/model.py

@@ -11,12 +11,10 @@ import joblib
 
 import numpy as np
 
-from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.feature_selection import VarianceThreshold
 from sklearn.preprocessing import StandardScaler
-from statsmodels.distributions.empirical_distribution import ECDF
 
+from .preprocess import get_feature_selection, get_ecdfs, create_quantiles
 from .utils import TASKS
 
 
@@ -45,8 +43,8 @@ class Tox21RFClassifier:
             )
             for task in self.tasks
         }
-        self.feature_selection = FeatureSelector()
-        self.ecdf = ECDFQuantileCreator(rdkit_desc_idxs)
+        self.feature_selection = None
+        self.ecdfs = None
         self.scaler = StandardScaler()
 
     def load_model(self, path: str) -> None:
@@ -61,9 +59,8 @@ class Tox21RFClassifier:
         self.scaler = model["scalers"]
         self.rdkit_desc_idxs = model["rdkit_desc_idxs"]
 
-        self.feature_selection.feature_selection = model["feature_selections"]
-        self.ecdf.ecdfs = model["ecdfs"]
-        self.ecdf.to_adapt = model["rdkit_desc_idxs"]
+        self.feature_selection = model["feature_selections"]
+        self.ecdfs = model["ecdfs"]
 
     def save_model(self, path: str) -> None:
         """Saves the model to a given path
@@ -76,26 +73,47 @@ class Tox21RFClassifier:
 
         model = {
             "models": self.models,
-            "feature_selections": self.feature_selection.feature_selection,
-            "ecdfs": self.ecdf.ecdfs,
+            "feature_selections": self.feature_selection,
+            "ecdfs": self.ecdfs,
             "scalers": self.scaler,
             "rdkit_desc_idxs": self.rdkit_desc_idxs,
         }
 
         joblib.dump(model, path)
 
-    def fit_preprocessing(self, X: np.ndarray) -> None:
+    def fit_preprocessing(self, X: np.ndarray, min_var=0.01, max_corr=0.95) -> None:
         X_ = X.copy()
 
-        X_ = self.ecdf.fit_transform(X_)
-        X_ = self.feature_selection.fit_transform(X_)
+        _, n_feat = X.shape
+
+        if self.rdkit_desc_idxs is None:
+            self.rdkit_desc_idxs = np.arange(n_feat)
+        else:
+            assert (
+                self.rdkit_desc_idxs < n_feat
+            ).all(), "passed to_adapt list contains more features than in X!"
+
+        self.ecdfs = get_ecdfs(X_[:, self.rdkit_desc_idxs])
+        X_[:, self.rdkit_desc_idxs] = create_quantiles(
+            X_[:, self.rdkit_desc_idxs], self.ecdfs
+        )
+
+        # get feature selection
+        self.feature_selection = get_feature_selection(
+            X_, min_var=min_var, max_corr=max_corr
+        )
+        X_ = X_[:, self.feature_selection]
+
+        # fit scaler
         X_ = self.scaler.fit(X_)
 
-    def preprocess(self, X: np.ndarray) -> None:
+    def _preprocess(self, X: np.ndarray) -> None:
         X_ = X.copy()
 
-        X_ = self.ecdf.transform(X_)
-        X_ = self.feature_selection.transform(X_)
+        X_[:, self.rdkit_desc_idxs] = create_quantiles(
+            X_[:, self.rdkit_desc_idxs], self.ecdfs
+        )
+        X_ = X_[:, self.feature_selection]
         X_ = self.scaler.transform(X_)
         return X_
 
@@ -109,8 +127,7 @@ class Tox21RFClassifier:
         """
         assert task in self.tasks, f"Unknown task: {task}"
 
-        X_ = self.preprocess(X)
-
+        X_ = self._preprocess(X)
         self.models[task].fit(X_, y)
 
     def predict(self, task: str, X: np.ndarray) -> np.ndarray:
@@ -128,73 +145,5 @@ class Tox21RFClassifier:
             len(X.shape) == 2
         ), f"Function expects 2D np.array. Current shape: {X.shape}"
 
-        X_ = self.preprocess(X)
+        X_ = self._preprocess(X)
         return self.models[task].predict_proba(X_)[:, 1]
-
-
-class FeatureSelector(BaseEstimator, TransformerMixin):
-    def __init__(self, min_var=0.01, max_corr=0.95):
-        self.min_var = min_var
-        self.max_corr = max_corr
-        self.feature_selection = None
-
-    def fit(self, X, y=None):
-        # select features with at least 0.01 variation
-        var_thresh = VarianceThreshold(threshold=self.min_var)
-        feature_selection = var_thresh.fit(X).get_support(
-            indices=True
-        )  # list containing selected feature indices
-
-        n_features_preselected = len(feature_selection)
-
-        # Remove highly correlated features
-        corr_matrix = np.corrcoef(X[:, feature_selection], rowvar=False)
-        upper_tri = np.triu(corr_matrix, k=1)
-        to_keep = np.ones((n_features_preselected,), dtype=bool)
-        for i in range(upper_tri.shape[0]):
-            for j in range(upper_tri.shape[1]):
-                if upper_tri[i, j] > self.max_corr:
-                    to_keep[j] = False
-
-        self.feature_selection = feature_selection[to_keep]
-
-        return self
-
-    def transform(self, X):
-        return X[:, self.feature_selection]
-
-
-class ECDFQuantileCreator(BaseEstimator, TransformerMixin):
-    def __init__(self, to_adapt=None):
-        self.to_adapt = to_adapt
-        self.ecdfs = None
-
-    def fit(self, X, y=None):
-        _, n_feat = X.shape
-
-        if self.to_adapt is None:
-            self.to_adapt = np.arange(n_feat)
-        else:
-            assert (
-                self.to_adapt < n_feat
-            ).all(), "passed to_adapt list contains more features than in X!"
-        selected_feat = X[:, self.to_adapt].copy()
-        ecdfs = []
-        for column in range(selected_feat.shape[1]):
-            raw_values = selected_feat[:, column].reshape(-1)
-            ecdfs.append(ECDF(raw_values))
-        self.ecdfs = ecdfs
-        return self
-
-    def transform(self, X):
-        selected_feat = X[:, self.to_adapt].copy()
-        quantiles = np.zeros_like(selected_feat)
-
-        for column in range(selected_feat.shape[1]):
-            raw_values = selected_feat[:, column].reshape(-1)
-            ecdf = self.ecdfs[column]
-            q = ecdf(raw_values)
-            quantiles[:, column] = q
-
-        X[:, self.to_adapt] = quantiles
-        return X

src/preprocess.py

@@ -13,6 +13,7 @@ import pandas as pd
 
 from datasets import load_dataset
 from sklearn.feature_selection import VarianceThreshold
+from statsmodels.distributions.empirical_distribution import ECDF
 
 from rdkit import Chem, DataStructs
 from rdkit.Chem import Descriptors, rdFingerprintGenerator, MACCSkeys
@@ -275,6 +276,14 @@ def create_descriptors(
     return features, clean_mol_mask
 
 
+def get_ecdfs(raw_features: np.ndarray, **kwargs) -> np.ndarray:
+    ecdfs = []
+    for column in range(raw_features.shape[1]):
+        raw_values = raw_features[:, column].reshape(-1)
+        ecdfs.append(ECDF(raw_values))
+    return ecdfs
+
+
 def get_feature_selection(
     raw_features: np.ndarray, min_var=0.01, max_corr=0.95, **kwargs
 ) -> np.ndarray:

train.py

@@ -88,7 +88,9 @@ def main(cfg):
         task_config=task_configs,
         rdkit_desc_idxs=rdkit_descr_idxs,
     )
-    model.fit_preprocessing(data)
+    model.fit_preprocessing(
+        data, min_var=cfg["feature_minvar"], max_corr=cfg["feature_maxcorr"]
+    )
 
     logger.info("Start training.")
     for i, task in enumerate(model.tasks):