Logistic SOSS&Pathfinder&BPS

bps.jl

using Soss
using ProgressMeter
using LinearAlgebra
using Revise
using ZigZagBoomerang
const ZZB = ZigZagBoomerang
using LinearAlgebra
const ∅ = nothing
using DelimitedFiles
using Random
using ForwardDiff
using ForwardDiff: Dual
using Pathfinder
using Pathfinder.PDMats
using MappedArrays

model = @model (N, C, c) begin
    α ~ Soss.Normal(0,1)
    cc ~ Soss.Normal(0,1) |> iid(C)
    y ~ For(1:N) do i
        v = α + cc[c[i]]
        Soss.Bernoulli(logistic(v))
    end
    return y
end

N = 1000
C = 10
data = (;c=rand(1:C, N), y=rand(Bool, N))
condmodel = model(;N,C,c=data.c) | (;y=data.y);


function make_grads(my_model, y)        
    post = my_model | (;y)
    as_post = as(post)
    obj(θ) = -logdensityof(post, transform(as_post, θ))
    ℓ(θ) = -obj(θ)
    @inline function dneglogp(t, x, v) # two directional derivatives
        f(t) = obj(x + t*v)
        u = ForwardDiff.derivative(f, Dual{:hSrkahPmmC}(0.0, 1.0))
        u.value, u.partials[]
    end
    
    #gconfig = ForwardDiff.GradientConfig(obj, rand(d), ForwardDiff.Chunk{25}())
    function ∇neglogp!(y, t, x)
        #ForwardDiff.gradient!(y, obj, x, gconfig)
        ForwardDiff.gradient!(y, obj, x)
        return
    end
    post, ℓ, dneglogp, ∇neglogp!
end

post, ℓ, dneglogp, ∇neglogp! = make_grads(model(;N,C,c=data.c), data.y)  
# Try things out


d = 11 # number of parameters 
dneglogp(2.4, randn(d), randn(d));
#∇neglogp!(randn(d), 2.1, randn(d));

t0 = 0.0;
n = 2000
c = 1.0 # initial guess for the bound

init_scale=1;
@time pf_result = pathfinder(ℓ; dim=d, init_scale);
x0 = pf_result.fit_distribution.μ
M = pf_result.fit_distribution.Σ
v0 = PDMats.unwhiten(M, normalize!(randn(length(x0))));
MAP = pf_result.optim_solution; # MAP, could be useful for control variates

# define BouncyParticle sampler (has two relevant parameters) 
Z = BouncyParticle(missing, # graphical structure 
    MAP, # MAP estimate, unused
    1.0, # momentum refreshment rate and sample saving rate 
    0.9, # momentum correlation / only gradually change momentum in refreshment/momentum update
    M, # metric (PDMat compatible object for momentum covariance)
    missing # legacy
); 

sampler = ZZB.NotFactSampler(Z, (dneglogp, ∇neglogp!), ZZB.LocalBound(c), t0 => (x0, v0), ZZB.Rng(ZZB.Seed()), (),
(;  adapt=true, # adapt bound c
    subsample=true, # keep only samples at refreshment times
));


using TupleVectors: chainvec
using Soss.MeasureTheory: transform


function collect_sampler(t, sampler, n; progress=true, progress_stops=20)
    if progress
        prg = Progress(progress_stops, 1)
    else
        prg = missing
    end
    stops = ismissing(prg) ? 0 : max(prg.n - 1, 0) # allow one stop for cleanup
    nstop = n/stops

    x1 = transform(t, sampler.u0[2][1])
    tv = chainvec(x1, n)
    ϕ = iterate(sampler)
    j = 1
    local state
    while ϕ !== nothing && j < n
        j += 1
        val, state = ϕ
        tv[j] = transform(t, val[2])
        ϕ = iterate(sampler, state)
        if j > nstop
            nstop += n/stops
            next!(prg) 
        end 
    end
    ismissing(prg) || ProgressMeter.finish!(prg)
    tv, (;uT=state[1], acc=state[3][1], total=state[3][2], bound=state[4].c)
end
collect_sampler(as(post), sampler, 10; progress=false);

elapsed_time = @elapsed @time begin
    global bps_samples, info 
    bps_samples, info = collect_sampler(as(post), sampler, n; progress=true)
end

 
using MCMCChains
bps_chain = MCMCChains.Chains([bps_samples.α bps_samples.cc.data'])
bps_chain = setinfo(bps_chain, (;start_time=0.0, stop_time=elapsed_time));
ess_bps = MCMCChains.ess_rhat(bps_chain).nt.ess_per_sec;

μ̂1 = round.(mean(bps_chain).nt[:mean], sigdigits=4)
println("μ̂ (BPS) = ", μ̂1)

@show info.bound
@show round(info.acc/info.total, sigdigits=2)

bps_chain

Project.toml

[deps]
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
Pathfinder = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
SampleChainsDynamicHMC = "6d9fd711-e8b2-4778-9c70-c1dfb499d4c4"
Soss = "8ce77f84-9b61-11e8-39ff-d17a774bf41c"
TupleVectors = "615932cf-77b6-4358-adcd-5b7eba981d7e"
UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
ZigZagBoomerang = "36347407-b186-4a6a-8c98-4f4567861712"