diff --git a/docs/make.jl b/docs/make.jl index cee8b8f..2d0fc5c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -13,7 +13,7 @@ makedocs(; ), pages = [ "Home" => "index.md", -# "Example: Gaussian Mixture" => "example_1.md", + # "Example: Gaussian Mixture" => "example_1.md", #"Reference" => "reference.md", ], ) diff --git a/examples/example_n1.jl b/examples/example_n1.jl index 91240fd..04246ec 100644 --- a/examples/example_n1.jl +++ b/examples/example_n1.jl @@ -45,7 +45,7 @@ res = sample( # You can also use Sequential Monte Carlo (SMC) to infer posterior parameters: -ressmc = smc(prior, cost, nparticles=500, epstol=0.01) +ressmc = smc(prior, cost, nparticles = 500, epstol = 0.01) @show ressmc diff --git a/examples/example_n2.jl b/examples/example_n2.jl index 91bfbf5..6bee68d 100644 --- a/examples/example_n2.jl +++ b/examples/example_n2.jl @@ -9,8 +9,8 @@ function model(P, N) d1 = @. r1 * σ_1 + μ_1 d2 = @. r1 * σ_2 + μ_2 - ps = @. (1 + sign(r2 - prob))/2 - @. (d1+ps*(d2-d1)) + ps = @. (1 + sign(r2 - prob)) / 2 + @. (d1 + ps * (d2 - d1)) end @@ -56,5 +56,5 @@ approx_density = ApproxPosterior(prior, D, 0.032) @show res # In this case, it is best to apply SMC, as it leads to tighter CI's and lower computational costs -@time res = smc(prior, D, verbose=false, nparticles=100, alpha=0.95) +@time res = smc(prior, D, verbose = false, nparticles = 100, alpha = 0.95) @show res.P diff --git a/src/KissABC.jl b/src/KissABC.jl index 210bc88..3ff45e6 100644 --- a/src/KissABC.jl +++ b/src/KissABC.jl @@ -55,8 +55,9 @@ function AbstractMCMC.step( particles[i] = op(float, unconditional_sample(rng, model)) logdensity[i] = loglike(model, push_p(model, particles[i])) retrys -= 1 - retrys < 0 && - error("Prior leads to ∞ costs too often, tune the prior or increase `retry_sampling`.") + retrys < 0 && error( + "Prior leads to ∞ costs too often, tune the prior or increase `retry_sampling`.", + ) end end diff --git a/src/smc.jl b/src/smc.jl index 262995d..cf94ef1 100644 --- a/src/smc.jl +++ b/src/smc.jl @@ -1,6 +1,6 @@ macro cthreads(condition::Symbol, loop) #does not work well because of #15276, but seems to work on Julia v0.7 return esc(quote - if $condition + if $condition Threads.@threads $loop else $loop @@ -110,41 +110,38 @@ function smc( r_epstol >= 0 || error("r_epstol must be >= 0") mcmc_tol >= 0 || error("mcmc_tol must be >= 0") max_stretch > 1 || error("max_stretch must be > 1") - Np=length(prior) - min_nparticles = ceil( - Int, - 3 * Np / (min(alpha, min_r_ess)), - ) + Np = length(prior) + min_nparticles = ceil(Int, 3 * Np / (min(alpha, min_r_ess))) nparticles >= min_nparticles || error("nparticles must be >= $min_nparticles.") θs = [op(float, Particle(rand(rng, prior))) for i = 1:nparticles] - Xs = parallel ? - fetch.([ - Threads.@spawn cost(push_p(prior, θs[$i].x)) for i = 1:nparticles]) : + Xs = + parallel ? + fetch.([Threads.@spawn cost(push_p(prior, θs[$i].x)) for i = 1:nparticles]) : [cost(push_p(prior, θs[i].x)) for i = 1:nparticles] lπs = [logpdf(prior, push_p(prior, θs[i].x)) for i = 1:nparticles] α = alpha ϵ = Inf - alive = fill(true,nparticles) + alive = fill(true, nparticles) iteration = 0 # Step 1 - adaptive threshold while true iteration += 1 ϵv = ϵ - ϵ = quantile(Xs[alive],α) - flag=false + ϵ = quantile(Xs[alive], α) + flag = false if ϵ > minimum(Xs[alive]) alive = Xs .< ϵ else alive = Xs .<= ϵ - flag=true + flag = true end ESS = sum(alive) verbose && @show iteration, ϵ, ESS # Step 2 - Resampling - if α*ESS <= nparticles * min_r_ess + if α * ESS <= nparticles * min_r_ess idxalive = (1:nparticles)[alive] - idx=repeat(idxalive,ceil(Int,nparticles/length(idxalive)))[1:nparticles] + idx = repeat(idxalive, ceil(Int, nparticles / length(idxalive)))[1:nparticles] θs = θs[idx] Xs = Xs[idx] lπs = lπs[idx] @@ -157,51 +154,55 @@ function smc( retry_N = 1 + mcmc_retrys for r = 1:retry_N - new_p = map(1:nparticles) do i - a = b = i - alive[i] || return (nothing,nothing,nothing) - while a==i; a = rand(rng,1:nparticles); end - while b==i || b==a; b = rand(rng,1:nparticles); end - W = op(*, op(-, θs[b], θs[a]), max_stretch*randn(rng)/sqrt(Np)) - (log(rand(rng)), op(+, θs[i], W), 0.0) + new_p = map(1:nparticles) do i + a = b = i + alive[i] || return (nothing, nothing, nothing) + while a == i + a = rand(rng, 1:nparticles) end - @cthreads parallel for i = 1:nparticles # non-ideal parallelism - alive[i] || continue - lprob, θp, logcorr = new_p[i] - isnothing(lprob) && continue - lπp = logpdf(prior, push_p(prior, θp.x)) - lπp < 0 && (!isfinite(lπp)) && continue - lM = min(lπp - lπs[i] + logcorr, 0.0) - if lprob < lM - Xp = cost(push_p(prior, θp.x)) - if flag - Xp > ϵ && continue - else - Xp >= ϵ && continue - end - θs[i] = θp - Xs[i] = Xp - lπs[i] = lπp - if parallel - Threads.atomic_add!(accepted, 1) - else - accepted += 1 - end + while b == i || b == a + b = rand(rng, 1:nparticles) + end + W = op(*, op(-, θs[b], θs[a]), max_stretch * randn(rng) / sqrt(Np)) + (log(rand(rng)), op(+, θs[i], W), 0.0) + end + @cthreads parallel for i = 1:nparticles # non-ideal parallelism + alive[i] || continue + lprob, θp, logcorr = new_p[i] + isnothing(lprob) && continue + lπp = logpdf(prior, push_p(prior, θp.x)) + lπp < 0 && (!isfinite(lπp)) && continue + lM = min(lπp - lπs[i] + logcorr, 0.0) + if lprob < lM + Xp = cost(push_p(prior, θp.x)) + if flag + Xp > ϵ && continue + else + Xp >= ϵ && continue + end + θs[i] = θp + Xs[i] = Xp + lπs[i] = lπp + if parallel + Threads.atomic_add!(accepted, 1) + else + accepted += 1 end end + end accepted[] >= mcmc_tol * nparticles && break end - if 2*abs(ϵv - ϵ) < r_epstol * (abs(ϵv)+abs(ϵ)) || + if 2 * abs(ϵv - ϵ) < r_epstol * (abs(ϵv) + abs(ϵ)) || ϵ <= epstol || accepted[] < mcmc_tol * nparticles - break + break end end θs = [push_p(prior, θs[i].x) for i = 1:nparticles][alive] l = length(prior) P = map(x -> Particles(x), getindex.(θs, i) for i = 1:l) - length(P)==1 && (P=first(P)) + length(P) == 1 && (P = first(P)) (P = P, C = Xs, ϵ = ϵ) end @@ -272,71 +273,91 @@ Particles(sigmapoints(mean(R.P),cov(R.P))) -function pfilter(prior, cost, N; rng=Random.GLOBAL_RNG, q=0.7, eff_tol = 0.1, epstol=-Inf, max_iters = Inf, proposal_width=0.75, verbose=false, parallel=false) - lowN=4*length(prior) - if N*q<=lowN - N=ceil(Int,(lowN+1)/q) +function pfilter( + prior, + cost, + N; + rng = Random.GLOBAL_RNG, + q = 0.7, + eff_tol = 0.1, + epstol = -Inf, + max_iters = Inf, + proposal_width = 0.75, + verbose = false, + parallel = false, +) + lowN = 4 * length(prior) + if N * q <= lowN + N = ceil(Int, (lowN + 1) / q) end - sample=[op(float, Particle(rand(rng, prior))) for i = 1:N] - logπ = [logpdf(prior, push_p(prior,sample[i].x)) for i = 1:N] - C = fill(cost(sample[1].x),N) + sample = [op(float, Particle(rand(rng, prior))) for i = 1:N] + logπ = [logpdf(prior, push_p(prior, sample[i].x)) for i = 1:N] + C = fill(cost(sample[1].x), N) @cthreads parallel for i = 1:N - trng=rng - parallel && (trng=Random.default_rng(Threads.threadid());) + trng = rng + parallel && (trng = Random.default_rng(Threads.threadid())) if isfinite(logπ[i]) C[i] = cost(sample[i].x) end while (!isfinite(C[i])) || (!isfinite(logπ[i])) - sample[i]=op(float, Particle(rand(trng, prior))) - logπ[i] = logpdf(prior, push_p(prior,sample[i].x)) + sample[i] = op(float, Particle(rand(trng, prior))) + logπ[i] = logpdf(prior, push_p(prior, sample[i].x)) C[i] = cost(sample[i].x) end end iters = 0 while true - iters += 1 - ϵ = quantile(C,q) - filter_bad=C .> ϵ - idxok=(1:N)[.!filter_bad] - idxbad=(1:N)[filter_bad] - nreps= Threads.Atomic{Int}(0) + iters += 1 + ϵ = quantile(C, q) + filter_bad = C .> ϵ + idxok = (1:N)[.!filter_bad] + idxbad = (1:N)[filter_bad] + nreps = Threads.Atomic{Int}(0) @cthreads parallel for i in idxbad - trng=rng - parallel && (trng=Random.default_rng(Threads.threadid());) - localreps=0 + trng = rng + parallel && (trng = Random.default_rng(Threads.threadid())) + localreps = 0 @label resample - b=c=d=rand(trng,idxok) - while c==b; c=rand(trng,idxok); end - while d==b || d==c; d=rand(trng,idxok); end - p=op(+,sample[b],op(*,op(-,sample[d],sample[c]), randn(trng)*proposal_width)) + b = c = d = rand(trng, idxok) + while c == b + c = rand(trng, idxok) + end + while d == b || d == c + d = rand(trng, idxok) + end + p = op( + +, + sample[b], + op(*, op(-, sample[d], sample[c]), randn(trng) * proposal_width), + ) localreps += 1 - ll = logpdf(prior, push_p(prior,p.x)) - if log(rand(trng)) > min(0.0,ll-logπ[i]) + ll = logpdf(prior, push_p(prior, p.x)) + if log(rand(trng)) > min(0.0, ll - logπ[i]) @goto resample end - Cp=cost(p.x) + Cp = cost(p.x) if Cp > ϵ @goto resample end C[i] = Cp sample[i] = p logπ[i] = ll - Threads.atomic_add!(nreps,localreps) + Threads.atomic_add!(nreps, localreps) end - eff=length(idxbad)/nreps[] + eff = length(idxbad) / nreps[] verbose && @show iters, ϵ, eff - eff max_iters && break end θs = [push_p(prior, sample[i].x) for i = 1:N] l = length(prior) P = map(x -> Particles(x), getindex.(θs, i) for i = 1:l) - length(P)==1 && (P=first(P)) - (P=P, C=Particles(C)) + length(P) == 1 && (P = first(P)) + (P = P, C = Particles(C)) end @@ -344,64 +365,76 @@ export pfilter -function ABCDE(prior, cost, ϵ_target; nparticles=50, generations=20, α=0, parallel=false, earlystop=false, verbose=true, rng=Random.GLOBAL_RNG, proposal_width=1.0) - @assert 0<=α<1 "α must be in 0 <= α < 1." - θs =[op(float, Particle(rand(rng, prior))) for i = 1:nparticles] +function ABCDE( + prior, + cost, + ϵ_target; + nparticles = 50, + generations = 20, + α = 0, + parallel = false, + earlystop = false, + verbose = true, + rng = Random.GLOBAL_RNG, + proposal_width = 1.0, +) + @assert 0 <= α < 1 "α must be in 0 <= α < 1." + θs = [op(float, Particle(rand(rng, prior))) for i = 1:nparticles] - logπ = [logpdf(prior, push_p(prior,θs[i].x)) for i = 1:nparticles] - Δs = fill(cost(θs[1].x),nparticles) + logπ = [logpdf(prior, push_p(prior, θs[i].x)) for i = 1:nparticles] + Δs = fill(cost(θs[1].x), nparticles) @cthreads parallel for i = 1:nparticles - trng=rng - parallel && (trng=Random.default_rng(Threads.threadid());) + trng = rng + parallel && (trng = Random.default_rng(Threads.threadid())) if isfinite(logπ[i]) Δs[i] = cost(θs[i].x) end while (!isfinite(Δs[i])) || (!isfinite(logπ[i])) - θs[i]=op(float, Particle(rand(trng, prior))) - logπ[i] = logpdf(prior, push_p(prior,θs[i].x)) + θs[i] = op(float, Particle(rand(trng, prior))) + logπ[i] = logpdf(prior, push_p(prior, θs[i].x)) Δs[i] = cost(θs[i].x) end end - nsims = zeros(Int,nparticles) - γ = proposal_width*2.38/sqrt(2*length(prior)) - iters=0 - complete=1-sum(Δs.>ϵ_target)/nparticles - while iters ϵ_target) / nparticles + while iters < generations + iters += 1 nθs = identity.(θs) nΔs = identity.(Δs) - nlogπ=identity.(logπ) + nlogπ = identity.(logπ) ϵ_l, ϵ_h = extrema(Δs) if earlystop - ϵ_h<=ϵ_target && break + ϵ_h <= ϵ_target && break end - ϵ_pop = max(ϵ_target,ϵ_l + α * (ϵ_h - ϵ_l)) - @cthreads parallel for i in 1:nparticles + ϵ_pop = max(ϵ_target, ϵ_l + α * (ϵ_h - ϵ_l)) + @cthreads parallel for i = 1:nparticles if earlystop Δs[i] <= ϵ_target && continue end - trng=rng - parallel && (trng=Random.default_rng(Threads.threadid());) + trng = rng + parallel && (trng = Random.default_rng(Threads.threadid())) s = i ϵ = ifelse(Δs[i] <= ϵ_target, ϵ_target, ϵ_pop) if Δs[i] > ϵ - s=rand(trng,(1:nparticles)[Δs .<= Δs[i]]) + s = rand(trng, (1:nparticles)[Δs.<=Δs[i]]) end a = s while a == s - a = rand(trng,1:nparticles) + a = rand(trng, 1:nparticles) end b = a while b == a || b == s - b = rand(trng,1:nparticles) + b = rand(trng, 1:nparticles) end - θp = op(+,θs[s],op(*,op(-,θs[a],θs[b]), γ)) - lπ= logpdf(prior, push_p(prior,θp.x)) + θp = op(+, θs[s], op(*, op(-, θs[a], θs[b]), γ)) + lπ = logpdf(prior, push_p(prior, θp.x)) w_prior = lπ - logπ[i] - log(rand(trng)) > min(0,w_prior) && continue - nsims[i]+=1 + log(rand(trng)) > min(0, w_prior) && continue + nsims[i] += 1 dp = cost(θp.x) if dp <= max(ϵ, Δs[i]) nΔs[i] = dp @@ -414,20 +447,22 @@ function ABCDE(prior, cost, ϵ_target; nparticles=50, generations=20, α=0, para logπ = nlogπ ncomplete = 1 - sum(Δs .> ϵ_target) / nparticles if verbose && (ncomplete != complete || complete >= (nparticles - 1) / nparticles) - @info "Finished run:" completion=ncomplete nsim = sum(nsims) range_ϵ = extrema(Δs) + @info "Finished run:" completion = ncomplete nsim = sum(nsims) range_ϵ = + extrema(Δs) end - complete=ncomplete + complete = ncomplete end - conv=maximum(Δs) <= ϵ_target + conv = maximum(Δs) <= ϵ_target if verbose - @info "End:" completion = complete converged = conv nsim = sum(nsims) range_ϵ = extrema(Δs) + @info "End:" completion = complete converged = conv nsim = sum(nsims) range_ϵ = + extrema(Δs) end θs = [push_p(prior, θs[i].x) for i = 1:nparticles] l = length(prior) P = map(x -> Particles(x), getindex.(θs, i) for i = 1:l) - length(P)==1 && (P=first(P)) - (P=P, C=Particles(Δs), reached_ϵ=conv) + length(P) == 1 && (P = first(P)) + (P = P, C = Particles(Δs), reached_ϵ = conv) end -export ABCDE \ No newline at end of file +export ABCDE diff --git a/test/runtests.jl b/test/runtests.jl index 3fe44e0..be90d69 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -127,7 +127,7 @@ end modelabc = ApproxPosterior(prior, cost, 0.1) sim = sample(modelabc, AIS(50), 100, discard_initial = 50000, progress = false) @test all(sim .≈ params) - @test all(smc(prior, cost, verbose=false,min_r_ess=0.55).P .≈ params) + @test all(smc(prior, cost, verbose = false, min_r_ess = 0.55).P .≈ params) end @testset "Classical Mixture Model 0.1N+N" begin @@ -162,7 +162,17 @@ end discard_initial = 5000, progress = false, ) - ressmc = smc(prior, cost, nparticles = 2000, alpha = 0.9, epstol = 0.01,mcmc_retrys=500,mcmc_tol=0.9,verbose=false).P + ressmc = + smc( + prior, + cost, + nparticles = 2000, + alpha = 0.9, + epstol = 0.01, + mcmc_retrys = 500, + mcmc_tol = 0.9, + verbose = false, + ).P testst(alg, r) = begin m = mean(abs, st(r) - st_n) println(":", alg, ": testing m = ", m) @@ -241,13 +251,32 @@ end pp = Factored(Normal(0, 5), Normal(0, 5)) cc((x, y)) = 50 * (x + randn() * 0.01 - y^2)^2 + (y - 1 + randn() * 0.01)^2 - R = smc(pp, cc, verbose = false, alpha = 0.9, nparticles = 500, epstol = 0.01, parallel=true).P + R = + smc( + pp, + cc, + verbose = false, + alpha = 0.9, + nparticles = 500, + epstol = 0.01, + parallel = true, + ).P @test R[1] ≈ 1 @test R[2] ≈ 1 - cc2((x, y)) = rand((50 * (x + randn() * 0.01 - y^2)^2 + (y - 1 + randn() * 0.01)^2,Inf)) + cc2((x, y)) = + rand((50 * (x + randn() * 0.01 - y^2)^2 + (y - 1 + randn() * 0.01)^2, Inf)) - R = smc(pp, cc2, verbose = false, alpha = 0.9, nparticles = 1000, epstol = 0.01, parallel=true).P + R = + smc( + pp, + cc2, + verbose = false, + alpha = 0.9, + nparticles = 1000, + epstol = 0.01, + parallel = true, + ).P @test R[1] ≈ 1 @test R[2] ≈ 1