Optimisation rework

Author: penelopeysm

Project.toml

@@ -15,6 +15,7 @@ BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -25,7 +26,6 @@ Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-NamedArrays = "86f7a689-2022-50b4-a561-43c23ac3c673"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
@@ -58,20 +58,20 @@ BangBang = "0.4.2"
 Bijectors = "0.14, 0.15"
 Compat = "4.15.0"
 DataStructures = "0.18, 0.19"
+DifferentiationInterface = "0.7"
 Distributions = "0.25.77"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.39.1"
+DynamicPPL = "0.39.8"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3, 1"
 Libtask = "0.9.3"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 MCMCChains = "5, 6, 7"
-NamedArrays = "0.9, 0.10"
 Optimization = "3, 4, 5"
-OptimizationOptimJL = "0.1, 0.2, 0.3, 0.4"
+OptimizationOptimJL = "0.1 - 0.4"
 OrderedCollections = "1"
 Printf = "1"
 Random = "1"

src/Turing.jl

@@ -11,7 +11,6 @@ using AdvancedVI: AdvancedVI
 using DynamicPPL: DynamicPPL
 import DynamicPPL: NoDist, NamedDist
 using LogDensityProblems: LogDensityProblems
-using NamedArrays: NamedArrays
 using Accessors: Accessors
 using StatsAPI: StatsAPI
 using StatsBase: StatsBase
@@ -45,6 +44,7 @@ end
 # Random probability measures.
 include("stdlib/distributions.jl")
 include("stdlib/RandomMeasures.jl")
+include("init_strategy.jl")
 include("mcmc/Inference.jl")  # inference algorithms
 using .Inference
 include("variational/Variational.jl")

src/init_strategy.jl

@@ -0,0 +1,27 @@
+using AbstractPPL: VarName
+using DynamicPPL: DynamicPPL
+
+# This function is shared by both MCMC and optimisation, so has to exist outside of both.
+"""
+    _convert_initial_params(initial_params)
+
+Convert `initial_params` to a `DynamicPPl.AbstractInitStrategy` if it is not already one, or
+throw a useful error message.
+"""
+_convert_initial_params(initial_params::DynamicPPL.AbstractInitStrategy) = initial_params
+function _convert_initial_params(nt::NamedTuple)
+    @info "Using a NamedTuple for `initial_params` will be deprecated in a future release. Please use `InitFromParams(namedtuple)` instead."
+    return DynamicPPL.InitFromParams(nt)
+end
+function _convert_initial_params(d::AbstractDict{<:VarName})
+    @info "Using a Dict for `initial_params` will be deprecated in a future release. Please use `InitFromParams(dict)` instead."
+    return DynamicPPL.InitFromParams(d)
+end
+function _convert_initial_params(::AbstractVector{<:Real})
+    errmsg = "`initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or ideally a `DynamicPPL.AbstractInitStrategy`. Using a vector of parameters for `initial_params` is no longer supported. Please see https://turinglang.org/docs/usage/sampling-options/#specifying-initial-parameters for details on how to update your code."
+    throw(ArgumentError(errmsg))
+end
+function _convert_initial_params(@nospecialize(_::Any))
+    errmsg = "`initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or a `DynamicPPL.AbstractInitStrategy`."
+    throw(ArgumentError(errmsg))
+end

src/mcmc/abstractmcmc.jl

@@ -19,30 +19,6 @@ parameters for sampling are chosen if not specified by the user. By default, thi
 """
 init_strategy(::AbstractSampler) = DynamicPPL.InitFromPrior()
 
-"""
-    _convert_initial_params(initial_params)
-
-Convert `initial_params` to a `DynamicPPl.AbstractInitStrategy` if it is not already one, or
-throw a useful error message.
-"""
-_convert_initial_params(initial_params::DynamicPPL.AbstractInitStrategy) = initial_params
-function _convert_initial_params(nt::NamedTuple)
-    @info "Using a NamedTuple for `initial_params` will be deprecated in a future release. Please use `InitFromParams(namedtuple)` instead."
-    return DynamicPPL.InitFromParams(nt)
-end
-function _convert_initial_params(d::AbstractDict{<:VarName})
-    @info "Using a Dict for `initial_params` will be deprecated in a future release. Please use `InitFromParams(dict)` instead."
-    return DynamicPPL.InitFromParams(d)
-end
-function _convert_initial_params(::AbstractVector{<:Real})
-    errmsg = "`initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or ideally a `DynamicPPL.AbstractInitStrategy`. Using a vector of parameters for `initial_params` is no longer supported. Please see https://turinglang.org/docs/usage/sampling-options/#specifying-initial-parameters for details on how to update your code."
-    throw(ArgumentError(errmsg))
-end
-function _convert_initial_params(@nospecialize(_::Any))
-    errmsg = "`initial_params` must be a `NamedTuple`, an `AbstractDict{<:VarName}`, or a `DynamicPPL.AbstractInitStrategy`."
-    throw(ArgumentError(errmsg))
-end
-
 """
     default_varinfo(rng, model, sampler)
 
@@ -89,7 +65,7 @@ function AbstractMCMC.sample(
         model,
         spl,
         N;
-        initial_params=_convert_initial_params(initial_params),
+        initial_params=Turing._convert_initial_params(initial_params),
         chain_type,
         kwargs...,
     )
@@ -134,7 +110,7 @@ function AbstractMCMC.sample(
         n_chains;
         chain_type,
         check_model=false, # no need to check again
-        initial_params=map(_convert_initial_params, initial_params),
+        initial_params=map(Turing._convert_initial_params, initial_params),
         kwargs...,
     )
 end

src/mcmc/emcee.jl

@@ -39,7 +39,7 @@ function Turing.Inference.init_strategy(spl::Emcee)
     return fill(DynamicPPL.InitFromPrior(), _get_n_walkers(spl))
 end
 # We also have to explicitly allow this or else it will error...
-function Turing.Inference._convert_initial_params(
+function Turing._convert_initial_params(
     x::AbstractVector{<:DynamicPPL.AbstractInitStrategy}
 )
     return x

src/mcmc/repeat_sampler.jl

@@ -115,7 +115,7 @@ function AbstractMCMC.sample(
         model,
         sampler,
         N;
-        initial_params=_convert_initial_params(initial_params),
+        initial_params=Turing._convert_initial_params(initial_params),
         chain_type=chain_type,
         progress=progress,
         kwargs...,
@@ -143,7 +143,7 @@ function AbstractMCMC.sample(
         ensemble,
         N,
         n_chains;
-        initial_params=map(_convert_initial_params, initial_params),
+        initial_params=map(Turing._convert_initial_params, initial_params),
         chain_type=chain_type,
         progress=progress,
         kwargs...,