Repository: JuliaStats/MixedModels.jl (julia)

Issue:
Add `fitlog` kwarg for logging the parameter vector during optimization
As discussed in [some visualization work](https://github.com/RePsychLing/SMLP2021/pull/14#issuecomment-897900060), we should add a `fitlog=nothing` kwarg to the `fit` and `fit!` methods. In the default case, nothing changes. In the case that `fitlog` is a collection supporting `push!`, then we `push!(fitlog, SVector{length(val)}(val)` in the `obj` function being optimized.

Pull request description:
~Addresses (but does not close) #549 (because we want to change `OptSummary` later to be smarter about this).~

Closes #549. There is a minor inefficiency in that `initial`, `final`, `finitial`, `fmin` store information duplicated in the `fitlog`. Not duplicating this information (by cheating and using properties that redirect to the relevant field entries) turns out to make the bookkeeping a lot messier, so I've let this inefficiency remain (which also keeps better compatibility). I use `Vector` inside of the log instead of `SVector` because the changing size of the parameter vector in GLMM creates problems. 

It also seems that I might have found another case where julialang/julia#40048 isn't as fixed as we had hoped -- I get a big stackoverflow type inference dump as a Heisenbug when stepping through the `contra` testset, BUT it continues afterwards with the correct result.

Patch:
diff --git a/NEWS.md b/NEWS.md
index b1b46f26f..f05ac08c0 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,11 +1,13 @@
 MixedModels v4.1.0 Release Notes
 ========================
 * Add support for specifying a fixed value of `σ`, the residual standard deviation,
-  in `LinearMixedModel`. `fit` takes a keyword-argument `σ`. For models constructed
-  separately, `σ` can be specified by setting `optsum.sigma`. [#551]
-  struct. [#551]
-
-
+  in `LinearMixedModel`. `fit` takes a keyword-argument `σ`. `fit!` does not expose `σ`,
+  but `σ` can be changed after model construction by setting `optsum.sigma`. [#551]
+* Add support for logging the non-linear optimizer's steps via a `thin`
+  keyword-argument for `fit` and `fit!`. The default behavior is 'maximal' thinning,
+  such that only the initial and final values are stored. `OptSummary` has a new field
+  `fitlog` that contains the aforementioned log as a  vector of tuples of parameter and
+  objective values.[#552]
 
 MixedModels v4.0.0 Release Notes
 ========================
@@ -275,3 +277,4 @@ Package dependencies
 [#537]: https://github.com/JuliaStats/MixedModels.jl/issues/537
 [#539]: https://github.com/JuliaStats/MixedModels.jl/issues/539
 [#551]: https://github.com/JuliaStats/MixedModels.jl/issues/551
+[#552]: https://github.com/JuliaStats/MixedModels.jl/issues/552
diff --git a/Project.toml b/Project.toml
index dfa10dad2..2638f4d41 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "MixedModels"
 uuid = "ff71e718-51f3-5ec2-a782-8ffcbfa3c316"
 author = ["Phillip Alday <me@phillipalday.com>", "Douglas Bates <dmbates@gmail.com>", "Jose Bayoan Santiago Calderon <jbs3hp@virginia.edu>"]
-version = "4.0.0"
+version = "4.1.0"
 
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
diff --git a/docs/src/optimization.md b/docs/src/optimization.md
index 6e38f1f1b..d6ae7bb77 100644
--- a/docs/src/optimization.md
+++ b/docs/src/optimization.md
@@ -189,6 +189,15 @@ fm2.optsum
 DisplayAs.Text(ans) # hide
 ```
 
+More detailed information about the intermediate steps of the nonlinear optimizer can be obtained the `fitlog` field.
+By default, `fitlog` contains entries for only the initial and final steps, but additional information about every nth step can be obtained with the `thin` keyword-argument to `fit`, `fit!` and `refit!`:
+
+```@example Main
+refit!(fm2; thin=1)
+fm2.optsum.fitlog[1:10]
+DisplayAs.Text(ans) # hide
+```
+
 ## A blocked Cholesky factor
 
 A `LinearMixedModel` object contains two blocked matrices; a symmetric matrix `A` (only the lower triangle is stored) and a lower-triangular `L` which is the lower Cholesky factor of the updated and inflated `A`.
@@ -220,18 +229,30 @@ To modify the optimization process the input fields can be changed after constru
 
 Suppose, for example, that the user wishes to try a [Nelder-Mead](https://en.wikipedia.org/wiki/Nelder%E2%80%93Mead_method) optimization method instead of the default [`BOBYQA`](https://en.wikipedia.org/wiki/BOBYQA) (Bounded Optimization BY Quadratic Approximation) method.
 ```@example Main
-fm2 = LinearMixedModel(@formula(reaction ~ 1+days+(1+days|subj)), sleepstudy);
-fm2.optsum.optimizer = :LN_NELDERMEAD;
-fit!(fm2)
-fm2.optsum
+fm2nm = LinearMixedModel(@formula(reaction ~ 1+days+(1+days|subj)), sleepstudy);
+fm2nm.optsum.optimizer = :LN_NELDERMEAD;
+fit!(fm2nm; thin=1)
+fm2nm.optsum
 DisplayAs.Text(ans) # hide
 ```
 
 The parameter estimates are quite similar to those using `:LN_BOBYQA` but at the expense of 140 functions evaluations for `:LN_NELDERMEAD` versus 57 for `:LN_BOBYQA`.
+When plotting the progress of the individual fits, it becomes obvious that `:LN_BOBYQA` has fully converged by the time `:LN_NELDERMEAD` begins to approach the optimum.
+
+```@example Main
+using Gadfly
+nm = fm2nm.optsum.fitlog
+bob = fm2.optsum.fitlog
+convdf = DataFrame(algorithm=[repeat(["NelderMead"], length(nm));
+                           repeat(["BOBYQA"], length(bob))],
+                   objective=[last.(nm); last.(bob)],
+                   step=[1:length(nm); 1:length(bob)])
+plot(convdf, x=:step, y=:objective, color=:algorithm, Geom.line)
+```
 
 Run time can be constrained with  `maxfeval` and `maxtime`.
 
-See the documentation for the [`NLopt`](https://github.com/JuliaOpt/NLopt.jl) package for details about the various settings. 
+See the documentation for the [`NLopt`](https://github.com/JuliaOpt/NLopt.jl) package for details about the various settings.
 
 ### Convergence to singular covariance matrices
 
diff --git a/src/generalizedlinearmixedmodel.jl b/src/generalizedlinearmixedmodel.jl
index c65690d0b..ae3540451 100644
--- a/src/generalizedlinearmixedmodel.jl
+++ b/src/generalizedlinearmixedmodel.jl
@@ -174,6 +174,7 @@ function fit(
     fast::Bool=false,
     nAGQ::Integer=1,
     progress::Bool=true,
+    thin::Int=typemax(Int),
 )
     return fit(
         GeneralizedLinearMixedModel,
@@ -181,13 +182,14 @@ function fit(
         columntable(tbl),
         d,
         l;
-        wts=wts,
-        offset=offset,
-        contrasts=contrasts,
-        verbose=verbose,
-        fast=fast,
-        nAGQ=nAGQ,
-        progress=progress,
+        wts,
+        offset,
+        contrasts,
+        verbose,
+        fast,
+        nAGQ,
+        progress,
+        thin,
     )
 end
 
@@ -204,15 +206,17 @@ function fit(
     fast::Bool=false,
     nAGQ::Integer=1,
     progress::Bool=true,
+    thin::Int=typemax(Int),
 )
     return fit!(
         GeneralizedLinearMixedModel(
             f, tbl, d, l; wts=wts, offset=offset, contrasts=contrasts
         );
-        verbose=verbose,
-        fast=fast,
-        nAGQ=nAGQ,
-        progress=progress,
+        verbose,
+        fast,
+        nAGQ,
+        progress,
+        thin,
     )
 end
 
@@ -230,6 +234,7 @@ function fit(
     fast::Bool=false,
     nAGQ::Integer=1,
     progress::Bool=true,
+    thin::Int=typemax(Int),
 )
     return fit(
         GeneralizedLinearMixedModel,
@@ -237,18 +242,21 @@ function fit(
         tbl,
         d,
         l;
-        wts=wts,
-        contrasts=contrasts,
-        offset=offset,
-        verbose=verbose,
-        fast=fast,
-        nAGQ=nAGQ,
-        progress=progress,
+        wts,
+        contrasts,
+        offset,
+        verbose,
+        fast,
+        nAGQ,
+        progress,
+        thin,
     )
 end
 
 """
-    fit!(m::GeneralizedLinearMixedModel[, verbose=false, fast=false, nAGQ=1, progress=true])
+    fit!(m::GeneralizedLinearMixedModel; fast=false, nAGQ=1,
+                                         verbose=false, progress=true,
+                                         thin::Int=1)
 
 Optimize the objective function for `m`.
 
@@ -261,6 +269,8 @@ during the iterations to minimize the deviance.  There is a delay before this di
 and it may not be shown at all for models that are optimized quickly.
 
 If `verbose` is `true`, then both the intermediate results of both the nonlinear optimization and PIRLS are also displayed on standard output.
+
+At every `thin`th iteration  is recorded in `fitlog`, optimization progress is saved in `m.optsum.fitlog`.
 """
 function fit!(
     m::GeneralizedLinearMixedModel{T};
@@ -268,6 +278,7 @@ function fit!(
     fast::Bool=false,
     nAGQ::Integer=1,
     progress::Bool=true,
+    thin::Int=typemax(Int),
 ) where {T}
     β = m.β
     lm = m.LMM
@@ -284,16 +295,23 @@ function fit!(
     end
     setpar! = fast ? setθ! : setβθ!
     prog = ProgressUnknown("Minimizing"; showspeed=true)
+    # start from zero for the initial call to obj before optimization
+    iter = 0
+    fitlog = optsum.fitlog
     function obj(x, g)
         isempty(g) || throw(ArgumentError("g should be empty for this objective"))
         val = deviance(pirls!(setpar!(m, x), fast, verbose), nAGQ)
+        iszero(rem(iter, thin)) && push!(fitlog, (copy(x), val))
         verbose && println(round(val; digits=5), " ", x)
         progress && ProgressMeter.next!(prog; showvalues=[(:objective, val)])
+        iter += 1
         return val
     end
     opt = Opt(optsum)
     NLopt.min_objective!(opt, obj)
     optsum.finitial = obj(optsum.initial, T[])
+    empty!(fitlog)
+    push!(fitlog, (copy(optsum.initial), optsum.finitial))
     fmin, xmin, ret = NLopt.optimize(opt, copyto!(optsum.final, optsum.initial))
     ProgressMeter.finish!(prog)
     ## check if very small parameter values bounded below by zero can be set to zero
@@ -303,10 +321,14 @@ function fit!(
             xmin_[i] = zero(T)
         end
     end
+    loglength = length(fitlog)
     if xmin ≠ xmin_
         if (zeroobj = obj(xmin_, T[])) ≤ (fmin + 1.e-5)
             fmin = zeroobj
             copyto!(xmin, xmin_)
+        elseif length(fitlog) > loglength
+            # remove unused extra log entry
+            pop!(fitlog)
         end
     end
     ## ensure that the parameter values saved in m are xmin
diff --git a/src/linearmixedmodel.jl b/src/linearmixedmodel.jl
index b1270dcdc..59bd6aa22 100644
--- a/src/linearmixedmodel.jl
+++ b/src/linearmixedmodel.jl
@@ -183,9 +183,18 @@ function fit(
     progress::Bool=true,
     REML::Bool=false,
     σ=nothing,
+    thin=typemax(Int),
 )
     return fit(
-        LinearMixedModel, f, Tables.columntable(tbl); wts, contrasts, progress, REML, σ
+        LinearMixedModel,
+        f,
+        Tables.columntable(tbl);
+        wts,
+        contrasts,
+        progress,
+        REML,
+        σ,
+        thin,
     )
 end
 
@@ -198,8 +207,9 @@ function fit(
     progress=true,
     REML=false,
     σ=nothing,
+    thin=typemax(Int),
 )
-    return fit!(LinearMixedModel(f, tbl; contrasts, wts, σ); progress, REML)
+    return fit!(LinearMixedModel(f, tbl; contrasts, wts, σ); progress, REML, thin)
 end
 
 function _offseterr()
@@ -220,11 +230,12 @@ function fit(
     REML::Bool=false,
     offset=[],
     σ=nothing,
+    thin=typemax(Int),
 )
     return if !isempty(offset)
         _offseterr()
     else
-        fit(LinearMixedModel, f, tbl; wts, contrasts, progress, REML, σ)
+        fit(LinearMixedModel, f, tbl; wts, contrasts, progress, REML, σ, thin)
     end
 end
 
@@ -240,13 +251,17 @@ function fit(
     REML::Bool=false,
     offset=[],
     σ=nothing,
-    fast::Bool=false,
-    nAGQ::Integer=1,
+    thin=typemax(Int),
+    fast=nothing,
+    nAGQ=nothing,
 )
     return if !isempty(offset)
         _offseterr()
     else
-        fit(LinearMixedModel, f, tbl; wts, contrasts, progress, REML, σ)
+        if !isnothing(fast) || !isnothing(nAGQ)
+            @warn "fast and nAGQ arguments are ignored when fitting a LinearMixedModel"
+        end
+        fit(LinearMixedModel, f, tbl; wts, contrasts, progress, REML, σ, thin)
     end
 end
 
@@ -398,13 +413,24 @@ function feL(m::LinearMixedModel)
 end
 
 """
-    fit!(m::LinearMixedModel[; progress::Bool=true, REML::Bool=false])
+    fit!(m::LinearMixedModel; progress::Bool=true, REML::Bool=false,
+                              σ::Union{Real, Nothing}=nothing,
+                              thin::Int=typemax(Int))
 
 Optimize the objective of a `LinearMixedModel`.  When `progress` is `true` a
 `ProgressMeter.ProgressUnknown` display is shown during the optimization of the
 objective, if the optimization takes more than one second or so.
+
+At every `thin`th iteration  is recorded in `fitlog`, optimization progress is
+saved in `m.optsum.fitlog`.
 """
-function fit!(m::LinearMixedModel{T}; progress::Bool=true, REML::Bool=false) where {T}
+function fit!(
+    m::LinearMixedModel{T};
+    progress::Bool=true,
+    REML::Bool=false,
+    σ::Union{Real,Nothing}=nothing,
+    thin::Int=typemax(Int),
+) where {T}
     optsum = m.optsum
     # this doesn't matter for LMM, but it does for GLMM, so let's be consistent
     if optsum.feval > 0
@@ -413,14 +439,21 @@ function fit!(m::LinearMixedModel{T}; progress::Bool=true, REML::Bool=false) whe
     opt = Opt(optsum)
     optsum.REML = REML
     prog = ProgressUnknown("Minimizing"; showspeed=true)
+    # start from zero for the initial call to obj before optimization
+    iter = 0
+    fitlog = optsum.fitlog
     function obj(x, g)
         isempty(g) || throw(ArgumentError("g should be empty for this objective"))
         val = objective(updateL!(setθ!(m, x)))
+        iszero(rem(iter, thin)) && push!(fitlog, (copy(x), val))
         progress && ProgressMeter.next!(prog; showvalues=[(:objective, val)])
+        iter += 1
         return val
     end
     NLopt.min_objective!(opt, obj)
     optsum.finitial = obj(optsum.initial, T[])
+    empty!(fitlog)
+    push!(fitlog, (copy(optsum.initial), optsum.finitial))
     fmin, xmin, ret = NLopt.optimize!(opt, copyto!(optsum.final, optsum.initial))
     ProgressMeter.finish!(prog)
     ## check if small non-negative parameter values can be set to zero
@@ -431,10 +464,14 @@ function fit!(m::LinearMixedModel{T}; progress::Bool=true, REML::Bool=false) whe
             xmin_[i] = zero(T)
         end
     end
-    if xmin_ ≠ xmin
+    loglength = length(fitlog)
+    if xmin ≠ xmin_
         if (zeroobj = obj(xmin_, T[])) ≤ (fmin + 1.e-5)
             fmin = zeroobj
             copyto!(xmin, xmin_)
+        elseif length(fitlog) > loglength
+            # remove unused extra log entry
+            pop!(fitlog)
         end
     end
     ## ensure that the parameter values saved in m are xmin
@@ -826,7 +863,7 @@ end
 
 Read, check, and restore the `optsum` field from a JSON stream or filename.
 """
-function restoreoptsum!(m::LinearMixedModel, io::IO)
+function restoreoptsum!(m::LinearMixedModel{T}, io::IO) where {T}
     dict = JSON3.read(io)
     ops = m.optsum
     okay =
@@ -852,6 +889,13 @@ function restoreoptsum!(m::LinearMixedModel, io::IO)
     ops.returnvalue = Symbol(dict.returnvalue)
     # provides compatibility with fits saved before the introduction of fixed sigma
     ops.sigma = get(dict, :sigma, nothing)
+    fitlog = get(dict, :fitlog, nothing)
+    ops.fitlog = if isnothing(fitlog)
+        # compat with fits saved before fitlog
+        [(ops.initial, ops.finitial, ops.final, ops.fmin)]
+    else
+        [(convert(Vector{T}, first(entry)), T(last(entry))) for entry in fitlog]
+    end
     return m
 end
 
diff --git a/src/optsummary.jl b/src/optsummary.jl
index a88d73388..f33c9eed4 100644
--- a/src/optsummary.jl
+++ b/src/optsummary.jl
@@ -4,8 +4,8 @@
 Summary of an `NLopt` optimization
 
 # Fields
-
 * `initial`: a copy of the initial parameter values in the optimization
+* `finitial`: the initial value of the objective
 * `lowerbd`: lower bounds on the parameter values
 * `ftol_rel`: as in NLopt
 * `ftol_abs`: as in NLopt
@@ -22,8 +22,14 @@ Summary of an `NLopt` optimization
 * `nAGQ`: number of adaptive Gauss-Hermite quadrature points in deviance evaluation for GLMMs
 * `REML`: use the REML criterion for LMM fits
 * `sigma`: a priori value for the residual standard deviation for LMM
+* `fitlog`: A vector of tuples of parameter and objectives values from steps in the optimization
+
+The latter four fields are MixedModels functionality and not related directly to the `NLopt` package or algorithms.
 
-The latter three fields are model characteristics and not related directly to the `NLopt` package or algorithms.
+!!! note
+    The internal storage of the parameter values within `fitlog` may change in
+    the future to use a different subtype of `AbstractVector` (e.g., `StaticArrays.SVector`)
+    for each snapshot without being considered a breaking change.
 """
 mutable struct OptSummary{T<:AbstractFloat}
     initial::Vector{T}
@@ -36,15 +42,17 @@ mutable struct OptSummary{T<:AbstractFloat}
     initial_step::Vector{T}
     maxfeval::Int
     maxtime::T
+    feval::Int
     final::Vector{T}
     fmin::T
-    feval::Int
     optimizer::Symbol
     returnvalue::Symbol
     nAGQ::Integer           # don't really belong here but I needed a place to store them
     REML::Bool
     sigma::Union{T,Nothing}
+    fitlog::Vector{Tuple{Vector{T},T}} # not SVector because we would need to parameterize on size (which breaks GLMM)
 end
+
 function OptSummary(
     initial::Vector{T},
     lowerbd::Vector{T},
@@ -52,10 +60,13 @@ function OptSummary(
     ftol_rel::T=zero(T),
     ftol_abs::T=zero(T),
     xtol_rel::T=zero(T),
+    xtol_abs::Vector{T}=zero(initial) .+ 1e-10,
     initial_step::Vector{T}=T[],
     maxfeval=-1,
     maxtime=T(-1),
 ) where {T<:AbstractFloat}
+    fitlog = [(initial, T(Inf))]
+
     return OptSummary(
         initial,
         lowerbd,
@@ -63,18 +74,19 @@ function OptSummary(
         ftol_rel,
         ftol_abs,
         xtol_rel,
-        zero(initial),
+        xtol_abs,
         initial_step,
         maxfeval,
         maxtime,
+        -1,
         copy(initial),
         T(Inf),
-        -1,
         optimizer,
         :FAILURE,
         1,
         false,
         nothing,
+        fitlog,
     )
 end
 


Test patch:
diff --git a/test/pirls.jl b/test/pirls.jl
index f306945c5..0cd7c9fc9 100644
--- a/test/pirls.jl
+++ b/test/pirls.jl
@@ -13,7 +13,11 @@ include("modelcache.jl")
 
 @testset "contra" begin
     contra = dataset(:contra)
-    gm0 = fit(MixedModel, only(gfms[:contra]), contra, Bernoulli(); fast=true, progress=false)
+    thin = 5
+    gm0 = fit(MixedModel, only(gfms[:contra]), contra, Bernoulli(); fast=true, progress=false, thin)
+    fitlog = gm0.optsum.fitlog
+    @test length(fitlog) == (div(gm0.optsum.feval, thin) + 1) # for the initial value
+    @test first(fitlog) == (gm0.optsum.initial, gm0.optsum.finitial)
     @test gm0.lowerbd == zeros(1)
     @test isapprox(gm0.θ, [0.5720734451352923], atol=0.001)
     @test !issingular(gm0)
diff --git a/test/pls.jl b/test/pls.jl
index a65aeae20..21e8ef443 100644
--- a/test/pls.jl
+++ b/test/pls.jl
@@ -130,7 +130,14 @@ end
     @test isa(vc, Matrix{Float64})
     @test only(vc) ≈ 375.7167775 rtol=1.e-6
     # since we're caching the fits, we should get it back to being correctly fitted
-    refit!(fm1; REML=false, progress=false)
+    # we also take this opportunity to test fitlog
+    @testset "fitlog" begin
+        fitlog = fm1.optsum.fitlog
+        thin = 2
+        refit!(fm1; REML=false, progress=false, thin)
+        @test length(fitlog) == (div(fm1.optsum.feval, thin) + 1) # for the initial value
+        @test first(fitlog) == (fm1.optsum.initial, fm1.optsum.finitial)
+    end
 end
 
 @testset "Dyestuff2" begin