Matrix completion

Low-rank matrix completion with SoftImpute and probabilistic PCA.

Low-rank matrix completion imputers.

Both algorithms are implemented directly on NumPy so the package does not depend on unmaintained third-party solvers.

SoftImputeImputer

SoftImputeImputer(max_iter: int = 100, init_fill_method: str = 'zero', shrinkage_value: float | None = None, convergence_threshold: float = 0.001, on_error: OnError = None)

Bases: BaseImputer

Impute missing values by low-rank matrix completion (SoftImpute).

Missing entries are initialised with init_fill_method and then iteratively replaced by a low-rank SVD reconstruction whose singular values are soft-thresholded, which is equivalent to nuclear-norm regularisation. Columns with no observed values are left untouched.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> from imputation_methods import SoftImputeImputer
>>> df = pd.DataFrame({"a": [1.0, np.nan, 3.0], "b": [4.0, 5.0, np.nan]})
>>> bool(SoftImputeImputer().impute(df).notna().all().all())
True

References

Mazumder, R., Hastie, T., & Tibshirani, R. (2010). Spectral regularization algorithms for learning large incomplete matrices. Journal of Machine Learning Research, 11, 2287-2322.

Initialize the imputer.

Parameters:

Name	Type	Description	Default
max_iter	int	Maximum number of SVD iterations.	100
init_fill_method	str	How to initialise missing entries before solving: "zero", "mean", "median" or "min".	'zero'
shrinkage_value	float | None	Amount subtracted from each singular value. Defaults to 1/50 of the largest singular value of the initial fill.	None
convergence_threshold	float	Stop when the relative change of the imputed entries between iterations falls below this value.	0.001
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 an argument is out of range.

Source code in src/imputation_methods/matrix.py

@renamed_parameters(max_iters="max_iter")
def __init__(
    self,
    max_iter: int = 100,
    init_fill_method: str = "zero",
    shrinkage_value: float | None = None,
    convergence_threshold: float = 1e-3,
    on_error: OnError = None,
) -> None:
    """Initialize the imputer.

    Args:
        max_iter: Maximum number of SVD iterations.
        init_fill_method: How to initialise missing entries before solving:
            ``"zero"``, ``"mean"``, ``"median"`` or ``"min"``.
        shrinkage_value: Amount subtracted from each singular value. Defaults
            to 1/50 of the largest singular value of the initial fill.
        convergence_threshold: Stop when the relative change of the imputed
            entries between iterations falls below this value.
        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 an argument is out of range.
    """
    if init_fill_method not in _INIT_FILL_METHODS:
        raise ValueError(
            f"init_fill_method must be one of {_INIT_FILL_METHODS}, "
            f"got {init_fill_method!r}"
        )
    if max_iter < 1:
        raise ValueError(f"max_iter must be >= 1, got {max_iter}")
    if shrinkage_value is not None and shrinkage_value < 0:
        raise ValueError(f"shrinkage_value must be >= 0, got {shrinkage_value}")
    self.max_iter = max_iter
    self.init_fill_method = init_fill_method
    self.shrinkage_value = shrinkage_value
    self.convergence_threshold = convergence_threshold
    self.on_error = check_on_error(on_error)

impute

impute(df: DataFrame) -> DataFrame

Fill missing values using a low-rank matrix approximation.

Parameters:

Name	Type	Description	Default
df	DataFrame	Dataframe with potential NaN values.	required

Returns:

Type	Description
DataFrame	Dataframe with missing entries imputed by SoftImpute.

Source code in src/imputation_methods/matrix.py

@preserve_dtypes
def impute(self, df: pd.DataFrame) -> pd.DataFrame:
    """Fill missing values using a low-rank matrix approximation.

    Args:
        df: Dataframe with potential NaN values.

    Returns:
        Dataframe with missing entries imputed by SoftImpute.
    """
    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

    try:
        completed = _soft_impute(
            result[modelled].to_numpy(dtype=float, na_value=np.nan),
            shrinkage_value=self.shrinkage_value,
            max_iter=self.max_iter,
            convergence_threshold=self.convergence_threshold,
            init_fill_method=self.init_fill_method,
        )
    except 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)

    result[modelled] = completed
    return result

PPCAImputer

PPCAImputer(n_components: int | None = 1, min_obs: int = 1, max_iter: int = 500, tol: float = 1e-06, on_error: OnError = None)

Bases: BaseImputer

Impute missing values with probabilistic PCA (PPCA).

Columns are standardised, then a PPCA model is fitted by EM while the missing entries are repeatedly replaced by their expected value under the model. The model is the maximum-likelihood PPCA of Tipping & Bishop; no priors are placed on the loadings.

Examples:

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

References

Tipping, M. E., & Bishop, C. M. (1999). Probabilistic principal component analysis. Journal of the Royal Statistical Society: Series B, 61(3), 611-622.

Initialize the imputer.

Parameters:

Name	Type	Description	Default
n_components	int | None	Number of latent dimensions. None uses 2. The value is capped at n_columns - 1 so the model stays low-rank.	1
min_obs	int	Minimum number of observed values a column needs to take part in the model. Other columns are mean-imputed.	1
max_iter	int	Maximum number of EM iterations.	500
tol	float	Stop when the relative change of the imputed entries between iterations falls below this value.	1e-06
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 an argument is out of range.

Source code in src/imputation_methods/matrix.py

def __init__(
    self,
    n_components: int | None = 1,
    min_obs: int = 1,
    max_iter: int = 500,
    tol: float = 1e-6,
    on_error: OnError = None,
) -> None:
    """Initialize the imputer.

    Args:
        n_components: Number of latent dimensions. ``None`` uses 2. The value
            is capped at ``n_columns - 1`` so the model stays low-rank.
        min_obs: Minimum number of observed values a column needs to take
            part in the model. Other columns are mean-imputed.
        max_iter: Maximum number of EM iterations.
        tol: Stop when the relative change of the imputed entries between
            iterations falls below this value.
        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 an argument is out of range.
    """
    if n_components is not None and n_components < 1:
        raise ValueError(f"n_components must be >= 1, got {n_components}")
    if min_obs < 1:
        raise ValueError(f"min_obs must be >= 1, got {min_obs}")
    if max_iter < 1:
        raise ValueError(f"max_iter must be >= 1, got {max_iter}")
    self.n_components = n_components
    self.min_obs = min_obs
    self.max_iter = max_iter
    self.tol = tol
    self.on_error = check_on_error(on_error)

impute

impute(df: DataFrame) -> DataFrame

Fill missing values using a probabilistic PCA model.

Parameters:

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

Returns:

Type	Description
DataFrame	Dataframe with missing entries filled via probabilistic PCA. A
DataFrame	single-column dataframe is returned unchanged, since PPCA needs at
DataFrame	least two columns.

Source code in src/imputation_methods/matrix.py

@preserve_dtypes
def impute(self, df: pd.DataFrame) -> pd.DataFrame:
    """Fill missing values using a probabilistic PCA model.

    Args:
        df: Dataframe with missing values.

    Returns:
        Dataframe with missing entries filled via probabilistic PCA. A
        single-column dataframe is returned unchanged, since PPCA needs at
        least two columns.
    """
    df = self._ensure_numeric(df)
    if df.shape[1] == 1:
        logger.warning("PPCA needs at least two columns; returning input copy")
        return df.copy()

    result = df.astype(float)
    enough_obs = result.notna().sum() >= self.min_obs
    modelled = result.columns[enough_obs]
    if len(modelled) < 2:
        logger.warning(
            "Fewer than two columns have %d observations; using mean imputation",
            self.min_obs,
        )
        return MeanImputer().impute(df)

    values = result[modelled].to_numpy(dtype=float, na_value=np.nan)
    means = np.nanmean(values, axis=0)
    stds = np.nanstd(values, axis=0)
    stds[~(stds > 0)] = 1.0
    requested = 2 if self.n_components is None else self.n_components
    n_components = min(requested, len(modelled) - 1)

    try:
        completed = _ppca_impute(
            (values - means) / stds,
            n_components=n_components,
            max_iter=self.max_iter,
            tol=self.tol,
        )
    except 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)

    result[modelled] = completed * stds + means
    excluded = result.columns[~enough_obs]
    if len(excluded) > 0:
        result[excluded] = MeanImputer().impute(result[excluded])
    return result