From 74869f4b825f9b99c75a2fc4bc7638acb4a4da25 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Thu, 27 Jul 2023 18:27:01 +0100 Subject: [PATCH] Bump bijectors compat (#2052) * CompatHelper: bump compat for Bijectors to 0.13, (keep existing compat) * Update Project.toml * Replacement for #2039 (#2040) * Fix testset for external samplers * Update abstractmcmc.jl * Update test/contrib/inference/abstractmcmc.jl Co-authored-by: Tor Erlend Fjelde * Update test/contrib/inference/abstractmcmc.jl Co-authored-by: Tor Erlend Fjelde * Update FillArrays compat to 1.4.1 (#2035) * Update FillArrays compat to 1.4.0 * Update test compat * Try to enable ReverseDiff tests * Update Project.toml * Update Project.toml * Bump version * Revert dependencies on FillArrays (#2042) * Update Project.toml * Update Project.toml * Fix redundant definition of `getstats` (#2044) * Fix redundant definition of `getstats` * Update Inference.jl * Revert "Update Inference.jl" This reverts commit e4f51c24fa7450d625d18b21ca3a273bb2d736d0. * Bump version --------- Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> * Transfer some test utility function into DynamicPPL (#2049) * Update OptimInterface.jl * Only run optimisation tests in numerical stage. * fix function lookup after moving functions --------- Co-authored-by: Xianda Sun * Move Optim support to extension (#2051) * Move Optim support to extension * More imports * Update Project.toml --------- Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> --------- Co-authored-by: CompatHelper Julia Co-authored-by: haris organtzidis Co-authored-by: Tor Erlend Fjelde Co-authored-by: David Widmann Co-authored-by: Xianda Sun Co-authored-by: Cameron Pfiffer --- Project.toml | 16 +++- .../TuringOptimExt.jl | 92 ++++++++++--------- src/Turing.jl | 41 +++++---- src/contrib/inference/abstractmcmc.jl | 1 - test/Project.toml | 4 +- test/contrib/inference/abstractmcmc.jl | 10 +- test/essential/ad.jl | 9 +- test/modes/OptimInterface.jl | 41 +-------- 8 files changed, 97 insertions(+), 117 deletions(-) rename src/modes/OptimInterface.jl => ext/TuringOptimExt.jl (69%) diff --git a/Project.toml b/Project.toml index a4cc4686a..f4fb14669 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Turing" uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" -version = "0.26.4" +version = "0.27" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" @@ -16,7 +16,6 @@ DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" EllipticalSliceSampling = "cad2338a-1db2-11e9-3401-43bc07c9ede2" -FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -44,20 +43,20 @@ AdvancedMH = "0.6.8, 0.7" AdvancedPS = "0.4" AdvancedVI = "0.2" BangBang = "0.3" -Bijectors = "0.12" +Bijectors = "0.13.2" DataStructures = "0.18" Distributions = "0.23.3, 0.24, 0.25" DistributionsAD = "0.6" DocStringExtensions = "0.8, 0.9" DynamicPPL = "0.23" EllipticalSliceSampling = "0.5, 1" -FillArrays = "=1.0.0" ForwardDiff = "0.10.3" Libtask = "0.7, 0.8" LogDensityProblems = "2" LogDensityProblemsAD = "1.4" MCMCChains = "5, 6" NamedArrays = "0.9" +Optim = "1" Reexport = "0.2, 1" Requires = "0.5, 1.0" SciMLBase = "1.37.1" @@ -68,3 +67,12 @@ StatsBase = "0.32, 0.33, 0.34" StatsFuns = "0.8, 0.9, 1" Tracker = "0.2.3" julia = "1.7" + +[weakdeps] +Optim = "429524aa-4258-5aef-a3af-852621145aeb" + +[extensions] +TuringOptimExt = "Optim" + +[extras] +Optim = "429524aa-4258-5aef-a3af-852621145aeb" diff --git a/src/modes/OptimInterface.jl b/ext/TuringOptimExt.jl similarity index 69% rename from src/modes/OptimInterface.jl rename to ext/TuringOptimExt.jl index f477fedb9..714cea202 100644 --- a/src/modes/OptimInterface.jl +++ b/ext/TuringOptimExt.jl @@ -1,14 +1,14 @@ -using Setfield -using DynamicPPL: DefaultContext, LikelihoodContext -using DynamicPPL: DynamicPPL -import .Optim -import .Optim: optimize -import ..ForwardDiff -import NamedArrays -import StatsBase -import Printf -import StatsAPI - +module TuringOptimExt + +if isdefined(Base, :get_extension) + import Turing + import Turing: Distributions, DynamicPPL, ForwardDiff, NamedArrays, Printf, Setfield, Statistics, StatsAPI, StatsBase + import Optim +else + import ..Turing + import ..Turing: Distributions, DynamicPPL, ForwardDiff, NamedArrays, Printf, Setfield, Statistics, StatsAPI, StatsBase + import ..Optim +end """ ModeResult{ @@ -23,7 +23,7 @@ A wrapper struct to store various results from a MAP or MLE estimation. struct ModeResult{ V<:NamedArrays.NamedArray, O<:Optim.MultivariateOptimizationResults, - M<:OptimLogDensity + M<:Turing.OptimLogDensity } <: StatsBase.StatisticalModel "A vector with the resulting point estimates." values::V @@ -57,10 +57,10 @@ function StatsBase.coeftable(m::ModeResult; level::Real=0.95) estimates = m.values.array[:, 1] stderrors = StatsBase.stderror(m) zscore = estimates ./ stderrors - p = map(z -> StatsAPI.pvalue(Normal(), z; tail=:both), zscore) + p = map(z -> StatsAPI.pvalue(Distributions.Normal(), z; tail=:both), zscore) # Confidence interval (CI) - q = quantile(Normal(), (1 + level) / 2) + q = Statistics.quantile(Distributions.Normal(), (1 + level) / 2) ci_low = estimates .- q .* stderrors ci_high = estimates .+ q .* stderrors @@ -80,7 +80,7 @@ function StatsBase.informationmatrix(m::ModeResult; hessian_function=ForwardDiff # Hessian is computed with respect to the untransformed parameters. linked = DynamicPPL.istrans(m.f.varinfo) if linked - @set! m.f.varinfo = invlink!!(m.f.varinfo, m.f.model) + Setfield.@set! m.f.varinfo = DynamicPPL.invlink!!(m.f.varinfo, m.f.model) end # Calculate the Hessian. @@ -90,7 +90,7 @@ function StatsBase.informationmatrix(m::ModeResult; hessian_function=ForwardDiff # Link it back if we invlinked it. if linked - @set! m.f.varinfo = link!!(m.f.varinfo, m.f.model) + Setfield.@set! m.f.varinfo = DynamicPPL.link!!(m.f.varinfo, m.f.model) end return NamedArrays.NamedArray(info, (varnames, varnames)) @@ -126,18 +126,18 @@ mle = optimize(model, MLE()) mle = optimize(model, MLE(), NelderMead()) ``` """ -function Optim.optimize(model::Model, ::MLE, options::Optim.Options=Optim.Options(); kwargs...) +function Optim.optimize(model::DynamicPPL.Model, ::Turing.MLE, options::Optim.Options=Optim.Options(); kwargs...) return _mle_optimize(model, options; kwargs...) end -function Optim.optimize(model::Model, ::MLE, init_vals::AbstractArray, options::Optim.Options=Optim.Options(); kwargs...) +function Optim.optimize(model::DynamicPPL.Model, ::Turing.MLE, init_vals::AbstractArray, options::Optim.Options=Optim.Options(); kwargs...) return _mle_optimize(model, init_vals, options; kwargs...) end -function Optim.optimize(model::Model, ::MLE, optimizer::Optim.AbstractOptimizer, options::Optim.Options=Optim.Options(); kwargs...) +function Optim.optimize(model::DynamicPPL.Model, ::Turing.MLE, optimizer::Optim.AbstractOptimizer, options::Optim.Options=Optim.Options(); kwargs...) return _mle_optimize(model, optimizer, options; kwargs...) end function Optim.optimize( - model::Model, - ::MLE, + model::DynamicPPL.Model, + ::Turing.MLE, init_vals::AbstractArray, optimizer::Optim.AbstractOptimizer, options::Optim.Options=Optim.Options(); @@ -146,9 +146,9 @@ function Optim.optimize( return _mle_optimize(model, init_vals, optimizer, options; kwargs...) end -function _mle_optimize(model::Model, args...; kwargs...) - ctx = OptimizationContext(DynamicPPL.LikelihoodContext()) - return _optimize(model, OptimLogDensity(model, ctx), args...; kwargs...) +function _mle_optimize(model::DynamicPPL.Model, args...; kwargs...) + ctx = Turing.OptimizationContext(DynamicPPL.LikelihoodContext()) + return _optimize(model, Turing.OptimLogDensity(model, ctx), args...; kwargs...) end """ @@ -172,18 +172,18 @@ map_est = optimize(model, MAP(), NelderMead()) ``` """ -function Optim.optimize(model::Model, ::MAP, options::Optim.Options=Optim.Options(); kwargs...) +function Optim.optimize(model::DynamicPPL.Model, ::Turing.MAP, options::Optim.Options=Optim.Options(); kwargs...) return _map_optimize(model, options; kwargs...) end -function Optim.optimize(model::Model, ::MAP, init_vals::AbstractArray, options::Optim.Options=Optim.Options(); kwargs...) +function Optim.optimize(model::DynamicPPL.Model, ::Turing.MAP, init_vals::AbstractArray, options::Optim.Options=Optim.Options(); kwargs...) return _map_optimize(model, init_vals, options; kwargs...) end -function Optim.optimize(model::Model, ::MAP, optimizer::Optim.AbstractOptimizer, options::Optim.Options=Optim.Options(); kwargs...) +function Optim.optimize(model::DynamicPPL.Model, ::Turing.MAP, optimizer::Optim.AbstractOptimizer, options::Optim.Options=Optim.Options(); kwargs...) return _map_optimize(model, optimizer, options; kwargs...) end function Optim.optimize( - model::Model, - ::MAP, + model::DynamicPPL.Model, + ::Turing.MAP, init_vals::AbstractArray, optimizer::Optim.AbstractOptimizer, options::Optim.Options=Optim.Options(); @@ -192,9 +192,9 @@ function Optim.optimize( return _map_optimize(model, init_vals, optimizer, options; kwargs...) end -function _map_optimize(model::Model, args...; kwargs...) - ctx = OptimizationContext(DynamicPPL.DefaultContext()) - return _optimize(model, OptimLogDensity(model, ctx), args...; kwargs...) +function _map_optimize(model::DynamicPPL.Model, args...; kwargs...) + ctx = Turing.OptimizationContext(DynamicPPL.DefaultContext()) + return _optimize(model, Turing.OptimLogDensity(model, ctx), args...; kwargs...) end """ @@ -203,8 +203,8 @@ end Estimate a mode, i.e., compute a MLE or MAP estimate. """ function _optimize( - model::Model, - f::OptimLogDensity, + model::DynamicPPL.Model, + f::Turing.OptimLogDensity, optimizer::Optim.AbstractOptimizer=Optim.LBFGS(), args...; kwargs... @@ -213,8 +213,8 @@ function _optimize( end function _optimize( - model::Model, - f::OptimLogDensity, + model::DynamicPPL.Model, + f::Turing.OptimLogDensity, options::Optim.Options=Optim.Options(), args...; kwargs... @@ -223,8 +223,8 @@ function _optimize( end function _optimize( - model::Model, - f::OptimLogDensity, + model::DynamicPPL.Model, + f::Turing.OptimLogDensity, init_vals::AbstractArray=DynamicPPL.getparams(f), options::Optim.Options=Optim.Options(), args...; @@ -234,8 +234,8 @@ function _optimize( end function _optimize( - model::Model, - f::OptimLogDensity, + model::DynamicPPL.Model, + f::Turing.OptimLogDensity, init_vals::AbstractArray=DynamicPPL.getparams(f), optimizer::Optim.AbstractOptimizer=Optim.LBFGS(), options::Optim.Options=Optim.Options(), @@ -244,8 +244,8 @@ function _optimize( ) # Convert the initial values, since it is assumed that users provide them # in the constrained space. - @set! f.varinfo = DynamicPPL.unflatten(f.varinfo, init_vals) - @set! f.varinfo = DynamicPPL.link!!(f.varinfo, model) + Setfield.@set! f.varinfo = DynamicPPL.unflatten(f.varinfo, init_vals) + Setfield.@set! f.varinfo = DynamicPPL.link!!(f.varinfo, model) init_vals = DynamicPPL.getparams(f) # Optimize! @@ -258,10 +258,10 @@ function _optimize( # Get the VarInfo at the MLE/MAP point, and run the model to ensure # correct dimensionality. - @set! f.varinfo = DynamicPPL.unflatten(f.varinfo, M.minimizer) - @set! f.varinfo = invlink!!(f.varinfo, model) + Setfield.@set! f.varinfo = DynamicPPL.unflatten(f.varinfo, M.minimizer) + Setfield.@set! f.varinfo = DynamicPPL.invlink!!(f.varinfo, model) vals = DynamicPPL.getparams(f) - @set! f.varinfo = link!!(f.varinfo, model) + Setfield.@set! f.varinfo = DynamicPPL.link!!(f.varinfo, model) # Make one transition to get the parameter names. ts = [Turing.Inference.Transition( @@ -275,3 +275,5 @@ function _optimize( return ModeResult(vmat, M, -M.minimum, f) end + +end # module diff --git a/src/Turing.jl b/src/Turing.jl index fc29cb4e9..b4ebef19d 100644 --- a/src/Turing.jl +++ b/src/Turing.jl @@ -11,6 +11,12 @@ import AdvancedVI using DynamicPPL: DynamicPPL, LogDensityFunction import DynamicPPL: getspace, NoDist, NamedDist import LogDensityProblems +import NamedArrays +import Setfield +import StatsAPI +import StatsBase + +import Printf import Random const PROGRESS = Ref(true) @@ -48,26 +54,9 @@ using .Inference include("variational/VariationalInference.jl") using .Variational -@init @require DynamicHMC="bbc10e6e-7c05-544b-b16e-64fede858acb" begin - @eval Inference begin - import ..DynamicHMC - - if isdefined(DynamicHMC, :mcmc_with_warmup) - include("contrib/inference/dynamichmc.jl") - else - error("Please update DynamicHMC, v1.x is no longer supported") - end - end -end - include("modes/ModeEstimation.jl") using .ModeEstimation -@init @require Optim="429524aa-4258-5aef-a3af-852621145aeb" @eval begin - include("modes/OptimInterface.jl") - export optimize -end - ########### # Exports # ########### @@ -146,4 +135,22 @@ export @model, # modelling optim_objective, optim_function, optim_problem + +function __init__() + @static if !isdefined(Base, :get_extension) + @require Optim="429524aa-4258-5aef-a3af-852621145aeb" include("../ext/TuringOptimExt.jl") + end + @require DynamicHMC="bbc10e6e-7c05-544b-b16e-64fede858acb" begin + @eval Inference begin + import ..DynamicHMC + + if isdefined(DynamicHMC, :mcmc_with_warmup) + include("contrib/inference/dynamichmc.jl") + else + error("Please update DynamicHMC, v1.x is no longer supported") + end + end + end +end + end diff --git a/src/contrib/inference/abstractmcmc.jl b/src/contrib/inference/abstractmcmc.jl index 19411ac99..0f158adfc 100644 --- a/src/contrib/inference/abstractmcmc.jl +++ b/src/contrib/inference/abstractmcmc.jl @@ -19,7 +19,6 @@ getparams(transition::AdvancedHMC.Transition) = transition.z.θ getstats(transition::AdvancedHMC.Transition) = transition.stat getparams(transition::AdvancedMH.Transition) = transition.params -getstats(transition) = NamedTuple() getvarinfo(f::DynamicPPL.LogDensityFunction) = f.varinfo getvarinfo(f::LogDensityProblemsAD.ADGradientWrapper) = getvarinfo(parent(f)) diff --git a/test/Project.toml b/test/Project.toml index cec3c0d4b..470d50afe 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,7 +8,6 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" -FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -42,14 +41,13 @@ Distributions = "0.25" DistributionsAD = "0.6.3" DynamicHMC = "2.1.6, 3.0" DynamicPPL = "0.23" -FillArrays = "=1.0.0" FiniteDifferences = "0.10.8, 0.11, 0.12" ForwardDiff = "0.10.12 - 0.10.32, 0.10" LogDensityProblems = "2" LogDensityProblemsAD = "1.4" MCMCChains = "5, 6" NamedArrays = "0.9.4" -Optim = "0.22, 1.0" +Optim = "1" Optimization = "3.5" OptimizationOptimJL = "0.1" PDMats = "0.10, 0.11" diff --git a/test/contrib/inference/abstractmcmc.jl b/test/contrib/inference/abstractmcmc.jl index 691130635..ca7e1b3ec 100644 --- a/test/contrib/inference/abstractmcmc.jl +++ b/test/contrib/inference/abstractmcmc.jl @@ -41,7 +41,7 @@ function initialize_mh(model) end @testset "External samplers" begin - @testset "AdvancedHMC.jl" begin + @turing_testset "AdvancedHMC.jl" begin for model in DynamicPPL.TestUtils.DEMO_MODELS # Need some functionality to initialize the sampler. # TODO: Remove this once the constructors in the respective packages become "lazy". @@ -52,12 +52,13 @@ end 5_000; nadapts=1_000, discard_initial=1_000, - rtol=0.2 + rtol=0.2, + sampler_name="AdvancedHMC" ) end end - @testset "AdvancedMH.jl" begin + @turing_testset "AdvancedMH.jl" begin for model in DynamicPPL.TestUtils.DEMO_MODELS # Need some functionality to initialize the sampler. # TODO: Remove this once the constructors in the respective packages become "lazy". @@ -68,7 +69,8 @@ end 10_000; discard_initial=1_000, thinning=10, - rtol=0.2 + rtol=0.2, + sampler_name="AdvancedMH" ) end end diff --git a/test/essential/ad.jl b/test/essential/ad.jl index 0a60a3e0f..8a0241a83 100644 --- a/test/essential/ad.jl +++ b/test/essential/ad.jl @@ -95,15 +95,14 @@ sample(dir(), HMC(0.01, 1), 1000) Turing.setrdcache(false) end - # FIXME: For some reasons PDMatDistribution AD tests fail with ReverseDiff @testset "PDMatDistribution AD" begin @model function wishart() theta ~ Wishart(4, Matrix{Float64}(I, 4, 4)) end Turing.setadbackend(:tracker) sample(wishart(), HMC(0.01, 1), 1000); - #Turing.setadbackend(:reversediff) - #sample(wishart(), HMC(0.01, 1), 1000); + Turing.setadbackend(:reversediff) + sample(wishart(), HMC(0.01, 1), 1000); Turing.setadbackend(:zygote) sample(wishart(), HMC(0.01, 1), 1000); @@ -112,8 +111,8 @@ end Turing.setadbackend(:tracker) sample(invwishart(), HMC(0.01, 1), 1000); - #Turing.setadbackend(:reversediff) - #sample(invwishart(), HMC(0.01, 1), 1000); + Turing.setadbackend(:reversediff) + sample(invwishart(), HMC(0.01, 1), 1000); Turing.setadbackend(:zygote) sample(invwishart(), HMC(0.01, 1), 1000); end diff --git a/test/modes/OptimInterface.jl b/test/modes/OptimInterface.jl index ea873ffee..2418037a4 100644 --- a/test/modes/OptimInterface.jl +++ b/test/modes/OptimInterface.jl @@ -1,38 +1,3 @@ -# TODO: Remove these once the equivalent is present in `DynamicPPL.TestUtils. -function likelihood_optima(::DynamicPPL.TestUtils.UnivariateAssumeDemoModels) - return (s=1/16, m=7/4) -end -function posterior_optima(::DynamicPPL.TestUtils.UnivariateAssumeDemoModels) - # TODO: Figure out exact for `s`. - return (s=0.907407, m=7/6) -end - -function likelihood_optima(model::DynamicPPL.TestUtils.MultivariateAssumeDemoModels) - # Get some containers to fill. - vals = Random.rand(model) - - # NOTE: These are "as close to zero as we can get". - vals.s[1] = 1e-32 - vals.s[2] = 1e-32 - - vals.m[1] = 1.5 - vals.m[2] = 2.0 - - return vals -end -function posterior_optima(model::DynamicPPL.TestUtils.MultivariateAssumeDemoModels) - # Get some containers to fill. - vals = Random.rand(model) - - # TODO: Figure out exact for `s[1]`. - vals.s[1] = 0.890625 - vals.s[2] = 1 - vals.m[1] = 3/4 - vals.m[2] = 1 - - return vals -end - # Used for testing how well it works with nested contexts. struct OverrideContext{C,T1,T2} <: DynamicPPL.AbstractContext context::C @@ -57,7 +22,7 @@ function DynamicPPL.tilde_observe(context::OverrideContext, right, left, vi) return context.loglikelihood_weight, vi end -@testset "OptimInterface.jl" begin +@numerical_testset "OptimInterface.jl" begin @testset "MLE" begin Random.seed!(222) true_value = [0.0625, 1.75] @@ -157,7 +122,7 @@ end # FIXME: Some models doesn't work for Tracker and ReverseDiff. if Turing.Essential.ADBACKEND[] === :forwarddiff @testset "MAP for $(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS - result_true = posterior_optima(model) + result_true = DynamicPPL.TestUtils.posterior_optima(model) @testset "$(nameof(typeof(optimizer)))" for optimizer in [LBFGS(), NelderMead()] result = optimize(model, MAP(), optimizer) @@ -188,7 +153,7 @@ end DynamicPPL.TestUtils.demo_dot_assume_matrix_dot_observe_matrix, ] @testset "MLE for $(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS - result_true = likelihood_optima(model) + result_true = DynamicPPL.TestUtils.likelihood_optima(model) # `NelderMead` seems to struggle with convergence here, so we exclude it. @testset "$(nameof(typeof(optimizer)))" for optimizer in [LBFGS(),]