From f58dd6abe2c92d630346e7d25eb19cc9480c9120 Mon Sep 17 00:00:00 2001
From: Luis Barroso-Luque <lbluque@berkeley.edu>
Date: Mon, 24 Jul 2023 18:09:36 -0700
Subject: [PATCH] DOC: more small edits to examples

---
 examples/plot_adaptive.py    | 18 ++++++++++--------
 examples/plot_line_search.py | 17 -----------------
 examples/plot_stepwise.py    | 19 ++++++++++---------
 3 files changed, 20 insertions(+), 34 deletions(-)

diff --git a/examples/plot_adaptive.py b/examples/plot_adaptive.py
index 0b0036c..11e39f4 100644
--- a/examples/plot_adaptive.py
+++ b/examples/plot_adaptive.py
@@ -79,6 +79,16 @@
 print(f"    train rmse: {alasso_train['rmse']:.3f}")
 print(f"    test rmse: {alasso_test['rmse']:.3f}")
 
+# plot model coefficients
+fig, ax = plt.subplots()
+ax.plot(coef, "o", label="True coefficients")
+ax.plot(lasso_cv.best_estimator_.coef_, "o", label="Lasso", alpha=0.5)
+ax.plot(alasso_cv.best_estimator_.coef_, "o", label="Adaptive Lasso", alpha=0.5)
+ax.set_xlabel("covariate index")
+ax.set_ylabel("coefficient value")
+ax.legend()
+fig.show()
+
 # plot predicted values
 fig, ax = plt.subplots()
 ax.plot(y_test, lasso_cv.predict(X_test), "o", label="lasso", alpha=0.5)
@@ -89,11 +99,3 @@
 ax.legend()
 fig.show()
 
-# plot model coefficients
-fig, ax = plt.subplots()
-ax.plot(coef, "o", label="True coefficients")
-ax.plot(lasso_cv.best_estimator_.coef_, "o", label="Lasso", alpha=0.5)
-ax.plot(alasso_cv.best_estimator_.coef_, "o", label="Adaptive Lasso", alpha=0.5)
-ax.set_xlabel("covariate index")
-ax.set_ylabel("coefficient value")
-fig.show()
diff --git a/examples/plot_line_search.py b/examples/plot_line_search.py
index eb6aa40..c25a3d8 100644
--- a/examples/plot_line_search.py
+++ b/examples/plot_line_search.py
@@ -68,20 +68,3 @@
 print(f"    test r2: {l2l0_test['r2']:.3f}")
 print(f"    train rmse: {l2l0_train['rmse']:.3f}")
 print(f"    test rmse: {l2l0_test['rmse']:.3f}")
-
-# plot predicted values
-fig, ax = plt.subplots()
-ax.plot(y_test, l2l0_cv.predict(X_test), "o", label="L2L0 line search", alpha=0.5)
-ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], "k--")
-ax.set_xlabel("true values")
-ax.set_ylabel("predicted values")
-ax.legend()
-fig.show()
-
-# plot model coefficients
-fig, ax = plt.subplots()
-ax.plot(coef, "o", label="True coefficients")
-ax.plot(l2l0_cv.best_estimator_.coef_, "o", label="L2L0 line search", alpha=0.5)
-ax.set_xlabel("covariate index")
-ax.set_ylabel("coefficient value")
-fig.show()
diff --git a/examples/plot_stepwise.py b/examples/plot_stepwise.py
index 0eeedef..387269d 100644
--- a/examples/plot_stepwise.py
+++ b/examples/plot_stepwise.py
@@ -83,6 +83,16 @@
 print(f"    train rmse: {stepwise_train['rmse']:.3f}")
 print(f"    test rmse: {stepwise_test['rmse']:.3f}")
 
+# plot model coefficients
+fig, ax = plt.subplots()
+ax.plot(coef, "o", label="True coefficients")
+ax.plot(stepwise.coef_[[0, 1, 2, 99]], "o", label="Stepwise (ridge)", alpha=0.5)
+ax.plot(stepwise.coef_[range(3, 99)], "o", label="Stepwise (lasso)", alpha=0.5)
+ax.set_xlabel("covariate index")
+ax.set_ylabel("coefficient value")
+ax.legend()
+fig.show()
+
 # plot predicted values
 fig, ax = plt.subplots()
 ax.plot(y_test, stepwise.predict(X_test), "o", label="Stepwise", alpha=0.5)
@@ -90,12 +100,3 @@
 ax.set_xlabel("true values")
 ax.set_ylabel("predicted values")
 ax.legend()
-fig.show()
-
-# plot model coefficients
-fig, ax = plt.subplots()
-ax.plot(coef, "o", label="True coefficients")
-ax.plot(stepwise.coef_, "o", label="Stepwise", alpha=0.5)
-ax.set_xlabel("covariate index")
-ax.set_ylabel("coefficient value")
-fig.show()