Reduce usage of sampler

Project.toml

@@ -47,7 +47,7 @@ DataStructures = "0.18"
 Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
-DynamicPPL = "0.21"
+DynamicPPL = "0.21.5"
 EllipticalSliceSampling = "0.5, 1"
 ForwardDiff = "0.10.3"
 Libtask = "0.7, 0.8"

src/Turing.jl

@@ -8,6 +8,7 @@ using Libtask
 using Tracker: Tracker
 
 import AdvancedVI
+using DynamicPPL: DynamicPPL, LogDensityFunction
 import DynamicPPL: getspace, NoDist, NamedDist
 import LogDensityProblems
 import Random
@@ -26,26 +27,6 @@ function setprogress!(progress::Bool)
     return progress
 end
 
-# Log density function
-struct LogDensityFunction{V,M,S,C}
-    varinfo::V
-    model::M
-    sampler::S
-    context::C
-end
-
-function (f::LogDensityFunction)(θ::AbstractVector)
-    vi_new = DynamicPPL.unflatten(f.varinfo, f.sampler, θ)
-    return getlogp(last(DynamicPPL.evaluate!!(f.model, vi_new, f.sampler, f.context)))
-end
-
-# LogDensityProblems interface
-LogDensityProblems.logdensity(f::LogDensityFunction, θ::AbstractVector) = f(θ)
-LogDensityProblems.dimension(f::LogDensityFunction) = length(f.varinfo[f.sampler])
-function LogDensityProblems.capabilities(::Type{<:LogDensityFunction})
-    return LogDensityProblems.LogDensityOrder{0}()
-end
-
 # Standard tag: Improves stacktraces
 # Ref: https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
 struct TuringTag end
@@ -154,6 +135,7 @@ export  @model,                 # modelling
         generated_quantities,
         logprior,
         logjoint,
+        LogDensityFunction,
 
         constrained_space,            # optimisation interface
         MAP,

src/essential/ad.jl

@@ -77,13 +77,18 @@ Find the autodifferentiation backend of the algorithm `alg`.
 """
 getADbackend(spl::Sampler) = getADbackend(spl.alg)
 getADbackend(::SampleFromPrior) = ADBackend()()
+getADbackend(ctx::DynamicPPL.SamplingContext) = getADbackend(ctx.sampler)
+getADbackend(ctx::DynamicPPL.AbstractContext) = getADbackend(DynamicPPL.NodeTrait(ctx), ctx)
+
+getADbackend(::DynamicPPL.IsLeaf, ctx::DynamicPPL.AbstractContext) = ADBackend()()
+getADbackend(::DynamicPPL.IsParent, ctx::DynamicPPL.AbstractContext) = getADbackend(DynamicPPL.childcontext(ctx))
 
 function LogDensityProblemsAD.ADgradient(ℓ::Turing.LogDensityFunction)
-    return LogDensityProblemsAD.ADgradient(getADbackend(ℓ.sampler), ℓ)
+    return LogDensityProblemsAD.ADgradient(getADbackend(ℓ.context), ℓ)
 end
 
 function LogDensityProblemsAD.ADgradient(ad::ForwardDiffAD, ℓ::Turing.LogDensityFunction)
-    θ = ℓ.varinfo[ℓ.sampler]
+    θ = DynamicPPL.getparams(ℓ)
     f = Base.Fix1(LogDensityProblems.logdensity, ℓ)
 
     # Define configuration for ForwardDiff.

src/inference/emcee.jl

@@ -74,7 +74,9 @@ function AbstractMCMC.step(
 )
     # Generate a log joint function.
     vi = state.vi
-    densitymodel = AMH.DensityModel(Turing.LogDensityFunction(vi, model, SampleFromPrior(), DynamicPPL.DefaultContext()))
+    densitymodel = AMH.DensityModel(
+        Base.Fix1(LogDensityProblems.logdensity, Turing.LogDensityFunction(model, vi))
+    )
 
     # Compute the next states.
     states = last(AbstractMCMC.step(rng, densitymodel, spl.alg.ensemble, state.states))

src/inference/ess.jl

@@ -124,10 +124,10 @@ end
 Distributions.mean(p::ESSPrior) = p.μ
 
 # Evaluate log-likelihood of proposals
-const ESSLogLikelihood{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo} = Turing.LogDensityFunction{V,M,S,DynamicPPL.DefaultContext()}
+const ESSLogLikelihood{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo} = Turing.LogDensityFunction{V,M,<:DynamicPPL.SamplingContext{<:S}}
 
 function (ℓ::ESSLogLikelihood)(f::AbstractVector)
-    sampler = ℓ.sampler
+    sampler = DynamicPPL.getsampler(ℓ)
     varinfo = setindex!!(ℓ.varinfo, f, sampler)
     varinfo = last(DynamicPPL.evaluate!!(ℓ.model, varinfo, sampler))
     return getlogp(varinfo)

src/inference/mh.jl

@@ -246,11 +246,11 @@ A log density function for the MH sampler.
 
 This variant uses the  `set_namedtuple!` function to update the `VarInfo`.
 """
-const MHLogDensityFunction{M<:Model,S<:Sampler{<:MH},V<:AbstractVarInfo} = Turing.LogDensityFunction{V,M,S,DynamicPPL.DefaultContext}
+const MHLogDensityFunction{M<:Model,S<:Sampler{<:MH},V<:AbstractVarInfo} = Turing.LogDensityFunction{V,M,<:DynamicPPL.SamplingContext{<:S}}
 
-function (f::MHLogDensityFunction)(x::NamedTuple)
+function LogDensityProblems.logdensity(f::MHLogDensityFunction, x::NamedTuple)
     # TODO: Make this work with immutable `f.varinfo` too.
-    sampler = f.sampler
+    sampler = DynamicPPL.getsampler(f)
     vi = f.varinfo
 
     x_old, lj_old = vi[sampler], getlogp(vi)
@@ -374,7 +374,9 @@ function propose!!(
     prev_trans = AMH.Transition(vt, getlogp(vi))
 
     # Make a new transition.
-    densitymodel = AMH.DensityModel(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
+    densitymodel = AMH.DensityModel(
+        Base.Fix1(LogDensityProblems.logdensity, Turing.LogDensityFunction(vi, model, DynamicPPL.SamplingContext(rng, spl)))
+    )
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
 
     # TODO: Make this compatible with immutable `VarInfo`.
@@ -400,7 +402,9 @@ function propose!!(
     prev_trans = AMH.Transition(vals, getlogp(vi))
 
     # Make a new transition.
-    densitymodel = AMH.DensityModel(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
+    densitymodel = AMH.DensityModel(
+        Base.Fix1(LogDensityProblems.logdensity, Turing.LogDensityFunction(vi, model, DynamicPPL.SamplingContext(rng, spl)))
+    )
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
 
     return setlogp!!(DynamicPPL.unflatten(vi, spl, trans.params), trans.lp)

src/modes/ModeEstimation.jl

@@ -46,45 +46,42 @@ DynamicPPL.childcontext(context::OptimizationContext) = context.context
 DynamicPPL.setchildcontext(::OptimizationContext, child) = OptimizationContext(child)
 
 # assume
-function DynamicPPL.tilde_assume(rng::Random.AbstractRNG, ctx::OptimizationContext, spl, dist, vn, vi)
-    return DynamicPPL.tilde_assume(ctx, spl, dist, vn, vi)
-end
-
-function DynamicPPL.tilde_assume(ctx::OptimizationContext{<:LikelihoodContext}, spl, dist, vn, vi)
-    r = vi[vn]
+function DynamicPPL.tilde_assume(ctx::OptimizationContext{<:LikelihoodContext}, dist, vn, vi)
+    r = vi[vn, dist]
     return r, 0, vi
 end
 
-function DynamicPPL.tilde_assume(ctx::OptimizationContext, spl, dist, vn, vi)
-    r = vi[vn]
+function DynamicPPL.tilde_assume(ctx::OptimizationContext, dist, vn, vi)
+    r = vi[vn, dist]
     return r, Distributions.logpdf(dist, r), vi
 end
 
 # dot assume
-function DynamicPPL.dot_tilde_assume(rng::Random.AbstractRNG, ctx::OptimizationContext, sampler, right, left, vns, vi)
-    return DynamicPPL.dot_tilde_assume(ctx, sampler, right, left, vns, vi)
-end
-
-function DynamicPPL.dot_tilde_assume(ctx::OptimizationContext{<:LikelihoodContext}, sampler::SampleFromPrior, right, left, vns, vi)
+function DynamicPPL.dot_tilde_assume(ctx::OptimizationContext{<:LikelihoodContext}, right, left, vns, vi)
     # Values should be set and we're using `SampleFromPrior`, hence the `rng` argument shouldn't
     # affect anything.
-    r = DynamicPPL.get_and_set_val!(Random.GLOBAL_RNG, vi, vns, right, sampler)
+    # TODO: Stop using `get_and_set_val!`.
+    r = DynamicPPL.get_and_set_val!(Random.default_rng(), vi, vns, right, SampleFromPrior())
     return r, 0, vi
 end
 
-function DynamicPPL.dot_tilde_assume(ctx::OptimizationContext, sampler::SampleFromPrior, right, left, vns, vi)
+_loglikelihood(dist::Distribution, x) = loglikelihood(dist, x)
+_loglikelihood(dists::AbstractArray{<:Distribution}, x) = loglikelihood(arraydist(dists), x)
+
+function DynamicPPL.dot_tilde_assume(ctx::OptimizationContext, right, left, vns, vi)
     # Values should be set and we're using `SampleFromPrior`, hence the `rng` argument shouldn't
     # affect anything.
-    r = DynamicPPL.get_and_set_val!(Random.GLOBAL_RNG, vi, vns, right, sampler)
-    return r, loglikelihood(right, r), vi
+    # TODO: Stop using `get_and_set_val!`.
+    r = DynamicPPL.get_and_set_val!(Random.default_rng(), vi, vns, right, SampleFromPrior())
+    return r, _loglikelihood(right, r), vi
 end
 
 """
     OptimLogDensity{M<:Model,C<:Context,V<:VarInfo}
 
 A struct that stores the negative log density function of a `DynamicPPL` model.
 """
-const OptimLogDensity{M<:Model,C<:OptimizationContext,V<:VarInfo} = Turing.LogDensityFunction{V,M,DynamicPPL.SampleFromPrior,C}
+const OptimLogDensity{M<:Model,C<:OptimizationContext,V<:VarInfo} = Turing.LogDensityFunction{V,M,C}
 
 """
     OptimLogDensity(model::Model, context::OptimizationContext)
@@ -93,21 +90,23 @@ Create a callable `OptimLogDensity` struct that evaluates a model using the give
 """
 function OptimLogDensity(model::Model, context::OptimizationContext)
     init = VarInfo(model)
-    return Turing.LogDensityFunction(init, model, DynamicPPL.SampleFromPrior(), context)
+    return Turing.LogDensityFunction(init, model, context)
 end
 
 """
-    (f::OptimLogDensity)(z)
+    LogDensityProblems.logdensity(f::OptimLogDensity, z)
 
 Evaluate the negative log joint (with `DefaultContext`) or log likelihood (with `LikelihoodContext`)
 at the array `z`.
 """
 function (f::OptimLogDensity)(z::AbstractVector)
-    sampler = f.sampler
-    varinfo = DynamicPPL.unflatten(f.varinfo, sampler, z)
-    return -getlogp(last(DynamicPPL.evaluate!!(f.model, varinfo, sampler, f.context)))
+    varinfo = DynamicPPL.unflatten(f.varinfo, z)
+    return -getlogp(last(DynamicPPL.evaluate!!(f.model, varinfo, f.context)))
 end
 
+# NOTE: This seems a bit weird IMO since this is the _negative_ log-likelihood.
+LogDensityProblems.logdensity(f::OptimLogDensity, z::AbstractVector) = f(z)
+
 function (f::OptimLogDensity)(F, G, z)
     if G !== nothing
         # Calculate negative log joint and its gradient.
@@ -127,7 +126,7 @@ function (f::OptimLogDensity)(F, G, z)
 
     # Only negative log joint requested but no gradient.
     if F !== nothing
-        return f(z)
+        return LogDensityProblems.logdensity(f, z)
     end
 
     return nothing
@@ -140,50 +139,44 @@ end
 #################################################
 
 function transform!!(f::OptimLogDensity)
-    spl = f.sampler
-
     ## Check link status of vi in OptimLogDensity
-    linked = DynamicPPL.islinked(f.varinfo, spl)
+    linked = DynamicPPL.istrans(f.varinfo)
 
     ## transform into constrained or unconstrained space depending on current state of vi
     @set! f.varinfo = if !linked
-        DynamicPPL.link!!(f.varinfo, spl, f.model)
+        DynamicPPL.link!!(f.varinfo, f.model)
     else
-        DynamicPPL.invlink!!(f.varinfo, spl, f.model)
+        DynamicPPL.invlink!!(f.varinfo, f.model)
     end
 
     return f
 end
 
 function transform!!(p::AbstractArray, vi::DynamicPPL.VarInfo, model::DynamicPPL.Model, ::constrained_space{true})
-    spl = DynamicPPL.SampleFromPrior()
-
-    linked = DynamicPPL.islinked(vi, spl)
+    linked = DynamicPPL.istrans(vi)
 
     !linked && return identity(p)  # TODO: why do we do `identity` here?
-    vi = DynamicPPL.setindex!!(vi, p, spl)
-    vi = DynamicPPL.invlink!!(vi, spl, model)
-    p .= vi[spl]
+    vi = DynamicPPL.unflatten(vi, p)
+    vi = DynamicPPL.invlink!!(vi, model)
+    p .= vi[:]
 
     # If linking mutated, we need to link once more.
-    linked && DynamicPPL.link!!(vi, spl, model)
+    linked && DynamicPPL.link!!(vi, model)
 
     return p
 end
 
 function transform!!(p::AbstractArray, vi::DynamicPPL.VarInfo, model::DynamicPPL.Model, ::constrained_space{false})
-    spl = DynamicPPL.SampleFromPrior()
-
-    linked = DynamicPPL.islinked(vi, spl)
+    linked = DynamicPPL.istrans(vi)
     if linked
-        vi = DynamicPPL.invlink!!(vi, spl, model)
+        vi = DynamicPPL.invlink!!(vi, model)
     end
-    vi = DynamicPPL.setindex!!(vi, p, spl)
-    vi = DynamicPPL.link!!(vi, spl, model)
-    p .= vi[spl]
+    vi = DynamicPPL.unflatten(vi, p)
+    vi = DynamicPPL.link!!(vi, model)
+    p .= vi[:]
 
     # If linking mutated, we need to link once more.
-    !linked && DynamicPPL.invlink!!(vi, spl, model)
+    !linked && DynamicPPL.invlink!!(vi, model)
 
     return p
 end
@@ -208,26 +201,26 @@ end
 
 function (t::AbstractTransform)(p::AbstractArray)
     return transform(p, t.vi, t.model, t.space)
-end 
+end
 
 function (t::Init)()
     return t.vi[DynamicPPL.SampleFromPrior()]
 end 
 
 function get_parameter_bounds(model::DynamicPPL.Model)
     vi = DynamicPPL.VarInfo(model)
-    spl = DynamicPPL.SampleFromPrior()
 
     ## Check link status of vi
-    linked = DynamicPPL.islinked(vi, spl) 
+    linked = DynamicPPL.istrans(vi)
 
     ## transform into unconstrained
     if !linked
-        vi = DynamicPPL.link!!(vi, spl, model)
+        vi = DynamicPPL.link!!(vi, model)
     end
-
-    lb = transform(fill(-Inf,length(vi[DynamicPPL.SampleFromPrior()])), vi, model, constrained_space{true}())
-    ub = transform(fill(Inf,length(vi[DynamicPPL.SampleFromPrior()])), vi, model, constrained_space{true}())
+
+    d = length(vi[:])
+    lb = transform(fill(-Inf, d), vi, model, constrained_space{true}())
+    ub = transform(fill(Inf, d), vi, model, constrained_space{true}())
 
     return lb, ub
 end