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]
