def residual_error(X_train, X_test, y_train, y_test, reg="linear"):
"""
Plot the residual error of the Regresssion model for the input data,
and return the fitted Regression model.
-------------------------------------------------------------------
# Parameters
# X_train,X_test,y_train,y_test (np.arrays): Given X, a 2-D array of Data,
and y, an array of target data, we can use:
sklearn.model_selection.train_test_split(X,y)
to obtain X_train, X_test, y_train, and y_test.
# reg (string): Whether the regression model is linear or logistical (default="linear").
"""
if reg.lower() == "linear":
reg = LinearRegression()
reg.fit(X_train, y_train)
elif reg.lower() == "logistic":
reg = LogisticRegression()
reg.fit(X_train, y_train)
## setting plot style
plt.style.use('fivethirtyeight')
## plotting residual errors in training data
plt.scatter(reg.predict(X_train),
reg.predict(X_train) - y_train,
color="green",
s=10,
label='Train data')
## plotting residual errors in test data
plt.scatter(reg.predict(X_test),
reg.predict(X_test) - y_test,
color="blue",
s=10,
label='Test data')
## plotting line for zero residual error
plt.hlines(y=0, xmin=0, xmax=50, linewidth=2)
## plotting legend
plt.legend(loc='upper right')
## plot title
plt.title("Residual errors")
return reg
