Update optimization to adopt new API from ProximalAlgorithms/ProximalCore

Author: qntwrsm

Project.toml

@@ -4,6 +4,7 @@ authors = ["Quint Wiersma <qnt.wrsm@gmail.com> and contributors"]
 version = "0.1.0"
 
 [deps]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -23,6 +24,8 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
 
 [compat]
+DifferentiationInterface = "0.6.53"
+FiniteDiff = "2.27.0"
 julia = "1.9"
 
 [extras]

src/DynamicFactorModels.jl

@@ -21,7 +21,7 @@ using Distributions: MvNormal
 
 using StatsAPI: StatisticalModel
 
-using MultivariateStats: PCA, projection, predict
+using MultivariateStats: PCA, fit, projection, predict
 
 using IrrationalConstants: log2π
 
@@ -30,9 +30,10 @@ using Distances: evaluate, Chebyshev
 using LogExpFunctions: logistic, logit
 using Optim
 using FiniteDiff
+using DifferentiationInterface: AutoFiniteDiff
 using ProximalCore
 using ProximalOperators: NormL1, NormL21, TotalVariation1D
-using ProximalAlgorithms: FastForwardBackward
+using ProximalAlgorithms
 
 using ProgressMeter
 

src/fit/solver.jl

@@ -79,14 +79,9 @@ function update_loadings!(F::AbstractUnrestrictedFactorProcess, y::AbstractMatri
 
         return (0.5 * dot(ΩΛ, Λ * Eff) - dot(ΩΛ, Eyf)) / length(V)
     end
-    function gradient!(∇::AbstractMatrix, Λ::AbstractMatrix)
-        ∇ .= Σ \ (Λ * Eff - Eyf)
-        ∇ ./= length(V)
-
-        return nothing
-    end
-    f = ObjectiveGradientWrapper(objective, gradient!)
-    ffb = FastForwardBackward(maxit = 100, tol = 1e-2)
+    gradient(Λ::AbstractMatrix) = (Σ \ (Λ * Eff - Eyf)) ./ length(V)
+    f = ObjectiveGradientWrapper(objective, gradient)
+    ffb = ProximalAlgorithms.FastForwardBackward(maxit = 1000, tol = 1e-4)
     (solution, _) = ffb(x0 = loadings(F), f = f, g = regularizer)
     loadings(F) .= solution
 
@@ -220,14 +215,9 @@ function update!(μ::Exogenous, y::AbstractMatrix, Σ::AbstractMatrix, regulariz
 
         return (0.5 * dot(Ωβ, β * XX) - dot(Ωβ, yX)) / size(regressors(μ), 2)
     end
-    function gradient!(∇::AbstractMatrix, β::AbstractMatrix)
-        ∇ .= Σ \ (β * XX - yX)
-        ∇ ./= size(regressors(μ), 2)
-
-        return nothing
-    end
-    f = ObjectiveGradientWrapper(objective, gradient!)
-    ffb = FastForwardBackward(maxit = 100, tol = 1e-2)
+    gradient(β::AbstractMatrix) = (Σ \ (β * XX - yX)) ./ size(regressors(μ), 2)
+    f = ObjectiveGradientWrapper(objective, gradient)
+    ffb = ProximalAlgorithms.FastForwardBackward(maxit = 1000, tol = 1e-4)
     (solution, _) = ffb(x0 = slopes(μ), f = f, g = regularizer)
     slopes(μ) .= solution
 
@@ -274,7 +264,7 @@ function update!(ε::SpatialAutoregression, e::AbstractMatrix, Λ::AbstractMatri
 
         return -logdet(G) + 0.5 * dot(Ω, Eee) / size(e, 2)
     end
-    opt = optimize(objective, logit.((spatial(ε) .+ offset) ./ scale), ConjugateGradient(),
+    opt = optimize(objective, logit.((spatial(ε) .+ offset) ./ scale), LBFGS(),
                    Optim.Options(g_tol = 1e-4))
     spatial(ε) .= scale .* logistic.(Optim.minimizer(opt)) .- offset
 
@@ -302,9 +292,8 @@ function update!(ε::SpatialAutoregression, e::AbstractMatrix, Λ::AbstractMatri
 
         return -logdet(G) + 0.5 * dot(Ω, Eee) / size(e, 2)
     end
-    cache = FiniteDiff.GradientCache(copy(spatial(ε)), copy(spatial(ε)))
-    f = ObjectiveWrapper(objective, cache)
-    ffb = FastForwardBackward(maxit = 100, tol = 1e-2)
+    f = ProximalAlgorithms.AutoDifferentiable(objective, AutoFiniteDiff())
+    ffb = ProximalAlgorithms.FastForwardBackward(maxit = 1000, tol = 1e-4)
     (solution, _) = ffb(x0 = logit.((spatial(ε) .+ offset) ./ scale), f = f,
                         g = regularizer)
     spatial(ε) .= scale .* logistic.(solution) .- offset
@@ -333,7 +322,7 @@ function update!(ε::SpatialMovingAverage, e::AbstractMatrix, Λ::AbstractMatrix
 
         return logdet(G) + 0.5 * tr(Σ \ Eee) / size(e, 2)
     end
-    opt = optimize(objective, logit.((spatial(ε) .+ offset) ./ scale), ConjugateGradient(),
+    opt = optimize(objective, logit.((spatial(ε) .+ offset) ./ scale), LBFGS(),
                    Optim.Options(g_tol = 1e-4))
     spatial(ε) .= scale .* logistic.(Optim.minimizer(opt)) .- offset
 
@@ -361,9 +350,8 @@ function update!(ε::SpatialMovingAverage, e::AbstractMatrix, Λ::AbstractMatrix
 
         return logdet(G) + 0.5 * tr(Σ \ Eee) / size(e, 2)
     end
-    cache = FiniteDiff.GradientCache(copy(spatial(ε)), copy(spatial(ε)))
-    f = ObjectiveWrapper(objective, cache)
-    ffb = FastForwardBackward(maxit = 100, tol = 1e-2)
+    f = ProximalAlgorithms.AutoDifferentiable(objective, AutoFiniteDiff())
+    ffb = ProximalAlgorithms.FastForwardBackward(maxit = 1000, tol = 1e-4)
     (solution, _) = ffb(x0 = logit.((spatial(ε) .+ offset) ./ scale), f = f,
                         g = regularizer)
     spatial(ε) .= scale .* logistic.(solution) .- offset

src/fit/utilities.jl

@@ -9,24 +9,6 @@ utilities.jl
 @date: 2023/12/08
 =#
 
-"""
-    ObjectiveWrapper
-
-Wrapper for objective functions for gradient dispatching to finit differences.
-"""
-struct ObjectiveWrapper{F, C}
-    f::F
-    cache::C
-end
-
-(f::ObjectiveWrapper)(x) = f.f(x)
-
-function ProximalCore.gradient!(y, f::ObjectiveWrapper, x)
-    FiniteDiff.finite_difference_gradient!(y, f, x, f.cache)
-
-    return f(x)
-end
-
 """
     ObjectiveGradientWrapper
 
@@ -35,16 +17,12 @@ function.
 """
 struct ObjectiveGradientWrapper{F, G}
     f::F
-    g!::G
+    g::G
 end
 
 (f::ObjectiveGradientWrapper)(x) = f.f(x)
 
-function ProximalCore.gradient!(y, f::ObjectiveGradientWrapper, x)
-    f.g!(y, x)
-
-    return f(x)
-end
+ProximalAlgorithms.value_and_gradient(f::ObjectiveGradientWrapper, x) = (f(x), f.g(x))
 
 """
     init!(model, method)
@@ -423,7 +401,7 @@ function objective(model::DynamicFactorModel, regularizer::NamedTuple)
 end
 
 function dof(model::DynamicFactorModel)
-    R = size(process(model))
+    R = nfactors(model)
 
     # factor component
     process(model) isa AbstractUnrestrictedFactorProcess &&