refactor optimization to make it easy to add new backends (#802)

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

* fix crossref

* NEWS and version bump

* OptSummary show specialization

* update PRIMA backend to use new infrastructure

* test other PRIMA optimizers

* nlopt for GLMM

* prima  for GLMM

* use scaling

* updated show methods

* make thinning a noop

* BlueStyle

* docs update

Author: palday

NEWS.md

@@ -1,3 +1,9 @@
+MixedModels v4.30.0 Release Notes
+==============================
+- Refactor calls to backend optimizer to make it easier to add and use different optimization backends.
+  The structure of `OptSummary` has been accordingly expanded and `prfit!` has been updated to use this new structure. [#802]
+- Make the `thin` argument to `fit!` a no-op. It complicated several bits of logic without having any real performance benefit in the majority of cases. This argument has been replaced with a `fitlog::Bool=false` that determines whether a log is kept.[#802]
+
 MixedModels v4.29.1 Release Notes
 ==============================
 - Populate `optsum` in `prfit!` call. [#801]
@@ -595,3 +601,4 @@ Package dependencies
 [#795]: https://github.com/JuliaStats/MixedModels.jl/issues/795
 [#799]: https://github.com/JuliaStats/MixedModels.jl/issues/799
 [#801]: https://github.com/JuliaStats/MixedModels.jl/issues/801
+[#802]: https://github.com/JuliaStats/MixedModels.jl/issues/802

Project.toml

@@ -1,7 +1,7 @@
 name = "MixedModels"
 uuid = "ff71e718-51f3-5ec2-a782-8ffcbfa3c316"
 author = ["Phillip Alday <[email protected]>", "Douglas Bates <[email protected]>"]
-version = "4.29.1"
+version = "4.30.0"
 
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"

docs/src/optimization.md

@@ -190,11 +190,11 @@ 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!`:
+By default, `fitlog` is not populated, but passing the keyword argument `fitlog=true` to `fit!` and `refit!` will result in it being populated with the values obtained at each step of optimization:
 
 ```@example Main
-refit!(fm2; thin=1)
-fm2.optsum.fitlog[1:10]
+refit!(fm2; fitlog=true)
+first(fm2.optsum.fitlog, 5)
 DisplayAs.Text(ans) # hide
 ```
 

ext/MixedModelsPRIMAExt.jl

@@ -1,27 +1,41 @@
 module MixedModelsPRIMAExt
 
-using MixedModels: MixedModels, LinearMixedModel, objective!
-using MixedModels: ProgressMeter, ProgressUnknown
+using MixedModels
+using MixedModels: Statistics
+using MixedModels: ProgressMeter, ProgressUnknown, objective!, _objective!
+using LinearAlgebra: PosDefException
 using PRIMA: PRIMA
 
+function __init__()
+    push!(MixedModels.OPTIMIZATION_BACKENDS, :prima)
+    return nothing
+end
+
+const PRIMABackend = Val{:prima}
+
 function MixedModels.prfit!(m::LinearMixedModel;
-                            progress::Bool=true,
-                            REML::Bool=m.optsum.REML,
-                            σ::Union{Real,Nothing}=m.optsum.sigma,
-                            thin::Int=1)
-    optsum = m.optsum
-    copyto!(optsum.final, optsum.initial)
-    optsum.REML = REML
-    optsum.sigma = σ
-    optsum.finitial = objective!(m, optsum.initial)
+                            kwargs...)
+
+    MixedModels.unfit!(m)
+    m.optsum.optimizer = :bobyqa
+    m.optsum.backend = :prima
 
+    return fit!(m; kwargs...)
+end
+
+prima_optimizer!(::Val{:bobyqa}, args...; kwargs...) = PRIMA.bobyqa!(args...; kwargs...)
+prima_optimizer!(::Val{:cobyla}, args...; kwargs...) = PRIMA.cobyla!(args...; kwargs...)
+prima_optimizer!(::Val{:lincoa}, args...; kwargs...) = PRIMA.lincoa!(args...; kwargs...)
+
+function MixedModels.optimize!(m::LinearMixedModel, ::PRIMABackend;
+                               progress::Bool=true, fitlog::Bool=false, kwargs...)
+    optsum = m.optsum
     prog = ProgressUnknown(; desc="Minimizing", showspeed=true)
-    # start from zero for the initial call to obj before optimization
-    iter = 0
-    fitlog = empty!(optsum.fitlog)
+    fitlog && empty!(optsum.fitlog)
+
     function obj(x)
-        iter += 1
-        val = if isone(iter) && x == optsum.initial
+        val = if x == optsum.initial
+            # fast path since we've already evaluated the initial value
             optsum.finitial
         else
             try
@@ -37,19 +51,84 @@ function MixedModels.prfit!(m::LinearMixedModel;
             end
         end
         progress && ProgressMeter.next!(prog; showvalues=[(:objective, val)])
-        if isone(iter) || iszero(rem(iter, thin))
-            push!(fitlog, (copy(x), val))
+        fitlog && push!(optsum.fitlog, (copy(x), val))
+        return val
+    end
+
+    maxfun = optsum.maxfeval > 0 ? optsum.maxfeval : 500 * length(optsum.initial)
+    info = prima_optimizer!(Val(optsum.optimizer), obj, optsum.final;
+                            xl=optsum.lowerbd, maxfun,
+                            optsum.rhoend, optsum.rhobeg)
+    ProgressMeter.finish!(prog)
+    optsum.feval = info.nf
+    optsum.fmin = info.fx
+    optsum.returnvalue = Symbol(info.status)
+    _check_prima_return(info)
+    return optsum.final, optsum.fmin
+end
+
+function MixedModels.optimize!(m::GeneralizedLinearMixedModel, ::PRIMABackend;
+                               progress::Bool=true, fitlog::Bool=false,
+                               fast::Bool=false, verbose::Bool=false, nAGQ=1,
+                               kwargs...)
+    optsum = m.optsum
+    prog = ProgressUnknown(; desc="Minimizing", showspeed=true)
+    fitlog && empty!(opstum.fitlog)
+
+    function obj(x)
+        val = try
+            _objective!(m, x, Val(fast); verbose, nAGQ)
+        catch ex
+            # this allows us to recover from models where e.g. the link isn't
+            # as constraining as it should be
+            ex isa Union{PosDefException,DomainError} || rethrow()
+            x == optsum.initial && rethrow()
+            m.optsum.finitial
         end
+        fitlog && push!(optsum.fitlog, (copy(x), val))
+        verbose && println(round(val; digits=5), " ", x)
+        progress && ProgressMeter.next!(prog; showvalues=[(:objective, val)])
         return val
     end
 
+    optsum.finitial = _objective!(m, optsum.initial, Val(fast); verbose, nAGQ)
+    maxfun = optsum.maxfeval > 0 ? optsum.maxfeval : 500 * length(optsum.initial)
+    scale = if fast
+        nothing
+    else
+        # scale by the standard deviation of the columns of the fixef model matrix
+        # when including the fixef in the nonlinear opt
+        sc = [map(std, eachcol(modelmatrix(m))); fill(1, length(m.θ))]
+        for (i, x) in enumerate(sc)
+            # for nearly constant things, e.g. intercept, we don't want to scale to zero...
+            # also, since we're scaling the _parameters_ and not the data,
+            # we need to invert the scale
+            sc[i] = ifelse(iszero(x), one(x), inv(x))
+        end
+        sc
+    end
+    info = prima_optimizer!(Val(optsum.optimizer), obj, optsum.final;
+                            xl=optsum.lowerbd, maxfun,
+                            optsum.rhoend, optsum.rhobeg,
+                            scale)
     ProgressMeter.finish!(prog)
-    info = PRIMA.bobyqa!(obj, optsum.final; xl=m.optsum.lowerbd)
+
     optsum.feval = info.nf
     optsum.fmin = info.fx
     optsum.returnvalue = Symbol(info.status)
-    optsum.optimizer = :PRIMA_BOBYQA
-    return m
+    _check_prima_return(info)
+
+    return optsum.final, optsum.fmin
+end
+
+function _check_prima_return(info::PRIMA.Info)
+    if !PRIMA.issuccess(info)
+        @warn "PRIMA optimization failure: $(info.status)\n$(PRIMA.reason(info))"
+    end
+
+    return nothing
 end
 
+MixedModels.opt_params(::PRIMABackend) = (:rhobeg, :rhoend, :maxfeval)
+
 end # module

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/bootstrap.jl

@@ -234,7 +234,7 @@ function parametricbootstrap(
     end
     β = convert(Vector{T}, β)
     θ = convert(Vector{T}, θ)
-    # scratch -- note that this is the length of the unpivoted coef vector 
+    # scratch -- note that this is the length of the unpivoted coef vector
     βsc = coef(morig)
     θsc = zeros(ftype, length(θ))
     p = length(βsc)
@@ -254,7 +254,7 @@ function parametricbootstrap(
     )
     samp = replicate(n; progress) do
         simulate!(rng, m; β, σ, θ)
-        refit!(m; progress=false)
+        refit!(m; progress=false, fitlog=false)
         (
             objective=ftype.(m.objective),
             σ=ismissing(m.σ) ? missing : ftype(m.σ),

src/generalizedlinearmixedmodel.jl

@@ -219,7 +219,8 @@ 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`.
+The `thin` argument is ignored: it had no impact on the final model fit and the logic around
+thinning the `fitlog` was needlessly complicated for a trivial performance gain.
 
 By default, the starting values for model fitting are taken from a (non mixed,
 i.e. marginal ) GLM fit. Experience with larger datasets (many thousands of
@@ -247,7 +248,10 @@ function StatsAPI.fit!(
     nAGQ::Integer=1,
     progress::Bool=true,
     thin::Int=typemax(Int),
+    fitlog::Bool=false,
     init_from_lmm=Set(),
+    backend::Symbol=m.optsum.backend,
+    optimizer::Symbol=m.optsum.optimizer,
 ) where {T}
     β = copy(m.β)
     θ = copy(m.θ)
@@ -277,60 +281,33 @@ function StatsAPI.fit!(
         optsum.initial = vcat(β, lm.optsum.final)
         optsum.final = copy(optsum.initial)
     end
-    setpar! = fast ? setθ! : setβθ!
-    prog = ProgressUnknown(; desc="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 = try
-            deviance(pirls!(setpar!(m, x), fast, verbose), nAGQ)
-        catch ex
-            # this allows us to recover from models where e.g. the link isn't
-            # as constraining as it should be
-            ex isa Union{PosDefException,DomainError} || rethrow()
-            iter == 1 && rethrow()
-            m.optsum.finitial
-        end
-        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)
+
+    optsum.backend = backend
+    optsum.optimizer = optimizer
+
+    xmin, fmin = optimize!(m; progress, fitlog, fast, verbose, nAGQ)
+
     ## check if very small parameter values bounded below by zero can be set to zero
     xmin_ = copy(xmin)
     for i in eachindex(xmin_)
         if iszero(optsum.lowerbd[i]) && zero(T) < xmin_[i] < optsum.xtol_zero_abs
             xmin_[i] = zero(T)
         end
     end
-    loglength = length(fitlog)
     if xmin ≠ xmin_
-        if (zeroobj = obj(xmin_, T[])) ≤ (fmin + optsum.ftol_zero_abs)
+        if (zeroobj = objective!(m, xmin_; nAGQ, fast, verbose)) ≤
+            (fmin + optsum.ftol_zero_abs)
             fmin = zeroobj
             copyto!(xmin, xmin_)
-        elseif length(fitlog) > loglength
-            # remove unused extra log entry
-            pop!(fitlog)
+            fitlog && push!(optsum.fitlog, (copy(xmin), fmin))
         end
     end
+
     ## ensure that the parameter values saved in m are xmin
-    pirls!(setpar!(m, xmin), fast, verbose)
-    optsum.nAGQ = nAGQ
-    optsum.feval = opt.numevals
+    objective!(m, xmin; fast, verbose, nAGQ)
     optsum.final = xmin
     optsum.fmin = fmin
-    optsum.returnvalue = ret
-    _check_nlopt_return(ret)
+    optsum.nAGQ = nAGQ
     return m
 end
 
@@ -540,6 +517,30 @@ function StatsAPI.loglikelihood(m::GeneralizedLinearMixedModel{T}) where {T}
     return accum - (mapreduce(u -> sum(abs2, u), +, m.u) + logdet(m)) / 2
 end
 
+# Base.Fix1 doesn't forward kwargs
+function objective!(m::GeneralizedLinearMixedModel; fast=false, kwargs...)
+    return x -> _objective!(m, x, Val(fast); kwargs...)
+end
+
+function objective!(m::GeneralizedLinearMixedModel{T}, x; fast=false, kwargs...) where {T}
+    return _objective!(m, x, Val(fast); kwargs...)
+end
+
+# normally, it doesn't make sense to move a simple branch to dispatch
+# HOWEVER, this winds up getting called in optimization a lot and
+# moving this to a realization here allows us to avoid dynamic dispatch on setθ! / setθβ!
+function _objective!(
+    m::GeneralizedLinearMixedModel{T}, x, ::Val{true}; nAGQ=1, verbose=false
+) where {T}
+    return deviance(pirls!(setθ!(m, x), true, verbose), nAGQ)
+end
+
+function _objective!(
+    m::GeneralizedLinearMixedModel{T}, x, ::Val{false}; nAGQ=1, verbose=false
+) where {T}
+    return deviance(pirls!(setβθ!(m, x), false, verbose), nAGQ)
+end
+
 function Base.propertynames(m::GeneralizedLinearMixedModel, private::Bool=false)
     return (
         :A,