update Optimisation code to not use LogDensityProblemsAD

Author: penelopeysm

ext/TuringOptimExt.jl

@@ -1,14 +1,8 @@
 module TuringOptimExt
 
-if isdefined(Base, :get_extension)
-    using Turing: Turing
-    import Turing: DynamicPPL, NamedArrays, Accessors, Optimisation
-    using Optim: Optim
-else
-    import ..Turing
-    import ..Turing: DynamicPPL, NamedArrays, Accessors, Optimisation
-    import ..Optim
-end
+using Turing: Turing
+import Turing: DynamicPPL, NamedArrays, Accessors, Optimisation
+using Optim: Optim
 
 ####################
 # Optim.jl methods #
@@ -42,7 +36,7 @@ function Optim.optimize(
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -65,7 +59,7 @@ function Optim.optimize(
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
 end
 function Optim.optimize(
@@ -112,7 +106,7 @@ function Optim.optimize(
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -135,7 +129,7 @@ function Optim.optimize(
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
 end
 function Optim.optimize(
@@ -162,17 +156,20 @@ Estimate a mode, i.e., compute a MLE or MAP estimate.
 function _optimize(
     model::DynamicPPL.Model,
     f::Optimisation.OptimLogDensity,
-    init_vals::AbstractArray=DynamicPPL.getparams(f),
+    init_vals::AbstractArray=DynamicPPL.getparams(f.ldf),
     optimizer::Optim.AbstractOptimizer=Optim.LBFGS(),
     options::Optim.Options=Optim.Options(),
     args...;
     kwargs...,
 )
     # Convert the initial values, since it is assumed that users provide them
     # in the constrained space.
-    f = Accessors.@set f.varinfo = DynamicPPL.unflatten(f.varinfo, init_vals)
-    f = Accessors.@set f.varinfo = DynamicPPL.link(f.varinfo, model)
-    init_vals = DynamicPPL.getparams(f)
+    # TODO(penelopeysm): As with in src/optimisation/Optimisation.jl, unclear
+    # whether initialisation is really necessary at all
+    vi = DynamicPPL.unflatten(f.ldf.varinfo, init_vals)
+    vi = DynamicPPL.link(vi, f.ldf.model)
+    f = Optimisation.OptimLogDensity(f.ldf.model, vi, f.ldf.context; adtype=f.ldf.adtype)
+    init_vals = DynamicPPL.getparams(f.ldf)
 
     # Optimize!
     M = Optim.optimize(Optim.only_fg!(f), init_vals, optimizer, options, args...; kwargs...)
@@ -186,12 +183,16 @@ function _optimize(
     end
 
     # Get the optimum in unconstrained space. `getparams` does the invlinking.
-    f = Accessors.@set f.varinfo = DynamicPPL.unflatten(f.varinfo, M.minimizer)
-    vns_vals_iter = Turing.Inference.getparams(model, f.varinfo)
+    vi = f.ldf.varinfo
+    vi_optimum = DynamicPPL.unflatten(vi, M.minimizer)
+    logdensity_optimum = Optimisation.OptimLogDensity(
+        f.ldf.model, vi_optimum, f.ldf.context
+    )
+    vns_vals_iter = Turing.Inference.getparams(model, vi_optimum)
     varnames = map(Symbol ∘ first, vns_vals_iter)
     vals = map(last, vns_vals_iter)
     vmat = NamedArrays.NamedArray(vals, varnames)
-    return Optimisation.ModeResult(vmat, M, -M.minimum, f)
+    return Optimisation.ModeResult(vmat, M, -M.minimum, logdensity_optimum)
 end
 
 end # module

src/optimisation/Optimisation.jl

@@ -4,7 +4,6 @@ using ..Turing
 using NamedArrays: NamedArrays
 using DynamicPPL: DynamicPPL
 using LogDensityProblems: LogDensityProblems
-using LogDensityProblemsAD: LogDensityProblemsAD
 using Optimization: Optimization
 using OptimizationOptimJL: OptimizationOptimJL
 using Random: Random
@@ -95,22 +94,40 @@ function DynamicPPL.tilde_observe(
 end
 
 """
-    OptimLogDensity{M<:DynamicPPL.Model,C<:Context,V<:DynamicPPL.VarInfo}
+    OptimLogDensity{M<:DynamicPPL.Model,V<:DynamicPPL.VarInfo,C<:OptimizationContext,AD<:ADTypes.AbstractADType}
 
 A struct that stores the negative log density function of a `DynamicPPL` model.
+
+TODO(penelopeysm): It _doesn't_ really store the negative, does it? It's more like we
+overrode logdensity to give the negative logdensity.
 """
-const OptimLogDensity{M<:DynamicPPL.Model,C<:OptimizationContext,V<:DynamicPPL.VarInfo,AD} = Turing.LogDensityFunction{
-    M,V,C,AD
+struct OptimLogDensity{
+    M<:DynamicPPL.Model,
+    V<:DynamicPPL.VarInfo,
+    C<:OptimizationContext,
+    AD<:ADTypes.AbstractADType,
 }
+    ldf::Turing.LogDensityFunction{M,V,C,AD}
+end
 
-"""
-    OptimLogDensity(model::DynamicPPL.Model, context::OptimizationContext)
+function OptimLogDensity(
+    model::DynamicPPL.Model,
+    vi::DynamicPPL.VarInfo,
+    ctx::OptimizationContext;
+    adtype::Union{Nothing,ADTypes.AbstractADType}=AutoForwardDiff(),
+)
+    return OptimLogDensity(Turing.LogDensityFunction(model, vi, ctx; adtype=adtype))
+end
 
-Create a callable `OptimLogDensity` struct that evaluates a model using the given `context`.
-"""
-function OptimLogDensity(model::DynamicPPL.Model, context::OptimizationContext)
-    init = DynamicPPL.VarInfo(model)
-    return Turing.LogDensityFunction(model, init, context)
+# No varinfo
+function OptimLogDensity(
+    model::DynamicPPL.Model,
+    ctx::OptimizationContext;
+    adtype::Union{Nothing,ADTypes.AbstractADType}=AutoForwardDiff(),
+)
+    return OptimLogDensity(
+        Turing.LogDensityFunction(model, DynamicPPL.VarInfo(model), ctx; adtype=adtype)
+    )
 end
 
 """
@@ -123,40 +140,30 @@ depends on the context of `f`.
 Any second argument is ignored. The two-argument method only exists to match interface the
 required by Optimization.jl.
 """
-function (f::OptimLogDensity)(z::AbstractVector)
-    varinfo = DynamicPPL.unflatten(f.varinfo, z)
-    return -DynamicPPL.getlogp(last(DynamicPPL.evaluate!!(f.model, varinfo, f.context)))
-end
-
+(f::OptimLogDensity)(z::AbstractVector) = -LogDensityProblems.logdensity(f.ldf, z)
 (f::OptimLogDensity)(z, _) = f(z)
 
-# NOTE: This seems a bit weird IMO since this is the _negative_ log-likelihood.
-LogDensityProblems.logdensity(f::OptimLogDensity, z::AbstractVector) = f(z)
-
 # NOTE: The format of this function is dictated by Optim. The first argument sets whether to
 # compute the function value, the second whether to compute the gradient (and stores the
 # gradient). The last one is the actual argument of the objective function.
 function (f::OptimLogDensity)(F, G, z)
     if G !== nothing
-        # Calculate negative log joint and its gradient.
-        # TODO: Make OptimLogDensity already an LogDensityProblems.ADgradient? Allow to
-        # specify AD?
-        ℓ = LogDensityProblemsAD.ADgradient(f)
-        neglogp, ∇neglogp = LogDensityProblems.logdensity_and_gradient(ℓ, z)
+        # Calculate log joint and its gradient.
+        logp, ∇logp = LogDensityProblems.logdensity_and_gradient(f.ldf, z)
 
-        # Save the gradient to the pre-allocated array.
-        copyto!(G, ∇neglogp)
+        # Save the negative gradient to the pre-allocated array.
+        copyto!(G, -∇logp)
 
         # If F is something, the negative log joint is requested as well.
         # We have already computed it as a by-product above and hence return it directly.
         if F !== nothing
-            return neglogp
+            return -logp
         end
     end
 
     # Only negative log joint requested but no gradient.
     if F !== nothing
-        return LogDensityProblems.logdensity(f, z)
+        return -LogDensityProblems.logdensity(f.ldf, z)
     end
 
     return nothing
@@ -232,9 +239,11 @@ function StatsBase.informationmatrix(
 
     # Convert the values to their unconstrained states to make sure the
     # Hessian is computed with respect to the untransformed parameters.
-    linked = DynamicPPL.istrans(m.f.varinfo)
+    linked = DynamicPPL.istrans(m.f.ldf.varinfo)
     if linked
-        m = Accessors.@set m.f.varinfo = DynamicPPL.invlink!!(m.f.varinfo, m.f.model)
+        new_vi = DynamicPPL.invlink!!(m.f.ldf.varinfo, m.f.ldf.model)
+        new_f = OptimLogDensity(m.f.ldf.model, new_vi, m.f.ldf.context)
+        m = Accessors.@set m.f = new_f
     end
 
     # Calculate the Hessian, which is the information matrix because the negative of the log
@@ -244,7 +253,9 @@ function StatsBase.informationmatrix(
 
     # Link it back if we invlinked it.
     if linked
-        m = Accessors.@set m.f.varinfo = DynamicPPL.link!!(m.f.varinfo, m.f.model)
+        new_vi = DynamicPPL.link!!(m.f.ldf.varinfo, m.f.ldf.model)
+        new_f = OptimLogDensity(m.f.ldf.model, new_vi, m.f.ldf.context)
+        m = Accessors.@set m.f = new_f
     end
 
     return NamedArrays.NamedArray(info, (varnames, varnames))
@@ -265,7 +276,7 @@ Return the values of all the variables with the symbol(s) `var_symbol` in the mo
 argument should be either a `Symbol` or a vector of `Symbol`s.
 """
 function Base.get(m::ModeResult, var_symbols::AbstractVector{Symbol})
-    log_density = m.f
+    log_density = m.f.ldf
     # Get all the variable names in the model. This is the same as the list of keys in
     # m.values, but they are more convenient to filter when they are VarNames rather than
     # Symbols.
@@ -297,9 +308,9 @@ richer format of `ModeResult`. It also takes care of transforming them back to t
 parameter space in case the optimization was done in a transformed space.
 """
 function ModeResult(log_density::OptimLogDensity, solution::SciMLBase.OptimizationSolution)
-    varinfo_new = DynamicPPL.unflatten(log_density.varinfo, solution.u)
+    varinfo_new = DynamicPPL.unflatten(log_density.ldf.varinfo, solution.u)
     # `getparams` performs invlinking if needed
-    vns_vals_iter = Turing.Inference.getparams(log_density.model, varinfo_new)
+    vns_vals_iter = Turing.Inference.getparams(log_density.ldf.model, varinfo_new)
     syms = map(Symbol ∘ first, vns_vals_iter)
     vals = map(last, vns_vals_iter)
     return ModeResult(
@@ -383,12 +394,15 @@ end
     OptimizationProblem(log_density::OptimLogDensity, adtype, constraints)
 
 Create an `OptimizationProblem` for the objective function defined by `log_density`.
+
+Note that the adtype parameter here overrides any adtype parameter the
+OptimLogDensity was constructed with.
 """
 function Optimization.OptimizationProblem(log_density::OptimLogDensity, adtype, constraints)
     # Note that OptimLogDensity is a callable that evaluates the model with given
     # parameters. Hence we can use it in the objective function as below.
     f = Optimization.OptimizationFunction(log_density, adtype; cons=constraints.cons)
-    initial_params = log_density.varinfo[:]
+    initial_params = log_density.ldf.varinfo[:]
     prob = if !has_constraints(constraints)
         Optimization.OptimizationProblem(f, initial_params)
     else
@@ -454,28 +468,34 @@ function estimate_mode(
     end
 
     # Create an OptimLogDensity object that can be used to evaluate the objective function,
-    # i.e. the negative log density. Set its VarInfo to the initial parameters.
-    log_density = let
-        inner_context = if estimator isa MAP
-            DynamicPPL.DefaultContext()
-        else
-            DynamicPPL.LikelihoodContext()
-        end
-        ctx = OptimizationContext(inner_context)
-        ld = OptimLogDensity(model, ctx)
-        Accessors.@set ld.varinfo = DynamicPPL.unflatten(ld.varinfo, initial_params)
+    # i.e. the negative log density.
+    inner_context = if estimator isa MAP
+        DynamicPPL.DefaultContext()
+    else
+        DynamicPPL.LikelihoodContext()
     end
+    ctx = OptimizationContext(inner_context)
 
+    # Set its VarInfo to the initial parameters.
+    # TODO(penelopeysm): Unclear if this is really needed? Any time that logp is calculated
+    # (using `LogDensityProblems.logdensity(ldf, x)`) the parameters in the
+    # varinfo are completely ignored. The parameters only matter if you are calling evaluate!!
+    # directly on the fields of the LogDensityFunction
+    vi = DynamicPPL.VarInfo(model)
+    vi = DynamicPPL.unflatten(vi, initial_params)
+
+    # Link the varinfo if needed.
     # TODO(mhauru) We currently couple together the questions of whether the user specified
     # bounds/constraints and whether we transform the objective function to an
     # unconstrained space. These should be separate concerns, but for that we need to
     # implement getting the bounds of the prior distributions.
     optimise_in_unconstrained_space = !has_constraints(constraints)
     if optimise_in_unconstrained_space
-        transformed_varinfo = DynamicPPL.link(log_density.varinfo, log_density.model)
-        log_density = Accessors.@set log_density.varinfo = transformed_varinfo
+        vi = DynamicPPL.link(vi, model)
     end
 
+    log_density = OptimLogDensity(model, vi, ctx)
+
     prob = Optimization.OptimizationProblem(log_density, adtype, constraints)
     solution = Optimization.solve(prob, solver; kwargs...)
     # TODO(mhauru) We return a ModeResult for compatibility with the older Optim.jl

test/ext/OptimInterface.jl

@@ -143,7 +143,6 @@ using Turing
         DynamicPPL.TestUtils.demo_assume_multivariate_observe_literal,
         DynamicPPL.TestUtils.demo_dot_assume_observe_submodel,
         DynamicPPL.TestUtils.demo_dot_assume_observe_matrix_index,
-        DynamicPPL.TestUtils.demo_dot_assume_matrix_dot_observe_matrix,
         DynamicPPL.TestUtils.demo_assume_submodel_observe_index_literal,
         DynamicPPL.TestUtils.demo_dot_assume_observe_index,
         DynamicPPL.TestUtils.demo_dot_assume_observe_index_literal,

test/optimisation/Optimisation.jl

@@ -545,7 +545,6 @@ using Turing
         DynamicPPL.TestUtils.demo_assume_multivariate_observe_literal,
         DynamicPPL.TestUtils.demo_dot_assume_observe_submodel,
         DynamicPPL.TestUtils.demo_dot_assume_observe_matrix_index,
-        DynamicPPL.TestUtils.demo_dot_assume_matrix_dot_observe_matrix,
         DynamicPPL.TestUtils.demo_assume_submodel_observe_index_literal,
         DynamicPPL.TestUtils.demo_dot_assume_observe_index,
         DynamicPPL.TestUtils.demo_dot_assume_observe_index_literal,