Neural networks

Neural-network imputers: an autoencoder and generative adversarial imputation (GAIN).

Neural-network imputers.

AutoencoderImputer

AutoencoderImputer(hidden_layer_sizes: tuple[int, ...] = (10,), max_iter: int = 200, random_state: int | None = None, on_error: OnError = None)

Bases: BaseImputer

Impute missing values using a simple autoencoder.

Initialize the imputer.

Parameters:

Name	Type	Description	Default
hidden_layer_sizes	tuple[int, ...]	Architecture of the MLPRegressor used as the autoencoder.	(10,)
max_iter	int	Maximum training iterations.	200
random_state	int | None	Random seed controlling network initialization.	None
on_error	OnError	What to do if the model can't be fitted: "raise" an :class:~imputation_methods.ImputationError, or "fallback" to use mean imputation. The default, None, falls back with a FutureWarning; it will change to "raise" in 1.0.0.	None

Source code in src/imputation_methods/neural.py

def __init__(
    self,
    hidden_layer_sizes: tuple[int, ...] = (10,),
    max_iter: int = 200,
    random_state: int | None = None,
    on_error: OnError = None,
) -> None:
    """Initialize the imputer.

    Args:
        hidden_layer_sizes: Architecture of the ``MLPRegressor`` used as
            the autoencoder.
        max_iter: Maximum training iterations.
        random_state: Random seed controlling network initialization.
        on_error: What to do if the model can't be fitted: ``"raise"`` an
            :class:`~imputation_methods.ImputationError`, or ``"fallback"``
            to use mean imputation. The default, ``None``, falls back with a
            ``FutureWarning``; it will change to ``"raise"`` in 1.0.0.
    """
    self.hidden_layer_sizes = hidden_layer_sizes
    self.max_iter = max_iter
    self.random_state = random_state
    self._model = MLPRegressor(
        hidden_layer_sizes=hidden_layer_sizes,
        activation="relu",
        max_iter=max_iter,
        random_state=random_state,
    )
    self.on_error = check_on_error(on_error)

impute

impute(df: DataFrame) -> DataFrame

Fill missing values using an autoencoder reconstruction.

Parameters:

Name	Type	Description	Default
df	DataFrame	Dataframe with missing values.	required

Returns:

Type	Description
DataFrame	Dataframe with imputed values predicted by the autoencoder. Columns
DataFrame	with no observed values are left as NaN and don't enter the model.

Raises:

Type	Description
ImputationError	If autoencoder training or prediction fails unexpectedly, or fails and on_error is "raise".

Source code in src/imputation_methods/neural.py

@preserve_dtypes
def impute(self, df: pd.DataFrame) -> pd.DataFrame:
    """Fill missing values using an autoencoder reconstruction.

    Args:
        df: Dataframe with missing values.

    Returns:
        Dataframe with imputed values predicted by the autoencoder. Columns
        with no observed values are left as NaN and don't enter the model.

    Raises:
        ImputationError: If autoencoder training or prediction fails
            unexpectedly, or fails and ``on_error`` is ``"raise"``.
    """
    df = self._ensure_numeric(df)
    modelled = df.columns[df.notna().any()]
    if modelled.empty:
        return df.copy()
    try:
        filled = df[modelled].fillna(df[modelled].mean())
        target = filled.to_numpy()
        # scikit-learn expects a 1-D target when there is a single column.
        self._model.fit(filled, target.ravel() if target.shape[1] == 1 else target)
        reconstructed = pd.DataFrame(
            self._model.predict(filled).reshape(len(df), len(modelled)),
            columns=modelled,
            index=df.index,
        )
        result = df.copy()
        result[modelled] = df[modelled].where(df[modelled].notna(), reconstructed)
        return result
    except (ValueError, np.linalg.LinAlgError) as e:
        raise_or_fall_back(
            self.on_error,
            imputer=type(self).__name__,
            error=e,
            fallback="mean imputation",
        )
        return MeanImputer().impute(df)
    except Exception as e:
        raise ImputationError(f"{type(self).__name__} failed: {e}") from e

GAINImputer

GAINImputer(batch_size: int = 128, hint_rate: float = 0.9, alpha: float = 100.0, max_iter: int = 10000, learning_rate: float = 0.001, random_state: int | None = None)

Bases: BaseImputer

Impute missing values with Generative Adversarial Imputation Nets (GAIN).

A generator network fills in the missing entries, while a discriminator tries to tell which entries were observed and which were imputed. A hint vector reveals part of the missingness mask to the discriminator, and a reconstruction loss on the observed entries keeps the generator faithful to the data. Columns are min-max scaled to [0, 1] for training. Both networks have two hidden layers as wide as the number of columns and are trained with Adam, implemented directly on NumPy.

GAIN needs a reasonable amount of data to train; on small datasets simpler methods such as :class:~imputation_methods.MICEImputer are usually more accurate. Columns with no observed values are left untouched.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> from imputation_methods import GAINImputer
>>> df = pd.DataFrame(
...     {"a": [1.0, 2.0, np.nan, 4.0], "b": [2.0, np.nan, 6.0, 8.0]}
... )
>>> imputed = GAINImputer(max_iter=100, random_state=0).impute(df)
>>> bool(imputed.notna().all().all())
True

References

Yoon, J., Jordon, J., & van der Schaar, M. (2018). GAIN: Missing data imputation using generative adversarial nets. ICML, 5689-5698.

Initialize the imputer.

The defaults are those of the reference implementation.

Parameters:

Name	Type	Description	Default
batch_size	int	Rows per training step (capped at the number of rows).	128
hint_rate	float	Probability that each entry of the missingness mask is revealed to the discriminator.	0.9
alpha	float	Weight of the reconstruction loss on observed entries.	100.0
max_iter	int	Number of training steps.	10000
learning_rate	float	Adam learning rate for both networks.	0.001
random_state	int | None	Seed for initialisation, batching, noise and hints.	None

Raises:

Type	Description
ValueError	If an argument is out of range.

Source code in src/imputation_methods/neural.py

@renamed_parameters(iterations="max_iter")
def __init__(
    self,
    batch_size: int = 128,
    hint_rate: float = 0.9,
    alpha: float = 100.0,
    max_iter: int = 10000,
    learning_rate: float = 0.001,
    random_state: int | None = None,
) -> None:
    """Initialize the imputer.

    The defaults are those of the reference implementation.

    Args:
        batch_size: Rows per training step (capped at the number of rows).
        hint_rate: Probability that each entry of the missingness mask is
            revealed to the discriminator.
        alpha: Weight of the reconstruction loss on observed entries.
        max_iter: Number of training steps.
        learning_rate: Adam learning rate for both networks.
        random_state: Seed for initialisation, batching, noise and hints.

    Raises:
        ValueError: If an argument is out of range.
    """
    if batch_size < 1:
        raise ValueError(f"batch_size must be >= 1, got {batch_size}")
    if not 0.0 <= hint_rate <= 1.0:
        raise ValueError(f"hint_rate must be in [0, 1], got {hint_rate}")
    if alpha < 0:
        raise ValueError(f"alpha must be >= 0, got {alpha}")
    if max_iter < 1:
        raise ValueError(f"max_iter must be >= 1, got {max_iter}")
    if learning_rate <= 0:
        raise ValueError(f"learning_rate must be > 0, got {learning_rate}")
    self.batch_size = batch_size
    self.hint_rate = hint_rate
    self.alpha = alpha
    self.max_iter = max_iter
    self.learning_rate = learning_rate
    self.random_state = random_state

impute

impute(df: DataFrame) -> DataFrame

Train GAIN on df and fill its missing values.

Parameters:

Name	Type	Description	Default
df	DataFrame	Dataframe with missing values.	required

Returns:

Type	Description
DataFrame	Imputed dataframe. Observed values are unchanged.

Source code in src/imputation_methods/neural.py

@preserve_dtypes
def impute(self, df: pd.DataFrame) -> pd.DataFrame:
    """Train GAIN on ``df`` and fill its missing values.

    Args:
        df: Dataframe with missing values.

    Returns:
        Imputed dataframe. Observed values are unchanged.
    """
    df = self._ensure_numeric(df)
    result = df.astype(float)
    modelled = result.columns[result.notna().any()]
    if modelled.empty or not result[modelled].isna().any().any():
        return result

    values = result[modelled].to_numpy(dtype=float, na_value=np.nan)
    observed = ~np.isnan(values)
    minimum = np.nanmin(values, axis=0)
    # As in the reference implementation, the small offset keeps constant
    # columns at their value instead of letting the generator move them.
    value_range = np.nanmax(values, axis=0) - minimum + 1e-6
    scaled = np.where(observed, (values - minimum) / value_range, 0.0)

    completed = _gain_impute(
        scaled,
        observed.astype(float),
        batch_size=self.batch_size,
        hint_rate=self.hint_rate,
        alpha=self.alpha,
        max_iter=self.max_iter,
        learning_rate=self.learning_rate,
        rng=np.random.default_rng(self.random_state),
    )
    result[modelled] = np.where(observed, values, completed * value_range + minimum)
    return result