[docs] fix lasso_regression example (#486)

Co-authored-by: Soederlind <paul.soderlind@unisg.ch>

docs/examples_literate/general_examples/lasso_regression.jl

@@ -8,9 +8,6 @@
 
 using DelimitedFiles, LinearAlgebra, Statistics, Plots, Convex, SCS
 
-import MathOptInterface
-const MOI = MathOptInterface
-
 # # Loading Data
 #
 # We use the diabetes data from Efron et al, downloaded from https://web.stanford.edu/~hastie/StatLearnSparsity_files/DATA/diabetes.html and then converted from a tab to a comma delimited file.
@@ -31,10 +28,10 @@ xNames = header[1:end-1];
 # (a)  The regression is $Y = Xb + u$,
 # where $Y$ and $u$ are $T \times 1$, $X$ is $T \times K$, and $b$ is the $K$-vector of regression coefficients.
 #
-# (b) We want to minimize $(Y-Xb)'(Y-Xb) + \gamma \sum |b_i| + \lambda \sum b_i^2$.
+# (b) We want to minimize $(Y-Xb)'(Y-Xb)/T + \gamma \sum |b_i| + \lambda \sum b_i^2$.
 #
 # (c) We can equally well minimise $b'Qb - 2c'b + \gamma \sum |b_i| + \lambda \sum b_i^2$,
-# where $Q = X'X$ and $c=X'Y$
+# where $Q = X'X/T$ and $c=X'Y/T$.
 #
 # (d) Lasso: $\gamma>0,\lambda=0$; Ridge: $\gamma=0,\lambda>0$; elastic net: $\gamma>0,\lambda>0$.
 
@@ -46,36 +43,40 @@ Do Lasso (set γ>0,λ=0), ridge (set γ=0,λ>0) or elastic net regression (set 
 
 ## Input
 - `Y::Vector`:     T-vector with the response (dependent) variable
-- `Y::VecOrMat`:   TxK matrix of covariates (regressors)
+- `X::VecOrMat`:   TxK matrix of covariates (regressors)
 - `γ::Number`:     penalty on sum(abs.(b))
 - `λ::Number`:     penalty on sum(b.^2)
 
 """
-function LassoEN(Y, X, γ, λ = 0.0)
-    K = size(X, 2)
+function LassoEN(Y, X, γ, λ = 0)
+    (T, K) = (size(X, 1), size(X, 2))
 
     b_ls = X \ Y                    #LS estimate of weights, no restrictions
 
-    Q = X'X
-    c = X'Y                      #c'b = Y'X*b
+    Q = X'X / T
+    c = X'Y / T                      #c'b = Y'X*b
 
     b = Variable(K)              #define variables to optimize over
     L1 = quadform(b, Q)            #b'Q*b
     L2 = dot(c, b)                 #c'b
     L3 = norm(b, 1)                #sum(|b|)
     L4 = sumsquares(b)            #sum(b^2)
 
-    Sol = minimize(L1 - 2 * L2 + γ * L3 + λ * L4)      #u'u + γ*sum(|b|) + λsum(b^2), where u = Y-Xb
+    if λ > 0
+        Sol = minimize(L1 - 2 * L2 + γ * L3 + λ * L4)      #u'u/T + γ*sum(|b|) + λ*sum(b^2), where u = Y-Xb
+    else
+        Sol = minimize(L1 - 2 * L2 + γ * L3)               #u'u/T + γ*sum(|b|) where u = Y-Xb
+    end
     solve!(Sol, SCS.Optimizer; silent_solver = true)
-    Sol.status == MOI.OPTIMAL ? b_i = vec(evaluate(b)) : b_i = NaN
+    Sol.status == Convex.MOI.OPTIMAL ? b_i = vec(evaluate(b)) : b_i = NaN
 
     return b_i, b_ls
 end
 
 # The next cell makes a Lasso regression for a single value of γ.
 
 K = size(X, 2)
-γ = 100
+γ = 0.25
 
 (b, b_ls) = LassoEN(Y, X, γ)
 
@@ -90,7 +91,7 @@ display([["" "OLS" "Lasso"]; xNames b_ls b])
 # Remark: it would be quicker to put this loop inside the `LassoEN()` function so as to not recreate `L1`-`L4`.
 
 nγ = 101
-γM = range(0; stop = 600, length = nγ)             #different γ values
+γM = range(0; stop = 1.5, length = nγ)             #different γ values
 
 bLasso = fill(NaN, size(X, 2), nγ)       #results for γM[i] are in bLasso[:,i]
 for i in 1:nγ
@@ -101,20 +102,20 @@ end
 #-
 
 plot(
-    log.(γM),
+    log10.(γM),
     bLasso',
     title = "Lasso regression coefficients",
-    xlabel = "log(γ)",
+    xlabel = "log10(γ)",
     label = permutedims(xNames),
     size = (600, 400),
 )
 
 # # Ridge Regression
 #
-# We use the same function to do a ridge regression. Alternatively, do `b = inv(X'X + λ*I)*X'Y`.
+# We use the same function to do a ridge regression. Alternatively, do `b = inv(X'X/T + λ*I)*X'Y/T`.
 
 nλ = 101
-λM = range(0; stop = 3000, length = nλ)
+λM = range(0; stop = 7.5, length = nλ)
 
 bRidge = fill(NaN, size(X, 2), nλ)
 for i in 1:nλ
@@ -125,17 +126,17 @@ end
 #-
 
 plot(
-    log.(λM),
+    log10.(λM),
     bRidge',
     title = "Ridge regression coefficients",
-    xlabel = "log(λ)",
+    xlabel = "log10(λ)",
     label = permutedims(xNames),
     size = (600, 400),
 )
 
 # # Elastic Net Regression
 
-λ = 200
+λ = 0.5
 println("redo the Lasso regression, but with λ=$λ: an elastic net regression")
 
 bEN = fill(NaN, size(X, 2), nγ)
@@ -147,10 +148,10 @@ end
 #-
 
 plot(
-    log.(γM),
+    log10.(γM),
     bEN',
     title = "Elastic Net regression coefficients",
-    xlabel = "log(γ)",
+    xlabel = "log10(γ)",
     label = permutedims(xNames),
     size = (600, 400),
 )