def test_pipeline_parameter_False():
"""Assert that no transformers used."""
atom = ATOMClassifier(X10_nan, y10, random_state=1)
atom.impute(strat_num="median")
atom.prune(max_sigma=2)
X = atom.transform(X10_nan, pipeline=False) # Use None
assert isinstance(X, list) # X is unchanged
def test_pipeline_parameter_True():
"""Assert that all transformers are used."""
atom = ATOMClassifier(X10_nan, y10, random_state=1)
atom.impute(strat_num="median")
atom.prune(max_sigma=2)
X = atom.transform(X10_nan, pipeline=True) # Use both transformers
assert len(X) < 10
def test_pipeline_parameter_None():
"""Assert that only some transformers are used."""
atom = ATOMClassifier(X10_nan, y10, random_state=1)
atom.impute(strat_num="median")
atom.prune(max_sigma=2)
X = atom.transform(X10_nan, pipeline=None) # Only use imputer
assert len(X) == 10
def test_getitem():
"""Assert that atom is subscriptable."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.impute()
assert atom[1].__class__.__name__ == "Imputer"
assert atom["mean radius"].equals(atom.dataset["mean radius"])
def test_branch_setter_from_valid():
"""Assert that we can create a new branch, not from the current one."""
atom = ATOMClassifier(X10_nan, y10, random_state=1)
atom.branch = "branch_2"
atom.impute()
atom.branch = "branch_3_from_master"
assert atom.branch.name == "branch_3"
assert atom.n_nans > 0
def test_vote_branch_transformation(pipeline):
"""Assert that the branches transform every estimator only once."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.impute()
atom.branch = "branch_2"
atom.encode()
atom.run(models=["Tree", "LGB"])
atom.voting()
assert isinstance(atom.vote.predict(X_bin, pipeline=pipeline), np.ndarray)
def test_plot_pipeline(show_params):
"""Assert that the plot_pipeline method work as intended."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.impute()
atom.prune()
atom.feature_selection("univariate", n_features=10)
atom.successive_halving(["Tree", "AdaB"])
pytest.raises(ValueError, atom.plot_pipeline, branch="invalid")
atom.plot_pipeline(show_params=show_params,
title="Pipeline plot",
display=False)
def test_vote_prediction_methods():
"""Assert that the prediction methods work as intended."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.run(models=["Tree"])
atom.branch = "branch_2"
atom.impute(strat_num="mean", strat_cat="most_frequent")
atom.run(["LGB"])
atom.voting(models=["Tree", "LGB"])
pytest.raises(AttributeError, atom.vote.decision_function, X_bin)
assert isinstance(atom.vote.predict(X_bin), np.ndarray)
assert isinstance(atom.vote.predict_proba(X_bin), np.ndarray)
assert isinstance(atom.vote.score(X_bin, y_bin), np.float64)
def test_status_method():
"""Assert that the status method prints the estimators without errors."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.impute()
atom.branch.status()
assert str(atom.branch).endswith("min_frac_cols: None\n --> Models: None")
def test_iter():
"""Assert that we can iterate over atom's pipeline."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.impute()
assert [item for item in atom][1] == atom.pipeline[1]
def test_impute():
"""Assert that the impute method imputes all missing values."""
atom = ATOMClassifier(X10_nan, y10, random_state=1)
atom.impute()
assert atom.dataset.isna().sum().sum() == 0