For a 0.10.4 release

Project.toml

@@ -1,6 +1,6 @@
 name = "MLJLinearModels"
 uuid = "6ee0df7b-362f-4a72-a706-9e79364fb692"
-version = "0.10.3"
+version = "0.10.4"
 authors = ["Thibaut Lienart <tlienart@me.com>"]
 
 [deps]

src/fit/proxgrad.jl

@@ -21,6 +21,8 @@ function _fit(glr::GLR, solver::ProxGrad, X, y, scratch)
     # functions
     _f = if solver.gram
         smooth_gram_objective(glr, X, y, n)
+    elseif c == 0
+        smooth_objective!(glr, X, y, scratch)
     else
         smooth_objective(glr, X, y; c=c)
     end

src/glr/d_l2loss.jl

@@ -69,10 +69,20 @@ function smooth_fg!(glr::GLR{L2Loss,<:ENR}, X, y, scratch)
         apply_Xt!(g, X, r)
         g .+= λ .* θ
         glr.fit_intercept && (glr.penalize_intercept || (g[end] -= λ * θ[end]))
-        return glr.loss(r) + get_l2(glr.penalty)(view_θ(glr, θ))
+        return l2loss!(r) + get_l2(glr.penalty)(view_θ(glr, θ))
     end
 end
 
+# residuals (Xθ-y) formed in scratch.n2, |y-Xθ|₂² = |Xθ-y|₂²
+function smooth_objective!(glr::GLR{L2Loss,<:ENR}, X, y, scratch)
+    J = smooth_objective(glr, size(X, 1))
+    r = scratch.n2
+    return θ -> (get_residuals!(r, X, θ, y); l2loss!(r) + J.penalty(view_θ(glr, θ)))
+end
+
+# L2Loss()(r) overwriting r; sums the squares as lp(r, 2) does so values match
+l2loss!(r) = sum(r .= abs2.(r)) / 2
+
 function smooth_gram_fg!(glr::GLR{L2Loss,<:ENR}, XX, Xy, n)
     λ = get_penalty_scale_l2(glr, n)
     (g, θ) -> begin

src/glr/utils.jl

@@ -33,6 +33,19 @@ smooth_objective(glr::GLR, X, y; c::Int=0) =
 """
 $SIGNATURES
 
+Return a function computing the smooth part of the objective at a given
+evaluation point `θ`, as `smooth_objective(glr, X, y)` for a single-output
+model but writing `X*θ` in `scratch.n2`.
+"""
+function smooth_objective!(glr::GLR, X, y, scratch)
+    J  = smooth_objective(glr, size(X, 1))
+    Xθ = scratch.n2
+    return θ -> (apply_X!(Xθ, X, θ); J(y, Xθ, view_θ(glr, θ)))
+end
+
+"""
+$SIGNATURES
+
 Return the smooth part of the objective function of a GLR.
 """
 smooth_objective(glr::GLR{<:SmoothLoss,<:ENR}, n) = glr.loss + get_l2(glr.penalty) * ifelse(glr.scale_penalty_with_samples, n, 1.)

test/fit/ols-ridge-lasso-elnet.jl

@@ -161,3 +161,40 @@ end
     @test isapprox(θ_fista, θ_gram_explicit, rtol=1e-5)
     @test isapprox(θ_gram_explicit, θ_gram_implicit; rtol=1e-5)
 end
+
+@testset "proxgrad in-place objective" begin
+    n, p = 20_000, 10
+    ((X, _, _), (_, y1, _)) = generate_continuous(n, p; seed=5)
+    yb = sign.(y1)
+    # (model, target, check allocations); scaled elnet does not converge here,
+    # see https://github.com/JuliaAI/MLJLinearModels.jl/issues/187
+    for (glr, t, check_alloc) in (
+            (LassoRegression(0.5), y1, true),
+            (LassoRegression(0.5; fit_intercept=false), y1, true),
+            (ElasticNetRegression(0.3, 0.5), y1, false),
+            (ElasticNetRegression(0.3, 0.5; scale_penalty_with_samples=false), y1, true),
+            (LogisticRegression(0.0, 0.2; penalty=:l1), yb, false),
+        )
+        s  = R.scratch(X; i=glr.fit_intercept)
+        f  = R.smooth_objective(glr, X, t)
+        f! = R.smooth_objective!(glr, X, t, s)
+        r  = StableRNG(1)
+        for _ in 1:5
+            θ = randn(r, p + Int(glr.fit_intercept))
+            @test f!(θ) == f(θ)
+        end
+        check_alloc || continue
+        # no length-n allocation in an objective or gradient call
+        fg! = R.smooth_fg!(glr, X, t, s)
+        θ   = randn(r, p + Int(glr.fit_intercept))
+        g   = similar(θ)
+        f!(θ); fg!(g, θ)
+        @test (@allocated f!(θ)) < 8n
+        @test (@allocated fg!(g, θ)) < 8n
+        # a fit allocates the 3 length-n scratch vectors and nothing else of length n
+        for solver in (FISTA(), ISTA())
+            fit(glr, X, t; solver=solver)
+            @test (@allocated fit(glr, X, t; solver=solver)) < 4 * 8n
+        end
+    end
+end