86 lines
2.4 KiB
Python
86 lines
2.4 KiB
Python
from sklearn.utils.extmath import randomized_svd
|
|
import numpy as np
|
|
from functools import partial
|
|
from tqdm import tqdm
|
|
from sklearn.base import BaseEstimator, TransformerMixin
|
|
|
|
|
|
def np_svd(M, k):
|
|
u, s, v = np.linalg.svd(M, full_matrices=False)
|
|
return u[:, :k], s[:k], v[:k, :]
|
|
|
|
|
|
def inverse_svd(u, s, v):
|
|
return np.dot(u * s, v)
|
|
|
|
|
|
def em_svd(X, n_iter, svd_f):
|
|
nan_mask = np.isnan(X)
|
|
|
|
# Fill NaNs in the columns first with means
|
|
mu = np.nanmean(X, axis=0)
|
|
X[nan_mask] = (mu[None, :] * np.ones(X.shape))[nan_mask]
|
|
|
|
# Iteratively replace the NaN indexes with an SVD approximation
|
|
# in an Expectation Maximization manner.
|
|
for _ in range(n_iter):
|
|
M_approx = inverse_svd(*svd_f(X))
|
|
|
|
if np.allclose(X[nan_mask], M_approx[nan_mask], rtol=0.01, atol=0.01):
|
|
break
|
|
X[nan_mask] = M_approx[nan_mask]
|
|
|
|
return X
|
|
|
|
|
|
class SVDImputer(BaseEstimator, TransformerMixin):
|
|
def __init__(
|
|
self,
|
|
n_rows,
|
|
k=4,
|
|
n_iter_fit=20,
|
|
n_iter_impute=5,
|
|
svd_method="randomized",
|
|
block_size=200,
|
|
):
|
|
self.k = k
|
|
self.svd_method = svd_method
|
|
self.svd = (
|
|
partial(np_svd, k=k)
|
|
if svd_method == "numpy"
|
|
else partial(randomized_svd, n_components=k)
|
|
)
|
|
self.n_rows = int(n_rows)
|
|
self.n_iter_fit = int(n_iter_fit)
|
|
self.n_iter_impute = n_iter_impute
|
|
self.block_size = int(block_size)
|
|
self.nan_filled_X = None
|
|
|
|
def fit(self, X, y=None):
|
|
assert self.n_rows < X.shape[0]
|
|
X = np.array(X, dtype=np.float32)
|
|
|
|
# Sort from least missing to most missing.
|
|
idx = np.isnan(X).sum(axis=1).argsort()
|
|
# Use only the least missing rows for fitting.
|
|
X = X[idx[: self.n_rows], :]
|
|
|
|
self.nan_filled_X = em_svd(X, self.n_iter_fit, self.svd)
|
|
|
|
def transform(self, X, y=None):
|
|
X = np.array(X, dtype=np.float32)
|
|
|
|
for i in tqdm(range(X.shape[0] // self.block_size + 1)):
|
|
start = i * self.block_size
|
|
stop = start + self.block_size
|
|
|
|
block = X[start:stop, :]
|
|
M = np.concatenate([block, self.nan_filled_X])
|
|
imputed = em_svd(M, n_iter=self.n_iter_impute, svd_f=self.svd)
|
|
X[start:stop, :] = imputed[: block.shape[0], :]
|
|
return X
|
|
|
|
def fit_transform(self, X, y=None, **fit_params):
|
|
self.fit(X)
|
|
return self.transform(X)
|