Use proposed samples from AdvancedHMC.jl for tuners in case of HMC Proposal

Author: Micki-D

ext/BATAdvancedHMCExt.jl

@@ -23,7 +23,7 @@ using BAT: MCMCChainState, HMCState, HamiltonianMC, HMCProposalState, MCMCChainS
 using BAT: MCMCBasicStats, push!, reweight_relative!
 using BAT: RAMTuning
 using BAT: MCMCProposalTunerState, MCMCTransformTunerState, NoMCMCTempering, NoMCMCTransformTuning
-using BAT: _current_sample_idx, _proposed_sample_idx, _current_sample_z_idx, _proposed_sample_z_idx, _cleanup_samples, current_sample_z, proposed_sample
+using BAT: _current_sample_idx, _proposed_sample_idx, _current_sample_z_idx, _proposed_sample_z_idx, _cleanup_samples, current_sample_z, proposed_sample_z, proposed_sample
 using BAT: AbstractTransformTarget, NoAdaptiveTransform, TriangularAffineTransform, valgrad_func
 using BAT: RNGPartition, get_rng, set_rng!
 using BAT: mcmc_step!!, nsamples, nsteps, samples_available, eff_acceptance_ratio
@@ -37,6 +37,8 @@ using BAT: AHMCSampleID, AHMCSampleIDVector
 using BAT: HMCMetric, DiagEuclideanMetric, UnitEuclideanMetric, DenseEuclideanMetric
 using BAT: HMCTuning, MassMatrixAdaptor, StepSizeAdaptor, NaiveHMCTuning, StanHMCTuning
 
+using ChangesOfVariables: with_logabsdet_jacobian
+
 using LinearAlgebra: cholesky
 
 using MeasureBase: pullback 

ext/ahmc_impl/ahmc_sampler_impl.jl

@@ -69,7 +69,7 @@ function BAT.next_cycle!(mc_state::HMCState)
     mc_state.nsamples = 0
     mc_state.stepno = 0
 
-    #reset_rng_counters!(mc_state)
+    reset_rng_counters!(mc_state)
 
     resize!(mc_state.samples, 1)
 
@@ -119,16 +119,15 @@ function BAT.mcmc_propose!!(mc_state::HMCState)
     z_phase = AdvancedHMC.phasepoint(hamiltonian, vec(z_current[:]), rand(rng, hamiltonian.metric, hamiltonian.kinetic))
     # Note: `RiemannianKinetic` requires an additional position argument, but including this causes issues. So only support the other kinetics.
 
-    proposal.transition = AdvancedHMC.transition(rng, τ, hamiltonian, z_phase)
-    z_proposed[:] = proposal.transition.z.θ
-    x_proposed[:] = f_transform(z_proposed)
-    
-    proposed_log_posterior = logdensityof(target, x_proposed)
-    samples.logd[proposed_x_idx] = proposed_log_posterior
+    proposal.transition, z_proposed_hmc, p_accept = _bat_transition(rng, τ, hamiltonian, z_phase)
+    accepted = z_current[:] != proposal.transition.z.θ
+    z_proposed[:] = accepted ? proposal.transition.z.θ : z_proposed_hmc
+
+    p_accept = AdvancedHMC.stat(proposal.transition).acceptance_rate
 
-    accepted = z_current != z_proposed
-    tstat = AdvancedHMC.stat(proposal.transition)
-    p_accept = accepted ? tstat.acceptance_rate : 0.0
+    x_proposed[:] = f_transform(z_proposed)    
+    logd_x_proposed = logdensityof(target, x_proposed)
+    samples.logd[proposed_x_idx] = logd_x_proposed
 
     return mc_state, accepted, p_accept
 end
@@ -142,7 +141,7 @@ function BAT._accept_reject!(mc_state::HMCState, accepted::Bool, p_accept::Float
         samples.info.sampletype[current] = ACCEPTED_SAMPLE
         samples.info.sampletype[proposed] = CURRENT_SAMPLE
         mc_state.nsamples += 1
-
+        
         tstat = AdvancedHMC.stat(proposal.transition)
         samples.info.hamiltonian_energy[proposed] = tstat.hamiltonian_energy
         # ToDo: Handle proposal-dependent tstat (only NUTS has tree_depth):
@@ -176,3 +175,79 @@ function BAT.set_mc_state_transform!!(mc_state::HMCState, f_transform_new::Funct
     mc_state_new = @set mc_state_new.f_transform = f_transform_new
     return mc_state_new
 end
+
+
+# Copied from AdvancedHMC.jl, but also return proposed point
+function _bat_transition(
+    rng::AbstractRNG,
+    τ::AdvancedHMC.Trajectory{TS,I,TC},
+    h::AdvancedHMC.Hamiltonian,
+    z0::AdvancedHMC.PhasePoint,
+) where {
+    TS<:AdvancedHMC.AbstractTrajectorySampler,
+    I<:AdvancedHMC.AbstractIntegrator,
+    TC<:AdvancedHMC.DynamicTerminationCriterion,
+}
+    H0 = AdvancedHMC.energy(z0)
+    tree = AdvancedHMC.BinaryTree(
+        z0,
+        z0,
+        AdvancedHMC.TurnStatistic(τ.termination_criterion, z0),
+        zero(H0),
+        zero(Int),
+        zero(H0),
+    )
+    sampler = TS(rng, z0)
+    termination = AdvancedHMC.Termination(false, false)
+    zcand = z0
+    proposed_zs = Vector[]
+
+    j = 0
+    while !AdvancedHMC.isterminated(termination) && j < τ.termination_criterion.max_depth
+        v = rand(rng, [-1, 1])
+        if v == -1
+            tree′, sampler′, termination′ =
+            AdvancedHMC.build_tree(rng, τ, h, tree.zleft, sampler, v, j, H0)
+            treeleft, treeright = tree′, tree
+        else
+            tree′, sampler′, termination′ =
+            AdvancedHMC.build_tree(rng, τ, h, tree.zright, sampler, v, j, H0)
+            treeleft, treeright = tree, tree′
+        end
+        if !AdvancedHMC.isterminated(termination′)
+            j = j + 1
+            if AdvancedHMC.mh_accept(rng, sampler, sampler′)
+                zcand = sampler′.zcand
+            end
+        end
+        push!(proposed_zs, sampler′.zcand.θ)
+
+        tree = AdvancedHMC.combine(treeleft, treeright)
+        sampler = AdvancedHMC.combine(zcand, sampler, sampler′)
+        termination =
+            termination *
+            termination′ *
+            AdvancedHMC.isterminated(τ.termination_criterion, h, tree, treeleft, treeright)
+    end
+
+    H = AdvancedHMC.energy(zcand)
+    tstat = AdvancedHMC.merge(
+        (
+            n_steps = tree.nα,
+            is_accept = true,
+            acceptance_rate = tree.sum_α / tree.nα,
+            log_density = zcand.ℓπ.value,
+            hamiltonian_energy = H,
+            hamiltonian_energy_error = H - H0,
+            max_hamiltonian_energy_error = tree.ΔH_max,
+            tree_depth = j,
+            numerical_error = termination.numerical,
+        ),
+        AdvancedHMC.stat(τ.integrator),
+    )
+
+    z_proposed = proposed_zs[end]    
+    p_accept = tstat.acceptance_rate
+
+    return AdvancedHMC.Transition(zcand, tstat), z_proposed, p_accept
+end

ext/ahmc_impl/ahmc_stan_tuner_impl.jl

@@ -55,8 +55,9 @@ function BAT.mcmc_tune_post_step!!(
     is_in_window =  stan_state.i >= stan_state.window_start && stan_state.i <= stan_state.window_end
     is_window_end = stan_state.i in stan_state.window_splits
 
-    # What to append?
-    is_in_window && BAT.push!(stats, proposed_sample(chain_state))
+    if is_in_window
+        BAT.push!(stats, proposed_sample(chain_state))
+    end
 
     if is_window_end 
         A = chain_state.f_transform.A

src/samplers/mcmc/mcmc_algorithm.jl

@@ -280,7 +280,7 @@ function mcmc_iterate!!(
     while (
         (nsteps(mcmc_state) - start_nsteps) < max_nsteps &&
         (time() - start_time) < max_time
-    ) 
+    )
         mcmc_state = mcmc_step!!(mcmc_state)
 
         if !isnothing(output)

src/samplers/mcmc/mcmc_sample.jl

@@ -88,23 +88,22 @@ function bat_sample_impl(m::BATMeasure, samplingalg::TransformedMCMC, context::B
     if !samplingalg.store_burnin
         chain_outputs .= DensitySampleVector.(mcmc_states)
     end
-
+    
     mcmc_states = mcmc_burnin!(
         samplingalg.store_burnin ? chain_outputs : nothing,
         mcmc_states,
         samplingalg,
         samplingalg.store_burnin ? samplingalg.callback : nop_func
     )
-
+    
     next_cycle!.(mcmc_states)
-
+    
     mcmc_states = mcmc_iterate!!(
         chain_outputs,
         mcmc_states;
         max_nsteps = samplingalg.nsteps,
         nonzero_weights = samplingalg.nonzero_weights
-    )
-
+    )    
     samples_transformed = DensitySampleVector(first(mcmc_states))
     isempty(chain_outputs) || append!.(Ref(samples_transformed), chain_outputs)
 

src/samplers/mcmc/mcmc_state.jl

@@ -145,7 +145,7 @@ end
 function mcmc_step!!(mcmc_state::MCMCState)
     _cleanup_samples(mcmc_state)
 
-    #reset_rng_counters!(mcmc_state)
+    reset_rng_counters!(mcmc_state)
 
     chain_state = mcmc_state.chain_state
 
@@ -156,7 +156,7 @@ function mcmc_step!!(mcmc_state::MCMCState)
     resize!(samples, size(samples, 1) + 1)
 
     samples.info[lastindex(samples)] = _get_sample_id(proposal, chain_state.info.id, chain_state.info.cycle, chain_state.stepno, PROPOSED_SAMPLE)[1]
-    
+
     chain_state, accepted, p_accept = mcmc_propose!!(chain_state)
 
     mcmc_state_new = mcmc_tune_post_step!!(mcmc_state, p_accept)
@@ -219,7 +219,7 @@ function next_cycle!(chain_state::MCMCChainState)
     chain_state.nsamples = 0
     chain_state.stepno = 0
 
-    #reset_rng_counters!(chain_state)
+    reset_rng_counters!(chain_state)
 
     resize!(chain_state.samples, 1)
 
@@ -277,7 +277,6 @@ end
 
 function mcmc_update_z_position!!(mc_state::MCMCChainState)
     f_transform = mc_state.f_transform
-    sample_z = mc_state.sample_z
 
     current_sample_x = current_sample(mc_state)
     proposed_sample_x = proposed_sample(mc_state)

src/samplers/mcmc/mcmc_tuning/mcmc_ram_tuner.jl

@@ -76,34 +76,28 @@ function mcmc_tune_post_step!!(
     mc_state::MCMCChainState,
     p_accept::Real,
 )
-    # TODO: MD: Discuss; apparently the RandomWalk sampler wants the trafo to be tuned even if p_accept = 0. If not, the burnin does not converge.
-    if iszero(p_accept) && !(mc_state isa MHChainState)
-        return mc_state, tuner_state
-    end
-
     (; f_transform, sample_z) = mc_state
     (; target_acceptance, gamma) = tuner_state.tuning
     b = f_transform.b
-    
+
     tuner_state_new = @set tuner_state.nsteps = tuner_state.nsteps + 1
-    
+
     n_dims = size(sample_z.v[1], 1)
     η = min(1, n_dims * tuner_state.nsteps^(-gamma))
 
     s_L = f_transform.A
 
     u = sample_z.v[2] - sample_z.v[1] # proposed - current
-    M = s_L * (I + η * (p_accept - target_acceptance) * (u * u') / norm(u)^2 ) * s_L'
-
+    M = s_L * (I + η * (p_accept - target_acceptance) * (u * u') / norm(u)^2) * s_L'
     new_s_L = oftype(s_L, cholesky(Positive, M).L)
-    
+       
     x = mc_state.samples[_proposed_sample_idx(mc_state)] # proposed in x-space
     mean_update_rate = η / 10 # heuristic
     α = mean_update_rate * p_accept
     new_b = oftype(b, (1- α) * b + α * x.v)
 
     f_transform_new = MulAdd(new_s_L, new_b)
-    
+
     mc_state_new = set_mc_state_transform!!(mc_state, f_transform_new)
     mc_state_new = mcmc_update_z_position!!(mc_state_new)
 

src/samplers/mcmc/mh_sampler.jl

@@ -144,8 +144,6 @@ function _accept_reject!(mc_state::MHChainState, accepted::Bool, p_accept::Float
         samples.info.sampletype[proposed] = CURRENT_SAMPLE
 
         mc_state.nsamples += 1
-
-        mc_state.sample_z[1] = deepcopy(proposed_sample_z(mc_state))
     else
         samples.info.sampletype[proposed] = REJECTED_SAMPLE
     end