changed src structure; added ensemble simulation

src/MultilevelChainSampler.jl

-module MultilevelChainSampler
-
-using AbstractMCMC: AbstractMCMC, AbstractModel
-using BangBang
-using Distributions
-
-using Random, StatsBase
-using Primes, HaltonSequences
-
-using Graphs
-using DataFrames
-
-import Base: getindex, length, size, eltype, convert, rand, push!
-import Distributions: logpdf
-import StatsBase: sample, autocov
 
-export logdensity
-export LogDensity, SampledLogDensity
-export MultilevelLogDensity, MultilevelSampledLogDensity
-
-export propose
-export RandomWalk, CyclicWalk
-export ErdosRenyi
-
-export sample
-export MetropolisHastings, ChristenFox, LykkegaardScheichl
-
-export nchains, info
-export levels, is_multilevel
-export states
-export autocov, logpart, empirical_cdf
-export RejectionChains
+module MultilevelChainSampler
+using Reexport
 
+include("sampling/Samplers.jl")
+@reexport using .Samplers
 
-include("chains/all.jl")
-include("models/all.jl")
-include("proposals/all.jl")
-include("algos/all.jl")
+include("powergrid/PowerGrids.jl")
+@reexport using .PowerGrids
 
 end
 

src/powergrid/PowerGrids.jl

+module PowerGrids
+    using Statistics
+    using Graphs
+    using OrdinaryDiffEq
+    using NonlinearSolve
+    using Primes, HaltonSequences
+
+    include("dynamics.jl")
+    include("stability.jl")
+
+    export swing_dynamics!
+    export synchronous_state
+    export stability_problem
+end
\ No newline at end of file

src/powergrid/dynamics.jl

+
+# differential equation
+function swing_dynamics!(dₜx, x, p, t)
+
+    # unpack
+    G,P,K,D = p.grid, p.power, p.coupling, p.damping 
+    N = nv(G); ϕ = x[1:N]; ω = x[N+1:2N]
+
+    # second-order, frequency is derivative of phase
+    dₜx[1:N] .= ω;
+
+    # swing equation
+    L = laplacian_matrix(G)
+    s = sin.( L * ϕ )
+    dₜx[N+1:2N] .= @. P - D * ω + K * s 
+end
+
+
+function sync_rhs(x,p)
+    N = length(x)÷2
+    
+    # time derivative at t = 0
+    dₜx = similar(x)
+    swing_dynamics!(dₜx, x, p, 0.0)
+    
+    # ignore constant angle shift 
+    dₜx[1:N] .-= mean(dₜx[1:N])
+
+    return dₜx[2:2N]
+end
+
+# steady state 
+function synchronous_state(p)
+    nlp = NonlinearProblem(sync_rhs, zeros(2*nv(p.grid)), p)
+    sync = solve(nlp, NewtonRaphson()).u
+    return sync
+end 
+

src/powergrid/stability.jl

+
+
+"""
+
+`freqamp` frequency perturbation amplitude
+`n` number of perturbations 
+
+Example
+≡≡≡≡≡≡≡≡
+
+julia> 
+    J = 1       # only node 1 
+    J = [1,2,3] # set of nodes 
+
+"""
+function stability_problem(p, J; freqamp = 1.0, t_final=120.0, tolerance = 0.1)
+
+    # default ode problem
+    u_sync = synchronous_state(p)
+    ode_prob = ODEProblem(swing_dynamics!, u_sync, (0, t_final), p)
+
+    # low discrepancy sequence for unit inverval
+    N = nv(p.grid)
+    m = length(J)
+    P = primes(4*m^2+1)[1:2m]
+    H = Halton.(P)
+
+    # perturb intial condition 
+    function _bstab_init(prob, i, repeat)
+        u0 = u_sync[:]
+        Z = getindex.(H, i) .* 2 .- 1
+        u0[J]      .= Z[1:m] .* pi
+        u0[J .+ N] .= Z[1+m:2m] .* freqamp
+        prob.u0 .= u0
+        prob
+    end
+
+    # check convergence to sync
+    function _bstab_out(sol, i) 
+        u = sol.u[end];
+        Δu = maximum( abs.( u[1+N:end] .- u_sync[1+N:end] ) )
+        S = (Δu < tolerance)
+        return (S, false)
+    end
+
+    return prob = EnsembleProblem(ode_prob; 
+        output_func = _bstab_out,
+        prob_func = _bstab_init
+    )    
+end
+

src/sampling/Samplers.jl

+module Samplers
+
+using AbstractMCMC: AbstractMCMC, AbstractModel
+using BangBang
+using Distributions
+
+using Random, StatsBase
+using Primes, HaltonSequences
+
+using Graphs
+using DataFrames
+
+import Base: getindex, length, size, eltype, convert, rand, push!
+import Distributions: logpdf
+import StatsBase: sample, autocov
+
+export logdensity
+export LogDensity, SampledLogDensity
+export MultilevelLogDensity, MultilevelSampledLogDensity
+
+export propose
+export RandomWalk, CyclicWalk
+export ErdosRenyi
+
+export sample
+export MetropolisHastings, ChristenFox, LykkegaardScheichl
+
+export nchains, info
+export levels, is_multilevel
+export states
+export autocov, logpart, empirical_cdf
+export RejectionChains
+
+
+include("chains/all.jl")
+include("models/all.jl")
+include("proposals/all.jl")
+include("algos/all.jl")
+
+end

src/sampling/algos/all.jl

+abstract type RejectionBasedSampler <: AbstractMCMC.AbstractSampler end
+
+# Ignore chain length by default
+function AbstractMCMC.save!!(samples, sample, iterations::Integer, model::AbstractModel, sampler::RejectionBasedSampler, ::Integer; kwargs...)
+    AbstractMCMC.save!!(samples, sample, iterations, model, sampler; kwargs...)
+end
+
+# Size hint sample container chain length 
+function AbstractMCMC.samples(sample, model::AbstractModel, sampler::RejectionBasedSampler, N::Integer; kwargs...)
+    samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
+    T = typeof(samples)
+    for (k,t) in zip(T.parameters[1], T.types)
+        if t <: AbstractVector
+            sizehint!(getfield(samples, k), N) 
+        end
+    end
+    return samples
+end
+
+function AbstractMCMC.bundle_samples(samples, m::AbstractModel, s::RejectionBasedSampler, state, chain_type::Type; kwargs... )
+    AbstractMCMC.bundle_samples(samples, m, s, state, RejectionChains; kwargs... )
+end
+
+function _chain_info(samples, model::AbstractModel, sampler::RejectionBasedSampler)
+    return Dict{Symbol, Vector}()
+end
+
+function AbstractMCMC.bundle_samples(samples, m::AbstractModel, s::RejectionBasedSampler, state, chain_type::Type{<:AbstractRejectionChains}; kwargs... )
+    x = AbstractMCMC.bundle_samples(samples, m, s, state, NamedTuple; kwargs...)
+    c = RejectionChains( x.states, x.logprobs, x.rejections )
+    data = _chain_info(samples, m, s)
+    for k in keys(data)
+        c.info[k] = [data[k]]
+    end
+    return c
+end
+
+
+
+include("metropolis_hastings.jl")
+include("christen_fox.jl")
+include("lykkegard_scheichl.jl")

src/sampling/algos/christen_fox.jl

+"""
+ Delayed Acceptance algorithm
+
+ If `saveproxies == true` save log-density lower levels. 
+ This is particularly useful for Multilevel integration.
+
+ References
+    
+    * Christen, J.A. and Fox, C. (2005). "Markov chain Monte Carlo Using an Approximation"
+      Journal of Computational and Graphical Statistics
+
+""" 
+struct ChristenFox{saveproxies, P <: AbstractProposal} <: RejectionBasedSampler
+    proposal :: P
+end 
+ChristenFox(proposal::AbstractProposal, saveproxies::Bool=false) = ChristenFox{saveproxies, typeof(proposal)}(proposal)
+
+## Save samples only for highest level
+function AbstractMCMC.samples(sample, model::AbstractMultilevelModel, ::ChristenFox{false, P}; kwargs...) where {P}
+    (level, x, f_x) = sample
+    return (; 
+        states = typeof(x)[],
+        logprobs = typeof(f_x)[],
+        rejections = Vector{Int}[], 
+    )
+end
+function AbstractMCMC.save!!(samples, sample, ::Integer, model::AbstractMultilevelModel, ::ChristenFox{false, P}; kwargs... ) where {P}
+    (level, x, f_x) = sample
+    
+    # accept  (at hightest level) 
+    if level == length(model)                
+        push!(samples.states, x)
+        push!(samples.logprobs, f_x)
+        push!(samples.rejections, zeros(Int, length(model)))
+
+    # reject
+    else samples.rejections[end][level+1] += 1 end
+    
+    return samples
+end
+
+
+## Save each level
+
+function AbstractMCMC.samples(sample, model::AbstractMultilevelModel, ::ChristenFox{true, P}; kwargs...) where {P}
+    (level, x, f_x) = sample
+    return (; 
+        states = typeof(x)[],
+        logprobs = typeof(f_x)[],
+        rejections = Vector{Int}[], 
+        current = Ref{Int}(0)  # reference to current top level state
+    )
+end
+function AbstractMCMC.save!!(samples, sample, ::Integer, model::AbstractMultilevelModel, ::ChristenFox{true, P}; kwargs... ) where {P}
+    (level, x, f_x) = sample
+
+    # accept 
+    if level == length(model)        
+        push!(samples.states, x)
+        push!(samples.logprobs, f_x)
+        push!(samples.rejections, zeros(Int, length(model)))
+
+        samples.current[] = length(samples.states) # reset current to accepted 
+    else
+        samples.rejections[samples.current[]][level+1] += 1 # update rejection counter 
+        
+        # store promoted
+        if level > 0 
+            push!(samples.states, x)
+            push!(samples.logprobs, f_x)
+            push!(samples.rejections, zeros(Int, length(model)))
+        end
+    end 
+    return samples
+end
+
+
+
+## General chain stepping
+
+function AbstractMCMC.step(rng::AbstractRNG, model::AbstractModel, sampler::ChristenFox; x0=nothing, f0=nothing, kwargs...)
+    
+    # Initialize states
+    _resample = isnothing(x0)  
+    x = _resample ? rand(rng, sampler.proposal) : x0
+    
+    # Initialize logprobs
+    _recompute = _resample || isnothing(f0)
+    f_x = _recompute ?  [ logdensity(model, x; level=l) for l=1:length(model) ] : f0
+
+    return (length(model), x, f_x), (x, f_x)
+end
+
+
+function AbstractMCMC.step(rng::AbstractRNG, model::AbstractMultilevelModel, sampler::ChristenFox, state; kwargs...)
+
+    # Load old state
+    x, f_x = state
+
+    # Propose
+    y = propose(rng, sampler.proposal, x)
+    f_y = [ logdensity(model, y; level=1) ]
+    q = logpratio(sampler.proposal, x, y)
+    
+    # Promotion probability
+    A_1 = min( f_y[1] - f_x[1] + q, 0)
+    accept = log(rand(rng)) < A_1
+    if !accept return (0, x, f_x), (x, f_x) end # completly reject, never promoted
+    
+    # Promotion loop
+    for l = 2:length(model)
+        push!(f_y, logdensity(model, y; level=l) )
+
+        # Next promotion
+        A_l = f_y[l] - f_x[l] + f_y[l-1] - f_x[l-1]   
+        accept = log(rand(rng)) < A_l 
+        if !accept return (l-1, y, f_y), (x, f_x) end # rejected at level l, promoted to l-1 
+    end
+
+    # Accept at highest level
+    return (length(model), y, f_y), (y, f_y) 
+end
+
+
+## Chain stepping specialized for Sampled-Based LogDensities, uses cached evaluation 
+
+function AbstractMCMC.step(rng::AbstractRNG, model::MultilevelSampledLogDensity, sampler::ChristenFox; x0=nothing, f0=nothing, kwargs...)
+    
+    # Initialize states
+    _resample = isnothing(x0)  
+    x = _resample ? rand(rng, sampler.proposal) : x0
+    
+    # Initialize logprobs
+    _recompute = _resample || isnothing(f0)
+    if _recompute 
+
+        # Recursively cached
+        f_x = [ logdensity(model, x; level=1) ]
+        sizehint!(f_x, length(model))
+        for l=2:length(model)
+            push!(f_x, logdensity(model, x; level=l, cache=f_x[end]))
+        end 
+    else f_x = f0 end
+
+    return (length(model), x, f_x), (x, f_x)
+end
+
+
+function AbstractMCMC.step(rng::AbstractRNG, model::MultilevelSampledLogDensity, sampler::ChristenFox, state; kwargs...)
+    # Load old state
+    x, f_x = state
+
+    # Propose
+    y = propose(rng, sampler.proposal, x)
+    f_y = [ logdensity(model, y; level=1) ]
+    q = logpratio(sampler.proposal, x, y)
+    
+    # Promotion probability
+    A_1 = min( f_y[1] - f_x[1] + q, 0)
+    accept = log(rand(rng)) < A_1
+    if !accept return (0, x, f_x), (x, f_x) end # completly reject, never promoted
+    
+    # Promotion loop
+    for l = 2:length(model)
+        push!(f_y, logdensity(model, y; level=l, cache=f_y[end]) ) # > Only this line changes!
+
+        # Next promotion
+        A_l = f_y[l] - f_x[l] + f_y[l-1] - f_x[l-1]   
+        accept = log(rand(rng)) < A_l 
+        if !accept return (l-1, y, f_y), (x, f_x) end # rejected at level l, promoted to l-1 
+    end
+
+    # Accept at highest level
+    return (length(model), y, f_y), (y, f_y) 
+end
+
+
+function AbstractMCMC.bundle_samples(samples, ::AbstractMultilevelModel, ::ChristenFox, state, chain_type::Type{<:NamedTuple}; kwargs...)
+    return (; states = samples.states, logprobs = samples.logprobs, rejections = samples.rejections)
+end
+
+function _chain_info(samples, model::AbstractMultilevelModel, ::ChristenFox)
+    n_total = sum( 1 .+ sum.(samples.rejections))
+    n_reject = sum(samples.rejections)
+    info = Dict(
+        :chain_length   => n_total,
+        :rejection_rate => n_reject ./ n_total,
+    )
+
+    if model isa MultilevelSampledLogDensity
+        evals_per_level = n_total .- cumsum([0, n_reject[1:end-1]... ])
+        costs_per_level = model.nlevels .- [0, model.nlevels[1:end-1]...] 
+        info[:total_costs] = sum( evals_per_level .* costs_per_level )
+    end
+
+    return info
+end

src/sampling/algos/lykkegard_scheichl.jl

+"""
+ Recursive Delayed Acceptance algorithm
+
+
+ References
+    
+    * Mikkel B. Lykkegaard, T. Dodwell, C. Fox, Grigorios Mingas, Robert Scheichl (2022). "Multilevel Delayed Acceptance MCMC"
+    SIAM/ASA J. Uncertain. Quantification
+
+""" 
+
+
+struct LykkegaardScheichl{saveproxies, P <: AbstractProposal, N} <: RejectionBasedSampler
+    proposal :: P
+    sublen :: N # sub-chain length (may be a distribution)
+end 
+LykkegaardScheichl(p,s=1,saveproxies::Bool=true) = LykkegaardScheichl{saveproxies, typeof(p), typeof(s)}(p,s)
+
+subchainlength(rng::AbstractRNG, s::LykkegaardScheichl{saveproxies,P,<:Distribution}) where {saveproxies,P} = rand(rng, s.sublen)
+subchainlength(rng::AbstractRNG, s::LykkegaardScheichl{saveproxies,P,<:Integer}) where {saveproxies,P} = s.sublen 
+
+
+## Only save top level
+
+function AbstractMCMC.samples(sample, mode::AbstractMultilevelModel, sampler::LykkegaardScheichl{false,P,N}; L=length(model), kwargs...) where {P,N}
+    accept, x, f_x, r_x, Y = sample 
+    return (; 
+        states = typeof(x)[], 
+        logprobs = typeof(f_x)[], 
+        rejections = typeof(r_x)[], 
+    )
+end
+
+function AbstractMCMC.save!!(samples, sample, ::Integer, model::AbstractMultilevelModel, sampler::LykkegaardScheichl{false, P,N}) where {P,N}
+    accept, x, f_x, r_x, Y = sample
+    
+    if accept # save new entry
+        push!(samples.states, x)
+        push!(samples.logprobs, f_x)
+        push!(samples.rejections, r_x )
+
+    else # update rejection counter
+        samples.rejections[end][L] += 1
+    end
+    return samples
+
+end
+
+## Save proxies
+
+function AbstractMCMC.samples(sample, model::AbstractMultilevelModel, sampler::LykkegaardScheichl{true,P,N}; L=length(model), kwargs...) where {P,N}
+    accept, x, f_x, r_x, Y = sample 
+    return (; 
+        states = typeof(x)[], 
+        logprobs = typeof(f_x)[], 
+        rejections = typeof(r_x)[], 
+
+        current=Ref{Int}(0) # reference to current top level state
+    )
+end
+
+
+function AbstractMCMC.save!!(samples, sample, iter::Integer, model::AbstractMultilevelModel, sampler::LykkegaardScheichl{true, P,N}; L=length(model), kwargs... ) where {P,N}
+    accept, x, f_x, r_x, Y = sample
+    
+    # Save proxies
+    append!(samples.states, Y.states)
+    append!(samples.logprobs, Y.logprobs)
+    append!(samples.rejections, Y.rejections)
+
+    if accept
+        push!(samples.states, x)
+        push!(samples.logprobs, f_x)
+        push!(samples.rejections, r_x)
+
+        samples.current[] = length(samples.states)  # update reference
+    else
+        samples.rejections[samples.current[]][L] += 1 # reference, 
+    end
+
+    return samples
+end 
+
+# Initialize chain 
+function AbstractMCMC.step(rng::AbstractRNG, model::AbstractMultilevelModel, sampler::LykkegaardScheichl; x0=nothing, f0=nothing, L=length(model), kwargs...) 
+    x = !isnothing(x0) ? x0 : rand(rng, sampler.proposal)
+    if !isnothing(x0) && !isnothing(f0) 
+        if (f0 isa Number) && (L == 1) 
+            f_x = [f0] 
+        elseif (f0 isa Vector)      
+            f_x = f0[1:L] 
+        end
+    else
+        f_x = [ logdensity(model, x; level=l) for l=1:L ]
+    end
+    Y = (; states=typeof(x)[], logprobs=typeof(f_x)[], rejections=Vector{Int}[])
+    return (true, x, f_x, zeros(Int, L), Y), (x, f_x)
+end
+
+function AbstractMCMC.step(rng::AbstractRNG, model::MultilevelSampledLogDensity, sampler::LykkegaardScheichl; x0=nothing, f0=nothing, L=length(model), kwargs...)
+    x = !isnothing(x0) ? x0 : rand(rng, sampler.proposal)
+    if !isnothing(x0) && !isnothing(f0) 
+        if (f0 isa Number && L == 1) 
+            f_x = [f0] 
+        elseif (f0 isa Vector)      
+            f_x = f0[1:L] 
+        end
+    else
+        f_x = [ logdensity(model, x; level=1) ]
+        sizehint!(f_x, L)
+        for l=2:L
+            push!(f_x, logdensity(model, x; level=l, cache=f_x[end]))
+        end 
+    end
+    Y = (; states=typeof(x)[], logprobs=typeof(f_x)[], rejections=Vector{Int}[])
+    return (true, x, f_x, zeros(Int, L), Y), (x, f_x)
+end
+ 
+
+function AbstractMCMC.step(rng::AbstractRNG, model::AbstractMultilevelModel, sampler::LykkegaardScheichl, state; L=length(model), kwargs...)
+    x, f_x = state
+
+    # Sample length at random
+    n = 1 + subchainlength(rng, sampler)
+    
+    if L == 2 # reduces to iterated MH 
+
+        # sample subchain
+        mh = MetropolisHastings(sampler.proposal)
+        model_1 = LogDensity(x->logdensity(model, x, level=1))
+        c = sample(rng, model_1, mh, n, chain_type=NamedTuple, x0=x,f0=f_x[1])
+        
+        # chain in between
+        Y = (; states     = c.states[2:end-1], 
+               logprobs   = map(x->[x],   c.logprobs[2:end-1]), 
+               rejections = map(x->[x,0], c.rejections[2:end-1])
+            )
+        
+        # end of chain
+        if length(Y.states) == 0
+            return (false, x, f_x, Vector{Int}[], Y), (x, f_x)
+        end
+
+        y = c.states[end]
+        f_y = [c.logprobs[end],   logdensity(model, y, level=2)]
+        r_y = [c.rejections[end], zeros(Int, L-1)... ]
+
+    else # Recursion 
+
+        # sample subchain
+        c = sample(rng, sampler, model, n, chain_type=(;), L=L-1, x0=x,f0=f_x[1:L-1], discard_initial=1)
+        
+        # chain in between
+        Y = (; states     = c.states[2:end-1], 
+               logprobs   = c.logprobs[2:end-1], 
+               rejections = map(x->[x,0], c.rejections[2:end-1])
+            )
+
+        # end of chain
+        if length(c.states) == 0
+            return (false, x, f_x, Vector{Int}[], Y), (x, f_x)
+        end
+        y = c.states[end]
+        f_y = [c.logprobs[end]...,   logdensity(model, y, level=L)] 
+        r_y = [c.rejections[end]..., 0 ]
+    end
+
+    # accept/reject step
+    A = min( f_y[L] - f_x[L] - f_y[L-1] + f_x[L-1], 0)
+    accept = log(rand(rng)) < A
+
+    if accept 
+        return (true, y, f_y, r_y, Y), (y, f_y)
+    else
+        return (false, x, f_x, Vector{Int}[], Y), (x, f_x)
+    end
+end
+
+function AbstractMCMC.bundle_samples(samples, model::AbstractMultilevelModel, sampler::LykkegaardScheichl, state, chain_type::Type{<:NamedTuple}; kwargs...) 
+    return (; states = samples.states, logprobs = samples.logprobs, rejections = samples.rejections )
+end
+
+function _chain_info(samples, model::AbstractMultilevelModel, sampler::LykkegaardScheichl) 
+    n_total = sum( 1 .+ sum.(samples.rejections))    
+    n_reject = sum(samples.rejections)
+    info = Dict(
+        :chain_length   => n_total,
+        :rejection_rate => n_reject ./ n_total,
+    )
+
+    if model isa MultilevelSampledLogDensity    # just copied from Christen-Fox sampler, is this valid ...?
+        evals_per_level = n_total .- cumsum([0, n_reject[1:end-1]... ])
+        costs_per_level = model.nlevels .- [0, model.nlevels[1:end-1]...] 
+        info[:total_costs] = sum( evals_per_level .* costs_per_level )
+    end
+
+    return info
+
+end
\ No newline at end of file

src/sampling/algos/metropolis_hastings.jl

+"""
+ Standard Metropolis-Hastings algorithm
+
+ References
+    
+    * Hastings, W.K. (1970). "Monte Carlo Sampling Methods Using Markov Chains and Their Applications". 
+      Biometrika, Volume 57, Issue 1
+
+""" 
+struct MetropolisHastings{P <: AbstractProposal} <: RejectionBasedSampler
+    proposal :: P
+end 
+
+# Initialize samples container
+function AbstractMCMC.samples(sample, model::AbstractModel, sampler::MetropolisHastings; kwargs...)
+    (accept, x, f_x) = sample
+    return (; 
+        states = typeof(x)[],
+        logprobs = typeof(f_x)[],
+        rejections = Int[], 
+    )
+end
+
+# Store sample to container
+function AbstractMCMC.save!!(samples, sample, ::Integer, ::AbstractModel, ::RejectionBasedSampler; kwargs... )
+    (accept, x, f_x) = sample
+
+    if sample[1] # accepted
+        push!(samples.states, x)
+        push!(samples.logprobs, f_x)
+        push!(samples.rejections, 0)
+    
+    else # rejected
+        samples.rejections[end] += 1 
+    end
+    return samples
+end
+
+
+function AbstractMCMC.step(rng::AbstractRNG, model::AbstractModel, sampler::MetropolisHastings; x0=nothing, f0=nothing, kwargs...)
+    
+    # Initialize states
+    _resample = isnothing(x0)  
+    x = _resample ? rand(rng, sampler.proposal) : x0
+    
+    # Initialize logprobs
+    _recompute = _resample || isnothing(f0)
+    f_x = _recompute ? logdensity(model, x) : f0
+
+    return (true, x, f_x), (x, f_x)
+end
+
+function AbstractMCMC.step(rng::AbstractRNG, model::AbstractModel, sampler::MetropolisHastings, state; kwargs...)
+    
+    # Load old state 
+    (x, f_x) = state
+
+    #propose new state  
+    y       = propose(rng, sampler.proposal, x)
+    f_y     = logdensity(model, y)
+    
+    q_ratio = logpratio(sampler.proposal, x, y)
+    logA = min( f_y - f_x + q_ratio, 0)
+    
+    # Accept / Reject step
+    if log(rand(rng)) < logA 
+        return (true, y, f_y), (y, f_y) # accept
+    else
+        return (false, x, f_x), (x, f_x) # reject
+    end
+end
+
+
+function AbstractMCMC.bundle_samples(samples, ::AbstractModel, ::MetropolisHastings, state, chain_type::Type{<:NamedTuple}; kwargs...)
+    return samples
+end
+
+function _chain_info(samples, model::AbstractModel, ::MetropolisHastings)
+    n_total = sum( 1 .+ samples.rejections)
+    n_reject = sum(samples.rejections)
+    info = Dict(
+        :chain_length   => n_total,
+        :rejection_rate => n_reject / n_total,
+    )
+
+    if model isa SampledLogDensity
+        info[:total_costs] = n_total .* length(model)
+    end
+
+    return info
+end

src/sampling/chains/all.jl

+abstract type AbstractRejectionChains <: AbstractMCMC.AbstractChains end
+
+include("rejection.jl")
+include("tables.jl")
+include("stats.jl")
+
+
+function Base.show(io::IO, c::AbstractRejectionChains)
+    print(io, "Chains (×$(nchains(c))) {", eltype(c), "} $(length(c)) elements. ")
+end
+
+function Base.display(c::AbstractRejectionChains)
+    println(c)
+    if is_multilevel(c)
+        ls = maximum.(levels(c))
+        if all( ls[2:end] .== ls[1] )
+            println(" proxy level L=", ls[1])
+        else
+            println(" proxy levels ", ls)
+        end
+    end
+    data = info(c)
+    for k in keys(data)
+        if nchains(c) == 1
+            println(" ", k, " : ", data[k][1])
+        else
+            m, s = mean(data[k]), std(data[k])
+            m = round.(m, sigdigits=5)
+            s = round.(s, sigdigits=3)
+            println(" ", k, " : ", m, " ± ", s)
+        end
+    end
+end
+

src/sampling/chains/estimator.jl

+abstract type AbstractEstimator{State} end
+
+
+struct ImportanceEstimator{S, P <: Real} <: AbstractEstimator{S}
+    states::Vector{S}
+    logprobs::Vector{Vector{P}}
+    rejections::Vector{Int}
+end
+
+#estimate(q::Q) = 
+
+levels = length.(logprobs)
+for l = 1:L
+s = states[(levels .== l)]
+mean(s)
+end

src/sampling/chains/rejection.jl

+# each chain may have different length
+struct RejectionChains{S,L,R} <: AbstractRejectionChains
+    samples::Vector{ 
+        Vector{ 
+            @NamedTuple begin
+                state::S
+                logprob::L
+                reject::R
+            end
+        }
+    }
+    info::Dict{Symbol, Vector}
+end
+
+
+# construct with empty info 
+function RejectionChains(samples::Vector{Vector{@NamedTuple{state::S, logprob::L, reject::R}}}) where {S,L,R}    
+    RejectionChains(samples, Dict{Symbol, Vector}())
+end
+
+
+# construct single chain
+function RejectionChains(s::Vector{S}, l::Vector{L}, r::Vector{R}) where {S,L,R}
+    samples = NamedTuple{ (:state, :logprob, :reject) }.([zip(s, l, r)...] )
+    return RejectionChains([samples])
+end
+
+# concatination method
+function AbstractMCMC.chainscat(c::RejectionChains ... )
+    infos = getfield.(c, :info)
+    lens = length.(c)
+    dict = Dict{Symbol, Vector}()
+    for k in union(keys.(infos) ... )
+        vals = [ haskey(infos[i],k) ? 
+                    infos[i][k] : repeat([missing], lens[i]) 
+                    for i = 1:length(c) 
+                ]
+        dict[Symbol(k)] = vcat(vals ... )
+    end
+    RejectionChains( vcat( getfield.(c, :samples) ... ), dict)
+end
+
+
+# Properties
+nchains(chains::RejectionChains) = length(chains.samples)
+length(chains::RejectionChains) = sum(length.(chains.samples)) # total number of unique samples (=storage costs)
+eltype(chains::RejectionChains{S,L,R}) where {S,L,R} = S       # state type
+
+# Get info on chain
+info(chains::RejectionChains) = chains.info
+
+# Get subset from list of chains samples
+function getindex(chains::RejectionChains, id::Integer)
+    samples = [chains.samples[id]]
+    info = Dict((k=>[chains.info[k][id]] for k in keys(chains.info)) ... ) 
+    RejectionChains(samples, info)
+end
+function getindex(chains::RejectionChains, ids::OrdinalRange) 
+    samples = chains.samples[ids]
+    info = Dict((k=>chains.info[k][ids] for k in keys(chains.info)) ... ) 
+    RejectionChains(samples, infos)
+end
+
+# Multilevel logprobs or rejection 
+is_multilevel(chains::RejectionChains{S,L,R}) where {S,L,R} = (L <: Union{AbstractVector, <:Tuple, <:NamedTuple})
+levels(chains::RejectionChains{S,L,R}) where {S, L <: Number, R} = ones(Int, nchains(chains))
+
+function levels(chains::RejectionChains{S,L,R}) where {S, L <: Union{AbstractVector, <:Tuple, <:NamedTuple}, R} 
+    return [
+        length.(getfield.(chains.samples[i], :logprob)) 
+        for i=1:nchains(chains)
+    ]
+end
+
+# get number of repetitions
+function repetitions(chains::RejectionChains)
+    return [ 
+        1 .+ sum.(getfield.(chains.samples[i], :reject)) 
+        for i=1:nchains(chains) 
+    ]
+end
+
+# Chain of values with repetition of where rejected 
+function states(chains::RejectionChains{S,L,R}) where {S,L,R}
+    V = Vector{S}[]
+    sizehint!(V, nchains(chains))
+    reps = repetitions(chains)
+    for i=1:nchains(chains)
+        x = chains.samples[i]
+        s = getfield.(x, :state)
+        v = vcat( fill.(s, reps[i]) ... )
+        push!(V, v)
+    end
+    return V # optionally: vcat(V ... )
+end

src/sampling/chains/stats.jl

+# Auto covariance
+#=
+function autocov(x::AbstractVector, lag::Int)
+    m = mean(x)
+    cov = (x[lag+1:end] .- m) .* (x[1:end-lag] .- m)
+    cor = (x .- m) .^ 2
+    return sum(cov)/sum(cor)
+end
+=#
+
+function autocov(c::RejectionChains)
+    X = states(c)
+    C = [ autocov(X[i]) for i=1:nchains(c) ]    
+    return C
+end
+
+function autocov(c::RejectionChains, lags::AbstractVector{<:AbstractVector})
+    X = states(c)
+    C = [ autocov(X[i], lags[i]) for i=1:nchains(c) ]    
+    return C
+end
+function autocov(c::RejectionChains, lags::AbstractVector{<:Integer})
+    X = states(c)
+    C = [ autocov(X[i], lags) for i=1:nchains(c) ]    
+    return C
+end
+
+
+# Log partition function, constant factor for normalize log density 
+function logpart(logp::AbstractVector)
+    return -log( sum( exp.(-logp) ) )
+end
+
+function logpart(c::RejectionChains)
+
+    # single level
+    if ! is_multilevel(c)
+        reps = repetitions(c)
+        return [ 
+            (
+                logprobs = getfield.(c.samples[i], :logprob);
+                logprobs = vcat( fill.(logprobs, reps[i])... );
+                logpart(logprobs)
+            )
+            for i=1:nchains(c)
+        ]
+
+    # multi level
+    else
+        reps = repetitions(c)
+        lvls = levels(c); maxlvl = maximum(lvls)
+        return [
+            [ # return one constant per level
+                (
+                    idx = lvls[i] <= l;
+                    logprobs = getindex.(getfield.(c.samples[i][idx], :logprob), l);
+                    logprobs = vcat( fill.(logprobs, reps[i][idx])... );
+                    logpart(logprobs)
+                )
+                for l = 1:maxlvl
+            ]
+            for i=1:nchains(c)
+        ]
+    end
+end
+
+function empirical_cdf(c::RejectionChains, x)
+    x = convert(eltype(c), x)
+
+    if ! is_multilevel(c)
+        reps = repetitions(c)
+        return [
+            (
+                s = getfield.(c.samples[i], :state);
+                sum( reps[i] .* [ all(X .<= x) for X in s] ) / sum(reps[i]) 
+            )
+            for i=1:nchains(c)
+        ]
+    else
+
+        reps = repetitions(c)
+        lvls = levels(c); maxlvl = maximum(lvls)
+        
+        cdfs = [ 
+            (
+                idx = lvls[i] .<= 1;
+                s = getfield.(c.samples[i], :state)[idx];
+                [ sum(reps[i][idx] .* [ all(X .<= x) for X in s]) / sum(reps[i][idx]) ]
+            )
+            for i=1:nchains(c)
+        ]
+        sizehint!.(cdfs, maxlvl)
+
+        logparts = logpart(c)
+        #Z = mean(logparts) # average over chain, otherwise for each chain seperately
+        for l=2:maxlvl
+            for i=1:nchains(c)
+                idx = lvls[i] <= l
+                s = getfield.(c.samples[i], :state)[idx]; # ?????
+                f = getfield.(c.samples[i], :logprob)[idx];
+                f0 = getindex.(f, l-1)
+                f1 = getindex.(f, l)
+                Z0,Z1 = logparts[i][l-1:l]
+                
+                _cdf = sum( 
+                    reps[i][idx] .* [ all(X .<= x) for x in s ] 
+                                 .* (1 .- exp.(f1 - f0 - Z1 + Z0)) 
+                ) / sum(reps[i][idx])
+                push!(cdfs[i], _cdf) 
+            end
+        end
+        return cdfs
+    end 
+end
\ No newline at end of file

src/sampling/chains/tables.jl

+
+
+# Write to table
+function convert(::Type{<:AbstractVector{<:DataFrame}}, chains::RejectionChains)
+    dfs = DataFrame[]; sizehint!(dfs,length(chains.samples))
+    for i=1:nchains(chains)
+        n = length(chains.samples[i])
+        df = DataFrame()
+        for k in (:state, :logprob, :reject)
+             df[:, k] = getfield.(chains.samples[i], k)
+        end
+        push!(dfs, df)
+    end
+    return dfs
+end
+
+# Read from table
+function convert(::Type{<:RejectionChains}, df::DataFrame)
+    s = df[:, :state] 
+    l = df[:, :logprob] 
+    r = df[:, :reject]
+    c = RejectionChains(s,l,r)
+    return c
+end
+
+function convert(::Type{<:RejectionChains}, df::Vector{<:DataFrame})
+    return AbstractMCMC.chainscat(convert.(RejectionChains, df) ... )
+end

src/sampling/models/all.jl

+abstract type AbstractLogDensity <: AbstractModel end
+
+struct LogDensity{L} <: AbstractLogDensity
+    density :: L
+end
+
+logdensity(model::LogDensity{<:Function}, x; kwargs...) = model.density(x)
+logdensity(model::LogDensity{<:Distribution}, x; kwargs...) = logpdf(model.density, x)
+logdensity(model::AbstractMCMC.LogDensityModel, x; kwargs...) = LogDensityProblems.logdensity(model.logdensity, x)
+
+include("samplebased.jl")
+include("multilevel.jl")
\ No newline at end of file

src/sampling/models/multilevel.jl

+abstract type AbstractMultilevelModel <: AbstractModel end
+
+struct MultilevelLogDensity{ L <: Tuple{Vararg{<:AbstractModel}}}  <: AbstractMultilevelModel
+    proxies :: L
+end
+MultilevelLogDensity(v::L... ) where {L <: AbstractModel} = MultilevelLogDensity(v)
+MultilevelLogDensity(v::Vector{ <: AbstractModel} ) = MultilevelLogDensity( tuple(v...) )
+
+length(m::MultilevelLogDensity) = length(m.proxies)
+logdensity(m::MultilevelLogDensity, x; level=length(m), kwargs...) = logdensity(m.proxies[level], x)
+
+function Base.show(io::IO, m::MultilevelLogDensity)
+    println(io, "MultilevelLogDensity ", length(m))
+    for i=1:length(m)
+        println(io, "  ", m.proxies[i])
+    end
+end
+
+struct MultilevelSampledLogDensity{X, F <: Function} <: AbstractMultilevelModel
+    density :: SampledLogDensity{X, F}
+    nlevels :: Vector{Int}
+    # @assert maximum(nlevels) <= length(density)
+end 
+MultilevelSampledLogDensity(f::Function, n::Vector{<:Integer}, d::Int=1, a=0, b=1) = MultilevelSampledLogDensity(SampledLogDensity(f, maximum(n), d, a, b), n)
+
+function Base.show(io::IO, m::MultilevelSampledLogDensity)
+    print(io, "MultilevelSampledLogDensity ", m.nlevels, " {", eltype(m.density), "} samples ")
+end
+
+length(m::MultilevelSampledLogDensity) = length(m.nlevels)
+logdensity(m::MultilevelSampledLogDensity, x; level=length(m), cache=nothing, kwargs...) = _logdensity(m, x, level, cache)
+
+_logdensity(m::MultilevelSampledLogDensity, x, level::Int, cache::Nothing) = logdensity(m.density, x; level=m.nlevels[level])
+_logdensity(m::MultilevelSampledLogDensity, x, level::Int, cache::Number)  = logdensity(m.density, x; level=m.nlevels[level], cache=(level-1, cache))
+
+
+
+

src/sampling/models/samplebased.jl

+
+struct SampledLogDensity{T, F <: Function} <: AbstractLogDensity
+    samples :: Vector{T}
+    func :: F
+end
+function SampledLogDensity(f::Function, n::Int=32, d::Int=1, a=0, b=1) 
+    p = shuffle!(primes(d^2+1)[1:d])
+    z = hcat((Halton(p[i])[1:n] for i=1:d) ... )
+    z = [Tuple(z[i,:]) for i=1:n]  
+    return SampledLogDensity(z, f)
+end
+eltype(::SampledLogDensity{T,F}) where {T,F} = T
+length(m::SampledLogDensity) = length(m.samples)
+
+function Base.show(io::IO, m::SampledLogDensity)
+    print(io, "SampledLogDensity ", length(m), " {", eltype(m), "} samples ")
+end
+
+logdensity(m::SampledLogDensity, x; level=length(m), cache=nothing) = _logdensity(m, x, level, cache)
+
+function _logdensity(m::SampledLogDensity, x, level::Int64, cache::Nothing)  
+    y = [m.func(x, z) for z=m.samples[1:level]] #TODO: parallelize evaluations of m.func !
+    return mean(y)
+end
+
+function _logdensity(m::SampledLogDensity, x, level::Int64, cache::Tuple{Int64, <:Number})  
+    if cache[1] <= level
+        y = [m.func(x, z) for z=m.samples[cache[1]+1:level]]
+        return mean(y) * length(y)/level + cache[2] * cache[1]/level
+    else
+        return logdensity(m, x; level, nothing) 
+    end
+end

src/sampling/proposals/all.jl

+abstract type AbstractProposal end
+rand(rng::AbstractRNG, p::AbstractProposal) = rand(Random.default_rng(), p)
+propose(rng::AbstractRNG, p::AbstractProposal, x=nothing) = propose(rng, p, x)
+propose(rng::AbstractRNG, p::AbstractProposal, x::Nothing) = rand(rng, p)
+logpratio(p::AbstractProposal, x, y) = 0 # symmetric proposal by default
+
+struct RandomWalk{I, S, F} <: AbstractProposal 
+    init :: I
+    step :: S
+    func :: F
+end
+RandomWalk(init::Distribution, step::Distribution) = RandomWalk(init, step, (.+))
+
+rand(rng::AbstractRNG, p::RandomWalk{<:Distribution, <:Any}) = rand(rng, p.init) 
+propose(rng::AbstractRNG, p::RandomWalk{<:Distribution, <:Distribution}, x) = p.func(x, rand(rng, p.step))
+
+
+function CyclicWalk(a=-1, b=1, s=.5) 
+    init = Uniform(a,b); step = Uniform(-s, s) 
+    func = (x,y) -> mod(x+y - a, b-a) + a
+    RandomWalk(init, step, func)
+end
+
+include("networks.jl")
\ No newline at end of file

src/sampling/proposals/networks.jl

+abstract type NetworkEnsemble <: AbstractProposal end
+
+struct ErdosRenyi <: NetworkEnsemble
+    n::Int64
+    p::Float64
+    s::Int64
+end
+ErdosRenyi(n::Int64) = ErdosRenyi(n, log(n)/n * (1 + exp(-rand())))
+ErdosRenyi(n::Int64, p::AbstractFloat) = ErdosRenyi(n, p, ceil(Int, log(n)))
+ErdosRenyi(n::Int64, m::Int64, s::Int64=ceil(Int, log(n))) = ErdosRenyi(n, m/(n*(n-1)/2), s)
+
+Base.rand(rng::AbstractRNG, e::ErdosRenyi) = erdos_renyi(e.n, e.p; rng)
+
+function propose(rng::AbstractRNG, e::ErdosRenyi, G::Graph)
+    G = copy(G)
+    for t = 1:e.s
+
+        # Choose a pair
+        u = rand(rng, 1:nv(G))
+        v = rand(rng, 1:nv(G)-1)
+        if (v>=u) v+=1 end
+
+        # Resample edge 
+        if has_edge(G, u, v)
+            if rand(rng) < 1 - e.p
+                rem_edge!(G, u, v)
+            end
+        else
+            if rand(rng) < e.p
+                add_edge!(G, u, v)
+            end
+        end
+    end
+    return G
+end
\ No newline at end of file