How to use fancyimpute - 10 common examples

def test_brr_like_sklearn():
    n = 10000
    d = 10
    sigma_sqr = 5
    X = np.random.randn(n, d)
    beta_true = np.random.random(d)
    y = np.dot(X, beta_true) + np.sqrt(sigma_sqr) * np.random.randn(n)
    X_tr = X[:n // 2, :]
    y_tr = y[:n // 2]
    X_ts = X[n // 2:, :]
    # prediction with my own bayesian ridge
    lambda_reg = 1
    brr = BayesianRidgeRegression(
        lambda_reg,
        add_ones=True,
        normalize_lambda=False)
    brr.fit(X_tr, y_tr)
    y_ts_brr = brr.predict(X_ts)

    # let's compare to scikit-learn's ridge regression
    rr = Ridge(lambda_reg)
    rr.fit(X_tr, y_tr)
    y_ts_rr = rr.predict(X_ts)

    assert np.mean(np.abs(y_ts_brr - y_ts_rr)) < 0.001, \
        "Predictions are different from sklearn's ridge regression."

def test_iterative_imputer_train_test_with_low_rank_random_matrix():
    XY_incomplete_train = XY_incomplete[:250]
    XY_incomplete_test = XY_incomplete[250:]
    XY_test = XY[250:]
    imputer = IterativeImputer(n_iter=50, random_state=0)
    imputer.fit(XY_incomplete_train)
    XY_completed_test = imputer.transform(XY_incomplete_test)
    _, missing_mae = reconstruction_error(
        XY_test,
        XY_completed_test,
        missing_mask,
        name="IterativeImputer Train/Test")
    assert missing_mae < 0.1, "Error too high with IterativeImputer train/test method!"

def test_br_like_sklearn():
    n = 10000
    d = 10
    sigma_sqr = 5
    X = np.random.randn(n, d)
    beta_true = np.random.random(d)
    y = np.dot(X, beta_true) + np.sqrt(sigma_sqr) * np.random.randn(n)
    X_tr = X[:n / 2, :]
    y_tr = y[:n / 2]
    X_ts = X[n / 2:, :]
    y_ts = y[n / 2:]

    # prediction with my own bayesian ridge
    br = BayesianRegression()
    br.fit(X_tr, y_tr)
    y_ts_br = br.predict(X_ts)
    br_error = np.mean(np.abs(y_ts_br - y_ts))

    # let's compare to scikit-learn's ridge regression
    rr = Ridge(1e-5)
    rr.fit(X_tr, y_tr)
    y_ts_rr = rr.predict(X_ts)
    rr_error = np.mean(np.abs(y_ts_rr - y_ts))

    assert br_error - rr_error < 0.1, \
        "Error is significantly worse than sklearn's ridge regression."

def test_soft_impute_with_low_rank_random_matrix():
    solver = SoftImpute()
    XY_completed = solver.fit_transform(XY_incomplete)
    _, missing_mae = reconstruction_error(
        XY,
        XY_completed,
        missing_mask,
        name="SoftImpute")
    assert missing_mae < 0.1, "Error too high!"

def test_mice_row_with_low_rank_random_matrix_approximate():
    mice = MICE(n_imputations=100, impute_type='pmm', n_nearest_columns=5)
    XY_completed = mice.complete(XY_incomplete)
    _, missing_mae = reconstruction_error(
        XY,
        XY_completed,
        missing_mask,
        name="MICE (impute_type=row)")
    assert missing_mae < 0.1, "Error too high with approximate PMM method!"

def test_mice_column_with_low_rank_random_matrix_approximate():
    mice = MICE(n_imputations=100, impute_type='col', n_nearest_columns=5)
    XY_completed = mice.complete(XY_incomplete)
    _, missing_mae = reconstruction_error(
        XY,
        XY_completed,
        missing_mask,
        name="MICE (impute_type=col)")
    assert missing_mae < 0.1, "Error too high with approximate column method!"

def test_mice_row_with_low_rank_random_matrix():
    mice = MICE(n_imputations=100, impute_type='pmm')
    XY_completed = mice.complete(XY_incomplete)
    _, missing_mae = reconstruction_error(
        XY,
        XY_completed,
        missing_mask,
        name="MICE (impute_type=row)")
    assert missing_mae < 0.1, "Error too high with PMM method!"

def test_mice_column_with_low_rank_random_matrix():
    mice = MICE(n_imputations=100, impute_type='col')
    XY_completed = mice.complete(XY_incomplete)
    _, missing_mae = reconstruction_error(
        XY,
        XY_completed,
        missing_mask,
        name="MICE (impute_type=col)")
    assert missing_mae < 0.1, "Error too high with column method!"

def test_iterative_imputer_with_low_rank_random_matrix_approximate():
    imputer = IterativeImputer(n_iter=50, n_nearest_features=5, random_state=0)
    XY_completed = imputer.fit_transform(XY_incomplete)
    _, missing_mae = reconstruction_error(
        XY,
        XY_completed,
        missing_mask,
        name="IterativeImputer with n_nearest_features=5")
    assert missing_mae < 0.1, "Error too high with IterativeImputer " \
                              "method using n_nearest_features=5!"

def test_iterative_imputer_as_mice_with_low_rank_random_matrix_approximate():
    n_imputations = 5
    XY_completed = []
    for i in range(n_imputations):
        imputer = IterativeImputer(n_iter=5, sample_posterior=True, random_state=i)
        XY_completed.append(imputer.fit_transform(XY_incomplete))
    _, missing_mae = reconstruction_error(
        XY,
        np.mean(XY_completed, axis=0),
        missing_mask,
        name="IterativeImputer as MICE")
    assert missing_mae < 0.1, "Error too high with IterativeImputer as MICE!"