Add support for GalacticOptim 3 and fix test errors (#1834)

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.21.3"
+version = "0.21.4"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/modes/ModeEstimation.jl

@@ -3,7 +3,7 @@ module ModeEstimation
 using ..Turing
 using Bijectors
 using Random
-using SciMLBase: OptimizationFunction, OptimizationProblem, AbstractADType
+using SciMLBase: OptimizationFunction, OptimizationProblem, AbstractADType, NoAD
 
 using DynamicPPL
 using DynamicPPL: Model, AbstractContext, VarInfo, VarName,
@@ -291,24 +291,47 @@ function optim_objective(model::DynamicPPL.Model, estimator::Union{MLE, MAP}; co
 end
 
 
-function optim_function(model::DynamicPPL.Model, estimator::Union{MLE, MAP}; constrained::Bool=true, autoad::Union{Nothing, AbstractADType}=nothing)
+function optim_function(
+    model::Model,
+    estimator::Union{MLE, MAP};
+    constrained::Bool=true,
+    autoad::Union{Nothing, AbstractADType}=NoAD(),
+)
+    if autoad === nothing
+        Base.depwarn("the use of `autoad=nothing` is deprecated, please use `autoad=SciMLBase.NoAD()`", :optim_function)
+    end
+
     obj, init, t = optim_objective(model, estimator; constrained=constrained)
 
-    l(x,p) = obj(x)
-    f = isa(autoad, AbstractADType) ? OptimizationFunction(l, autoad) : OptimizationFunction(l; grad = (G,x,p) -> obj(nothing, G, nothing, x), hess = (H,x,p) -> obj(nothing, nothing, H, x))
+    l(x, _) = obj(x)
+    f = if autoad isa AbstractADType && autoad !== NoAD()
+        OptimizationFunction(l, autoad)
+    else
+        OptimizationFunction(
+            l;
+            grad = (G,x,p) -> obj(nothing, G, nothing, x),
+            hess = (H,x,p) -> obj(nothing, nothing, H, x),
+        )
+    end
 
     return (func=f, init=init, transform = t)
 end
 
 
-function optim_problem(model::DynamicPPL.Model, estimator::Union{MAP, MLE}; constrained::Bool=true, init_theta=nothing, autoad::Union{Nothing, AbstractADType}=nothing, kwargs...)
-    f = optim_function(model, estimator; constrained=constrained, autoad=autoad)
-
-    init_theta = init_theta === nothing ? f.init() : f.init(init_theta)
+function optim_problem(
+    model::Model,
+    estimator::Union{MAP, MLE};
+    constrained::Bool=true,
+    init_theta=nothing,
+    autoad::Union{Nothing, AbstractADType}=NoAD(),
+    kwargs...,
+)
+    f, init, transform = optim_function(model, estimator; constrained=constrained, autoad=autoad)
 
-    prob = OptimizationProblem(f.func, init_theta, nothing; kwargs...)
+    u0 = init_theta === nothing ? init() : init(init_theta)
+    prob = OptimizationProblem(f, u0; kwargs...)
 
-    return (prob=prob, init=f.init, transform = f.transform)
+    return (; prob, init, transform)
 end
 
 end

test/Project.toml

@@ -12,6 +12,7 @@ DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GalacticOptim = "a75be94c-b780-496d-a8a9-0878b188d577"
+GalacticOptimJL = "9d3c5eb1-403b-401b-8c0f-c11105342e6b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
@@ -43,7 +44,8 @@ DynamicHMC = "2.1.6, 3.0"
 DynamicPPL = "0.19.1"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12"
-GalacticOptim = "0.4, 1, 2"
+GalacticOptim = "3"
+GalacticOptimJL = "0.1"
 MCMCChains = "5"
 Memoization = "0.1.4"
 NamedArrays = "0.9.4"

test/modes/ModeEstimation.jl

@@ -1,13 +1,11 @@
-
-
 @testset "ModeEstimation.jl" begin
     @testset "gdemo" begin
         @testset "MLE" begin
             Random.seed!(222)
             true_value = [0.0625, 1.75]
 
             f1 = optim_function(gdemo_default, MLE();constrained=false)
-            p1 = GalacticOptim.OptimizationProblem(f1.func, f1.init(true_value), nothing)
+            p1 = OptimizationProblem(f1.func, f1.init(true_value))
 
             p2 = optim_objective(gdemo_default, MLE();constrained=false)
 
@@ -39,7 +37,7 @@
             true_value = [49 / 54, 7 / 6]
 
             f1 = optim_function(gdemo_default, MAP();constrained=false)
-            p1 = GalacticOptim.OptimizationProblem(f1.func, f1.init(true_value), nothing)
+            p1 = OptimizationProblem(f1.func, f1.init(true_value))
 
             p2 = optim_objective(gdemo_default, MAP();constrained=false)
 
@@ -73,7 +71,7 @@
             ub = [2.0, 2.0]
 
             f1 = optim_function(gdemo_default, MLE();constrained=true)
-            p1 = GalacticOptim.OptimizationProblem(f1.func, f1.init(true_value), nothing; lb=lb, ub=ub)
+            p1 = GalacticOptim.OptimizationProblem(f1.func, f1.init(true_value); lb=lb, ub=ub)
 
             p2 = optim_objective(gdemo_default, MLE();constrained=true)
 
@@ -101,7 +99,7 @@
             ub = [2.0, 2.0]
 
             f1 = optim_function(gdemo_default, MAP();constrained=true)
-            p1 = GalacticOptim.OptimizationProblem(f1.func, f1.init(true_value), nothing; lb=lb, ub=ub)
+            p1 = GalacticOptim.OptimizationProblem(f1.func, f1.init(true_value); lb=lb, ub=ub)
 
             p2 = optim_objective(gdemo_default, MAP();constrained=true)
 

test/runtests.jl

@@ -7,6 +7,7 @@ using DistributionsAD
 using FiniteDifferences
 using ForwardDiff
 using GalacticOptim
+using GalacticOptimJL
 using MCMCChains
 using Memoization
 using NamedArrays