refactor fit!(::LinearMixedModel) to support a more flexible optimization backend

Author: palday

src/MixedModels.jl

@@ -205,6 +205,7 @@ include("grouping.jl")
 include("mimeshow.jl")
 include("serialization.jl")
 include("profile/profile.jl")
+include("nlopt.jl")
 include("prima.jl")
 
 # COV_EXCL_START

src/linearmixedmodel.jl

@@ -177,7 +177,6 @@ function LinearMixedModel(
     θ = foldl(vcat, getθ(c) for c in reterms)
     optsum = OptSummary(θ, lbd)
     optsum.sigma = isnothing(σ) ? nothing : T(σ)
-    fill!(optsum.xtol_abs, 1.0e-10)
     return LinearMixedModel(
         form,
         reterms,
@@ -459,55 +458,10 @@ function StatsAPI.fit!(
             ArgumentError("The response is constant and thus model fitting has failed")
         )
     end
-    opt = Opt(optsum)
     optsum.REML = REML
     optsum.sigma = σ
-    prog = ProgressUnknown(; desc="Minimizing", showspeed=true)
-    # start from zero for the initial call to obj before optimization
-    iter = 0
+    xmin, fmin = optimize!(m; progress, thin)
     fitlog = optsum.fitlog
-    function obj(x, g)
-        isempty(g) || throw(ArgumentError("g should be empty for this objective"))
-        iter += 1
-        val = if isone(iter) && x == optsum.initial
-            optsum.finitial
-        else
-            try
-                objective(updateL!(setθ!(m, x)))
-            catch ex
-                # This can happen when the optimizer drifts into an area where
-                # there isn't enough shrinkage. Why finitial? Generally, it will
-                # be the (near) worst case scenario value, so the optimizer won't
-                # view it as an optimum. Using Inf messes up the quadratic
-                # approximation in BOBYQA.
-                ex isa PosDefException || rethrow()
-                optsum.finitial
-            end
-        end
-        progress && ProgressMeter.next!(prog; showvalues=[(:objective, val)])
-        !isone(iter) && iszero(rem(iter, thin)) && push!(fitlog, (copy(x), val))
-        return val
-    end
-    NLopt.min_objective!(opt, obj)
-    try
-        # use explicit evaluation w/o calling opt to avoid confusing iteration count
-        optsum.finitial = objective(updateL!(setθ!(m, optsum.initial)))
-    catch ex
-        ex isa PosDefException || rethrow()
-        # give it one more try with a massive change in scaling
-        @info "Initial objective evaluation failed, rescaling initial guess and trying again."
-        @warn """Failure of the initial evaluation is often indicative of a model specification
-                 that is not well supported by the data and/or a poorly scaled model.
-              """
-        optsum.initial ./=
-            (isempty(m.sqrtwts) ? 1.0 : maximum(m.sqrtwts)^2) *
-            maximum(response(m))
-        optsum.finitial = objective(updateL!(setθ!(m, optsum.initial)))
-    end
-    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
     xmin_ = copy(xmin)
     lb = optsum.lowerbd
@@ -518,7 +472,7 @@ function StatsAPI.fit!(
     end
     loglength = length(fitlog)
     if xmin ≠ xmin_
-        if (zeroobj = obj(xmin_, T[])) ≤ (fmin + optsum.ftol_zero_abs)
+        if (zeroobj = objective!(m, xmin_)) ≤ (fmin + optsum.ftol_zero_abs)
             fmin = zeroobj
             copyto!(xmin, xmin_)
         elseif length(fitlog) > loglength
@@ -529,11 +483,8 @@ function StatsAPI.fit!(
     ## ensure that the parameter values saved in m are xmin
     updateL!(setθ!(m, xmin))
 
-    optsum.feval = opt.numevals
     optsum.final = xmin
     optsum.fmin = fmin
-    optsum.returnvalue = ret
-    _check_nlopt_return(ret)
     return m
 end
 
@@ -828,25 +779,6 @@ function objective(m::LinearMixedModel{T}) where {T}
     return isempty(wts) ? val : val - T(2.0) * sum(log, wts)
 end
 
-"""
-    objective!(m::LinearMixedModel, θ)
-    objective!(m::LinearMixedModel)
-
-Equivalent to `objective(updateL!(setθ!(m, θ)))`.
-
-When `m` has a single, scalar random-effects term, `θ` can be a scalar.
-
-The one-argument method curries and returns a single-argument function of `θ`.
-
-Note that these methods modify `m`.
-The calling function is responsible for restoring the optimal `θ`.
-"""
-function objective! end
-
-function objective!(m::LinearMixedModel)
-    return Base.Fix1(objective!, m)
-end
-
 function objective!(m::LinearMixedModel{T}, θ) where {T}
     return objective(updateL!(setθ!(m, θ)))
 end

src/mixedmodel.jl

@@ -117,6 +117,25 @@ StatsAPI.modelmatrix(m::MixedModel) = m.feterm.x
 
 StatsAPI.nobs(m::MixedModel) = length(m.y)
 
+"""
+    objective!(m::MixedModel, θ)
+    objective!(m::MixedModel)
+
+Equivalent to `objective(updateL!(setθ!(m, θ)))`.
+
+When `m` has a single, scalar random-effects term, `θ` can be a scalar.
+
+The one-argument method curries and returns a single-argument function of `θ`.
+
+Note that these methods modify `m`.
+The calling function is responsible for restoring the optimal `θ`.
+"""
+function objective! end
+
+function objective!(m::MixedModel)
+    return Base.Fix1(objective!, m)
+end
+
 StatsAPI.predict(m::MixedModel) = fitted(m)
 
 function retbl(mat, trm)
@@ -131,13 +150,13 @@ StatsAPI.adjr2(m::MixedModel) = r2(m)
 function StatsAPI.r2(m::MixedModel)
     @error (
         """There is no uniquely defined coefficient of determination for mixed models
-         that has all the properties of the corresponding value for classical 
+         that has all the properties of the corresponding value for classical
          linear models. The GLMM FAQ provides more detail:
-         
+
          https://bbolker.github.io/mixedmodels-misc/glmmFAQ.html#how-do-i-compute-a-coefficient-of-determination-r2-or-an-analogue-for-glmms
 
 
-         Alternatively, MixedModelsExtras provides a naive implementation, but 
+         Alternatively, MixedModelsExtras provides a naive implementation, but
          the warnings there and in the FAQ should be taken seriously!
          """
     )
@@ -148,7 +167,7 @@ end
     raneftables(m::MixedModel; uscale = false)
 
 Return the conditional means of the random effects as a `NamedTuple` of Tables.jl-compliant tables.
-	
+
 !!! note
     The API guarantee is only that the NamedTuple contains Tables.jl tables and not on the particular concrete type of each table.
 """

src/nlopt.jl

@@ -0,0 +1,98 @@
+push!(OPTIMIZATION_BACKENDS, :nlopt)
+
+const NLoptBackend = Val{:nlopt}
+
+function optimize!(m::LinearMixedModel, ::NLoptBackend; progress::Bool=true, thin::Int=tyepmax(Int))
+    optsum = m.optsum
+    opt = Opt(optsum)
+
+    function obj(x, g)
+        isempty(g) || throw(ArgumentError("g should be empty for this objective"))
+        iter += 1
+        val = if isone(iter) && x == optsum.initial
+            optsum.finitial
+        else
+            try
+                objective!(m, x)
+            catch ex
+                # This can happen when the optimizer drifts into an area where
+                # there isn't enough shrinkage. Why finitial? Generally, it will
+                # be the (near) worst case scenario value, so the optimizer won't
+                # view it as an optimum. Using Inf messes up the quadratic
+                # approximation in BOBYQA.
+                ex isa PosDefException || rethrow()
+                optsum.finitial
+            end
+        end
+        progress && ProgressMeter.next!(prog; showvalues=[(:objective, val)])
+        !isone(iter) && iszero(rem(iter, thin)) && push!(fitlog, (copy(x), val))
+        return val
+    end
+    NLopt.min_objective!(opt, obj)
+    prog = ProgressUnknown(; desc="Minimizing", showspeed=true)
+    # start from zero for the initial call to obj before optimization
+    iter = 0
+    fitlog = optsum.fitlog
+
+    try
+        # use explicit evaluation w/o calling opt to avoid confusing iteration count
+        optsum.finitial = objective!(m, optsum.initial)
+    catch ex
+        ex isa PosDefException || rethrow()
+        # give it one more try with a massive change in scaling
+        @info "Initial objective evaluation failed, rescaling initial guess and trying again."
+        @warn """Failure of the initial evaluation is often indicative of a model specification
+                that is not well supported by the data and/or a poorly scaled model.
+            """
+        optsum.initial ./=
+            (isempty(m.sqrtwts) ? 1.0 : maximum(m.sqrtwts)^2) *
+            maximum(response(m))
+        optsum.finitial = objective!(m, optsum.initial)
+    end
+    empty!(fitlog)
+    push!(fitlog, (copy(optsum.initial), optsum.finitial))
+    fmin, xmin, ret = NLopt.optimize!(opt, copyto!(optsum.final, optsum.initial))
+    ProgressMeter.finish!(prog)
+    optsum.feval = opt.numevals
+    optsum.returnvalue = ret
+    _check_nlopt_return(ret)
+    return xmin, fmin
+end
+
+function NLopt.Opt(optsum::OptSummary)
+    lb = optsum.lowerbd
+
+    opt = NLopt.Opt(optsum.optimizer, length(lb))
+    NLopt.ftol_rel!(opt, optsum.ftol_rel) # relative criterion on objective
+    NLopt.ftol_abs!(opt, optsum.ftol_abs) # absolute criterion on objective
+    NLopt.xtol_rel!(opt, optsum.xtol_rel) # relative criterion on parameter values
+    if length(optsum.xtol_abs) == length(lb)  # not true for fast=false optimization in GLMM
+        NLopt.xtol_abs!(opt, optsum.xtol_abs) # absolute criterion on parameter values
+    end
+    NLopt.lower_bounds!(opt, lb)
+    NLopt.maxeval!(opt, optsum.maxfeval)
+    NLopt.maxtime!(opt, optsum.maxtime)
+    if isempty(optsum.initial_step)
+        optsum.initial_step = NLopt.initial_step(opt, optsum.initial, similar(lb))
+    else
+        NLopt.initial_step!(opt, optsum.initial_step)
+    end
+    return opt
+end
+
+
+const _NLOPT_FAILURE_MODES = [
+    :FAILURE,
+    :INVALID_ARGS,
+    :OUT_OF_MEMORY,
+    :FORCED_STOP,
+    :MAXEVAL_REACHED,
+    :MAXTIME_REACHED,
+]
+
+function _check_nlopt_return(ret, failure_modes=_NLOPT_FAILURE_MODES)
+    ret == :ROUNDOFF_LIMITED && @warn("NLopt was roundoff limited")
+    if ret ∈ failure_modes
+        @warn("NLopt optimization failure: $ret")
+    end
+end