even more fixes

Author: penelopeysm

src/mcmc/external_sampler.jl

@@ -126,7 +126,13 @@ function make_updated_varinfo(
     return if ismissing(new_logp)
         last(DynamicPPL.evaluate!!(f.model, new_varinfo, f.context))
     else
-        DynamicPPL.setlogp!!(new_varinfo, new_logp)
+        # TODO(DPPL0.37/penelopeysm) This is obviously wrong. Note that we
+        # have the same problem here as in HMC in that the sampler doesn't
+        # tell us about how logp is broken down into prior and likelihood.
+        # We should probably just re-evaluate unconditionally. A bit
+        # unfortunate.
+        DynamicPPL.setlogprior!!(new_varinfo, 0.0)
+        DynamicPPL.setloglikelihood!!(new_varinfo, new_logp)
     end
 end
 

src/mcmc/is.jl

@@ -59,6 +59,6 @@ function DynamicPPL.assume(rng, ::Sampler{<:IS}, dist::Distribution, vn::VarName
         r = rand(rng, dist)
         vi = push!!(vi, vn, r, dist)
     end
-    vi = accumulate_assume!!(vi, r, 0.0, vn, dist)
+    vi = DynamicPPL.accumulate_assume!!(vi, r, 0.0, vn, dist)
     return r, vi
 end

src/mcmc/particle_mcmc.jl

@@ -429,6 +429,8 @@ function trace_local_rng_maybe(rng::Random.AbstractRNG)
     end
 end
 
+# TODO(DPPL0.37/penelopeysm) The whole tilde pipeline for particle MCMC needs to be
+# thoroughly fixed.
 function DynamicPPL.assume(
     rng, ::Sampler{<:Union{PG,SMC}}, dist::Distribution, vn::VarName, _vi::AbstractVarInfo
 )
@@ -442,13 +444,12 @@ function DynamicPPL.assume(
         DynamicPPL.unset_flag!(vi, vn, "del") # Reference particle parent
         r = rand(trng, dist)
         vi[vn] = DynamicPPL.tovec(r)
-        DynamicPPL.setorder!(vi, vn, DynamicPPL.get_num_produce(vi))
+        vi = DynamicPPL.setorder!!(vi, vn, DynamicPPL.get_num_produce(vi))
     else
         r = vi[vn]
     end
-    # TODO: Should we make this `zero(promote_type(eltype(dist), eltype(r)))` or something?
-    lp = 0
-    return r, lp, vi
+    # TODO: call accumulate_assume?!
+    return r, vi
 end
 
 # TODO(mhauru) Fix this.

src/optimisation/Optimisation.jl

@@ -611,7 +611,12 @@ function estimate_mode(
     ub=nothing,
     kwargs...,
 )
-    check_model && DynamicPPL.check_model(model; error_on_failure=true)
+    if check_model
+        spl_model = DynamicPPL.contextualize(
+            model, DynamicPPL.SamplingContext(model.context)
+        )
+        DynamicPPL.check_model(spl_model, VarInfo(); error_on_failure=true)
+    end
 
     constraints = ModeEstimationConstraints(lb, ub, cons, lcons, ucons)
     initial_params = generate_initial_params(model, initial_params, constraints)

test/mcmc/external_sampler.jl

@@ -20,7 +20,9 @@ function initialize_nuts(model::DynamicPPL.Model)
     linked_vi = DynamicPPL.link!!(vi, model)
 
     # Create a LogDensityFunction
-    f = DynamicPPL.LogDensityFunction(model, linked_vi; adtype=Turing.DEFAULT_ADTYPE)
+    f = DynamicPPL.LogDensityFunction(
+        model, DynamicPPL.getlogjoint, linked_vi; adtype=Turing.DEFAULT_ADTYPE
+    )
 
     # Choose parameter dimensionality and initial parameter value
     D = LogDensityProblems.dimension(f)