Nearest neighbors

Borrow values from similar rows.

Distance-based imputers that borrow values from similar rows.

KNNImputer

KNNImputer(n_neighbors: int = 5, on_error: OnError = None)

Bases: BaseImputer

Impute missing values using K-nearest neighbors.

Examples:

>>> import pandas as pd
>>> import numpy as np
>>> from imputation_methods import KNNImputer
>>> df = pd.DataFrame({"a": [1, 2, np.nan, 4], "b": [5, np.nan, 7, 8]})
>>> imputer = KNNImputer(n_neighbors=2)
>>> imputed = imputer.impute(df)
>>> assert not imputed.isna().any().any()

Initialize the imputer.

Parameters:

Name	Type	Description	Default
n_neighbors	int	Number of neighbors to consider.	5
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

Raises:

Type	Description
ValueError	If n_neighbors is not a positive integer.

Source code in src/imputation_methods/neighbors.py

@renamed_parameters(k="n_neighbors")
def __init__(self, n_neighbors: int = 5, on_error: OnError = None) -> None:
    """Initialize the imputer.

    Args:
        n_neighbors: Number of neighbors to consider.
        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.

    Raises:
        ValueError: If n_neighbors is not a positive integer.
    """
    if not isinstance(n_neighbors, int):
        raise TypeError(
            f"n_neighbors must be an integer, got {type(n_neighbors).__name__}"
        )
    if n_neighbors <= 0:
        raise ValueError(f"n_neighbors must be positive, got {n_neighbors}")
    self.n_neighbors = n_neighbors
    self._imputer = _SklearnKNNImputer(n_neighbors=n_neighbors)
    self.on_error = check_on_error(on_error)

impute

impute(df: DataFrame) -> DataFrame

Impute using the fitted KNN strategy.

Parameters:

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

Returns:

Type	Description
DataFrame	Imputed dataframe.

Raises:

Type	Description
ValueError	If dataframe is empty or has insufficient data for KNN.
RuntimeError	If KNN imputation fails.

Source code in src/imputation_methods/neighbors.py

@preserve_dtypes
def impute(self, df: pd.DataFrame) -> pd.DataFrame:
    """Impute using the fitted KNN strategy.

    Args:
        df: Dataframe with missing values.

    Returns:
        Imputed dataframe.

    Raises:
        ValueError: If dataframe is empty or has insufficient data for KNN.
        RuntimeError: If KNN imputation fails.
    """
    df = self._ensure_numeric(df)
    if not df.isna().any().any():
        return df.copy()

    try:
        return fit_transform_non_empty(self._imputer, df)
    except ValueError 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

RadiusNeighborsImputer

RadiusNeighborsImputer(radius: float = 1.0, weights: str = 'distance', metric: str = 'euclidean', on_error: OnError = None)

Bases: BaseImputer

Radius-based neighbors imputation using distance threshold.

Imputes using all neighbors within a specified radius rather than a fixed number of neighbors. Adaptive to local density.

Parameters:

Name	Type	Description	Default
radius	float	Distance threshold for neighbors. Default: 1.0	1.0
weights	str	Weight function ('uniform' or 'distance'). Default: 'distance'	'distance'
metric	str	Distance metric. Default: 'euclidean'	'euclidean'

Examples:

>>> import pandas as pd
>>> import numpy as np
>>> from imputation_methods import RadiusNeighborsImputer
>>> df = pd.DataFrame({
...     'a': [1, 2, np.nan, 4, 5],
...     'b': [2, 4, 6, np.nan, 10]
... })
>>> imputer = RadiusNeighborsImputer(radius=2.0)
>>> imputed = imputer.impute(df)

References

Radius-based neighborhood for adaptive local imputation.

Initialize the radius neighbors imputer.

Parameters:

Name	Type	Description	Default
radius	float	Distance threshold	1.0
weights	str	Weighting function	'distance'
metric	str	Distance metric	'euclidean'
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 for that column. The default, None, falls back with a FutureWarning; it will change to "raise" in 1.0.0.	None

Source code in src/imputation_methods/neighbors.py

def __init__(
    self,
    radius: float = 1.0,
    weights: str = "distance",
    metric: str = "euclidean",
    on_error: OnError = None,
) -> None:
    """Initialize the radius neighbors imputer.

    Args:
        radius: Distance threshold
        weights: Weighting function
        metric: Distance metric
        on_error: What to do if the model can't be fitted: ``"raise"`` an
            :class:`~imputation_methods.ImputationError`, or ``"fallback"``
            to use mean imputation for that column. The default, ``None``,
            falls back with a ``FutureWarning``; it will change to
            ``"raise"`` in 1.0.0.
    """
    self.radius = radius
    self.weights = weights
    self.metric = metric
    self.on_error = check_on_error(on_error)

impute

impute(df: DataFrame) -> DataFrame

Impute using radius neighbors.

Parameters:

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

Returns:

Type	Description
DataFrame	Imputed dataframe.

Source code in src/imputation_methods/neighbors.py

@preserve_dtypes
def impute(self, df: pd.DataFrame) -> pd.DataFrame:
    """Impute using radius neighbors.

    Args:
        df: Dataframe with missing values.

    Returns:
        Imputed dataframe.
    """
    df = self._ensure_numeric(df)
    result = df.copy()

    # For each column with missing values
    for column in result.columns:
        if result[column].isna().any():
            train_mask = ~result[column].isna()
            predict_mask = result[column].isna()

            if train_mask.sum() == 0:
                result[column] = result[column].fillna(result[column].mean())
                continue

            feature_cols = [c for c in result.columns if c != column]
            if len(feature_cols) == 0:
                result[column] = result[column].fillna(result[column].mean())
                continue

            X_train = result.loc[train_mask, feature_cols].fillna(0).values
            y_train = result.loc[train_mask, column].values
            X_predict = result.loc[predict_mask, feature_cols].fillna(0).values

            if len(X_train) > 0 and len(X_predict) > 0:
                try:
                    model = RadiusNeighborsRegressor(
                        radius=self.radius, weights=self.weights, metric=self.metric
                    )
                    model.fit(X_train, y_train)
                    predictions = model.predict(X_predict)
                    result.loc[predict_mask, column] = predictions
                except Exception as e:
                    raise_or_fall_back(
                        self.on_error,
                        imputer=type(self).__name__,
                        error=e,
                        fallback="mean imputation",
                        column=column,
                    )
                    result.loc[predict_mask, column] = result[column].mean()

    return result

LocalMeanImputer

LocalMeanImputer(n_neighbors: int = 5, distance_weight_power: float = 2.0)

Bases: BaseImputer

Local weighted mean imputation based on feature similarity.

Computes weighted average of similar observations, with weights decreasing by distance.

Parameters:

Name	Type	Description	Default
n_neighbors	int	Number of neighbors to consider. Default: 5	5
distance_weight_power	float	Power for distance weighting. Default: 2.0	2.0

Examples:

>>> import pandas as pd
>>> import numpy as np
>>> from imputation_methods import LocalMeanImputer
>>> df = pd.DataFrame({'a': [1, 2, np.nan, 4, 5]})
>>> imputer = LocalMeanImputer(n_neighbors=3)
>>> imputed = imputer.impute(df)

References

Locally weighted averaging for smooth imputation.

Initialize the local mean imputer.

Parameters:

Name	Type	Description	Default
n_neighbors	int	Number of neighbors	5
distance_weight_power	float	Power for weighting by distance	2.0

Source code in src/imputation_methods/neighbors.py

def __init__(
    self, n_neighbors: int = 5, distance_weight_power: float = 2.0
) -> None:
    """Initialize the local mean imputer.

    Args:
        n_neighbors: Number of neighbors
        distance_weight_power: Power for weighting by distance
    """
    self.n_neighbors = n_neighbors
    self.distance_weight_power = distance_weight_power

impute

impute(df: DataFrame) -> DataFrame

Impute using local weighted mean.

Parameters:

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

Returns:

Type	Description
DataFrame	Imputed dataframe.

Source code in src/imputation_methods/neighbors.py

@preserve_dtypes
def impute(self, df: pd.DataFrame) -> pd.DataFrame:
    """Impute using local weighted mean.

    Args:
        df: Dataframe with missing values.

    Returns:
        Imputed dataframe.
    """
    df = self._ensure_numeric(df)
    result = df.copy()

    for column in result.columns:
        if result[column].isna().any():
            if result[column].notna().sum() == 0:
                # No observed values to learn from: leave the column as NaN.
                continue
            for idx in result[result[column].isna()].index:
                # Get feature values for this row (excluding target column)
                feature_cols = [c for c in result.columns if c != column]
                if len(feature_cols) == 0:
                    result.loc[idx, column] = result[column].mean()
                    continue

                row_features = result.loc[idx, feature_cols].fillna(0).values

                # Find distances to all complete observations
                complete_mask = ~result[column].isna()
                complete_features = (
                    result.loc[complete_mask, feature_cols].fillna(0).values
                )
                complete_values = result.loc[complete_mask, column].to_numpy(
                    dtype=float
                )

                # Compute Euclidean distances
                distances = np.sqrt(
                    np.sum((complete_features - row_features) ** 2, axis=1)
                )

                # Get k nearest neighbors
                k = min(self.n_neighbors, len(distances))
                nearest_idx = np.argsort(distances)[:k]

                # Compute weights (inverse distance)
                nearest_distances = distances[nearest_idx]
                # Avoid division by zero
                nearest_distances = np.maximum(nearest_distances, 1e-10)
                weights = 1.0 / (nearest_distances**self.distance_weight_power)
                weights /= weights.sum()

                # Weighted average
                result.loc[idx, column] = np.sum(
                    weights * complete_values[nearest_idx]
                )

    return result