jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_multisurfstar_pipeline_missing_values():
    """Check: Data (Missing Values): MultiSURF* works in a sklearn pipeline"""
    np.random.seed(320931)

    clf = make_pipeline(MultiSURFstar(n_features_to_select=2),
                        Imputer(),
                        RandomForestClassifier(n_estimators=100, n_jobs=-1))

    assert np.mean(cross_val_score(clf, features_missing_values,
                                   labels_missing_values, cv=3, n_jobs=-1)) > 0.7
//  [python]
def test_multisurfstar_pipeline_missing_values():
    """Check: Data (Missing Values): MultiSURF* works in a sklearn pipeline"""
    np.random.seed(320931)

    clf = make_pipeline(MultiSURFstar(n_features_to_select=2),
                        SimpleImputer(),
                        RandomForestClassifier(n_estimators=100, n_jobs=-1))

    assert np.mean(cross_val_score(clf, features_missing_values,
                                   labels_missing_values, cv=3, n_jobs=-1)) > 0.7
//  [python]
def test_multisurf_pipeline_missing_values():
    """Check: Data (Missing Values): MultiSURF works in a sklearn pipeline"""
    np.random.seed(320931)

    clf = make_pipeline(MultiSURF(n_features_to_select=2),
                        Imputer(),
                        RandomForestClassifier(n_estimators=100, n_jobs=-1))

    assert np.mean(cross_val_score(clf, features_missing_values,
                                   labels_missing_values, cv=3, n_jobs=-1)) > 0.7
//  [python]
def test_multisurf_pipeline_missing_values():
    """Check: Data (Missing Values): MultiSURF works in a sklearn pipeline"""
    np.random.seed(320931)

    clf = make_pipeline(MultiSURF(n_features_to_select=2),
                        SimpleImputer(),
                        RandomForestClassifier(n_estimators=100, n_jobs=-1))

    assert np.mean(cross_val_score(clf, features_missing_values,
                                   labels_missing_values, cv=3, n_jobs=-1)) > 0.7
//  [python]
numeric_transformer = Imputer()
categorical_transformer = Pipeline(
    steps=[("imputer", Imputer(strategy="most_frequent")), ("encoder", OrdinalEncoder())]
)
//  [python]
numeric_transformer = SimpleImputer()
categorical_transformer = Pipeline(
    steps=[("imputer", SimpleImputer(strategy="most_frequent")), ("encoder", OrdinalEncoder())]
)
//  [python]
pipe1 = Pipeline(
    steps=[
        ("imputer", Imputer()),
        ("scaler", StandardScaler()),
        ("model", LogisticRegression)
    ]
)
# OR
pipe2 = make_pipeline(Imputer(), StandardScaler(), LogisticRegression())
//  [python]
pipe1 = Pipeline(
    steps=[
        ("imputer", SimpleImputer()),
        ("scaler", StandardScaler()),
        ("model", LogisticRegression)
    ]
)
# OR
pipe2 = make_pipeline(SimpleImputer(), StandardScaler(), LogisticRegression())
//  [python]
def num_pipeline_transformer(data):
    numerics = ['float64', 'int64']

    num_attrs = data.select_dtypes(include=numerics)

    num_pipeline = Pipeline([
        ('imputer', Imputer(strategy="median")),
        ('attrs_adder', CustomAttrAdder()),
        ('std_scaler', StandardScaler()),
        ])
    return num_attrs, num_pipeline
//  [python]
def num_pipeline_transformer(data):
    numerics = ['float64', 'int64']

    num_attrs = data.select_dtypes(include=numerics)

    num_pipeline = Pipeline([
        ('imputer', SimpleImputer(strategy="median")),
        ('attrs_adder', CustomAttrAdder()),
        ('std_scaler', StandardScaler()),
        ])
    return num_attrs, num_pipeline
//  [python]
def test_permutation_importance_mixed_types_pandas():
    pd = pytest.importorskip("pandas")
    rng = np.random.RandomState(42)
    n_repeats = 5

    # Last column is correlated with y
    X = pd.DataFrame({"col1": [1.0, 2.0, 3.0, np.nan], "col2": ["a", "b", "a", "b"]})
    y = np.array([0, 1, 0, 1])

    num_preprocess = make_pipeline(Imputer(), StandardScaler())
    preprocess = ColumnTransformer(
        [("num", num_preprocess, ["col1"]), ("cat", OneHotEncoder(), ["col2"])]
    )
    clf = make_pipeline(preprocess, LogisticRegression(solver="lbfgs"))
    clf.fit(X, y)

    result = permutation_importance(clf, X, y, n_repeats=n_repeats, random_state=rng)

    assert result.importances.shape == (X.shape[1], n_repeats)
    # the correlated feature with y is the last column and should
    # have the highest importance
    assert np.all(result.importances_mean[-1] > result.importances_mean[:-1])
//  [python]
def test_permutation_importance_mixed_types_pandas():
    pd = pytest.importorskip("pandas")
    rng = np.random.RandomState(42)
    n_repeats = 5

    # Last column is correlated with y
    X = pd.DataFrame({"col1": [1.0, 2.0, 3.0, np.nan], "col2": ["a", "b", "a", "b"]})
    y = np.array([0, 1, 0, 1])

    num_preprocess = make_pipeline(SimpleImputer(), StandardScaler())
    preprocess = ColumnTransformer(
        [("num", num_preprocess, ["col1"]), ("cat", OneHotEncoder(), ["col2"])]
    )
    clf = make_pipeline(preprocess, LogisticRegression(solver="lbfgs"))
    clf.fit(X, y)

    result = permutation_importance(clf, X, y, n_repeats=n_repeats, random_state=rng)

    assert result.importances.shape == (X.shape[1], n_repeats)
    # the correlated feature with y is the last column and should
    # have the highest importance
    assert np.all(result.importances_mean[-1] > result.importances_mean[:-1])
//  [python]
def test_sklearn_transformer_wrapper():
    check_estimator(SklearnTransformerWrapper(transformer=Imputer()))
//  [python]
def test_sklearn_transformer_wrapper():
    check_estimator(SklearnTransformerWrapper(transformer=SimpleImputer()))
//  [python]
df = pd.read_csv("mammographic_masses_converted.csv", names = ['BI-RADS', 'Age','Shape', 'Margin', 'Density', 'Severity'], na_values =["?"])
in_mean = Imputer(strategy='median')
//  [python]
df = pd.read_csv("mammographic_masses_converted.csv", names = ['BI-RADS', 'Age','Shape', 'Margin', 'Density', 'Severity'], na_values =["?"])
in_mean = SimpleImputer(strategy='median')
//  [python]
def test_permutation_importance_mixed_types():
    rng = np.random.RandomState(42)
    n_repeats = 4

    # Last column is correlated with y
    X = np.array([[1.0, 2.0, 3.0, np.nan], [2, 1, 2, 1]]).T
    y = np.array([0, 1, 0, 1])

    clf = make_pipeline(Imputer(), LogisticRegression(solver="lbfgs"))
    clf.fit(X, y)
    result = permutation_importance(clf, X, y, n_repeats=n_repeats, random_state=rng)

    assert result.importances.shape == (X.shape[1], n_repeats)

    # the correlated feature with y is the last column and should
    # have the highest importance
    assert np.all(result.importances_mean[-1] > result.importances_mean[:-1])

    # use another random state
    rng = np.random.RandomState(0)
    result2 = permutation_importance(clf, X, y, n_repeats=n_repeats, random_state=rng)
    assert result2.importances.shape == (X.shape[1], n_repeats)

    assert not np.allclose(result.importances, result2.importances)

    # the correlated feature with y is the last column and should
    # have the highest importance
    assert np.all(result2.importances_mean[-1] > result2.importances_mean[:-1])
//  [python]
def test_permutation_importance_mixed_types():
    rng = np.random.RandomState(42)
    n_repeats = 4

    # Last column is correlated with y
    X = np.array([[1.0, 2.0, 3.0, np.nan], [2, 1, 2, 1]]).T
    y = np.array([0, 1, 0, 1])

    clf = make_pipeline(SimpleImputer(), LogisticRegression(solver="lbfgs"))
    clf.fit(X, y)
    result = permutation_importance(clf, X, y, n_repeats=n_repeats, random_state=rng)

    assert result.importances.shape == (X.shape[1], n_repeats)

    # the correlated feature with y is the last column and should
    # have the highest importance
    assert np.all(result.importances_mean[-1] > result.importances_mean[:-1])

    # use another random state
    rng = np.random.RandomState(0)
    result2 = permutation_importance(clf, X, y, n_repeats=n_repeats, random_state=rng)
    assert result2.importances.shape == (X.shape[1], n_repeats)

    assert not np.allclose(result.importances, result2.importances)

    # the correlated feature with y is the last column and should
    # have the highest importance
    assert np.all(result2.importances_mean[-1] > result2.importances_mean[:-1])
//  [python]
def scale_impute(X: pd.DataFrame,
                 )->np.ndarray:
    """
    Performs standard scaling. Imputes mean for continuous data and most-frequent for categorical data.
    """
    pipe = pipeline.Pipeline(
        [("scaler", StandardScaler()), ("imputer", Imputer())]
    )
    # Categorical values are only booleans in the this dataset and none are missing
    continuous = [column for column in X.columns if len(np.unique(X[column].values)) > 2]
    categorical = list(set(X.columns) ^ set(continuous))

    if len(continuous) != 0:
        X_cont_ = pipe.fit_transform(X[continuous])
        X_cont = pd.DataFrame(X_cont_, columns=continuous, index=X.index)
    else:
        X_cont = pd.DataFrame(None, index=X.index)

    if len(categorical) != 0:
        categ_imputer = Imputer(strategy="most_frequent")
        X_cat_ = categ_imputer.fit_transform(X[categorical])
        X_cat = pd.DataFrame(X_cat_, columns=categorical, index=X.index)
    else:
        X_cat = pd.DataFrame(None, index=X.index)

    X_proc = pd.concat([X_cont, X_cat], axis=1, join='inner')
    X_proc = X_proc.reindex(sorted(X_proc.columns), axis=1)

    return X_proc.values
//  [python]
def scale_impute(X: pd.DataFrame,
                 )->np.ndarray:
    """
    Performs standard scaling. Imputes mean for continuous data and most-frequent for categorical data.
    """
    pipe = pipeline.Pipeline(
        [("scaler", StandardScaler()), ("imputer", SimpleImputer())]
    )
    # Categorical values are only booleans in the this dataset and none are missing
    continuous = [column for column in X.columns if len(np.unique(X[column].values)) > 2]
    categorical = list(set(X.columns) ^ set(continuous))

    if len(continuous) != 0:
        X_cont_ = pipe.fit_transform(X[continuous])
        X_cont = pd.DataFrame(X_cont_, columns=continuous, index=X.index)
    else:
        X_cont = pd.DataFrame(None, index=X.index)

    if len(categorical) != 0:
        categ_imputer = SimpleImputer(strategy="most_frequent")
        X_cat_ = categ_imputer.fit_transform(X[categorical])
        X_cat = pd.DataFrame(X_cat_, columns=categorical, index=X.index)
    else:
        X_cat = pd.DataFrame(None, index=X.index)

    X_proc = pd.concat([X_cont, X_cat], axis=1, join='inner')
    X_proc = X_proc.reindex(sorted(X_proc.columns), axis=1)

    return X_proc.values
//  [python]
def get_data_transformer_object(cls)->Pipeline:
        try:
            robust_scaler = RobustScaler()
            simple_imputer = Imputer(strategy="constant", fill_value=0)
            preprocessor = Pipeline(
                steps=[
                    ("Imputer", simple_imputer), #replace missing values with zero
                    ("RobustScaler", robust_scaler) #keep every feature in same range and handle outlier
                    ]
            )
            
            return preprocessor

        except Exception as e:
            raise SensorException(e, sys) from e
//  [python]
def get_data_transformer_object(cls)->Pipeline:
        try:
            robust_scaler = RobustScaler()
            simple_imputer = SimpleImputer(strategy="constant", fill_value=0)
            preprocessor = Pipeline(
                steps=[
                    ("Imputer", simple_imputer), #replace missing values with zero
                    ("RobustScaler", robust_scaler) #keep every feature in same range and handle outlier
                    ]
            )
            
            return preprocessor

        except Exception as e:
            raise SensorException(e, sys) from e
//  [python]
def test_with_pipe_and_skl_transformer_input_df(df_vartypes, input_features):
    pipe = Pipeline(
        [
            ("imputer", Imputer(strategy="constant")),
            ("transformer", MockTransformer()),
        ]
    )
    pipe.fit(df_vartypes)
    assert pipe.get_feature_names_out(input_features=input_features) == variables_str
//  [python]
def test_with_pipe_and_skl_transformer_input_df(df_vartypes, input_features):
    pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="constant")),
            ("transformer", MockTransformer()),
        ]
    )
    pipe.fit(df_vartypes)
    assert pipe.get_feature_names_out(input_features=input_features) == variables_str
//  [python]
mean_imputer = Imputer(missing_values='NaN', strategy='mean')
mode_imputer = Imputer(missing_values='NaN', strategy='most_frequent')
//  [python]
mean_imputer = SimpleImputer(missing_values=np.nan, strategy='mean')
mode_imputer = SimpleImputer(missing_values=np.nan, strategy='most_frequent')
//  [python]
def test_with_pipe_and_skl_transformer_input_array(df_vartypes, input_features):
    pipe = Pipeline(
        [
            ("imputer", Imputer(strategy="constant")),
            ("transformer", MockTransformer()),
        ]
    )
    pipe.fit(df_vartypes.to_numpy())

    if input_features is None:
        assert (
            pipe.get_feature_names_out(input_features=input_features) == variables_arr
        )
    else:
        assert pipe.get_feature_names_out(input_features=input_features) == list(
            input_features
        )
//  [python]
def test_with_pipe_and_skl_transformer_input_array(df_vartypes, input_features):
    pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="constant")),
            ("transformer", MockTransformer()),
        ]
    )
    pipe.fit(df_vartypes.to_numpy())

    if input_features is None:
        assert (
            pipe.get_feature_names_out(input_features=input_features) == variables_arr
        )
    else:
        assert pipe.get_feature_names_out(input_features=input_features) == list(
            input_features
        )
//  [python]
numeric_features = ['age', 'balance']
numeric_transformer = Pipeline(steps=[
    ('imputer', Imputer(strategy='median')),
    ('scaler', StandardScaler())])
//  [python]
numeric_features = ['age', 'balance']
numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())])
//  [python]
def test_new_feature_names_pipe_with_skl_transformer_and_array(
    df_vartypes, input_features
):
    df_vartypes = df_vartypes.drop(["dob"], axis=1)

    pipe = Pipeline(
        [
            ("imputer", Imputer(strategy="constant")),
            ("transformer", MockSelector()),
        ]
    )
    pipe.fit(df_vartypes.to_numpy())

    if input_features is not None:
        input_features = input_features[0:-1]
        features_out = input_features[2:]
        assert pipe.get_feature_names_out(input_features=input_features) == features_out
    else:
        features_out = ["x2", "x3"]
        assert pipe.get_feature_names_out(input_features=input_features) == features_out
//  [python]
def test_new_feature_names_pipe_with_skl_transformer_and_array(
    df_vartypes, input_features
):
    df_vartypes = df_vartypes.drop(["dob"], axis=1)

    pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="constant")),
            ("transformer", MockSelector()),
        ]
    )
    pipe.fit(df_vartypes.to_numpy())

    if input_features is not None:
        input_features = input_features[0:-1]
        features_out = input_features[2:]
        assert pipe.get_feature_names_out(input_features=input_features) == features_out
    else:
        features_out = ["x2", "x3"]
        assert pipe.get_feature_names_out(input_features=input_features) == features_out
//  [python]
def test_new_feature_names_pipe_with_skl_transformer_and_df(
    df_vartypes, features_in, input_features
):
    pipe = Pipeline(
        [
            ("imputer", Imputer(strategy="constant")),
            ("transformer", MockCreator(variables=features_in, drop_original=False)),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = list(df_vartypes.columns) + [f"{i}_plus" for i in features_in]
    assert pipe.get_feature_names_out(input_features=input_features) == features_out
    pipe = Pipeline(
        [
            ("imputer", Imputer(strategy="constant")),
            ("transformer", MockCreator(variables=features_in, drop_original=True)),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = [f for f in df_vartypes.columns if f not in features_in] + [
        f"{i}_plus" for i in features_in
    ]
    assert pipe.get_feature_names_out(input_features=input_features) == features_out
//  [python]
def test_new_feature_names_pipe_with_skl_transformer_and_df(
    df_vartypes, features_in, input_features
):
    pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="constant")),
            ("transformer", MockCreator(variables=features_in, drop_original=False)),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = list(df_vartypes.columns) + [f"{i}_plus" for i in features_in]
    assert pipe.get_feature_names_out(input_features=input_features) == features_out
    pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="constant")),
            ("transformer", MockCreator(variables=features_in, drop_original=True)),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = [f for f in df_vartypes.columns if f not in features_in] + [
        f"{i}_plus" for i in features_in
    ]
    assert pipe.get_feature_names_out(input_features=input_features) == features_out
//  [python]
num_imputers = []
num_imputers.append(("impute_mean", Imputer(), ["age_group", "days_last_event", "days_since_creation", "days_since_last_activity", "event_count", "gender", "order_count", "session_count", "total_amount", "total_item"]))
//  [python]
num_imputers = []
num_imputers.append(("impute_mean", SimpleImputer(), ["age_group", "days_last_event", "days_since_creation", "days_since_last_activity", "event_count", "gender", "order_count", "session_count", "total_amount", "total_item"]))
//  [python]
def test_remove_feature_names_pipe_with_skl_transformer_and_df(
    df_vartypes, input_features
):
    df_vartypes = df_vartypes.drop(["dob"], axis=1)
    if input_features is not None:
        input_features = input_features[0:-1]

    pipe = Pipeline(
        [
            ("transformer", MockSelector()),
            ("imputer", Imputer(strategy="constant")),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = list(df_vartypes.columns)[2:]
    assert all(
        pipe.get_feature_names_out(input_features=input_features) == features_out
    )

    pipe = Pipeline(
        [
            ("imputer", Imputer(strategy="constant")),
            ("transformer", MockSelector()),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = list(df_vartypes.columns)[2:]
    assert pipe.get_feature_names_out(input_features=input_features) == features_out
//  [python]
def test_remove_feature_names_pipe_with_skl_transformer_and_df(
    df_vartypes, input_features
):
    df_vartypes = df_vartypes.drop(["dob"], axis=1)
    if input_features is not None:
        input_features = input_features[0:-1]

    pipe = Pipeline(
        [
            ("transformer", MockSelector()),
            ("imputer", SimpleImputer(strategy="constant")),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = list(df_vartypes.columns)[2:]
    assert all(
        pipe.get_feature_names_out(input_features=input_features) == features_out
    )

    pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="constant")),
            ("transformer", MockSelector()),
        ]
    )
    pipe.fit(df_vartypes)
    features_out = list(df_vartypes.columns)[2:]
    assert pipe.get_feature_names_out(input_features=input_features) == features_out
//  [python]
def test_sklearn_imputer_allfeatures_with_constant(df_na):
    transformer = SklearnTransformerWrapper(
        transformer=Imputer(fill_value="missing", strategy="constant")
    )

    # transformed dataframe
    ref = df_na.copy()
    ref = ref.fillna("missing")

    dataframe_na_transformed = transformer.fit_transform(df_na)

    # transformed output
    assert all(dataframe_na_transformed.isna().sum() == 0)
    pd.testing.assert_frame_equal(ref, dataframe_na_transformed)
//  [python]
def test_sklearn_imputer_allfeatures_with_constant(df_na):
    transformer = SklearnTransformerWrapper(
        transformer=SimpleImputer(fill_value="missing", strategy="constant")
    )

    # transformed dataframe
    ref = df_na.copy()
    ref = ref.fillna("missing")

    dataframe_na_transformed = transformer.fit_transform(df_na)

    # transformed output
    assert all(dataframe_na_transformed.isna().sum() == 0)
    pd.testing.assert_frame_equal(ref, dataframe_na_transformed)
//  [python]
def test_wrap_simple_imputer(df_na):
    variables_to_impute = ["Age", "Marks"]
    na_variables_left_after_imputation = [
        col
        for col in df_na.loc[:, df_na.isna().any()].columns
        if col not in variables_to_impute
    ]

    transformer = SklearnTransformerWrapper(
        transformer=Imputer(fill_value=-999, strategy="constant"),
        variables=variables_to_impute,
    )

    # transformed dataframe
    ref = df_na.copy()
    ref[variables_to_impute] = ref[variables_to_impute].fillna(-999)

    dataframe_na_transformed = transformer.fit_transform(df_na)

    # transformed output
    assert all(
        dataframe_na_transformed[na_variables_left_after_imputation].isna().sum() != 0
    )
    assert all(dataframe_na_transformed[variables_to_impute].isna().sum() == 0)
    pd.testing.assert_frame_equal(ref, dataframe_na_transformed)
//  [python]
def test_wrap_simple_imputer(df_na):
    variables_to_impute = ["Age", "Marks"]
    na_variables_left_after_imputation = [
        col
        for col in df_na.loc[:, df_na.isna().any()].columns
        if col not in variables_to_impute
    ]

    transformer = SklearnTransformerWrapper(
        transformer=SimpleImputer(fill_value=-999, strategy="constant"),
        variables=variables_to_impute,
    )

    # transformed dataframe
    ref = df_na.copy()
    ref[variables_to_impute] = ref[variables_to_impute].fillna(-999)

    dataframe_na_transformed = transformer.fit_transform(df_na)

    # transformed output
    assert all(
        dataframe_na_transformed[na_variables_left_after_imputation].isna().sum() != 0
    )
    assert all(dataframe_na_transformed[variables_to_impute].isna().sum() == 0)
    pd.testing.assert_frame_equal(ref, dataframe_na_transformed)
//  [python]
def test_sklearn_imputer_object_with_constant(df_na):
    variables_to_impute = ["Name", "City"]
    na_variables_left_after_imputation = [
        col
        for col in df_na.loc[:, df_na.isna().any()].columns
        if col not in variables_to_impute
    ]

    transformer = SklearnTransformerWrapper(
        transformer=Imputer(strategy="constant", fill_value="missing"),
        variables=variables_to_impute,
    )

    # transformed dataframe
    ref = df_na.copy()
    ref[variables_to_impute] = ref[variables_to_impute].fillna("missing")

    dataframe_na_transformed = transformer.fit_transform(df_na)

    # transformed output
    assert all(
        dataframe_na_transformed[na_variables_left_after_imputation].isna().sum() != 0
    )
    assert all(dataframe_na_transformed[variables_to_impute].isna().sum() == 0)
    pd.testing.assert_frame_equal(ref, dataframe_na_transformed)
//  [python]
def test_sklearn_imputer_object_with_constant(df_na):
    variables_to_impute = ["Name", "City"]
    na_variables_left_after_imputation = [
        col
        for col in df_na.loc[:, df_na.isna().any()].columns
        if col not in variables_to_impute
    ]

    transformer = SklearnTransformerWrapper(
        transformer=SimpleImputer(fill_value="missing", strategy="constant"),
        variables=variables_to_impute,
    )

    # transformed dataframe
    ref = df_na.copy()
    ref[variables_to_impute] = ref[variables_to_impute].fillna("missing")

    dataframe_na_transformed = transformer.fit_transform(df_na)

    # transformed output
    assert all(
        dataframe_na_transformed[na_variables_left_after_imputation].isna().sum() != 0
    )
    assert all(dataframe_na_transformed[variables_to_impute].isna().sum() == 0)
    pd.testing.assert_frame_equal(ref, dataframe_na_transformed)
//  [python]
X = data[numeric_columns].drop(columns=["x_e_out [-]"])
y = data["x_e_out [-]"]

# Split the data into training and NaN sets
train_data = data.dropna(subset=["x_e_out [-]"])
nan_data = data[data["x_e_out [-]"].isna()]

# Impute missing values in features with mean
imputer = Imputer(strategy="mean")
//  [python]
X = data[numeric_columns].drop(columns=["x_e_out [-]"])
y = data["x_e_out [-]"]

# Split the data into training and NaN sets
train_data = data.dropna(subset=["x_e_out [-]"])
nan_data = data[data["x_e_out [-]"].isna()]

# Impute missing values in features with mean
imputer = SimpleImputer(strategy="mean")
//  [python]
def test_support_missing_values(MultiClassClassifier):
    # smoke test to check that pipeline OvR and OvO classifiers are letting
    # the validation of missing values to
    # the underlying pipeline or classifiers
    rng = np.random.RandomState(42)
    X, y = iris.data, iris.target
    X = np.copy(X)  # Copy to avoid that the original data is modified
    mask = rng.choice([1, 0], X.shape, p=[0.1, 0.9]).astype(bool)
    X[mask] = np.nan
    lr = make_pipeline(Imputer(), LogisticRegression(random_state=rng))

    MultiClassClassifier(lr).fit(X, y).score(X, y)
//  [python]
def test_support_missing_values(MultiClassClassifier):
    # smoke test to check that pipeline OvR and OvO classifiers are letting
    # the validation of missing values to
    # the underlying pipeline or classifiers
    rng = np.random.RandomState(42)
    X, y = iris.data, iris.target
    X = np.copy(X)  # Copy to avoid that the original data is modified
    mask = rng.choice([1, 0], X.shape, p=[0.1, 0.9]).astype(bool)
    X[mask] = np.nan
    lr = make_pipeline(SimpleImputer(), LogisticRegression(random_state=rng))

    MultiClassClassifier(lr).fit(X, y).score(X, y)
//  [python]
def create_pipeline(X, task_type):
    pipe_cfg = {
        'num_cols': X.dtypes[X.dtypes == 'int64'].index.values.tolist(),
        'cat_cols': X.dtypes[X.dtypes == 'object'].index.values.tolist(),
    }
    num_pipe = Pipeline([
        ('num_imputer', Imputer(strategy='median')),
        ('num_scaler', StandardScaler())
    ])
    cat_pipe = Pipeline([
        ('cat_imputer', Imputer(strategy='constant', fill_value='?')),
        ('cat_encoder', OneHotEncoder(handle_unknown='ignore', sparse=False))
    ])
    feat_pipe = ColumnTransformer([
        ('num_pipe', num_pipe, pipe_cfg['num_cols']),
        ('cat_pipe', cat_pipe, pipe_cfg['cat_cols'])
    ])

    if task_type == "classification":
        pipeline = Pipeline(
            steps=[('preprocessor', feat_pipe),
                   ('model', RandomForestClassifier(n_estimators=10,
                                                    max_depth=5))])
    else:
        pipeline = Pipeline(
            steps=[('preprocessor', feat_pipe),
                   ('model', RandomForestRegressor(n_estimators=10,
                                                   max_depth=5))])

    return pipeline
//  [python]
def create_pipeline(X, task_type):
    pipe_cfg = {
        'num_cols': X.dtypes[X.dtypes == 'int64'].index.values.tolist(),
        'cat_cols': X.dtypes[X.dtypes == 'object'].index.values.tolist(),
    }
    num_pipe = Pipeline([
        ('num_imputer', SimpleImputer(strategy='median')),
        ('num_scaler', StandardScaler())
    ])
    cat_pipe = Pipeline([
        ('cat_imputer', SimpleImputer(strategy='constant', fill_value='?')),
        ('cat_encoder', OneHotEncoder(handle_unknown='ignore', sparse=False))
    ])
    feat_pipe = ColumnTransformer([
        ('num_pipe', num_pipe, pipe_cfg['num_cols']),
        ('cat_pipe', cat_pipe, pipe_cfg['cat_cols'])
    ])

    if task_type == "classification":
        pipeline = Pipeline(
            steps=[('preprocessor', feat_pipe),
                   ('model', RandomForestClassifier(n_estimators=10,
                                                    max_depth=5))])
    else:
        pipeline = Pipeline(
            steps=[('preprocessor', feat_pipe),
                   ('model', RandomForestRegressor(n_estimators=10,
                                                   max_depth=5))])

    return pipeline
//  [python]
def test_pandas_indexing(input_passthrough):
    X, y = datasets.load_diabetes(return_X_y=True, as_frame=True)

    passcols = ["age", "sex", "bmi", "bp"]

    builder = DAGBuilder(infer_dataframe=True)
    if input_passthrough:
        builder.add_step("inp", "passthrough")

    preprocessing = (
        builder.add_step(
            "imp", s(), deps=["inp"] if input_passthrough else None
        )
        .add_step("vitals", "passthrough", deps={"imp": passcols})
        .add_step(
            "blood",
            PCA(n_components=2, random_state=0),
            deps={"imp": make_column_selector("s[0-9]+")},
        )
        .add_step("out", "passthrough", deps=["vitals", "blood"])
        .make_dag()
    )

    X_tr = preprocessing.fit_transform(X, y)
    assert isinstance(X_tr, pd.DataFrame)
    assert (X_tr.index == X.index).all()
    assert X_tr.columns.tolist() == [f"imp__{col}" for col in passcols] + [
        "blood__pca0",
        "blood__pca1",
    ]

    predictor = (
        DAGBuilder(infer_dataframe=True)
        .add_step("rf", RandomForestRegressor(random_state=0))
        .make_dag()
    )

    dag = preprocessing.join(
        predictor,
        edges=[("out", "rf")],
    )

    y_pred = dag.fit_predict(X, y)

    assert isinstance(y_pred, pd.Series)
    assert (y_pred.index == y.index).all()
    assert y_pred.name == dag.leaves_[0].name
//  [python]
def test_pandas_indexing(input_passthrough):
    X, y = datasets.load_diabetes(return_X_y=True, as_frame=True)

    passcols = ["age", "sex", "bmi", "bp"]

    builder = DAGBuilder(infer_dataframe=True)
    if input_passthrough:
        builder.add_step("inp", "passthrough")

    preprocessing = (
        builder.add_step(
            "imp", SimpleImputer(), deps=["inp"] if input_passthrough else None
        )
        .add_step("vitals", "passthrough", deps={"imp": passcols})
        .add_step(
            "blood",
            PCA(n_components=2, random_state=0),
            deps={"imp": make_column_selector("s[0-9]+")},
        )
        .add_step("out", "passthrough", deps=["vitals", "blood"])
        .make_dag()
    )

    X_tr = preprocessing.fit_transform(X, y)
    assert isinstance(X_tr, pd.DataFrame)
    assert (X_tr.index == X.index).all()
    assert X_tr.columns.tolist() == [f"imp__{col}" for col in passcols] + [
        "blood__pca0",
        "blood__pca1",
    ]

    predictor = (
        DAGBuilder(infer_dataframe=True)
        .add_step("rf", RandomForestRegressor(random_state=0))
        .make_dag()
    )

    dag = preprocessing.join(
        predictor,
        edges=[("out", "rf")],
    )

    y_pred = dag.fit_predict(X, y)

    assert isinstance(y_pred, pd.Series)
    assert (y_pred.index == y.index).all()
    assert y_pred.name == dag.leaves_[0].name
//  [python]
def test_get_visual_block_pipeline():
    pipe = Pipeline(
        [
            ("imputer", Imputer()),
            ("do_nothing", "passthrough"),
            ("do_nothing_more", None),
            ("classifier", LogisticRegression()),
        ]
    )
    est_html_info = _get_visual_block(pipe)
    assert est_html_info.kind == "serial"
    assert est_html_info.estimators == tuple(step[1] for step in pipe.steps)
    assert est_html_info.names == [
        "imputer: Imputer",
        "do_nothing: passthrough",
        "do_nothing_more: passthrough",
        "classifier: LogisticRegression",
    ]
    assert est_html_info.name_details == [str(est) for _, est in pipe.steps]
//  [python]
def test_get_visual_block_pipeline():
    pipe = Pipeline(
        [
            ("imputer", SimpleImputer()),
            ("do_nothing", "passthrough"),
            ("do_nothing_more", None),
            ("classifier", LogisticRegression()),
        ]
    )
    est_html_info = _get_visual_block(pipe)
    assert est_html_info.kind == "serial"
    assert est_html_info.estimators == tuple(step[1] for step in pipe.steps)
    assert est_html_info.names == [
        "imputer: SimpleImputer",
        "do_nothing: passthrough",
        "do_nothing_more: passthrough",
        "classifier: LogisticRegression",
    ]
    assert est_html_info.name_details == [str(est) for _, est in pipe.steps]
//  [python]
def imputers():
    return [IterativeImputer(tol=0.1), KNNImputer(), Imputer()]
//  [python]
def imputers():
    return [IterativeImputer(tol=0.1), KNNImputer(), SimpleImputer()]
//  [python]
def sparse_imputers():
    return [Imputer()]
//  [python]
def sparse_imputers():
    return [SimpleImputer()]
//  [python]
def test_get_visual_block_pipeline():
    pipe = Pipeline([
        ('imputer', Imputer()),
        ('do_nothing', 'passthrough'),
        ('do_nothing_more', None),
        ('classifier', LogisticRegression())
    ])
    est_html_info = _get_visual_block(pipe)
    assert est_html_info.kind == 'serial'
    assert est_html_info.estimators == tuple(step[1] for step in pipe.steps)
    assert est_html_info.names == ['imputer: Imputer',
                                   'do_nothing: passthrough',
                                   'do_nothing_more: passthrough',
                                   'classifier: LogisticRegression']
    assert est_html_info.name_details == [str(est) for _, est in pipe.steps]
//  [python]
def test_get_visual_block_pipeline():
    pipe = Pipeline([
        ('imputer', SimpleImputer()),
        ('do_nothing', 'passthrough'),
        ('do_nothing_more', None),
        ('classifier', LogisticRegression())
    ])
    est_html_info = _get_visual_block(pipe)
    assert est_html_info.kind == 'serial'
    assert est_html_info.estimators == tuple(step[1] for step in pipe.steps)
    assert est_html_info.names == ['imputer: SimpleImputer',
                                   'do_nothing: passthrough',
                                   'do_nothing_more: passthrough',
                                   'classifier: LogisticRegression']
    assert est_html_info.name_details == [str(est) for _, est in pipe.steps]
//  [python]
def __init__(self, metadata,
                 nbins=10, thresh_rare=0,
                 max_quantile = 1, anonymity_set_size=1,
                 drop_cols=None, quids=None):

        self.metadata = self._read_meta(metadata, drop_cols, quids)
        self.datatype = DataFrame

        self.histogram_size = nbins
        self.unique_threshold = thresh_rare
        self.quids = quids
        self.max_quantile = max_quantile
        self.anonymity_set_size = anonymity_set_size

        self.ImputerCat = Imputer(strategy='most_frequent')
        self.ImputerNum = Imputer(strategy='median')

        self.trained = False

        self.__name__ = f'SanitiserNHSk{self.anonymity_set_size}'
//  [python]
def __init__(self, metadata,
                 nbins=10, thresh_rare=0,
                 max_quantile = 1, anonymity_set_size=1,
                 drop_cols=None, quids=None):

        self.metadata = self._read_meta(metadata, drop_cols, quids)
        self.datatype = DataFrame

        self.histogram_size = nbins
        self.unique_threshold = thresh_rare
        self.quids = quids
        self.max_quantile = max_quantile
        self.anonymity_set_size = anonymity_set_size

        self.ImputerCat = SimpleImputer(strategy='most_frequent')
        self.ImputerNum = SimpleImputer(strategy='median')

        self.trained = False

        self.__name__ = f'SanitiserNHSk{self.anonymity_set_size}'
//  [python]
def __init__(self):
        self.age = Pipeline(
            [("scale", MinMaxScaler()),
            ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.fnlwgt = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.education_num = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.capital_net = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.hours_per_week = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.workclass = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.marital_status = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.occupation = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.relationship = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.race = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
             ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True)),
             ("scale", MinMaxScaler())]
        )
        self.sex = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
             ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True)),
             ("scale", MinMaxScaler())]
        )
        self.native_country = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
             ("input", Imputer(strategy="constant", fill_value=-1, add_indicator=True)),
             ("scale", MinMaxScaler())]
        )
//  [python]
def __init__(self):
        self.age = Pipeline(
            [("scale", MinMaxScaler()),
            ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.fnlwgt = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.education_num = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.capital_net = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.hours_per_week = Pipeline(
            [("scale", MinMaxScaler()),
             ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True))]
        )
        self.workclass = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.marital_status = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.occupation = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.relationship = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
            ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True)),
            ("scale", MinMaxScaler())]
        )
        self.race = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
             ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True)),
             ("scale", MinMaxScaler())]
        )
        self.sex = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
             ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True)),
             ("scale", MinMaxScaler())]
        )
        self.native_country = Pipeline(
            [("encode", OrdinalEncoder(dtype=np.int)),
             ("input", SimpleImputer(strategy="constant", fill_value=-1, add_indicator=True)),
             ("scale", MinMaxScaler())]
        )
//  [python]
def operate(self, input_datanode, target_fields=None):
        from sklearn.preprocessing import Imputer

        X, y = input_datanode.data
        self.target_fields = target_fields

        # Fetch the related fields.
        import pandas as pd
        if isinstance(X, pd.DataFrame):
            X = X.values
        X_input = X[:, target_fields]
        # Different imputation strategies applied for columns.
        if self.model is None:
            self.model = Imputer(strategy=self.params, copy=False)
            self.model.fit(X_input)
        new_X = self.model.transform(X_input).reshape(-1, 1)
        X_output = X.copy()
        X_output[:, target_fields] = new_X
        new_feature_types = input_datanode.feature_types.copy()
        output_datanode = DataNode((X_output, y), new_feature_types, input_datanode.task_type)
        output_datanode.trans_hist = input_datanode.trans_hist.copy()
        output_datanode.trans_hist.append(self.type)

        return output_datanode
//  [python]
def operate(self, input_datanode, target_fields=None):
        from sklearn.impute import SimpleImputer

        X, y = input_datanode.data
        self.target_fields = target_fields

        # Fetch the related fields.
        import pandas as pd
        if isinstance(X, pd.DataFrame):
            X = X.values
        X_input = X[:, target_fields]
        # Different imputation strategies applied for columns.
        if self.model is None:
            self.model = SimpleImputer(strategy=self.params, copy=False)
            self.model.fit(X_input)
        new_X = self.model.transform(X_input).reshape(-1, 1)
        X_output = X.copy()
        X_output[:, target_fields] = new_X
        new_feature_types = input_datanode.feature_types.copy()
        output_datanode = DataNode((X_output, y), new_feature_types, input_datanode.task_type)
        output_datanode.trans_hist = input_datanode.trans_hist.copy()
        output_datanode.trans_hist.append(self.type)

        return output_datanode
//  [python]
def test_simple_imputer_float_inputs(self):
        model = Imputer(strategy="mean", fill_value="nan")
        data = [[1, 2], [np.nan, 3], [7, 6]]
        model.fit(data)

        model_onnx = convert_sklearn(
            model,
            "scikit-learn simple imputer",
            [("input", FloatTensorType([None, 2]))],
            target_opset=TARGET_OPSET,
        )
        self.assertTrue(model_onnx.graph.node is not None)

        # should contain only node
        self.assertEqual(len(model_onnx.graph.node), 1)

        # last node should contain the Imputer
        outputs = model_onnx.graph.output
        self.assertEqual(len(outputs), 1)
        self.assertEqual(outputs[0].type.tensor_type.shape.dim[-1].dim_value, 2)
        dump_data_and_model(
            np.array(data, dtype=np.float32),
            model,
            model_onnx,
            basename="SklearnSimpleImputerMeanFloat32",
        )
//  [python]
def test_simple_imputer_float_inputs(self):
        model = SimpleImputer(strategy="mean", fill_value="nan")
        data = [[1, 2], [np.nan, 3], [7, 6]]
        model.fit(data)

        model_onnx = convert_sklearn(
            model,
            "scikit-learn simple imputer",
            [("input", FloatTensorType([None, 2]))],
            target_opset=TARGET_OPSET,
        )
        self.assertTrue(model_onnx.graph.node is not None)

        # should contain only node
        self.assertEqual(len(model_onnx.graph.node), 1)

        # last node should contain the Imputer
        outputs = model_onnx.graph.output
        self.assertEqual(len(outputs), 1)
        self.assertEqual(outputs[0].type.tensor_type.shape.dim[-1].dim_value, 2)
        dump_data_and_model(
            np.array(data, dtype=np.float32),
            model,
            model_onnx,
            basename="SklearnSimpleImputerMeanFloat32",
        )
//  [python]
def test_simple_imputer_float_inputs_int_mostf(self):
        model = Imputer(strategy="most_frequent", fill_value="nan")
        data = [[1, 2], [np.nan, 3], [7, 6], [8, np.nan]]
        model.fit(data)

        model_onnx = convert_sklearn(
            model,
            "scikit-learn simple imputer",
            [("input", Int64TensorType([None, 2]))],
            target_opset=TARGET_OPSET,
        )
        self.assertTrue(model_onnx.graph.node is not None)

        # should contain only node
        self.assertEqual(len(model_onnx.graph.node), 1)

        # last node should contain the Imputer
        outputs = model_onnx.graph.output
        self.assertEqual(len(outputs), 1)
        self._check_outputs_ints(model, model_onnx, data)
//  [python]
def test_simple_imputer_float_inputs_int_mostf(self):
        model = SimpleImputer(strategy="most_frequent", fill_value="nan")
        data = [[1, 2], [np.nan, 3], [7, 6], [8, np.nan]]
        model.fit(data)

        model_onnx = convert_sklearn(
            model,
            "scikit-learn simple imputer",
            [("input", Int64TensorType([None, 2]))],
            target_opset=TARGET_OPSET,
        )
        self.assertTrue(model_onnx.graph.node is not None)

        # should contain only node
        self.assertEqual(len(model_onnx.graph.node), 1)

        # last node should contain the Imputer
        outputs = model_onnx.graph.output
        self.assertEqual(len(outputs), 1)
        self._check_outputs_ints(model, model_onnx, data)
//  [python]
features_train = data_train.loc[:, data_train.columns != target_column_name]
labels_train = data_train[target_column_name]
features_test = data_test.loc[:, data_test.columns != target_column_name]
labels_test = data_test[target_column_name]

if args.library == 'pytorch' or args.library == 'sklearn':
	from sklearn.preprocessing import StandardScaler
	from sklearn.compose import ColumnTransformer
	from sklearn.pipeline import Pipeline
	from sklearn.preprocessing import StandardScaler, OneHotEncoder
	numeric_features = features_train.select_dtypes(
		include=[np.float64, np.int64]
	).columns
	numeric_transformer = Pipeline(steps=[
		('imputer', Imputer(strategy='median')),
		('scaler', StandardScaler())
	])
	categorical_features = features_train.select_dtypes(
		include=['category']
	).columns
	categorical_transformer = Pipeline(
		steps=[
			('imputer', Imputer(strategy='constant', fill_value='missing')),
			('onehot', OneHotEncoder(handle_unknown='ignore'))
	])
	preprocessor = ColumnTransformer(
		transformers=[
			('num', numeric_transformer, numeric_features),
			('cat', categorical_transformer, categorical_features)
	])
	features_train = preprocessor.fit_transform(features_train)
	features_test = preprocessor.transform(features_test)
	print(features_train.shape)