even more fixes (oh goodness when will this end)

Author: penelopeysm

src/mcmc/emcee.jl

@@ -53,10 +53,14 @@ function AbstractMCMC.step(
         length(initial_params) == n ||
             throw(ArgumentError("initial parameters have to be specified for each walker"))
         vis = map(vis, initial_params) do vi, init
+            # TODO(DPPL0.37/penelopeysm) This whole thing can be replaced with init!!
             vi = DynamicPPL.initialize_parameters!!(vi, init, model)
 
             # Update log joint probability.
-            last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromPrior()))
+            spl_model = DynamicPPL.contextualize(
+                model, DynamicPPL.SamplingContext(rng, SampleFromPrior(), model.context)
+            )
+            last(DynamicPPL.evaluate!!(spl_model, vi))
         end
     end
 
@@ -68,7 +72,7 @@ function AbstractMCMC.step(
         vis[1],
         map(vis) do vi
             vi = DynamicPPL.link!!(vi, model)
-            AMH.Transition(vi[:], getlogp(vi), false)
+            AMH.Transition(vi[:], DynamicPPL.getlogjoint(vi), false)
         end,
     )
 

src/mcmc/hmc.jl

@@ -236,12 +236,12 @@ function DynamicPPL.initialstep(
     if t.stat.is_accept
         vi = DynamicPPL.unflatten(vi, t.z.θ)
         # Re-evaluate to calculate log probability density.
-        # TODO(penelopeysm): This seems a little bit wasteful. The need for
-        # this stems from the fact that the HMC sampler doesn't keep track of
-        # prior and likelihood separately but rather a single log-joint, for
-        # which we have no way to decompose this back into prior and
-        # likelihood. I don't immediately see how to solve this without
-        # re-evaluating the model.
+        # TODO(penelopeysm): This seems a little bit wasteful. Unfortunately,
+        # even though `t.stat.log_density` contains some kind of logp, this
+        # doesn't track prior and likelihood separately but rather a single
+        # log-joint (and in linked space), so which we have no way to decompose
+        # this back into prior and likelihood. I don't immediately see how to
+        # solve this without re-evaluating the model.
         _, vi = DynamicPPL.evaluate!!(model, vi)
     else
         # Reset VarInfo back to its original state.
@@ -291,8 +291,9 @@ function AbstractMCMC.step(
     vi = state.vi
     if t.stat.is_accept
         vi = DynamicPPL.unflatten(vi, t.z.θ)
-        # TODO(mhauru) Is setloglikelihood! the right thing here?
-        vi = setloglikelihood!!(vi, t.stat.log_density)
+        # Re-evaluate to calculate log probability density.
+        # TODO(penelopeysm): This seems a little bit wasteful. See note above.
+        _, vi = DynamicPPL.evaluate!!(model, vi)
     end
 
     # Compute next transition and state.

src/mcmc/mh.jl

@@ -195,7 +195,7 @@ function LogDensityProblems.logdensity(f::LogDensityFunction, x::NamedTuple)
     vi = deepcopy(f.varinfo)
     set_namedtuple!(vi, x)
     vi_new = last(DynamicPPL.evaluate!!(f.model, vi, f.context))
-    lj = getlogp(vi_new)
+    lj = f.getlogdensity(vi_new)
     return lj
 end
 
@@ -304,7 +304,7 @@ function propose!!(
 
     # Create a sampler and the previous transition.
     mh_sampler = AMH.MetropolisHastings(dt)
-    prev_trans = AMH.Transition(vt, getlogp(vi), false)
+    prev_trans = AMH.Transition(vt, DynamicPPL.getlogjoint(vi), false)
 
     # Make a new transition.
     densitymodel = AMH.DensityModel(
@@ -339,7 +339,7 @@ function propose!!(
 
     # Create a sampler and the previous transition.
     mh_sampler = AMH.MetropolisHastings(spl.alg.proposals)
-    prev_trans = AMH.Transition(vals, getlogp(vi), false)
+    prev_trans = AMH.Transition(vals, DynamicPPL.getlogjoint(vi), false)
 
     # Make a new transition.
     densitymodel = AMH.DensityModel(