def test_delete():
"""Assert that models can be deleted."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("RF")
atom.rf.delete()
assert not atom.models
assert not atom.metric
def test_automl_regression(cls):
"""Assert that the automl method works for regression tasks."""
pl = Pipeline(steps=[('rbfsampler', RBFSampler(gamma=0.95, random_state=2)
), ('lassolarscv', LassoLarsCV(normalize=False))])
cls.return_value.fitted_pipeline_ = pl.fit(X_reg, y_reg)
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.automl(scoring="r2", random_state=2)
cls.assert_called_with(n_jobs=1, scoring="r2", verbosity=0, random_state=2)
assert atom.metric == "r2"
def test_canvas():
"""Assert that the canvas works."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("Tree")
with atom.canvas(1, 2, title="Title", display=False):
atom.plot_residuals(title="Residuals plot")
atom.plot_feature_importance(title="Feature importance plot")
def test_models_regression(model):
"""Assert that all models work with regression."""
atom = ATOMRegressor(X_reg, y_reg, test_size=0.24, random_state=1)
atom.run(
models=model,
metric="neg_mean_absolute_error",
n_calls=2,
n_initial_points=1,
bo_params={"base_estimator": "gbrt", "cv": 1},
)
assert not atom.errors
assert hasattr(atom, model) and hasattr(atom, model.lower())
def test_scoring_metric_is_none():
"""Assert that the scoring method works when metric is None."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
pytest.raises(NotFittedError, atom.scoring)
atom.run(["Tree", "RF"])
atom.run("LGB", bagging=5) # Test with and without bagging
atom.scoring()
def test_plot_prc():
"""Assert that the plot_prc method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("LGB")
pytest.raises(PermissionError, atom.plot_prc) # Task is not binary
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
pytest.raises(NotFittedError, atom.plot_prc)
atom.run(["XGB", "LGB"], metric="f1")
atom.plot_prc(display=False)
atom.lgb.plot_prc(display=False)
def test_plot_lift():
"""Assert that the plot_lift method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("Ridge")
pytest.raises(PermissionError, atom.plot_lift) # Task is not binary
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
pytest.raises(NotFittedError, atom.plot_lift)
atom.run(["Tree", "LGB", "PA"], metric="f1")
pytest.raises(AttributeError, atom.pa.plot_lift) # No predict_proba
atom.plot_lift(["Tree", "LGB"], display=False)
atom.lgb.plot_lift(display=False)
def test_plot_calibration():
"""Assert that the plot_calibration method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("Ridge")
pytest.raises(PermissionError, atom.plot_calibration) # Task is not binary
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
pytest.raises(NotFittedError, atom.plot_calibration)
atom.run(["Tree", "kSVM"], metric="f1")
pytest.raises(ValueError, atom.plot_calibration, n_bins=4)
atom.plot_calibration(display=False)
atom.tree.plot_calibration(display=False)
def test_plot_bo():
"""Assert that the plot_bo method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
pytest.raises(NotFittedError, atom.plot_bo)
atom.run("lasso", metric="max_error", n_calls=0)
pytest.raises(PermissionError, atom.plot_bo) # No BO in pipeline
atom.run(["lasso", "ridge"], metric="max_error", n_calls=10)
atom.plot_bo(display=False)
atom.lasso.plot_bo(display=False)
def test_plot_threshold(metric):
"""Assert that the plot_threshold method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("Ridge")
pytest.raises(PermissionError, atom.plot_threshold) # Task is not binary
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
pytest.raises(NotFittedError, atom.plot_threshold)
atom.run(["Tree", "LGB", "PA"], metric="f1")
pytest.raises(AttributeError, atom.pa.plot_threshold) # No predict_proba
pytest.raises(ValueError, atom.tree.plot_threshold, metric="unknown")
atom.plot_threshold(models=["Tree", "LGB"], display=False)
atom.lgb.plot_threshold(metric=metric, display=False)
def test_task_assignment():
"""Assert that the correct task is assigned."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
assert atom.task == "binary classification"
atom = ATOMClassifier(X_class, y_class, random_state=1)
assert atom.task == "multiclass classification"
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
assert atom.task == "regression"
def test_default_solver_from_task():
"""Assert that the solver is inferred from the task when a model is selected."""
# For classification tasks
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.feature_selection(strategy="rfe", solver="lgb", n_features=8)
assert type(atom.pipeline[0].solver).__name__ == "LGBMClassifier"
# For regression tasks
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.feature_selection(strategy="rfe", solver="lgb", n_features=25)
assert type(atom.pipeline[0].solver).__name__ == "LGBMRegressor"
def test_stacking_default_estimator():
"""Assert that a default estimator is provided per goal."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.run(["LR", "LGB"])
atom.stacking()
assert atom.stack.estimator.__class__.__name__ == "LogisticRegression"
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run(["Tree", "LGB"])
atom.stacking()
assert atom.stack.estimator.__class__.__name__ == "Ridge"
def test_default_solver_univariate():
"""Assert that the default solver is selected for strategy="univariate"."""
# For classification tasks
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.feature_selection(strategy="univariate", solver=None, n_features=8)
assert atom.pipeline[0].solver.__name__ == "f_classif"
# For regression tasks
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.feature_selection(strategy="univariate", solver=None, n_features=8)
assert atom.pipeline[0].solver.__name__ == "f_regression"
def test_plot_errors():
"""Assert that the plot_errors method work as intended."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.run("LR")
pytest.raises(PermissionError, atom.plot_errors) # Task is not regression
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
pytest.raises(NotFittedError, atom.plot_errors)
atom.run(["Tree", "Bag"], metric="MSE")
atom.plot_errors(display=False)
atom.tree.plot_errors(display=False)
def test_errors_are_updated():
"""Assert that the found exceptions are updated in the errors attribute."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
# Produce an error on one model (when n_initial_points > n_calls)
atom.run(["Tree", "LGB"], n_calls=(3, 2), n_initial_points=(2, 5))
assert list(atom.errors) == ["LGB"]
# Subsequent runs should remove the original model
atom.run(["Tree", "LGB"], n_calls=(5, 3), n_initial_points=(7, 1))
assert atom.models == "LGB"
def test_plot_learning_curve():
"""Assert that the plot_learning_curve method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
pytest.raises(NotFittedError, atom.plot_learning_curve)
atom.run("LGB")
atom.delete() # Clear the pipeline to allow ts
atom.train_sizing(["Tree", "LGB"], metric="max_error", bagging=4)
atom.plot_learning_curve(display=False)
atom.train_sizing(["Tree", "LGB"], metric="max_error")
atom.plot_learning_curve(display=False)
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)
def test_invalid_scoring():
"""Assert that an error is raised when the provided scoring is invalid."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("Tree", metric="mse")
pytest.raises(ValueError, atom.automl, scoring="r2")
def test_plot_confusion_matrix():
"""Assert that the plot_confusion_matrix method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run(["Ridge"])
pytest.raises(PermissionError,
atom.plot_confusion_matrix) # Task is not classif
# For binary classification tasks
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
pytest.raises(NotFittedError, atom.plot_confusion_matrix)
atom.run(["RF", "LGB"])
atom.plot_confusion_matrix(display=False)
atom.lgb.plot_confusion_matrix(normalize=True, display=False)
# For multiclass classification tasks
atom = ATOMClassifier(X_class, y_class, random_state=1)
atom.run(["RF", "LGB"])
pytest.raises(NotImplementedError, atom.plot_confusion_matrix)
atom.lgb.plot_confusion_matrix(normalize=True, display=False)
def test_scaling_is_passed():
"""Assert that the scaling is passed to the trainer."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.scale("minmax")
atom.run("LGB")
assert atom.dataset.equals(atom.lgb.dataset)
def test_plot_results(metric):
"""Assert that the plot_results method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
pytest.raises(NotFittedError, atom.plot_results)
# Without bagging
atom.run(["Tree", "LGB"], metric=metric, bagging=0)
atom.voting()
atom.plot_results(metric="me", display=False)
atom.tree.plot_results(display=False)
# With bagging
atom.run("Tree", metric=metric, bagging=3)
atom.plot_results(metric="me", display=False)
atom.tree.plot_results(display=False)
def test_goal_ATOMRegressor():
"""Assert that the goal is set correctly for ATOMRegressor."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
assert atom.goal == "regression"
def test_plot_successive_halving():
"""Assert that the plot_successive_halving method work as intended."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
pytest.raises(NotFittedError, atom.plot_successive_halving)
atom.run("LGB")
atom.delete() # Clear the pipeline to allow sh
atom.successive_halving(models=["Tree", "Bag", "RF", "LGB"], bagging=4)
pytest.raises(ValueError, atom.plot_successive_halving, models="unknown")
pytest.raises(ValueError, atom.plot_successive_halving, models="BR")
pytest.raises(ValueError, atom.plot_successive_halving, metric="unknown")
pytest.raises(ValueError, atom.plot_successive_halving, metric=-1)
pytest.raises(ValueError, atom.plot_successive_halving, metric=1)
pytest.raises(ValueError, atom.plot_successive_halving, metric="roc_auc")
atom.plot_successive_halving(display=False)
atom.successive_halving(models=["Tree", "Bag", "RF", "LGB"])
atom.plot_successive_halving(display=False)
def test_models_and_metric_are_updated():
"""Assert that the models and metric attributes are updated correctly."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run(["OLS", "Tree"], metric=get_scorer("max_error"))
assert atom.models == ["OLS", "Tree"]
assert atom.metric == "max_error"
def test_trainer_becomes_atom():
"""Assert that the parent trainer is converted to atom."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("Tree")
assert atom is atom.tree.T
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_passthrough_not_scaled():
"""Assert that the features are not scaled when models don't need it."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run(["Tree", "RF"])
atom.stacking(passthrough=True)
assert not check_scaling(atom.stack.X.iloc[:, 2:])
def test_calibrate_invalid_task():
"""Assert than an error is raised when task="regression"."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run("OLS")
pytest.raises(PermissionError, atom.calibrate)
def test_custom_models(model):
"""Assert that ATOM works with custom models."""
atom = ATOMRegressor(X_reg, y_reg, random_state=1)
atom.run(models=model, n_calls=2, n_initial_points=1)
assert atom.rfr.fullname == "RandomForestRegressor"
assert atom.rfr.estimator.get_params()["random_state"] == 1