Fix Hamiltonians in sghmc.jl / DynamicHMCExt

Author: penelopeysm

ext/TuringDynamicHMCExt.jl

@@ -35,8 +35,7 @@ State of the [`DynamicNUTS`](@ref) sampler.
 # Fields
 $(TYPEDFIELDS)
 """
-struct DynamicNUTSState{L,V<:DynamicPPL.AbstractVarInfo,C,M,S}
-    logdensity::L
+struct DynamicNUTSState{V<:DynamicPPL.AbstractVarInfo,C,M,S}
     vi::V
     "Cache of sample, log density, and gradient of log density evaluation."
     cache::C
@@ -48,30 +47,17 @@ function DynamicPPL.initialsampler(::DynamicPPL.Sampler{<:DynamicNUTS})
     return DynamicPPL.SampleFromUniform()
 end
 
-function DynamicPPL.initialstep(
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
-    model::DynamicPPL.Model,
-    spl::DynamicPPL.Sampler{<:DynamicNUTS},
-    vi::DynamicPPL.AbstractVarInfo;
+    ldf::DynamicPPL.LogDensityFunction,
+    spl::DynamicPPL.Sampler{<:DynamicNUTS};
     kwargs...,
 )
-    # 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)))
-    end
-
-    # Define log-density function.
-    ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
-    )
+    vi = ldf.varinfo
 
     # Perform initial step.
     results = DynamicHMC.mcmc_keep_warmup(
-        rng, ℓ, 0; initialization=(q=vi[:],), reporter=DynamicHMC.NoProgressReport()
+        rng, ldf, 0; initialization=(q=vi[:],), reporter=DynamicHMC.NoProgressReport()
     )
     steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
     Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)
@@ -81,32 +67,31 @@ function DynamicPPL.initialstep(
     vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create first sample and state.
-    sample = Turing.Inference.Transition(model, vi)
-    state = DynamicNUTSState(ℓ, vi, Q, steps.H.κ, steps.ϵ)
+    sample = Turing.Inference.Transition(ldf.model, vi)
+    state = DynamicNUTSState(vi, Q, steps.H.κ, steps.ϵ)
 
     return sample, state
 end
 
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
-    model::DynamicPPL.Model,
+    ldf::DynamicPPL.LogDensityFunction,
     spl::DynamicPPL.Sampler{<:DynamicNUTS},
     state::DynamicNUTSState;
     kwargs...,
 )
     # Compute next sample.
     vi = state.vi
-    ℓ = state.logdensity
-    steps = DynamicHMC.mcmc_steps(rng, spl.alg.sampler, state.metric, ℓ, state.stepsize)
+    steps = DynamicHMC.mcmc_steps(rng, spl.alg.sampler, state.metric, ldf, state.stepsize)
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
     # Update the variables.
     vi = DynamicPPL.unflatten(vi, Q.q)
     vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create next sample and state.
-    sample = Turing.Inference.Transition(model, vi)
-    newstate = DynamicNUTSState(ℓ, vi, Q, state.metric, state.stepsize)
+    sample = Turing.Inference.Transition(ldf.model, vi)
+    newstate = DynamicNUTSState(vi, Q, state.metric, state.stepsize)
 
     return sample, newstate
 end

src/mcmc/hmc.jl

@@ -1,7 +1,44 @@
+# InferenceAlgorithm interface
+
 abstract type Hamiltonian <: InferenceAlgorithm end
+
+DynamicPPL.initialsampler(::Sampler{<:Hamiltonian}) = SampleFromUniform()
+requires_unconstrained_space(::Hamiltonian) = true
+# TODO(penelopeysm): This is really quite dangerous code because it implicitly
+# assumes that any concrete type that subtypes `Hamiltonian` has an adtype
+# field.
+get_adtype(alg::Hamiltonian) = alg.adtype
+
 abstract type StaticHamiltonian <: Hamiltonian end
 abstract type AdaptiveHamiltonian <: Hamiltonian end
 
+function update_sample_kwargs(alg::AdaptiveHamiltonian, N::Integer, kwargs)
+    resume_from = get(kwargs, :resume_from, nothing)
+    nadapts = get(kwargs, :nadapts, alg.n_adapts)
+    discard_adapt = get(kwargs, :discard_adapt, true)
+    discard_initial = get(kwargs, :discard_initial, -1)
+
+    return if resume_from === nothing
+        # If `nadapts` is `-1`, then the user called a convenience constructor
+        # like `NUTS()` or `NUTS(0.65)`, and we should set a default for them.
+        if nadapts == -1
+            _nadapts = min(1000, N ÷ 2)  # Default to 1000 if not specified
+        else
+            _nadapts = nadapts
+        end
+        # If `discard_initial` is `-1`, then users did not specify the keyword argument.
+        if discard_initial == -1
+            _discard_initial = discard_adapt ? _nadapts : 0
+        else
+            _discard_initial = discard_initial
+        end
+
+        (nadapts=_nadapts, discard_initial=_discard_initial, kwargs...)
+    else
+        (nadapts=0, discard_adapt=false, discard_initial=0, kwargs...)
+    end
+end
+
 ###
 ### Sampler states
 ###
@@ -80,37 +117,6 @@ function HMC(
     return HMC(ϵ, n_leapfrog, metricT; adtype=adtype)
 end
 
-DynamicPPL.initialsampler(::Sampler{<:Hamiltonian}) = SampleFromUniform()
-
-get_adtype(alg::Hamiltonian) = alg.adtype
-
-function update_sample_kwargs(alg::AdaptiveHamiltonian, N::Integer, kwargs)
-    resume_from = get(kwargs, :resume_from, nothing)
-    nadapts = get(kwargs, :nadapts, alg.n_adapts)
-    discard_adapt = get(kwargs, :discard_adapt, true)
-    discard_initial = get(kwargs, :discard_initial, -1)
-
-    return if resume_from === nothing
-        # If `nadapts` is `-1`, then the user called a convenience constructor
-        # like `NUTS()` or `NUTS(0.65)`, and we should set a default for them.
-        if nadapts == -1
-            _nadapts = min(1000, N ÷ 2)  # Default to 1000 if not specified
-        else
-            _nadapts = nadapts
-        end
-        # If `discard_initial` is `-1`, then users did not specify the keyword argument.
-        if discard_initial == -1
-            _discard_initial = discard_adapt ? _nadapts : 0
-        else
-            _discard_initial = discard_initial
-        end
-
-        (nadapts=_nadapts, discard_initial=_discard_initial, kwargs...)
-    else
-        (nadapts=0, discard_adapt=false, discard_initial=0, kwargs...)
-    end
-end
-
 function find_initial_params(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
@@ -168,7 +174,7 @@ function AbstractMCMC.step(
     vi, z = if initial_params === nothing
         find_initial_params(rng, ldf.model, ldf.varinfo, hamiltonian)
     else
-        vi, AHMC.phasepoint(rng, theta, hamiltonian)
+        ldf.varinfo, AHMC.phasepoint(rng, theta, hamiltonian)
     end
     theta = vi[:]
 
@@ -425,9 +431,9 @@ function NUTS(; kwargs...)
     return NUTS(-1, 0.65; kwargs...)
 end
 
-for alg in (:HMC, :HMCDA, :NUTS)
-    @eval getmetricT(::$alg{<:Any,metricT}) where {metricT} = metricT
-end
+getmetricT(::HMC{<:Any,metricT}) where {metricT} = metricT
+getmetricT(::HMCDA{<:Any,metricT}) where {metricT} = metricT
+getmetricT(::NUTS{<:Any,metricT}) where {metricT} = metricT
 
 #####
 ##### HMC core functions

src/mcmc/sghmc.jl

@@ -45,50 +45,37 @@ function SGHMC(;
     return SGHMC(_learning_rate, _momentum_decay, adtype)
 end
 
-struct SGHMCState{L,V<:AbstractVarInfo,T<:AbstractVector{<:Real}}
-    logdensity::L
+struct SGHMCState{V<:AbstractVarInfo,T<:AbstractVector{<:Real}}
     vi::V
     velocity::T
 end
 
-function DynamicPPL.initialstep(
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
-    model::Model,
-    spl::Sampler{<:SGHMC},
-    vi::AbstractVarInfo;
+    ldf::DynamicPPL.LogDensityFunction,
+    spl::Sampler{<:SGHMC};
     kwargs...,
 )
-    # Transform the samples to unconstrained space and compute the joint log probability.
-    if !DynamicPPL.islinked(vi)
-        vi = DynamicPPL.link!!(vi, model)
-        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
-    end
+    vi = ldf.varinfo
 
     # Compute initial sample and state.
-    sample = Transition(model, vi)
-    ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
-    )
-    state = SGHMCState(ℓ, vi, zero(vi[:]))
+    sample = Transition(ldf.model, vi)
+    state = SGHMCState(vi, zero(vi[:]))
 
     return sample, state
 end
 
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
-    model::Model,
+    ldf::DynamicPPL.LogDensityFunction,
     spl::Sampler{<:SGHMC},
     state::SGHMCState;
     kwargs...,
 )
     # Compute gradient of log density.
-    ℓ = state.logdensity
     vi = state.vi
     θ = vi[:]
-    grad = last(LogDensityProblems.logdensity_and_gradient(ℓ, θ))
+    grad = last(LogDensityProblems.logdensity_and_gradient(ldf, θ))
 
     # Update latent variables and velocity according to
     # equation (15) of Chen et al. (2014)
@@ -100,11 +87,11 @@ function AbstractMCMC.step(
 
     # Save new variables and recompute log density.
     vi = DynamicPPL.unflatten(vi, θ)
-    vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
+    vi = last(DynamicPPL.evaluate!!(ldf.model, vi, DynamicPPL.SamplingContext(rng, spl)))
 
     # Compute next sample and state.
-    sample = Transition(model, vi)
-    newstate = SGHMCState(ℓ, vi, newv)
+    sample = Transition(ldf.model, vi)
+    newstate = SGHMCState(vi, newv)
 
     return sample, newstate
 end
@@ -208,57 +195,45 @@ metadata(t::SGLDTransition) = (lp=t.lp, SGLD_stepsize=t.stepsize)
 
 DynamicPPL.getlogp(t::SGLDTransition) = t.lp
 
-struct SGLDState{L,V<:AbstractVarInfo}
-    logdensity::L
+struct SGLDState{V<:AbstractVarInfo}
     vi::V
     step::Int
 end
 
-function DynamicPPL.initialstep(
+function AbstractMCMC.step(
     rng::Random.AbstractRNG,
-    model::Model,
-    spl::Sampler{<:SGLD},
-    vi::AbstractVarInfo;
+    ldf::DynamicPPL.LogDensityFunction,
+    spl::Sampler{<:SGLD};
     kwargs...,
 )
-    # Transform the samples to unconstrained space and compute the joint log probability.
-    if !DynamicPPL.islinked(vi)
-        vi = DynamicPPL.link!!(vi, model)
-        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
-    end
-
     # Create first sample and state.
-    sample = SGLDTransition(model, vi, zero(spl.alg.stepsize(0)))
-    ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
-    )
-    state = SGLDState(ℓ, vi, 1)
-
+    vi = ldf.varinfo
+    sample = SGLDTransition(ldf.model, vi, zero(spl.alg.stepsize(0)))
+    state = SGLDState(vi, 1)
     return sample, state
 end
 
 function AbstractMCMC.step(
-    rng::Random.AbstractRNG, model::Model, spl::Sampler{<:SGLD}, state::SGLDState; kwargs...
+    rng::Random.AbstractRNG,
+    ldf::LogDensityFunction,
+    spl::Sampler{<:SGLD},
+    state::SGLDState;
+    kwargs...,
 )
     # Perform gradient step.
-    ℓ = state.logdensity
     vi = state.vi
     θ = vi[:]
-    grad = last(LogDensityProblems.logdensity_and_gradient(ℓ, θ))
+    grad = last(LogDensityProblems.logdensity_and_gradient(ldf, θ))
     step = state.step
     stepsize = spl.alg.stepsize(step)
     θ .+= (stepsize / 2) .* grad .+ sqrt(stepsize) .* randn(rng, eltype(θ), length(θ))
 
     # Save new variables and recompute log density.
     vi = DynamicPPL.unflatten(vi, θ)
-    vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
+    vi = last(DynamicPPL.evaluate!!(ldf.model, vi, DynamicPPL.SamplingContext(rng, spl)))
 
     # Compute next sample and state.
-    sample = SGLDTransition(model, vi, stepsize)
-    newstate = SGLDState(ℓ, vi, state.step + 1)
-
+    sample = SGLDTransition(ldf.model, vi, stepsize)
+    newstate = SGLDState(vi, state.step + 1)
     return sample, newstate
 end