def test_stack_predictions_regression():
"""Assert that the prediction methods work for regression tasks."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.clean()
atom.run(models=["Tree"])
atom.branch = "branch_2"
atom.impute(strat_num="mean", strat_cat="most_frequent")
atom.run(["PA"])
atom.stacking(models=["Tree", "PA"], passthrough=True)
assert isinstance(atom.stack.predict(X_reg), np.ndarray)
assert isinstance(atom.stack.score(X_reg, y_reg), np.float64)
def test_plot_probabilities():
"""Assert that the plot_probabilities method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("Ridge")
pytest.raises(PermissionError,
atom.plot_probabilities) # Task is not classif
y = ["a" if i == 0 else "b" for i in y_bin]
atom = ATOMClassifier(X_bin, y, random_state=1)
atom.clean() # Encode the target column
pytest.raises(NotFittedError, atom.plot_probabilities)
atom.run(["Tree", "LGB", "PA"], metric="f1")
pytest.raises(AttributeError,
atom.pa.plot_probabilities) # No predict_proba
atom.plot_probabilities(models=["Tree", "LGB"], target="a", display=False)
atom.lgb.plot_probabilities(target="b", display=False)