plugin_ice.py

# stdlib
from typing import Any, List

# third party
import pandas as pd

# hyperimpute absolute
import hyperimpute.plugins.core.params as params
import hyperimpute.plugins.imputers.base as base
from hyperimpute.plugins.imputers.plugin_hyperimpute import plugin as base_model


class IterativeChainedEquationsPlugin(base.ImputerPlugin):
    """Imputation plugin for completing missing values using the Multivariate Iterative chained equations Imputation strategy.

    Method:
        Multivariate Iterative chained equations(MICE) methods model each feature with missing values as a function of other features in a round-robin fashion. For each step of the round-robin imputation, we use a BayesianRidge estimator, which does a regularized linear regression.

    Args:
        max_iter: int, default=500
            maximum number of imputation rounds to perform.
        random_state: int, default set to the current time.
            seed of the pseudo random number generator to use.

    Example:
        >>> import numpy as np
        >>> from hyperimpute.plugins.imputers import Imputers
        >>> plugin = Imputers().get("ice")
        >>> plugin.fit_transform([[1, 1, 1, 1], [np.nan, np.nan, np.nan, np.nan], [1, 2, 2, 1], [2, 2, 2, 2]])

    Reference: "mice: Multivariate Imputation by Chained Equations in R", Stef van Buuren, Karin Groothuis-Oudshoorn
    """

    def __init__(
        self,
        max_iter: int = 1000,
        initial_strategy: int = 0,
        imputation_order: int = 0,
        random_state: int = 0,
    ) -> None:
        super().__init__(random_state=random_state)

        self.max_iter = max_iter
        self.initial_strategy = initial_strategy
        self.imputation_order = imputation_order
        self.random_state = random_state

        self._model = base_model(
            classifier_seed=["logistic_regression"],
            regression_seed=["linear_regression"],
            imputation_order=imputation_order,
            baseline_imputer=initial_strategy,
            random_state=random_state,
            n_inner_iter=max_iter,
            class_threshold=5,
        )

    @staticmethod
    def name() -> str:
        return "ice"

    @staticmethod
    def hyperparameter_space(*args: Any, **kwargs: Any) -> List[params.Params]:
        return [
            params.Integer("max_iter", 100, 1000, 100),
            params.Integer(
                "initial_strategy",
                0,
                len(base_model.initial_strategy_vals) - 1,
            ),
            params.Integer(
                "imputation_order",
                0,
                len(base_model.imputation_order_vals) - 1,
            ),
        ]

    def _fit(
        self, X: pd.DataFrame, *args: Any, **kwargs: Any
    ) -> "IterativeChainedEquationsPlugin":
        self._model.fit(X, *args, **kwargs)

        return self

    def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
        return self._model.transform(X)


plugin = IterativeChainedEquationsPlugin