def test_gpipeline_regression(): gpipeline = GraphPipeline({ "PT": PassThrough(), "Ridge": Ridge() }, [("PT", "Ridge")]) X = dfX.loc[:, ["num1", "num2", "num3"]] gpipeline.fit(X, y) yhat = gpipeline.predict(X) yhat2 = gpipeline.models["Ridge"].predict(X) assert yhat.shape == y.shape assert (yhat == yhat2).all() with pytest.raises(AttributeError): gpipeline.predict_proba(X) with pytest.raises(AttributeError): gpipeline.predict_log_proba(X) assert gpipeline.get_feature_names_at_node("PT") == list(X.columns) assert gpipeline.get_input_features_at_node("PT") == list(X.columns) assert gpipeline.get_input_features_at_node("Ridge") == list(X.columns) with pytest.raises(ValueError): assert gpipeline.get_feature_names_at_node("DONTEXIST")
def test_gpipeline_classification(): gpipeline = GraphPipeline({"PT": PassThrough(), "Logit": LogisticRegression()}, [("PT", "Logit")]) gpipeline.fit(X, yc) yhat_proba = gpipeline.predict_proba(X) yhat_proba2 = gpipeline.models["Logit"].predict_proba(X) assert yhat_proba.shape == (X.shape[0], 2) assert (yhat_proba == yhat_proba2).all() assert list(gpipeline.classes_) == [0, 1]