Skip to content

Instantly share code, notes, and snippets.

@sazio
Created August 29, 2020 12:11
Show Gist options
  • Save sazio/cdc5ad045de12964d7e143f1488993b9 to your computer and use it in GitHub Desktop.
Save sazio/cdc5ad045de12964d7e143f1488993b9 to your computer and use it in GitHub Desktop.
import numpy as np
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
# max iter corresponds to imputation rounds
imp = IterativeImputer(max_iter=10, random_state=0)
X_train = [[1, 2], [3, 6], [4, 8], [np.nan, 3], [7, np.nan]]
imp.fit(X_train)
X_test = [[np.nan, 2], [6, np.nan], [np.nan, 6]]
# the model learns that the second feature is double the first
print(np.round(imp.transform(X_test)))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment