<jupyter_start><jupyter_text>Disable AdvancedHMC's NUTS logging<jupyter_code>using Logging
using LoggingExtras

function ignore_sampling_filter(log_args)
    !(occursin("sampling steps",log_args.message) || occursin("adapation steps",log_args.message))
end
logger = ActiveFilteredLogger(ignore_sampling_filter, global_logger())


if !(@isdefined old_logger) #do this only once
    old_logger = global_logger(logger)
end
function my_nuts(trace, selection, 
        n_postadapt_steps = 2,
        n_adapts = 1,
        initial_ϵ_reduce_fac = 10)
    
    n_NUTS_steps = n_postadapt_steps + n_adapts
    
    filtered_choices = get_selected(get_choices(trace), selection)
    cur_xy = to_array(filtered_choices, Float64)
    dimension = length(cur_xy)
    metric = DiagEuclideanMetric(dimension)
    
    retval_grad = nothing #accepts_output_grad(get_gen_fn(trace)) ? zero(get_retval(trace)) : nothing
    
    function update_xy(val)
        extra_constraints = from_array(filtered_choices, val)
        update(trace, (), (NoChange(),), extra_constraints)
    end
    
    function val_to_lp_plus_c(val)
        (new_trace, weight, discard, retdiff) = update_xy(val)
        weight
    end
    
    function val_to_grad(val)
        (new_trace, weight, discard, retdiff) = update_xy(val)
        (retval_grad_out, values_trie, gradient_trie) = choice_gradients(new_trace, selection, retval_grad)
        grad = [gradient_trie[:x], gradient_trie[:y]]
        (weight, grad)
    end
    
    # Define a Hamiltonian system, using metric defined globally above
    hamiltonian = Hamiltonian(metric, val_to_lp_plus_c, val_to_grad)
    
    # Define a leapfrog solver, with initial step size chosen heuristically
    initial_ϵ = find_good_stepsize(hamiltonian, cur_xy) ./ initial_ϵ_reduce_fac
    integrator = Leapfrog(initial_ϵ)
    
    # Define an HMC sampler, with the following components
    #   - multinomial sampling scheme,
    #   - generalised No-U-Turn criteria, and
    #   - windowed adaption for step-size and diagonal mass matrix
    proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
    adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))

    # Run the sampler to draw samples from the specified Gaussian, where
    #   - `samples` will store the samples
    #   - `stats` will store diagnostic statistics for each sample
    
    samples, stats = sample(hamiltonian, proposal, cur_xy, n_NUTS_steps, adaptor, n_adapts; progress=false)
    
    #println(samples[3])
    
    
    (new_trace, weight, discard, retdiff) = update_xy(samples[n_NUTS_steps])
    new_trace
end

iters = 200
show = 5
ρ = .99


samps = mcmc_inference(ρ, iters, my_nuts, select(:x,:y))
samps[(iters-show+1):iters,:]
println(cor(samps[1:iters-1,1],samps[2:iters,1])) #serial correlation; lower is better
println(ρ^4) #for comparison, gibbs would be ρ² for each step; ρ⁴ for two steps
positive<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup