DOC use notebook-style for plot_lasso_and_elasticnet
@@ -9,13 +9,15 @@ compared with the ground-truth.
"""
+# %%
+# Data Generation
+# ---------------------------------------------------
+
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score
-# #############################################################################
-# Generate some sparse data to play with
np.random.seed(42)
n_samples, n_features = 50, 100
@@ -35,8 +37,10 @@ n_samples = X.shape[0]
X_train, y_train = X[: n_samples // 2], y[: n_samples // 2]
X_test, y_test = X[n_samples // 2 :], y[n_samples // 2 :]
-# #############################################################################
+# %%
# Lasso
+# ---------------------------------------------------
+
from sklearn.linear_model import Lasso
alpha = 0.1
@@ -47,8 +51,10 @@ r2_score_lasso = r2_score(y_test, y_pred_lasso)
print(lasso)
print("r^2 on test data : %f" % r2_score_lasso)
-# #############################################################################
+# %%
# ElasticNet
+# ---------------------------------------------------
+
from sklearn.linear_model import ElasticNet
enet = ElasticNet(alpha=alpha, l1_ratio=0.7)
@@ -58,6 +64,11 @@ r2_score_enet = r2_score(y_test, y_pred_enet)
print(enet)
print("r^2 on test data : %f" % r2_score_enet)
+
+# %%
+# Plot
+# ---------------------------------------------------
+
m, s, _ = plt.stem(
np.where(enet.coef_)[0],
enet.coef_[enet.coef_ != 0],
