Remove context argument from LogDensityFunction

Author: penelopeysm

ext/TuringDynamicHMCExt.jl

@@ -58,16 +58,11 @@ function DynamicPPL.initialstep(
     # Ensure that initial sample is in unconstrained space.
     if !DynamicPPL.islinked(vi)
         vi = DynamicPPL.link!!(vi, model)
-        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
+        vi = last(DynamicPPL.evaluate!!(model, vi))
     end
 
     # Define log-density function.
-    ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
-    )
+    ℓ = DynamicPPL.LogDensityFunction(model, vi; adtype=spl.alg.adtype)
 
     # Perform initial step.
     results = DynamicHMC.mcmc_keep_warmup(

src/mcmc/ess.jl

@@ -49,8 +49,8 @@ function AbstractMCMC.step(
         rng,
         EllipticalSliceSampling.ESSModel(
             ESSPrior(model, spl, vi),
-            DynamicPPL.LogDensityFunction{:LogLikelihood}(
-                model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
+            ESSLikelihood(
+                DynamicPPL.LogDensityFunction(model, DynamicPPL.getloglikelihood, vi)
             ),
         ),
         EllipticalSliceSampling.ESS(),
@@ -63,7 +63,7 @@ function AbstractMCMC.step(
 
     return Transition(model, vi), vi
 end
-
+f
 # Prior distribution of considered random variable
 struct ESSPrior{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo,T}
     model::M
@@ -97,6 +97,10 @@ function Base.rand(rng::Random.AbstractRNG, p::ESSPrior)
     sampler = p.sampler
     varinfo = p.varinfo
     # TODO: Surely there's a better way of doing this now that we have `SamplingContext`?
+    # TODO(DPPL0.37/penelopeysm): This can be replaced with `init!!(p.model,
+    # p.varinfo, PriorInit())` after TuringLang/DynamicPPL.jl#984. The reason
+    # why we had to use the 'del' flag before this was because
+    # SampleFromPrior() wouldn't overwrite existing variables.
     vns = keys(varinfo)
     for vn in vns
         set_flag!(varinfo, vn, "del")
@@ -108,14 +112,9 @@ end
 # Mean of prior distribution
 Distributions.mean(p::ESSPrior) = p.μ
 
-function DynamicPPL.tilde_assume(
-    rng::Random.AbstractRNG, ctx::DefaultContext, ::Sampler{<:ESS}, right, vn, vi
-)
-    return DynamicPPL.tilde_assume(rng, ctx, SampleFromPrior(), right, vn, vi)
+# Evaluate log-likelihood of proposals
+struct ESSLogLikelihood{M<:Model,V<:AbstractVarInfo,AD<:ADTypes.AbstractADType}
+    ldf::DynamicPPL.LogDensityFunction{M,V,AD}
 end
 
-function DynamicPPL.tilde_observe!!(
-    ctx::DefaultContext, ::Sampler{<:ESS}, right, left, vn, vi
-)
-    return DynamicPPL.tilde_observe!!(ctx, SampleFromPrior(), right, left, vn, vi)
-end
+(ℓ::ESSLogLikelihood)(f::AbstractVector) = LogDensityProblems.logdensity(ℓ.ldf, f)

src/mcmc/external_sampler.jl

@@ -96,12 +96,12 @@ getlogp_external(::Any, ::Any) = missing
 getlogp_external(mh::AdvancedMH.Transition, ::AdvancedMH.Transition) = mh.lp
 getlogp_external(hmc::AdvancedHMC.Transition, ::AdvancedHMC.HMCState) = hmc.stat.log_density
 
-struct TuringState{S,V1<:AbstractVarInfo,M,V,C}
+struct TuringState{S,V1<:AbstractVarInfo,M,V}
     state::S
     # Note that this varinfo has the correct parameters and logp obtained from
     # the state, whereas `ldf.varinfo` will in general have junk inside it.
     varinfo::V1
-    ldf::DynamicPPL.LogDensityFunction{M,V,C}
+    ldf::DynamicPPL.LogDensityFunction{M,V}
 end
 
 varinfo(state::TuringState) = state.varinfo

src/mcmc/gibbs.jl

@@ -559,7 +559,7 @@ function setparams_varinfo!!(
     params::AbstractVarInfo,
 )
     logdensity = DynamicPPL.LogDensityFunction(
-        model, state.ldf.varinfo, state.ldf.context; adtype=sampler.alg.adtype
+        model, state.ldf.varinfo; adtype=sampler.alg.adtype
     )
     new_inner_state = setparams_varinfo!!(
         AbstractMCMC.LogDensityModel(logdensity), sampler, state.state, params

src/mcmc/hmc.jl

@@ -190,16 +190,7 @@ function DynamicPPL.initialstep(
     # Create a Hamiltonian.
     metricT = getmetricT(spl.alg)
     metric = metricT(length(theta))
-    ldf = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        # TODO(penelopeysm): Can we just use leafcontext(model.context)? Do we
-        # need to pass in the sampler? (In fact LogDensityFunction defaults to
-        # using leafcontext(model.context) so could we just remove the argument
-        # entirely?)
-        DynamicPPL.SamplingContext(rng, spl, DynamicPPL.leafcontext(model.context));
-        adtype=spl.alg.adtype,
-    )
+    ldf = DynamicPPL.LogDensityFunction(model, vi; adtype=spl.alg.adtype)
     lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)
     lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf)
     hamiltonian = AHMC.Hamiltonian(metric, lp_func, lp_grad_func)
@@ -305,16 +296,7 @@ end
 
 function get_hamiltonian(model, spl, vi, state, n)
     metric = gen_metric(n, spl, state)
-    ldf = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        # TODO(penelopeysm): Can we just use leafcontext(model.context)? Do we
-        # need to pass in the sampler? (In fact LogDensityFunction defaults to
-        # using leafcontext(model.context) so could we just remove the argument
-        # entirely?)
-        DynamicPPL.SamplingContext(spl, DynamicPPL.leafcontext(model.context));
-        adtype=spl.alg.adtype,
-    )
+    ldf = DynamicPPL.LogDensityFunction(model, vi; adtype=spl.alg.adtype)
     lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)
     lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf)
     return AHMC.Hamiltonian(metric, lp_func, lp_grad_func)

src/mcmc/mh.jl

@@ -189,7 +189,7 @@ 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} =
-    DynamicPPL.LogDensityFunction{M,V,<:DynamicPPL.SamplingContext{<:S},AD} where {AD}
+    DynamicPPL.LogDensityFunction{M,V,AD} where {AD}
 
 function LogDensityProblems.logdensity(f::MHLogDensityFunction, x::NamedTuple)
     vi = deepcopy(f.varinfo)

src/mcmc/particle_mcmc.jl

@@ -206,7 +206,9 @@ function DynamicPPL.initialstep(
     )
 
     # Perform particle sweep.
+    @info "Hello!"
     logevidence = AdvancedPS.sweep!(rng, particles, spl.alg.resampler, spl)
+    @info "Goodbye!"
 
     # Extract the first particle and its weight.
     particle = particles.vals[1]
@@ -222,6 +224,7 @@ end
 function AbstractMCMC.step(
     ::AbstractRNG, model::AbstractModel, spl::Sampler{<:SMC}, state::SMCState; kwargs...
 )
+    @info "helloooooo from step"
     # Extract the index of the current particle.
     index = state.particleindex
 

src/optimisation/Optimisation.jl

@@ -169,6 +169,9 @@ or
 OptimLogDensity(model; adtype=adtype)
 ```
 
+Here, `ctx` must be a context that contains an `OptimizationContext` as its
+leaf.
+
 If not specified, `adtype` defaults to `AutoForwardDiff()`.
 
 An OptimLogDensity does not, in itself, obey the LogDensityProblems interface.
@@ -189,24 +192,40 @@ optim_ld(z)  # returns -logp
 ```
 """
 struct OptimLogDensity{
-    M<:DynamicPPL.Model,
-    F<:Function,
-    V<:DynamicPPL.AbstractVarInfo,
-    C<:DynamicPPL.AbstractContext,
-    AD<:ADTypes.AbstractADType,
+    M<:DynamicPPL.Model,F<:Function,V<:DynamicPPL.AbstractVarInfo,AD<:ADTypes.AbstractADType
 }
-    ldf::DynamicPPL.LogDensityFunction{M,F,V,C,AD}
-end
+    ldf::DynamicPPL.LogDensityFunction{M,F,V,AD}
 
-function OptimLogDensity(
-    model::DynamicPPL.Model,
-    getlogdensity::Function,
-    vi::DynamicPPL.AbstractVarInfo=DynamicPPL.ldf_default_varinfo(model, getlogdensity);
-    adtype=AutoForwardDiff(),
-)
-    return OptimLogDensity(
-        DynamicPPL.LogDensityFunction(model, getlogdensity, vi; adtype=adtype)
+    # Inner constructors enforce that the model has an OptimizationContext as
+    # its leaf context.
+    function OptimLogDensity(
+        model::DynamicPPL.Model,
+        getlogdensity::Function,
+        vi::DynamicPPL.VarInfo,
+        ctx::OptimizationContext;
+        adtype::ADTypes.AbstractADType=Turing.DEFAULT_ADTYPE,
     )
+        new_context = DynamicPPL.setleafcontext(model, ctx)
+        new_model = contextualize(model, new_context)
+        return new{typeof(new_model),typeof(getlogdensity),typeof(vi),typeof(adtype)}(
+            DynamicPPL.LogDensityFunction(new_model, getlogdensity, vi; adtype=adtype)
+        )
+    end
+    function OptimLogDensity(
+        model::DynamicPPL.Model,
+        getlogdensity::Function,
+        ctx::OptimizationContext;
+        adtype::ADTypes.AbstractADType=Turing.DEFAULT_ADTYPE,
+    )
+        # No varinfo
+        return OptimLogDensity(
+            model,
+            getlogdensity,
+            DynamicPPL.ldf_default_varinfo(model, getlogdensity),
+            ctx;
+            adtype=adtype,
+        )
+    end
 end
 
 """

src/variational/VariationalInference.jl

@@ -17,12 +17,6 @@ export vi, q_locationscale, q_meanfield_gaussian, q_fullrank_gaussian
 
 include("deprecated.jl")
 
-function make_logdensity(model::DynamicPPL.Model)
-    weight = 1.0
-    ctx = DynamicPPL.MiniBatchContext(DynamicPPL.DefaultContext(), weight)
-    return DynamicPPL.LogDensityFunction(model, DynamicPPL.VarInfo(model), ctx)
-end
-
 """
     q_initialize_scale(
         [rng::Random.AbstractRNG,]
@@ -68,7 +62,7 @@ function q_initialize_scale(
     num_max_trials::Int=10,
     reduce_factor::Real=one(eltype(scale)) / 2,
 )
-    prob = make_logdensity(model)
+    prob = LogDensityFunction(model)
     ℓπ = Base.Fix1(LogDensityProblems.logdensity, prob)
     varinfo = DynamicPPL.VarInfo(model)
 
@@ -309,7 +303,7 @@ function vi(
 )
     return AdvancedVI.optimize(
         rng,
-        make_logdensity(model),
+        LogDensityFunction(model),
         objective,
         q,
         n_iterations;