Parameter validation with pydantic (#16)

* bugfixing

* bump version

* cleanup

Author: bcebere

setup.cfg

@@ -52,6 +52,7 @@ install_requires =
     miracle-imputation
     catboost
     xgboost
+    pydantic
     importlib-metadata; python_version<"3.8"
 
 

src/hyperimpute/plugins/core/base_plugin.py

@@ -9,6 +9,7 @@
 # third party
 from optuna.trial import Trial
 import pandas as pd
+from pydantic import validate_arguments
 
 # hyperimpute absolute
 import hyperimpute.logger as log
@@ -110,20 +111,24 @@ def subtype() -> str:
     def fqdn(cls) -> str:
         return cls.type() + "." + cls.subtype() + "." + cls.name()
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def fit_transform(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame:
         return pd.DataFrame(self.fit(X, *args, *kwargs).transform(X))
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def fit_predict(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame:
         return pd.DataFrame(self.fit(X, *args, *kwargs).predict(X))
 
-    def fit(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> "Plugin":
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
+    def fit(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> Any:
         X = cast.to_dataframe(X)
         return self._fit(X, *args, **kwargs)
 
     @abstractmethod
     def _fit(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> "Plugin":
         ...
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def transform(self, X: pd.DataFrame) -> pd.DataFrame:
         X = cast.to_dataframe(X)
         return pd.DataFrame(self._transform(X))
@@ -132,6 +137,7 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
     def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
         ...
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def predict(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame:
         X = cast.to_dataframe(X)
         return pd.DataFrame(self._predict(X, *args, *kwargs))

src/hyperimpute/plugins/core/params.py

@@ -5,9 +5,11 @@
 # third party
 import numpy as np
 from optuna.trial import Trial
+from pydantic import validate_arguments
 
 
 class Params(metaclass=ABCMeta):
+    @validate_arguments
     def __init__(self, name: str, bounds: Tuple[Any, Any]) -> None:
         self.name = name
         self.bounds = bounds
@@ -26,6 +28,7 @@ def sample_np(self) -> Any:
 
 
 class Categorical(Params):
+    @validate_arguments
     def __init__(self, name: str, choices: List[Any]) -> None:
         super().__init__(name, (min(choices), max(choices)))
         self.name = name
@@ -42,6 +45,7 @@ def sample_np(self) -> Any:
 
 
 class Float(Params):
+    @validate_arguments
     def __init__(self, name: str, low: float, high: float) -> None:
         low = float(low)
         high = float(high)
@@ -54,6 +58,7 @@ def __init__(self, name: str, low: float, high: float) -> None:
     def get(self) -> List[Any]:
         return [self.name, self.low, self.high]
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def sample(self, trial: Trial) -> float:
         return trial.suggest_float(self.name, self.low, self.high)
 
@@ -62,6 +67,7 @@ def sample_np(self) -> Any:
 
 
 class Integer(Params):
+    @validate_arguments
     def __init__(self, name: str, low: int, high: int, step: int = 1) -> None:
         self.low = low
         self.high = high
@@ -77,6 +83,7 @@ def __init__(self, name: str, low: int, high: int, step: int = 1) -> None:
     def get(self) -> List[Any]:
         return [self.name, self.low, self.high, self.step]
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def sample(self, trial: Trial) -> Any:
         return trial.suggest_int(self.name, self.low, self.high, self.step)
 

src/hyperimpute/plugins/imputers/plugin_hyperimpute.py

@@ -9,6 +9,7 @@
 import numpy as np
 import optuna
 import pandas as pd
+from pydantic import validate_arguments
 from sklearn.impute import MissingIndicator
 from sklearn.metrics import mean_squared_error
 from sklearn.preprocessing import LabelEncoder
@@ -61,6 +62,7 @@ def default(self, obj: Any) -> Any:
 
 
 class HyperbandOptimizer:
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def __init__(
         self,
         name: str,
@@ -112,6 +114,7 @@ def _hash_dict(self, name: str, dict_val: dict) -> str:
             {"name": name, "val": dict_val}, sort_keys=True, cls=NpEncoder
         )
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _sample_model(self, name: str, n: int) -> list:
         hashed = self._hash_dict(name, {})
         result: List[Tuple] = []
@@ -132,19 +135,22 @@ def _sample_model(self, name: str, n: int) -> list:
 
         return result
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _sample(self, n: int) -> list:
         results = []
         for name in self.seeds:
             results.extend(self._sample_model(name, n))
         return results
 
-    def _baseline(self, X: pd.DataFrame, y: pd.DataFrame) -> None:
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
+    def _baseline(self, X: pd.DataFrame, y: pd.Series) -> None:
         for seed in self.seeds:
             self._eval_params(seed, X, y, hyperparam_search_iterations=1)
         # TODO: balance methods
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _eval_params(
-        self, model_name: str, X: pd.DataFrame, y: pd.DataFrame, **params: Any
+        self, model_name: str, X: pd.DataFrame, y: pd.Series, **params: Any
     ) -> float:
         model = self.predictions.get(model_name, **params)
         for n_folds in [2, 1]:
@@ -171,9 +177,8 @@ def _eval_params(
 
         return score
 
-    def evaluate(
-        self, X: pd.DataFrame, y: pd.DataFrame
-    ) -> Tuple[PredictionPlugin, float]:
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
+    def evaluate(self, X: pd.DataFrame, y: pd.Series) -> Tuple[PredictionPlugin, float]:
         self._reset()
         self._baseline(X, y)
 
@@ -236,6 +241,7 @@ def evaluate(
 
 
 class BayesianOptimizer:
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def __init__(
         self,
         name: str,
@@ -272,11 +278,12 @@ def __init__(
         self.best_candidate = self.seeds[0]
         self.best_params: dict = {}
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def evaluate_plugin(
         self,
         plugin_name: str,
         X: pd.DataFrame,
-        y: pd.DataFrame,
+        y: pd.Series,
         prev_best_score: float,
     ) -> tuple:
         # BO evaluation for a single plugin
@@ -333,8 +340,9 @@ def objective(trial: optuna.Trial) -> float:
 
         return study.best_value, study.best_trial.params
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def evaluate(
-        self, X_train: pd.DataFrame, y_train: pd.DataFrame
+        self, X_train: pd.DataFrame, y_train: pd.Series
     ) -> Tuple[PredictionPlugin, float]:
         best_score = self.failure_score
 
@@ -366,6 +374,7 @@ def evaluate(
 
 
 class SimpleOptimizer:
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def __init__(
         self,
         name: str,
@@ -398,8 +407,9 @@ def __init__(
         for seed in self.seeds:
             self.model_best_score[seed] = -np.inf
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _eval_params(
-        self, model_name: str, X: pd.DataFrame, y: pd.DataFrame, **params: Any
+        self, model_name: str, X: pd.DataFrame, y: pd.Series, **params: Any
     ) -> float:
         model = self.predictions.get(
             model_name, random_state=self.random_state, **params
@@ -430,9 +440,8 @@ def _eval_params(
 
         return score
 
-    def evaluate(
-        self, X: pd.DataFrame, y: pd.DataFrame
-    ) -> Tuple[PredictionPlugin, float]:
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
+    def evaluate(self, X: pd.DataFrame, y: pd.Series) -> Tuple[PredictionPlugin, float]:
         for seed in self.seeds:
             self._eval_params(seed, X, y)
         log.info(
@@ -447,6 +456,7 @@ def evaluate(
 
 
 class IterativeErrorCorrection:
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def __init__(
         self,
         study: str,
@@ -530,10 +540,12 @@ def intersect_or_right(left: list, right: list) -> list:
             "regression_seed": reg,
         }
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _setup(self, X: pd.DataFrame) -> pd.DataFrame:
         # Encode the categorical columns
         # Reset internal caches
         X = pd.DataFrame(X).copy()
+        X.columns = X.columns.map(str)
 
         self.mask = self._missing_indicator(X)
 
@@ -590,25 +602,29 @@ def _setup(self, X: pd.DataFrame) -> pd.DataFrame:
 
         return X
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _tear_down(self, X: pd.DataFrame) -> pd.DataFrame:
         # Revert the encoding after processing the data
         for col in self.encoders:
             X[col] = self.encoders[col].inverse_transform(X[col].astype(int))
 
         return X
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _get_neighbors_for_col(self, col: str) -> list:
         covs = list(self.all_cols)
         covs.remove(col)
 
         return covs
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _is_same_type(self, lhs: str, rhs: str) -> bool:
         ltype = lhs in self.categorical_cols
         rtype = rhs in self.categorical_cols
 
         return ltype == rtype
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _check_similar(self, X: pd.DataFrame, col: str) -> Any:
         if not self.select_lazy:
             return None
@@ -630,6 +646,7 @@ def _check_similar(self, X: pd.DataFrame, col: str) -> Any:
 
         return None
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _optimize_model_for_column(self, X: pd.DataFrame, col: str) -> float:
         # BO evaluation for a single column
         if self.mask[col].sum() == 0:
@@ -666,6 +683,7 @@ def _optimize_model_for_column(self, X: pd.DataFrame, col: str) -> float:
 
         return score
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _optimize(self, X: pd.DataFrame) -> float:
         # BO evaluation to select the best models for each columns
         if self.select_model_by_iteration:
@@ -677,6 +695,7 @@ def _optimize(self, X: pd.DataFrame) -> float:
 
         return iteration_score
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _impute_single_column(
         self, X: pd.DataFrame, col: str, train: bool
     ) -> pd.DataFrame:
@@ -723,14 +742,17 @@ def _get_imputation_order(self) -> list:
             random.shuffle(self.imputation_order)
             return self.imputation_order
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _initial_imputation(self, X: pd.DataFrame) -> pd.DataFrame:
         # Use baseline imputer for initial values
         return self.baseline_imputer.fit_transform(X)
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _is_categorical(self, X: pd.DataFrame, col: str) -> bool:
         # Helper for filtering categorical columns
         return len(X[col].unique()) < self.class_threshold
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _missing_indicator(self, X: pd.DataFrame) -> pd.DataFrame:
         # Helper for generating missingness mask
         return pd.DataFrame(
@@ -739,6 +761,7 @@ def _missing_indicator(self, X: pd.DataFrame) -> pd.DataFrame:
             index=X.index,
         )
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def _fit_transform_inner_optimization(self, X: pd.DataFrame) -> pd.DataFrame:
         log.info("  > HyperImpute using inner optimization")
         best_obj_score = -10e10
@@ -780,6 +803,7 @@ def _fit_transform_inner_optimization(self, X: pd.DataFrame) -> pd.DataFrame:
 
         return X
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def fit_transform(self, X: pd.DataFrame) -> pd.DataFrame:
         # Run imputation
         X = self._setup(X)

src/hyperimpute/plugins/imputers/plugin_miracle.py

@@ -95,8 +95,7 @@ def _fit(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> "MiraclePlugin":
         return self
 
     def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
-        X = np.asarray(X)
-        missing_idxs = np.where(np.any(np.isnan(np.asarray(X)), axis=0))[0]
+        missing_idxs = np.where(np.any(np.isnan(X.values), axis=0))[0]
 
         _model = MIRACLE(
             num_inputs=X.shape[1],
@@ -117,7 +116,7 @@ def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
         seed_imputer = self._get_seed_imputer(self.seed_imputation)
         X_seed = seed_imputer.fit_transform(X)
 
-        return _model.fit(X, X_seed=X_seed)
+        return _model.fit(X.values, X_seed=X_seed.values)
 
     def save(self) -> bytes:
         return b""

src/hyperimpute/plugins/prediction/base.py

@@ -4,6 +4,7 @@
 
 # third party
 import pandas as pd
+from pydantic import validate_arguments
 
 # hyperimpute absolute
 import hyperimpute.plugins.core.base_plugin as plugin
@@ -37,12 +38,15 @@ def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
             "Prediction plugins do not implement the 'transform' method"
         )
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def score(self, X: pd.DataFrame, y: pd.DataFrame, metric: str = "aucroc") -> float:
         raise NotImplementedError(f"Score not implemented for {self.name()}")
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def explain(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame:
         raise NotImplementedError(f"Explainer not implemented for {self.name()}")
 
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def predict_proba(self, X: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame:
         X = cast.to_dataframe(X)
         return pd.DataFrame(self._predict_proba(X, *args, **kwargs))