diff --git a/Example/Using_Custom_VI/demonstarte.jl b/Example/Using_Custom_VI/demonstarte.jl new file mode 100644 index 00000000..6c3b3d9e --- /dev/null +++ b/Example/Using_Custom_VI/demonstarte.jl @@ -0,0 +1,57 @@ + +## It is not yet a package... +using Distributions +using AdvancedPS +using Libtask +using BenchmarkTools +using NamedTupleTools +include("AdvancedPS/Example/Using_Custom_VI/test_interface.jl") +# Define a short model. +# The syntax is rather simple. Observations need to be reported with report_observation. +# Transitions must be reported using report_transition. +# The trace contains the variables which we want to infer using particle gibbs. +# Thats all! +n = 3000 + +y = Vector{Float64}(undef,n-1) +for i =1:n-1 + y[i] = 0 +end + +function task_f(y) + var = initialize() + x = zeros(n,1) + r = vectorize(Normal(), rand(Normal())) + x[1,:] = update_var(var, 1, r) + report_transition!(var,0.0,0.0) + for i = 2:n + # Sampling + r = vectorize(Normal(),rand(Normal())) + x[i,:] = update_var(var, i, r) + logγ = logpdf(Normal(),x[i,1]) #γ(x_t|x_t-1) + logp = logpdf(Normal(),x[i,1]) # p(x_t|x_t-1) + report_transition!(var,logp,logγ) + #Proposal and Resampling + logpy = logpdf(Normal(x[i,1], 1.0), y[i-1]) + var = report_observation!(var,logpy) + end +end + + + + + +task = create_task(task_f, y) +model = PFModel(task) +tcontainer = Container(zeros(n,1),Vector{Bool}(undef,n),zeros(n),0) + + + +alg = AdvancedPS.SMCAlgorithm() +uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) +@btime sample(model, alg, uf, tcontainer, 100) + + +alg = AdvancedPS.PGAlgorithm(AdvancedPS.resample_systematic, 1.0, 10) +uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) +@btime chn2 =sample(model, alg, uf, tcontainer, 5) diff --git a/Example/Using_Custom_VI/test_interface.jl b/Example/Using_Custom_VI/test_interface.jl new file mode 100644 index 00000000..45e98465 --- /dev/null +++ b/Example/Using_Custom_VI/test_interface.jl @@ -0,0 +1,83 @@ +## Some function which make the model easier to define. + + +# A very light weight container for a state space model +# The state space is one dimensional +# Note that this is only for a very simple demonstration. + + + +# This is a very shallow container solely for testing puropose +mutable struct Container + x::Array{Float64,2} + marked::Vector{Bool} + produced_at::Vector{Int64} + num_produce::Float64 +end + +function Base.deepcopy(vi::Container) + return Container(deepcopy(vi.x),deepcopy(vi.marked),deepcopy(vi.produced_at),deepcopy(vi.num_produce)) +end + + +function set_retained_vns_del_by_spl!(container::Container) + for i in 1:length(container.marked) + if container.marked[i] + if container.produced_at[i] >= container.num_produce + container.marked[i] = false + end + end + end + container +end + + + +# This is important for initalizaiton +const initialize = current_trace + +function report_observation!(trace, logp::Float64) + produce(logp) + trace = current_trace() +end + +# logγ corresponds to the proposal distributoin we are sampling from. +function report_transition!(trace,logp::Float64,logγ::Float64) + trace.taskinfo.logp += logp - logγ + trace.taskinfo.logpseq += logp +end + +function update_var(trace, vn::Int64, r::Vector{Float64}) + if !trace.vi.marked[vn] + trace.vi.x[vn,:] = r + trace.vi.marked[vn] = true + trace.vi.produced_at[vn] = trace.vi.num_produce + return r + end + return trace.vi.x[vn,:] +end + + +# The reason for this is that we need to pass it! +function copy_container(vi::Container) + Container(deepcopy(vi.x),deepcopy(vi.marked),deepcopy(vi.produced_at),copy(vi.num_produce)) +end + +function create_task(f::Function, args...) + return CTask(() -> begin new_vi=f(args...); produce(Val{:done}); new_vi; end ) +end + +function tonamedtuple(vi::Container) + tnames = Tuple([Symbol("x$i") for i in 1:size(vi.x)[1]]) + tvalues = Tuple([vi.x[i,:] for i in 1:size(vi.x)[1]]) + return namedtuple(tnames, tvalues) +end +function Base.empty!(vi::Container) + for i in 1:length(vi.marked) + vi.marked[i] = false + end + vi +end +vectorize(d::UnivariateDistribution, r::Real) = [r] +vectorize(d::MultivariateDistribution, r::AbstractVector{<:Real}) = copy(r) +vectorize(d::MatrixDistribution, r::AbstractMatrix{<:Real}) = copy(vec(r)) diff --git a/Example/Using_Turing_VI/Turing_baseline.jl b/Example/Using_Turing_VI/Turing_baseline.jl new file mode 100644 index 00000000..1c9a8f5a --- /dev/null +++ b/Example/Using_Turing_VI/Turing_baseline.jl @@ -0,0 +1,25 @@ +using Turing +using BenchmarkTools +n = 500 + +y = Vector{Float64}(undef,n-1) +for i =1:n-1 + y[i] = 0 +end + +@model demo(y) = begin + x = TArray{Float64}(undef,n) + x[1] ~ Normal() + for i = 2:n + x[i] ~ Normal() + y[i-1] ~ Normal(x[i],1.0) + end +end + + + + +#@elapsed sample(demo(),PG(10),5) +chn = @btime sample(demo(y), SMC(), 100) + +chn.na.axes diff --git a/Example/Using_Turing_VI/demonstrate_with_Turing_VI.jl b/Example/Using_Turing_VI/demonstrate_with_Turing_VI.jl new file mode 100644 index 00000000..b17048da --- /dev/null +++ b/Example/Using_Turing_VI/demonstrate_with_Turing_VI.jl @@ -0,0 +1,75 @@ + +## It is not yet a package... +using Distributions +using AdvancedPS +using BenchmarkTools +using Libtask +include("AdvancedPS/Example/Using_Turing_VI/turing_interface.jl") +# Define a short model. +# The syntax is rather simple. Observations need to be reported with report_observation. +# Transitions must be reported using report_transition. +# The trace contains the variables which we want to infer using particle gibbs. +# Thats all! +n = 500 + +y = Vector{Float64}(undef,n-1) +for i =1:n-1 + y[i] = 0 +end + +function task_f(y) + var = initialize() + x = TArray{Float64}(undef,n) + vn = @varname x[1] + x[1] = update_var(var, vn, rand(Normal())) + report_transition!(var,0.0,0.0) + for i = 2:n + # Sampling + r = rand(Normal()) + vn = @varname x[i] + x[i] = update_var(var, vn, r) + logγ = logpdf(Normal(),x[i]) #γ(x_t|x_t-1) + logp = logγ # p(x_t|x_t-1) + report_transition!(var,logp,logγ) + #Proposal and Resampling + logpy = logpdf(Normal(x[i], 1.0), y[i-1]) + var = report_observation!(var,logpy) + end +end + + + + + + + +task = create_task(task_f, y) +model = PFModel(task) + + +################################################################# +# Get type stability!! # +################################################################# +alg = AdvancedPS.SMCAlgorithm() +uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) +untypedcontainer = VarInfo() +T = Trace{typeof(untypedcontainer),SMCTaskInfo{Float64}} +particles = T[ Trace(untypedcontainer, task, SMCTaskInfo(), uf.copy) for _ =1:1] +pc = ParticleContainer{typeof(particles[1])}(particles,zeros(1),0.0,0) +AdvancedPS.sample!(pc, alg, uf, nothing) +container = uf.empty!(TypedVarInfo(pc[1].vi)) +tcontainer = container + + + + +alg = AdvancedPS.SMCAlgorithm() +uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) +@btime sample(model, alg, uf, tcontainer, 100) + + + + +alg = AdvancedPS.PGAlgorithm(AdvancedPS.resample_systematic, 1.0, 10) +uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) +@elapsed chn2 =sample(model, alg, uf, tcontainer, 5) diff --git a/Example/Using_Turing_VI/turing_interface.jl b/Example/Using_Turing_VI/turing_interface.jl new file mode 100644 index 00000000..4ae8ef10 --- /dev/null +++ b/Example/Using_Turing_VI/turing_interface.jl @@ -0,0 +1,246 @@ + +using Turing.Core.RandomVariables +import Turing.Core: @varname +import Turing.Utilities: vectorize +using Turing +using AdvancedPS + +# This is important for initalizaiton +const initialize = AdvancedPS.current_trace +const TypedVarInfo = VarInfo{<:NamedTuple} +const Selector = Turing.Selector +const BASE_SELECTOR = Selector(:PS) + + + +function report_observation!(trace, logp::Float64) + produce(logp) + trace = AdvancedPS.current_trace() +end + +# logγ corresponds to the proposal distributoin we are sampling from. +function report_transition!(trace, logp::Float64, logγ::Float64) + trace.taskinfo.logp += logp - logγ + trace.taskinfo.logpseq += logp +end + + +# We obtain the new value for our variable +# If the distribution is not specified, we simply set it to be Normal +function update_var(trace, vn, val, dist= Normal()) + # check if the symbol is contained in the varinfo... + if haskey(trace.vi,vn) + if is_flagged(trace.vi, vn, "del") + unset_flag!(trace.vi, vn, "del") + trace.vi[vn] = vectorize(dist,val) + setgid!(trace.vi, BASE_SELECTOR, vn) + setorder!(trace.vi, vn, trace.vi.num_produce) + return val + else + updategid!(trace.vi, BASE_SELECTOR, vn) + val = trace.vi[vn] + end + else + #We do not specify the distribution... Thats why we set it to be Normal() + push!(trace.vi, vn, val, dist) + end + return val +end + +# The reason for this is that we need to pass it! +function create_task(f::Function, args...) + return CTask(() -> begin new_vi=f(args...); produce(Val{:done}); new_vi; end ) +end + + + +######################################################################### +# This is copied from turing compiler3.0, but we need to extract the # +# sampler form set_retained_vns_del_by_spl! # +######################################################################### + + +# Get all indices of variables belonging to SampleFromPrior: + +# if the gid/selector of a var is an empty Set, then that var is assumed to be assigned to + +# the SampleFromPrior sampler + +tonamedtuple(vi::TypedVarInfo) = Turing.tonamedtuple(vi) +tonamedtuple(vi::UntypedVarInfo) = tonamedtuple(TypedVarInfo(vi)) + + +""" + +`getidx(vi::UntypedVarInfo, vn::VarName)` + + + +Returns the index of `vn` in `vi.metadata.vns`. + +""" + +getidx(vi::UntypedVarInfo, vn::VarName) = vi.idcs[vn] + + +""" + +`getidx(vi::TypedVarInfo, vn::VarName{sym})` + + + +Returns the index of `vn` in `getfield(vi.metadata, sym).vns`. + +""" + +function getidx(vi::TypedVarInfo, vn::VarName{sym}) where sym + + getfield(vi.metadata, sym).idcs[vn] + +end + + + +""" + +`setgid!(vi::VarInfo, gid::Selector, vn::VarName)` + + +Adds `gid` to the set of sampler selectors associated with `vn` in `vi`. + +""" + +setgid!(vi::UntypedVarInfo, gid::Selector, vn::VarName) = push!(vi.gids[getidx(vi, vn)], gid) + +function setgid!(vi::TypedVarInfo, gid::Selector, vn::VarName{sym}) where sym + + push!(getfield(vi.metadata, sym).gids[getidx(vi, vn)], gid) + +end + +""" + +`updategid!(vi::VarInfo, vn::VarName, spl::Sampler)` + + + +If `vn` doesn't have a sampler selector linked and `vn`'s symbol is in the space of + +`spl`, this function will set `vn`'s `gid` to `Set([spl.selector])`. + +""" + +function updategid!(vi::AbstractVarInfo, sel::Selector, vn::VarName) + setgid!(vi, sel, vn) +end + + +@generated function _getidcs(metadata::NamedTuple{names}) where {names} + exprs = [] + for f in names + push!(exprs, :($f = findinds(metadata.$f))) + end + length(exprs) == 0 && return :(NamedTuple()) + return :($(exprs...),) +end + +# Get all indices of variables belonging to a given sampler + + +@inline function _getidcs(vi::UntypedVarInfo, s::Selector) + findinds(vi, s) +end + +@inline function _getidcs(vi::TypedVarInfo, s::Selector) + + return _getidcs(vi.metadata, s) +end + +# Get a NamedTuple for all the indices belonging to a given selector for each symbol + +@generated function _getidcs(metadata::NamedTuple{names}, s::Selector) where {names} + exprs = [] + # Iterate through each varname in metadata. + for f in names + # If the varname is in the sampler space + # or the sample space is empty (all variables) + # then return the indices for that variable. + push!(exprs, :($f = findinds(metadata.$f, s))) + end + length(exprs) == 0 && return :(NamedTuple()) + return :($(exprs...),) +end + +@inline function findinds(f_meta, s::Selector) + + # Get all the idcs of the vns in `space` and that belong to the selector `s` + return filter((i) -> + (s in f_meta.gids[i] || isempty(f_meta.gids[i])), 1:length(f_meta.gids)) +end + +@inline function findinds(f_meta) + # Get all the idcs of the vns + return filter((i) -> isempty(f_meta.gids[i]), 1:length(f_meta.gids)) +end + +""" + +`set_retained_vns_del_by_spl!(vi::VarInfo, spl::Sampler)` + + + +Sets the `"del"` flag of variables in `vi` with `order > vi.num_produce` to `true`. + +""" + +function set_retained_vns_del_by_spl!(vi::AbstractVarInfo) + return set_retained_vns_del_by_spl!(vi, BASE_SELECTOR) +end + + +function set_retained_vns_del_by_spl!(vi::UntypedVarInfo, sel::Selector) + # Get the indices of `vns` that belong to `spl` as a vector + gidcs = _getidcs(vi, sel) + if vi.num_produce == 0 + for i = length(gidcs):-1:1 + vi.flags["del"][gidcs[i]] = true + end + else + for i in 1:length(vi.orders) + if i in gidcs && vi.orders[i] > vi.num_produce + vi.flags["del"][i] = true + end + end + end + return nothing +end + +function set_retained_vns_del_by_spl!(vi::TypedVarInfo, sel::Selector) + # Get the indices of `vns` that belong to `spl` as a NamedTuple, one entry for each symbol + gidcs = _getidcs(vi, sel) + return _set_retained_vns_del_by_spl!(vi.metadata, gidcs, vi.num_produce) +end + +@generated function _set_retained_vns_del_by_spl!(metadata, gidcs::NamedTuple{names}, num_produce) where {names} + expr = Expr(:block) + for f in names + f_gidcs = :(gidcs.$f) + f_orders = :(metadata.$f.orders) + f_flags = :(metadata.$f.flags) + push!(expr.args, quote + # Set the flag for variables with symbol `f` + if num_produce == 0 + for i = length($f_gidcs):-1:1 + $f_flags["del"][$f_gidcs[i]] = true + end + else + for i in 1:length($f_orders) + if i in $f_gidcs && $f_orders[i] > num_produce + $f_flags["del"][i] = true + end + end + end + end) + end + return expr +end diff --git a/Project.toml b/Project.toml new file mode 100644 index 00000000..4342b747 --- /dev/null +++ b/Project.toml @@ -0,0 +1,3 @@ +name = "AdvancedPS" +authors = ["kongi "] +version = "0.1.0" diff --git a/Tests/Using_Turing_VI/numerical_tests.jl b/Tests/Using_Turing_VI/numerical_tests.jl new file mode 100644 index 00000000..abb5f878 --- /dev/null +++ b/Tests/Using_Turing_VI/numerical_tests.jl @@ -0,0 +1,159 @@ +using Random +using Test +using Distributions +using Turing +using AdvancedPS + +dir = splitdir(splitdir(pathof(AdvancedPS))[1])[1] + +include(dir*"/Example/Using_Turing_VI/turing_interface.jl") +include(dir*"/Tests/test_utils/AllUtils.jl") + +import Turing.Core: tonamedtuple +tonamedtuple(vi::UntypedVarInfo) = tonamedtuple(TypedVarInfo(vi)) + + +@testset "apf.jl" begin + @apf_testset "apf constructor" begin + N = 200 + f = aps_gdemo_default + task = create_task(f) + model = PFModel(task) + tcontainer = VarInfo() + + alg = AdvancedPS.SMCAlgorithm() + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + + alg = AdvancedPS.PGAlgorithm(5) + uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + + end + @numerical_testset "apf inference" begin + N = 2000 + f = aps_gdemo_default + task = create_task(f) + model = PFModel(task) + tcontainer = VarInfo() + + alg = AdvancedPS.SMCAlgorithm() + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + check_gdemo(caps1, atol = 0.1) + + alg = AdvancedPS.PGAlgorithm(5) + uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + check_gdemo(caps1, atol = 0.1) + end +end + + + + + +@testset "test_against_turing.jl" begin + + @numerical_testset "apf inference" begin + N = 5000 + + + ## Testset1 + y = [0 for i in 1:10] + + # We use the Turing model as baseline + chn_base = sample(large_demo(y),PG(20),N) + + ############################################# + # Using a Proposal # + ############################################# + f = large_demo_apf_proposal + task = create_task(f,y) + model = PFModel(task) + tcontainer = VarInfo() + + ##SMC + alg = AdvancedPS.SMCAlgorithm() + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, 5*N) + check_numerical(caps1,chn_base, 0.1, 0.2) + + # PG + alg = AdvancedPS.PGAlgorithm(20) + uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + check_numerical(caps1,chn_base,0.1, 0.2) + + + + ############################################# + # No Proposal # + ############################################# + f = large_demo_apf + task = create_task(f,y) + model = PFModel(task) + tcontainer = VarInfo() + + #SMC + alg = AdvancedPS.SMCAlgorithm() + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, 5*N) + check_numerical(caps1,chn_base, 0.1, 0.2) + + # PG + alg = AdvancedPS.PGAlgorithm(20) + uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + check_numerical(caps1,chn_base,0.1, 0.2) + + + + ## Testset2 + y = [-0.1*i for i in 1:10] + + # We use the Turing model as baseline + chn_base = sample(large_demo(y),PG(20),N) + + ############################################# + # Using a Proposal # + ############################################# + f = large_demo_apf_proposal + task = create_task(f,y) + model = PFModel(task) + tcontainer = VarInfo() + + #SMC + alg = AdvancedPS.SMCAlgorithm() + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, 5*N) + check_numerical(caps1,chn_base,0.1, 0.2) + + # PG + alg = AdvancedPS.PGAlgorithm(20) + uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + check_numerical(caps1,chn_base, 0.1, 0.2) + + + ############################################# + # No Proposal # + ############################################# + f = large_demo_apf + task = create_task(f,y) + model = PFModel(task) + tcontainer = VarInfo() + + #SMC + alg = AdvancedPS.SMCAlgorithm() + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, 5*N) + check_numerical(caps1,chn_base,0.21, 0.2) + + # PG + alg = AdvancedPS.PGAlgorithm(20) + uf = AdvancedPS.PGUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + caps1 = sample(model, alg, uf, tcontainer, N) + check_numerical(caps1,chn_base, 0.1, 0.2) + end +end diff --git a/Tests/test_resample.jl b/Tests/test_resample.jl new file mode 100644 index 00000000..54a000e7 --- /dev/null +++ b/Tests/test_resample.jl @@ -0,0 +1,19 @@ + + + +@apf_testset "resample.jl" begin + D = [0.3, 0.4, 0.3] + num_samples = Int(1e6) + resSystematic = AdvancedPS.resample_systematic(D, num_samples ) + resStratified = AdvancedPS.resample_stratified(D, num_samples ) + resMultinomial= AdvancedPS.resample_multinomial(D, num_samples ) + resResidual = AdvancedPS.resample_residual(D, num_samples ) + AdvancedPS.resample(D) + resSystematic2=AdvancedPS.resample(D, num_samples ) + + @test sum(resSystematic .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples + @test sum(resSystematic2 .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples + @test sum(resStratified .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples + @test sum(resMultinomial .== 2) ≈ (num_samples * 0.4) atol=1e-2*num_samples + @test sum(resResidual .== 2) ≈ (num_samples * 0.4) atol=1e-2*num_samples +end diff --git a/Tests/test_utils/AllUtils.jl b/Tests/test_utils/AllUtils.jl new file mode 100644 index 00000000..a850d6c2 --- /dev/null +++ b/Tests/test_utils/AllUtils.jl @@ -0,0 +1,9 @@ +using Turing, Random, Test + +# Import utility functions and reused models. +include("staging.jl") +include("numerical_tests.jl") +include("ad_utils.jl") +include("models.jl") +include("random_measure_utils.jl") +include("testing_functions.jl") diff --git a/Tests/test_utils/ad_utils.jl b/Tests/test_utils/ad_utils.jl new file mode 100644 index 00000000..1e5798d3 --- /dev/null +++ b/Tests/test_utils/ad_utils.jl @@ -0,0 +1,77 @@ +using Turing: gradient_logp_forward, gradient_logp_reverse +using Test + +function test_ad(f, at = 0.5; rtol = 1e-6, atol = 1e-6) + isarr = isa(at, AbstractArray) + reverse = Tracker.data(Tracker.gradient(f, at)[1]) + if isarr + forward = ForwardDiff.gradient(f, at) + @test isapprox(reverse, forward, rtol=rtol, atol=atol) + else + forward = ForwardDiff.derivative(f, at) + finite_diff = central_fdm(5,1)(f, at) + @test isapprox(reverse, forward, rtol=rtol, atol=atol) + @test isapprox(reverse, finite_diff, rtol=rtol, atol=atol) + end +end + +""" + test_reverse_mode_ad(forward, f, ȳ, x...; rtol=1e-6, atol=1e-6) + +Check that the reverse-mode sensitivities produced by an AD library are correct for `f` +at `x...`, given sensitivity `ȳ` w.r.t. `y = f(x...)` up to `rtol` and `atol`. +`forward` should be either `Tracker.forward` or `Zygote.forward`. +""" +function test_reverse_mode_ad(forward, f, ȳ, x...; rtol=1e-6, atol=1e-6) + + # Perform a regular forwards-pass. + y = f(x...) + + # Use tracker to compute reverse-mode sensitivities. + y_tracker, back = forward(f, x...) + x̄s_tracker = back(ȳ) + + # Use finite differencing to compute reverse-mode sensitivities. + x̄s_fdm = FDM.j′vp(central_fdm(5, 1), f, ȳ, x...) + + # Check that forwards-pass produces the correct answer. + @test isapprox(y, Tracker.data(y_tracker), atol=atol, rtol=rtol) + + # Check that reverse-mode sensitivities are correct. + @test all(zip(x̄s_tracker, x̄s_fdm)) do (x̄_tracker, x̄_fdm) + isapprox(Tracker.data(x̄_tracker), x̄_fdm; atol=atol, rtol=rtol) + end +end + +# See `test_reverse_mode_ad` for details. +function test_tracker_ad(f, ȳ, x...; rtol=1e-6, atol=1e-6) + return test_reverse_mode_ad(Tracker.forward, f, ȳ, x...; rtol=rtol, atol=atol) +end + +function test_model_ad(model, f, syms::Vector{Symbol}) + # Set up VI. + vi = Turing.VarInfo(model) + + # Collect symbols. + vnms = Vector(undef, length(syms)) + vnvals = Vector{Float64}() + for i in 1:length(syms) + s = syms[i] + vnms[i] = getfield(vi.metadata, s).vns[1] + + vals = getval(vi, vnms[i]) + for i in eachindex(vals) + push!(vnvals, vals[i]) + end + end + + spl = SampleFromPrior() + _, ∇E = gradient_logp_forward(vi[spl], vi, model) + grad_Turing = sort(∇E) + + # Call ForwardDiff's AD + grad_FWAD = sort(ForwardDiff.gradient(f, vec(vnvals))) + + # Compare result + @test grad_Turing ≈ grad_FWAD atol=1e-9 +end diff --git a/Tests/test_utils/models.jl b/Tests/test_utils/models.jl new file mode 100644 index 00000000..0327f1ea --- /dev/null +++ b/Tests/test_utils/models.jl @@ -0,0 +1,82 @@ + +@model gdemo_d() = begin + s ~ InverseGamma(2, 3) + m ~ Normal(0, sqrt(s)) + 1.5 ~ Normal(m, sqrt(s)) + 2.0 ~ Normal(m, sqrt(s)) + return s, m +end + +gdemo_default = gdemo_d() + +function gdemo_d_apf() + var = initialize() + r = rand(InverseGamma(2, 3)) + vn = @varname s + s = update_var(var, vn, r) + + r = rand(Normal(0, sqrt(s))) + vn = @varname m + m = update_var(var, vn, r) + + logp = logpdf(Normal(m, sqrt(s)), 1.5) + report_observation!(var,logp) + + logp = logpdf(Normal(m, sqrt(s)), 2.0) + report_observation!(var,logp) +end + + +aps_gdemo_default = gdemo_d_apf + +@model large_demo(y) = begin + x = TArray{Float64}(undef,11) + x[1] ~ Normal() + for i = 2:11 + x[i] ~ Normal(0.1*x[i-1]-0.1,0.5) + y[i-1] ~ Normal(x[i],0.3) + end +end + +function large_demo_apf(y) + var = initialize() + x = TArray{Float64}(undef,11) + vn = @varname x[1] + x[1] = update_var(var, vn, rand(Normal())) + # there is nothing to report since we are not using proposal sampling + report_transition!(var,0.0,0.0) + for i = 2:11 + # Sampling + r = rand( Normal(0.1*x[i-1]-0.1,0.5)) + vn = @varname x[i] + x[i] = update_var(var, vn, r) + logγ = logpdf( Normal(0.1*x[i-1]-0.1,0.5),x[i]) #γ(x_t|x_t-1) + logp = logγ # p(x_t|x_t-1) + report_transition!(var,logp,logγ) + #Proposal and Resampling + logpy = logpdf(Normal(x[i], 0.3), y[i-1]) + var = report_observation!(var,logpy) + end +end + +## Here, we even have a proposal distributin +function large_demo_apf_proposal(y) + var = initialize() + x = TArray{Float64}(undef,11) + vn = @varname x[1] + x[1] = update_var(var, vn, rand(Normal())) + # there is nothing to report since we are not using proposal sampling + report_transition!(var,0.0,0.0) + for i = 2:11 + # Sampling + r = rand(Normal(0.1*x[i-1],0.5)) + vn = @varname x[i] + x[i] = update_var(var, vn, r) + logγ = logpdf(Normal(0.1*x[i-1],0.5),x[i]) #γ(x_t|x_t-1) + logp = logpdf( Normal(0.1*x[i-1]-0.1,0.5),x[i]) # p(x_t|x_t-1) + report_transition!(var,logp,logγ) + #Proposal and Resampling + logpy = logpdf(Normal(x[i], 0.3), y[i-1]) + var = report_observation!(var,logpy) + end +end diff --git a/Tests/test_utils/numerical_tests.jl b/Tests/test_utils/numerical_tests.jl new file mode 100644 index 00000000..3e981fbd --- /dev/null +++ b/Tests/test_utils/numerical_tests.jl @@ -0,0 +1,85 @@ +function check_dist_numerical(dist, chn; mean_tol = 0.1, var_atol = 1.0, var_tol = 0.5) + @testset "numerical" begin + # Extract values. + chn_xs = chn[1:2:end, :x, :].value + + # Check means. + dist_mean = mean(dist) + mean_shape = size(dist_mean) + if !all(isnan, dist_mean) && !all(isinf, dist_mean) + chn_mean = Array(mean(chn_xs, dims=1)) + chn_mean = length(chn_mean) == 1 ? + chn_mean[1] : + reshape(chn_mean, mean_shape) + atol_m = length(chn_mean) > 1 ? + mean_tol * length(chn_mean) : + max(mean_tol, mean_tol * chn_mean) + @test chn_mean ≈ dist_mean atol=atol_m + end + + # Check variances. + # var() for Distributions.MatrixDistribution is not defined + if !(dist isa Distributions.MatrixDistribution) + # Variance + dist_var = var(dist) + var_shape = size(dist_var) + if !all(isnan, dist_var) && !all(isinf, dist_var) + chn_var = Array(var(chn_xs, dims=1)) + chn_var = length(chn_var) == 1 ? + chn_var[1] : + reshape(chn_var, var_shape) + atol_v = length(chn_mean) > 1 ? + mean_tol * length(chn_mean) : + max(mean_tol, mean_tol * chn_mean) + @test chn_mean ≈ dist_mean atol=atol_v + end + end + end +end + +# Helper function for numerical tests +function check_numerical(chain, + symbols::Vector, + exact_vals::Vector; + atol=0.2, + rtol=0.0) + for (sym, val) in zip(symbols, exact_vals) + E = val isa Real ? + mean(chain[sym].value) : + vec(mean(chain[sym].value, dims=[1])) + @info (symbol=sym, exact=val, evaluated=E) + @test E ≈ val atol=atol rtol=rtol + end +end + +# Compare two chains +function check_numerical(chain1, + chain2, + atol=0.2, + rtol=0.0) + for name in zip(chain2.name_map.parameters) + sym = Symbol(name[1]) + val = chain2[sym].value[1] isa Real ? + mean(chain2[sym].value) : + vec(mean(chain2[sym].value, dims=[1])) + E = val isa Real ? + mean(chain1[sym].value) : + vec(mean(chain1[sym].value, dims=[1])) + @info (symbol=sym, exact=val, evaluated=E) + @test E ≈ val atol=atol rtol=rtol + end +end + + + +# Wrapper function to quickly check gdemo accuracy. +function check_gdemo(chain; atol=0.2, rtol=0.0) + check_numerical(chain, [:s, :m], [49/24, 7/6], atol=atol, rtol=rtol) +end +# Wrapper function to check MoGtest. +function check_MoGtest_default(chain; atol=0.2, rtol=0.0) + check_numerical(chain, + [:z1, :z2, :z3, :z4, :mu1, :mu2], + [1.0, 1.0, 2.0, 2.0, 1.0, 4.0], + atol=atol, rtol=rtol) +end diff --git a/Tests/test_utils/random_measure_utils.jl b/Tests/test_utils/random_measure_utils.jl new file mode 100644 index 00000000..868bbd31 --- /dev/null +++ b/Tests/test_utils/random_measure_utils.jl @@ -0,0 +1,36 @@ +using SpecialFunctions + +function compute_log_joint(observations, partition, tau0, tau1, sigma, theta) + n = length(observations) + k = length(partition) + prob = k*log(sigma) + lgamma(theta) + lgamma(theta/sigma + k) - lgamma(theta/sigma) - lgamma(theta + n) + for cluster in partition + prob += lgamma(length(cluster) - sigma) - lgamma(1 - sigma) + prob += compute_log_conditonal_observations(observations, cluster, tau0, tau1) + end + prob +end + +function compute_log_conditonal_observations(observations, cluster, tau0, tau1) + nl = length(cluster) + prob = (nl/2)*log(tau1) - (nl/2)*log(2*pi) + 0.5*log(tau0) + 0.5*log(tau0+nl) + prob += -tau1/2*(sum(observations)) + 0.5*(tau0*mu_0+tau1*sum(observations[cluster]))^2/(tau0+nl*tau1) + prob +end + +# Test of similarity between distributions +function correct_posterior(empirical_probs, data, partitions, τ0, τ1, σ, θ) + true_log_probs = map(p -> compute_log_joint(data, p, τ0, τ1, σ, θ), partitions) + true_probs = exp.(true_log_probs) + true_probs /= sum(true_probs) + + empirical_probs /= sum(empirical_probs) + + # compare distribitions + # L2 + L2 = sum((empirical_probs - true_probs).^2) + + # Discrepancy + discr = maximum(abs.(empirical_probs - true_probs)) + return L2, discr +end diff --git a/Tests/test_utils/staging.jl b/Tests/test_utils/staging.jl new file mode 100644 index 00000000..68b0b29a --- /dev/null +++ b/Tests/test_utils/staging.jl @@ -0,0 +1,52 @@ +function get_stage() + # Appveyor uses "True" for non-Ubuntu images. + if get(ENV, "APPVEYOR", "") == "True" || get(ENV, "APPVEYOR", "") == "true" + return "nonnumeric" + end + + # Handle Travis and Github Actions specially. + if get(ENV, "TRAVIS", "") == "true" || get(ENV, "GITHUB_ACTIONS", "") == "true" + if "STAGE" in keys(ENV) + return ENV["STAGE"] + else + return "all" + end + end + + return "all" +end + +function do_test(stage_str) + stg = get_stage() + + # If the tests are being run by Appveyor, don't run + # any numerical tests. + if stg == "nonnumeric" + if stage_str == "numerical" + return false + else + return true + end + end + + # Otherwise run the regular testing procedure. + if stg == "all" || stg == stage_str + return true + end + + return false +end + +macro stage_testset(stage_string::String, args...) + if do_test(stage_string) + return esc(:(@testset($(args...)))) + end +end + +macro numerical_testset(args...) + esc(:(@stage_testset "numerical" $(args...))) +end + +macro apf_testset(args...) + esc(:(@stage_testset "test" $(args...))) +end diff --git a/Tests/test_utils/testing_functions.jl b/Tests/test_utils/testing_functions.jl new file mode 100644 index 00000000..809a62a6 --- /dev/null +++ b/Tests/test_utils/testing_functions.jl @@ -0,0 +1,73 @@ +using Turing + +GKernel(var) = (x) -> Normal(x, sqrt.(var)) + +varname(s::Symbol) = nothing, s +function varname(expr::Expr) + # Initialize an expression block to store the code for creating uid + local sym + @assert expr.head == :ref "expr needs to be an indexing expression, e.g. :(x[1])" + indexing_ex = Expr(:block) + # Add the initialization statement for uid + push!(indexing_ex.args, quote indexing_list = [] end) + # Initialize a local container for parsing and add the expr to it + to_eval = []; pushfirst!(to_eval, expr) + # Parse the expression and creating the code for creating uid + find_head = false + while length(to_eval) > 0 + evaling = popfirst!(to_eval) # get the current expression to deal with + if isa(evaling, Expr) && evaling.head == :ref && ~find_head + # Add all the indexing arguments to the left + pushfirst!(to_eval, "[", insdelim(evaling.args[2:end])..., "]") + # Add first argument depending on its type + # If it is an expression, it means it's a nested array calling + # Otherwise it's the symbol for the calling + if isa(evaling.args[1], Expr) + pushfirst!(to_eval, evaling.args[1]) + else + # push!(indexing_ex.args, quote pushfirst!(indexing_list, $(string(evaling.args[1]))) end) + push!(indexing_ex.args, quote sym = Symbol($(string(evaling.args[1]))) end) # store symbol in runtime + find_head = true + sym = evaling.args[1] # store symbol in compilation time + end + else + # Evaluting the concrete value of the indexing variable + push!(indexing_ex.args, quote push!(indexing_list, string($evaling)) end) + end + end + push!(indexing_ex.args, quote indexing = reduce(*, indexing_list) end) + return indexing_ex, sym +end + +genvn(sym::Symbol) = VarName(gensym(), sym, "", 1) +function genvn(expr::Expr) + ex, sym = varname(expr) + VarName(gensym(), sym, eval(ex), 1) +end + +function randr(vi::Turing.VarInfo, + vn::Turing.VarName, + dist::Distribution, + spl::Turing.Sampler, + count::Bool = false) + if ~haskey(vi, vn) + r = rand(dist) + Turing.push!(vi, vn, r, dist, spl) + return r + elseif is_flagged(vi, vn, "del") + unset_flag!(vi, vn, "del") + r = rand(dist) + Turing.RandomVariables.setval!(vi, Turing.vectorize(dist, r), vn) + return r + else + if count Turing.checkindex(vn, vi, spl) end + Turing.updategid!(vi, vn, spl) + return vi[vn] + end +end + + + +function insdelim(c, deli=",") + return reduce((e, res) -> append!(e, [res, deli]), c; init = [])[1:end-1] +end diff --git a/Tests/tests_container.jl b/Tests/tests_container.jl new file mode 100644 index 00000000..d91b3662 --- /dev/null +++ b/Tests/tests_container.jl @@ -0,0 +1,117 @@ +using Test +using AdvancedPS + + +using Random +using Test +using Distributions +using Turing +using AdvancedPS + +dir = splitdir(splitdir(pathof(AdvancedPS))[1])[1] + +include(dir*"/Example/Using_Turing_VI/turing_interface.jl") +include(dir*"/Tests/test_utils/AllUtils.jl") +import Turing.Core: tonamedtuple +tonamedtuple(vi::UntypedVarInfo) = tonamedtuple(TypedVarInfo(vi)) + + + + + + + + + +@testset "Particle.jl" begin + @apf_testset "copy particle container" begin + pc = ParticleContainer(Trace[], Float64[], 0.0, 0) + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, () -> (), () -> (), () ->()) + newpc = copy(pc, uf) + @test newpc.logE == pc.logE + @test newpc.logWs == pc.logWs + @test newpc.n_consume == pc.n_consume + @test typeof(pc) === typeof(newpc) + end + + @apf_testset "particle container" begin + n = Ref(0) + alg = AdvancedPS.PGAlgorithm(5) + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + + dist = Normal(0, 1) + + function fpc() + t = TArray(Float64, 1); + var = initialize() + t[1] = 0; + while true + vn = @varname x[n] + r = rand(dist) + update_var(var, vn, r) + n[] += 1 + produce(0) + var = current_trace() + r = rand(dist) + update_var(var, vn, r) + t[1] = 1 + t[1] + end + end + + task = create_task(fpc) + model = PFModel(task) + particles = [Trace(Turing.VarInfo(), fpc, PGTaskInfo(0.0, 0.0), deepcopy) for _ in 1:3] + pc = ParticleContainer(particles) + + @test weights(pc) == [1/3, 1/3, 1/3] + @test logZ(pc) ≈ log(3) + @test pc.logE ≈ log(1) + @test consume(pc) == log(1) + + Ws = weights(pc) + indx = alg.resampler(Ws, length(pc)) + resample!(pc, uf, indx, nothing) + @test weights(pc) == [1/3, 1/3, 1/3] + @test logZ(pc) ≈ log(3) + @test pc.logE ≈ log(1) + @test AdvancedPS.effectiveSampleSize(pc) == 3 + @test consume(pc) ≈ log(1) + Ws = weights(pc) + indx = alg.resampler(Ws, length(pc)) + resample!(pc, uf, indx, nothing) + @test consume(pc) ≈ log(1) + end + + + @apf_testset "trace" begin + n = Ref(0) + alg = AdvancedPS.PGAlgorithm(5) + uf = AdvancedPS.SMCUtilityFunctions(deepcopy, set_retained_vns_del_by_spl!, empty!, tonamedtuple) + dist = Normal(0, 1) + function f2() + t = TArray(Float64, 1); + var = initialize() + t[1] = 0; + while true + vn = @varname x[n] + r = rand(dist) + update_var(var, vn, r) + produce(t[1]); + var = current_trace() + vn = @varname x[n] + r = rand(dist) + update_var(var, vn, r) + t[1] = 1 + t[1] + end + end + task = create_task(fpc) + model = PFModel(task) + tr = Trace(Turing.VarInfo(), f2, PGTaskInfo(0.0, 0.0), deepcopy) + consume(tr); consume(tr) + a = AdvancedPS.fork(tr, deepcopy); + consume(a); consume(a) + @test consume(tr) == 2 + @test consume(a) == 4 + + end +end diff --git a/src/AdvancedPS.jl b/src/AdvancedPS.jl new file mode 100644 index 00000000..0d558a7d --- /dev/null +++ b/src/AdvancedPS.jl @@ -0,0 +1,89 @@ +### Attention this is a development package! It wount run. + + +module AdvancedPS + + + using Libtask + using StatsFuns: logsumexp, softmax! + using AbstractMCMC + import MCMCChains: Chains + import Base.copy + using Distributions + + abstract type AbstractTaskInfo end + abstract type AbstractParticleContainer end + abstract type AbstractTrace end + abstract type AbstractPFAlgorithm end + abstract type AbstractPFUtilitFunctions end + + abstract type AbstractPFTransition <: AbstractTransition end + abstract type AbstractPFSampler <: AbstractSampler end + + abstract type AbstractSMCUtilitFunctions <: AbstractPFUtilitFunctions end + abstract type AbstractPGASUtilityFunctions <: AbstractSMCUtilitFunctions end + + abstract type AbstractPFModel <: AbstractModel end + + export AbstractTaskInfo, + AbstractParticleContainer, + AbstractTrace, + AbstractPFAlgorithm, + AbstractPFUtilitFunctions, + AbstractPFTransition, + AbstractPFSampler, + AbstractSMCUtilitFunctions, + AbstractPGASUtilityFunctions, + AbstractPFModel + + include("Core/Algorithms/Algorithms.jl") + include("Core/Container/Trace.jl") + include("Core/Container/ParticleContainer.jl") + include("Core/Resample/resample.jl") + include("Core/Algorithms/Taskinfo.jl") + include("Core/Algorithms/sample.jl") + include("Core/Resample/resample_functions.jl") + include("Core/Utilities/UtilityFunctions.jl") + + export ParticleContainer, + Trace, + weights, + logZ, + current_trace, + extend!, + empty!, + resample!, + PGTaskInfo, + SMCTaskInfo, + PGASTaskInfo + sample!, + resample, + randcat, + resample_multinomial, + resample_residual, + resample_stratified, + resample_systematic, + SMCUtilityFunctions, + PGUtilityFunctions, + PGASUtilityFunctions, + SMCAlgorithm, + PGAlgorithm, + forkr + + include("Inference/Model.jl") + include("Inference/Transitions.jl") + include("Inference/Sampler.jl") + include("Inference/sample_init.jl") + include("Inference/step.jl") + include("Inference/Inference.jl") + + export PFModel, + sample, + bundle_samples, + sample_init!, + step!, + Sampler, + PFTransition, + transition_type + +end # module diff --git a/src/Core/Algorithms/Algorithms.jl b/src/Core/Algorithms/Algorithms.jl new file mode 100644 index 00000000..bc4d67d0 --- /dev/null +++ b/src/Core/Algorithms/Algorithms.jl @@ -0,0 +1,33 @@ + +struct SMCAlgorithm{RT} <: AbstractPFAlgorithm where RT<:AbstractFloat + resampler :: Function + resampler_threshold :: RT + +end + +function SMCAlgorithm() + SMCAlgorithm(resample_systematic, 0.5) +end + +struct PGAlgorithm{RT} <: AbstractPFAlgorithm where RT<:AbstractFloat + resampler :: Function + resampler_threshold :: RT + n :: Int64 + +end + +function PGAlgorithm(n::Int64) + PGAlgorithm(resample_systematic, 0.5, n) +end + + +struct PGASAlgorithm{RT} <: AbstractPFAlgorithm + resampler :: Function + resampler_threshold :: RT + merge_traj :: Function + n :: Int64 +end + +function PGASAlgorithm(merge_traj::Function, n::Int64) + SMCAlgorithm(resample_systematic, 0.5, merge_traj, n) +end diff --git a/src/Core/Algorithms/Taskinfo.jl b/src/Core/Algorithms/Taskinfo.jl new file mode 100644 index 00000000..e35c9a10 --- /dev/null +++ b/src/Core/Algorithms/Taskinfo.jl @@ -0,0 +1,43 @@ + + +# The idea behind the TaskInfo struct is that it +# allows to propagate information trough the computation. +# This is important because we do not want to + +mutable struct SMCTaskInfo{T} <: AbstractTaskInfo where {T <:AbstractFloat} + # This corresponds to p(yₜ | xₜ) p(xₜ | xₜ₋₁) / γ(xₜ | xₜ₋₁, yₜ) + # where γ is the porposal. + # We need this variable to compute the weights! + logp::T + # This corresponds to p(xₜ | xₜ₋₁) p(xₜ₋₁ | xₜ₋₂) ⋯ p(x₀) + # or |x_{0:t-1} for non markovian models, we need this to compute + # the ancestor weights. + logpseq::T +end + +SMCTaskInfo() = SMCTaskInfo(0.0, 0.0) + +const PGTaskInfo = SMCTaskInfo + + + +mutable struct PGASTaskInfo{T, F} <: AbstractTaskInfo where {T <: AbstractFloat} + # Same as above + logp::T + logpseq::T + # The ancestor weight + ancestor_weight::Float64 +end +function PGASTaskInfo(logp::Float64,logpseq::Float64) + PGASTaskInfo(logp, logpseq, 0.0) +end + +function Base.copy(info::PGTaskInfo) + PGTaskInfo(info.logp, info.logpseq) +end +function Base.copy(info::PGASTaskInfo) + PGASTaskInfo(info.logp, info.logpseq, info.ancestor_weight) +end + +reset_logp!(ti::AbstractTaskInfo) = (ti.logp = 0.0) +set_ancestor_weight!(ti::PGASTaskInfo, w::Float64) = (ti.ancestor_weight = w) diff --git a/src/Core/Algorithms/sample.jl b/src/Core/Algorithms/sample.jl new file mode 100644 index 00000000..14e39452 --- /dev/null +++ b/src/Core/Algorithms/sample.jl @@ -0,0 +1,29 @@ + + +## This is all we need for Turing! + +function sample!(pc::PC, alg::ALG, utility_functions::AbstractSMCUtilitFunctions, ref_traj::T) where { + PC <:ParticleContainer, + ALG <:Union{SMCAlgorithm, PGAlgorithm}, + T <:Union{Trace,Nothing} +} + n = length(pc.vals) + while consume(pc) != Val{:done} + ess = effectiveSampleSize(pc) + if ref_traj !== nothing || ess <= alg.resampler_threshold * length(pc) + # compute weights + Ws = weights(pc) + + # check that weights are not NaN + @assert !any(isnan, Ws) + # sample ancestor indices + # Ancestor trajectory is not sampled + ref_traj !== nothing ? nresamples = n-1 : nresamples = n + indx = alg.resampler(Ws, nresamples) + # We add ancestor trajectory to the path. + # For ancestor sampling, we would change n at this point. + ref_traj !== nothing ? push!(indx,n) : nothing + resample!(pc, utility_functions, indx, ref_traj) + end + end +end diff --git a/src/Core/Container/ParticleContainer.jl b/src/Core/Container/ParticleContainer.jl new file mode 100644 index 00000000..dcfda86e --- /dev/null +++ b/src/Core/Container/ParticleContainer.jl @@ -0,0 +1,105 @@ + +""" +Data structure for particle filters +- effectiveSampleSize(pc :: ParticleContainer) +- normalise!(pc::ParticleContainer) +- consume(pc::ParticleContainer): return incremental likelihood +""" +mutable struct ParticleContainer{T} <: AbstractParticleContainer where T<:Trace + vals::Vector{T} + # A named tuple with functions to manipulate the particles vi. + logWs::Vector{Float64} + # log model evidence + logE::Float64 + # helpful for rejuvenation steps, e.g. in SMC2 + n_consume::Int +end + + +ParticleContainer(particles::Vector{<:Trace}) = + ParticleContainer(particles, zeros(length(particles)), 0.0, 0) + + +Base.collect(pc :: ParticleContainer) = pc.vals # prev: Dict, now: Array +Base.length(pc :: ParticleContainer) = length(pc.vals) +# pc[i] returns the i'th particle +Base.@propagate_inbounds Base.getindex(pc::ParticleContainer, i::Int) = pc.vals[i] + + + +# registers a new x-particle in the container +function Base.push!(pc::ParticleContainer, p::Particle) + push!(pc.vals, p) + push!(pc.logWs, 0.0) + return pc +end + +# clears the container but keep params, logweight etc. +function Base.empty!(pc::ParticleContainer) + pc.vals = eltype(pc.vals)[] + pc.logWs = Float64[] + return pc +end + +# clones a theta-particle +function Base.copy(pc::ParticleContainer, utility_functions::AbstractPFUtilitFunctions) + # fork particles + vals = eltype(pc.vals)[fork(p, utility_functions.copy) for p in pc.vals] + # copy weights + logWs = copy(pc.logWs) + ParticleContainer( vals, logWs, pc.logE, pc.n_consume) +end + +# run particle filter for one step, return incremental likelihood +function Libtask.consume(pc::ParticleContainer) + # normalisation factor: 1/N + z1 = logZ(pc) + n = length(pc) + + particles = collect(pc) + num_done = 0 + for i=1:n + p = particles[i] + score = Libtask.consume(p) + if score isa Real + score += p.taskinfo.logp + reset_logp!(p.taskinfo) + increase_logweight!(pc, i, Float64(score)) + elseif score == Val{:done} + num_done += 1 + else + println(score) + error("[consume]: error in running particle filter.") + end + end + + if num_done == n + res = Val{:done} + elseif num_done != 0 + error("[consume]: mis-aligned execution traces, num_particles= $(n), num_done=$(num_done).") + else + # update incremental likelihoods + z2 = logZ(pc) + res = increase_logevidence!(pc, z2 - z1) + pc.n_consume += 1 + # res = increase_loglikelihood(pc, z2 - z1) + end + return res +end + + +# compute the normalized weights +weights(pc::ParticleContainer) = softmax!(copy(pc.logWs)) # There exists only softmax! + +# compute the log-likelihood estimate, ignoring constant term ``- \log num_particles`` +logZ(pc::ParticleContainer) = logsumexp(pc.logWs) + +# compute the effective sample size ``1 / ∑ wᵢ²``, where ``wᵢ```are the normalized weights +function effectiveSampleSize(pc :: ParticleContainer) + Ws = weights(pc) + return inv(sum(abs2, Ws)) +end + +increase_logweight!(pc::ParticleContainer, t::Int, logw::Float64) = (pc.logWs[t] += logw) + +increase_logevidence!(pc::ParticleContainer, logw::Float64) = (pc.logE += logw) diff --git a/src/Core/Container/Trace.jl b/src/Core/Container/Trace.jl new file mode 100644 index 00000000..a62bb127 --- /dev/null +++ b/src/Core/Container/Trace.jl @@ -0,0 +1,76 @@ +# Idea: We decouple tasks form the model by only allowing to pass a function. +# This way, we do no longer need to have the model and the sampler saved in the trace struct +# However, we still need to insantiate the VarInfo + +# TaskInfo stores additional information about the task +mutable struct Trace{Tvi, TInfo} <: AbstractTrace where {Tvi, TInfo <: AbstractTaskInfo} + vi::Tvi # Unfortunatley, we can not set force this to be a subtype of VarInfo... + task::Task + taskinfo::TInfo +end + +function Base.copy(trace::Trace{Tvi,TInfo}, copy_vi::Function) where {Tvi, TInfo <: AbstractTaskInfo} + return Trace(trace.vi, trace.task, trace.taskinfo, copy_vi) +end + +# The procedure passes a function which is specified by the model. + +function Trace(vi::Tvi, f::Function, taskinfo::TInfo, copy_vi::Function) where {Tvi,TInfo<:AbstractTaskInfo} + task = CTask( () -> begin res=f(); produce(Val{:done}); res; end ) + res = Trace(copy_vi(vi), task, copy(taskinfo)) + # CTask(()->f()); + if res.task.storage === nothing + res.task.storage = IdDict() + end + res.task.storage[:turing_trace] = res # create a backward reference in task_local_storage + return res +end + +## We need to design the task in the Turing wrapper. +function Trace(vi::Tvi, task::Task, taskinfo::TInfo, copy_vi::Function) where {Tvi,TInfo<:AbstractTaskInfo} + res = Trace(copy_vi(vi), Libtask.copy(task), copy(taskinfo)) + # CTask(()->f()); + if res.task.storage === nothing + res.task.storage = IdDict() + end + res.task.storage[:turing_trace] = res # create a backward reference in task_local_storage + return res +end + + +# NOTE: this function is called by `forkr` + +function Trace(vi::Tvi, f::Function, taskinfo::TInfo, copy_vi::Function) where {Tvi,TInfo<:AbstractTaskInfo} + # CTask(()->f()); + task = CTask( () -> begin res=f(); produce(Val{:done}); res; end ) + res = Trace(copy_vi(vi), task, copy(taskinfo)) + if res.task.storage === nothing + res.task.storage = IdDict() + end + res.task.storage[:turing_trace] = res # create a backward reference in task_local_storage + return res +end +# step to the next observe statement, return log likelihood +Libtask.consume(t::Trace) = (t.vi.num_produce += 1; consume(t.task)) + +# Task copying version of fork for Trace. +function fork(trace::Trace, copy_vi::Function, is_ref::Bool = false, set_retained_vns_del_by_spl!::Union{Function,Nothing} = nothing) + newtrace = copy(trace, copy_vi) + if is_ref + @assert set_retained_vns_del_by_spl! !== nothing "[AdvancedPS] set_retained_vns_del_by_spl! is not set." + set_retained_vns_del_by_spl!(newtrace.vi) + end + newtrace.task.storage[:turing_trace] = newtrace + return newtrace +end + + +function forkr(trace::Trace, copy_vi::Function) + newtrace = Trace(trace.vi, trace.task.code, trace.taskinfo, copy_vi) + newtrace.vi.num_produce = 0 + return newtrace +end + +current_trace() = current_task().storage[:turing_trace] + +const Particle = Trace diff --git a/src/Core/Resample/resample.jl b/src/Core/Resample/resample.jl new file mode 100644 index 00000000..d61c7922 --- /dev/null +++ b/src/Core/Resample/resample.jl @@ -0,0 +1,58 @@ + +function resample!( + pc :: ParticleContainer, + utility_functions :: AbstractSMCUtilitFunctions, + indx :: Vector{Int64}, + ref :: Union{Particle, Nothing} = nothing, + new_ref:: Union{Particle, Nothing} = nothing +) + + n = length(pc.vals) + # count number of children for each particle + num_children = zeros(Int, n) + + @inbounds for i in indx + num_children[i] += 1 + end + + + + # fork particles + particles = collect(pc) + children = similar(particles) + j = 0 + @inbounds for i in 1:n + ni = num_children[i] + + if ni > 0 + # fork first child + pi = particles[i] + isref = pi === ref + p = isref ? fork(pi, utility_functions.copy, isref, utility_functions.set_retained_vns_del_by_spl!) : pi + children[j += 1] = p + + # fork additional children + for _ in 2:ni + children[j += 1] = fork(p, utility_functions.copy, isref, utility_functions.set_retained_vns_del_by_spl!) + end + end + end + + if ref !== nothing + ancestor_idx = indx[n] + # Insert the retained particle. This is based on the replaying trick for efficiency + # reasons. If we implement PG using task copying, we need to store Nx * T particles! + # This is a rather effcient way of how to solve the ancestor problem. + if ancestor_idx == n + @inbounds children[n] = ref + elseif new_ref !== nothing + @inbounds children[n] = new_ref + end + end + + # replace particles and log weights in the container with new particles and weights + pc.vals = children + pc.logWs = zeros(n) + + pc +end diff --git a/src/Core/Resample/resample_functions.jl b/src/Core/Resample/resample_functions.jl new file mode 100644 index 00000000..06d7a6c4 --- /dev/null +++ b/src/Core/Resample/resample_functions.jl @@ -0,0 +1,163 @@ +### + +### Resampler Functions + +### + + +# Default resampling scheme + +function resample(w::AbstractVector{<:Real}, num_particles::Integer=length(w)) + return resample_systematic(w, num_particles) +end + + + +# More stable, faster version of rand(Categorical) + +function randcat(p::AbstractVector{T}) where T<:Real + r, s = rand(T), 1 + for j in eachindex(p) + r -= p[j] + if r <= zero(T) + s = j + break + end + end + return s +end + +function resample_multinomial(w::AbstractVector{<:Real}, num_particles::Integer) + return rand(Distributions.sampler(Categorical(w)), num_particles) +end + + + +function resample_residual(w::AbstractVector{<:Real}, num_particles::Integer) + M = length(w) + # "Repetition counts" (plus the random part, later on): + Ns = floor.(length(w) .* w) + # The "remainder" or "residual" count: + R = Int(sum(Ns)) + # The number of particles which will be drawn stocastically: + M_rdn = num_particles - R + # The modified weights: + Ws = (M .* w - floor.(M .* w)) / M_rdn + # Draw the deterministic part: + indx1, i = Array{Int}(undef, R), 1 + for j in 1:M + for k in 1:Ns[j] + indx1[i] = j + i += 1 + end + end + # And now draw the stocastic (Multinomial) part: + return append!(indx1, rand(Distributions.sampler(Categorical(w)), M_rdn)) +end + + + +""" + + resample_stratified(weights, n) + + + +Return a vector of `n` samples `x₁`, ..., `xₙ` from the numbers 1, ..., `length(weights)`, + +generated by stratified resampling. + + + +In stratified resampling `n` ordered random numbers `u₁`, ..., `uₙ` are generated, where + +``uₖ \\sim U[(k - 1) / n, k / n)``. Based on these numbers the samples `x₁`, ..., `xₙ` + +are selected according to the multinomial distribution defined by the normalized `weights`, + +i.e., `xᵢ = j` if and only if + +``uᵢ \\in [\\sum_{s=1}^{j-1} weights_{s}, \\sum_{s=1}^{j} weights_{s})``. + +""" + +function resample_stratified(weights::AbstractVector{<:Real}, n::Integer) + # check input + m = length(weights) + m > 0 || error("weight vector is empty") + # pre-calculations + @inbounds v = n * weights[1] + # generate all samples + samples = Array{Int}(undef, n) + sample = 1 + @inbounds for i in 1:n + # sample next `u` (scaled by `n`) + u = oftype(v, i - 1 + rand()) + # as long as we have not found the next sample + while v < u + # increase and check the sample + sample += 1 + sample > m && + error("sample could not be selected (are the weights normalized?)") + # update the cumulative sum of weights (scaled by `n`) + v += n * weights[sample] + end + # save the next sample + samples[i] = sample + end + return samples +end + + + +""" + + resample_systematic(weights, n) + + + +Return a vector of `n` samples `x₁`, ..., `xₙ` from the numbers 1, ..., `length(weights)`, + +generated by systematic resampling. + + + +In systematic resampling a random number ``u \\sim U[0, 1)`` is used to generate `n` ordered + +numbers `u₁`, ..., `uₙ` where ``uₖ = (u + k − 1) / n``. Based on these numbers the samples + +`x₁`, ..., `xₙ` are selected according to the multinomial distribution defined by the + +normalized `weights`, i.e., `xᵢ = j` if and only if + +``uᵢ \\in [\\sum_{s=1}^{j-1} weights_{s}, \\sum_{s=1}^{j} weights_{s})``. + +""" + +function resample_systematic(weights::AbstractVector{<:Real}, n::Integer) + # check input + m = length(weights) + m > 0 || error("weight vector is empty") + # pre-calculations + @inbounds v = n * weights[1] + u = oftype(v, rand()) + # find all samples + samples = Array{Int}(undef, n) + sample = 1 + @inbounds for i in 1:n + # as long as we have not found the next sample + while v < u + # increase and check the sample + sample += 1 + sample > m && + error("sample could not be selected (are the weights normalized?)") + # update the cumulative sum of weights (scaled by `n`) + v += n * weights[sample] + end + # save the next sample + samples[i] = sample + # update `u` + u += one(u) + end + return samples +end diff --git a/src/Core/Utilities/UtilityFunctions.jl b/src/Core/Utilities/UtilityFunctions.jl new file mode 100644 index 00000000..e30d209a --- /dev/null +++ b/src/Core/Utilities/UtilityFunctions.jl @@ -0,0 +1,22 @@ + + + +struct SMCUtilityFunctions<:AbstractSMCUtilitFunctions + copy :: Function + set_retained_vns_del_by_spl! :: Function + empty! :: Function + tonamedtuple :: Function +end + + + +const PGUtilityFunctions = SMCUtilityFunctions + +struct PGASUtilityFunctions{AP}<:AbstractPGASUtilityFunctions + copy :: Function + set_retained_vns_del_by_spl! :: Function + merge_traj :: Function + empty! :: Function + tonamedtuple :: Function + +end diff --git a/src/Inference/Inference.jl b/src/Inference/Inference.jl new file mode 100644 index 00000000..66fe383a --- /dev/null +++ b/src/Inference/Inference.jl @@ -0,0 +1,59 @@ + + +const PROGRESS = Ref(true) + +function AbstractMCMC.sample( + model::ModelType, + alg::AlgType, + uf::UF, + vi::C, + N::Integer; + kwargs... +) where { + C, + ModelType<:AbstractPFModel, + AlgType<: AbstractPFAlgorithm, + UF<:AbstractPFUtilitFunctions +} + return sample(model, Sampler(alg, uf, vi), N; progress=PROGRESS[], kwargs...) +end + +const ACVANCEDPS_INTERNAL_VARS = (internals = [ + "lp", + "weight", + "le", +],) + +function AbstractMCMC.bundle_samples( + rng::AbstractRNG, + model::AbstractPFModel, + spl::AbstractPFSampler, + N::Integer, + ts::Vector{T}; + kwargs... +) where {T<:AbstractPFTransition} + + # Convert transitions to array format. + # Also retrieve the variable names. + nms, vals = _params_to_array(ts) + + # Get the values of the extra parameters in each Transition struct. + + extra_params, extra_values = get_transition_extras(ts) + # Extract names & construct param array. + nms = string.(vcat(nms..., string.(extra_params)...)) + parray = hcat(vals, extra_values) + + + le = missing + # Set up the info tuple. + info = NamedTuple() + # Chain construction. + return Chains( + parray, + string.(nms), + deepcopy(ACVANCEDPS_INTERNAL_VARS); + evidence=le, + info=info + ) +end diff --git a/src/Inference/Model.jl b/src/Inference/Model.jl new file mode 100644 index 00000000..a6a74150 --- /dev/null +++ b/src/Inference/Model.jl @@ -0,0 +1,25 @@ + +### + +### The model + +### + + +struct PFModel <: AbstractPFModel + task::Task +end + +function PFModel(f::Function, args) + task = create_task(f,args) + PFModel(task) +end + + +function create_task(f::Function, args=nothing) + if args === nothing + return CTask(() -> begin new_vi=f(); produce(Val{:done}); new_vi; end ) + else + return CTask(() -> begin new_vi=f(args...); produce(Val{:done}); new_vi; end ) + end +end diff --git a/src/Inference/Sampler.jl b/src/Inference/Sampler.jl new file mode 100644 index 00000000..375c2a87 --- /dev/null +++ b/src/Inference/Sampler.jl @@ -0,0 +1,45 @@ + + +# The particle container is our state, + +mutable struct SMCSampler{PC, ALG, UF, C} <: AbstractPFSampler where { + PC<:ParticleContainer, + ALG<:SMCAlgorithm, + UF<:SMCUtilityFunctions +} + pc :: PC + alg :: ALG + uf :: UF + vi :: C +end + + +function Sampler(alg:: ALG, uf::UF, vi::C) where { + C, + ALG<: SMCAlgorithm, + UF<: SMCUtilityFunctions, +} + pc = ParticleContainer(Trace{typeof(vi),SMCTaskInfo{Float64}}[]) + SMCSampler(pc, alg, uf, uf.empty!(vi)) +end + + + + +mutable struct PGSampler{T, ALG, UF, C} <: AbstractPFSampler where { + T <:Particle, + ALG<:PGAlgorithm, + UF<:PGUtilityFunctions +} + alg :: ALG + uf :: UF + ref_traj :: Union{T, Nothing} + vi :: C +end + + +Sampler(alg:: ALG, uf::UF, vi::T) where { + T, + ALG<: PGAlgorithm, + UF<: PGUtilityFunctions +} = PGSampler{Trace{typeof(vi),PGTaskInfo{Float64}},typeof(alg),typeof(uf),typeof(vi)}(alg, uf, nothing, vi) diff --git a/src/Inference/Transitions.jl b/src/Inference/Transitions.jl new file mode 100644 index 00000000..7dd56b83 --- /dev/null +++ b/src/Inference/Transitions.jl @@ -0,0 +1,132 @@ + + + +struct PFTransition{T, F<:AbstractFloat} <: AbstractPFTransition + θ::T + lp::F + le::F + weight::F +end + +AbstractMCMC.transition_type(spl::S) where S<:AbstractPFSampler = PFTransition +function additional_parameters(::Type{<:PFTransition}) + return [:lp,:le, :weight] +end + + + +function get_transition_extras(ts::Vector{T}) where T<:AbstractTransition + # Get the extra field names from the sampler state type. + # This handles things like :lp or :weight. + extra_params = additional_parameters(T) + # Get the values of the extra parameters. + local extra_names + all_vals = [] + # Iterate through each transition. + for t in ts + extra_names = String[] + vals = [] + # Iterate through each of the additional field names + # in the struct. + for p in extra_params + # Check whether the field contains a NamedTuple, + # in which case we need to iterate through each + # key/value pair. + prop = getproperty(t, p) + if prop isa NamedTuple + for (k, v) in pairs(prop) + push!(extra_names, string(k)) + push!(vals, v) + end + else + push!(extra_names, string(p)) + push!(vals, prop) + end + end + push!(all_vals, vals) + end + # Convert the vector-of-vectors to a matrix. + valmat = [all_vals[i][j] for i in 1:length(ts), j in 1:length(all_vals[1])] + return extra_names, valmat +end + + + + + +function flatten(names, value :: AbstractArray, k :: String, v) + if isa(v, Number) + name = k + push!(value, v) + push!(names, name) + elseif isa(v, Array) + for i = eachindex(v) + if isa(v[i], Number) + name = string(ind2sub(size(v), i)) + name = replace(name, "(" => "["); + name = replace(name, ",)" => "]"); + name = replace(name, ")" => "]"); + name = k * name + isa(v[i], Nothing) && println(v, i, v[i]) + push!(value, v[i]) + push!(names, name) + elseif isa(v[i], AbstractArray) + name = k * string(ind2sub(size(v), i)) + flatten(names, value, name, v[i]) + else + error("Unknown var type: typeof($v[i])=$(typeof(v[i]))") + end + end + else + error("Unknown var type: typeof($v)=$(typeof(v))") + end + return +end + +function flatten_namedtuple(nt::NamedTuple{pnames}) where {pnames} + vals = Vector{Real}() + names = Vector{AbstractString}() + for k in pnames + v = nt[k] + if length(v) == 1 + flatten(names, vals, string(k), v) + else + for (vnval, vn) in zip(v[1], v[2]) + flatten(names, vals, vn, vnval) + end + end + end + return names, vals +end + + + + +function _params_to_array(ts::Vector{T}) where {T<:AbstractTransition} + names = Set{String}() + dicts = Vector{Dict{String, Any}}() + # Extract the parameter names and values from each transition. + for t in ts + nms, vs = flatten_namedtuple(t.θ) + push!(names, nms...) + # Convert the names and values to a single dictionary. + d = Dict{String, Any}() + for (k, v) in zip(nms, vs) + d[k] = v + end + push!(dicts, d) + end + # Convert the set to an ordered vector so the parameter ordering + # is deterministic. + ordered_names = collect(names) + vals = Matrix{Union{Real, Missing}}(undef, length(ts), length(ordered_names)) + # Place each element of all dicts into the returned value matrix. + for i in eachindex(dicts) + for (j, key) in enumerate(ordered_names) + vals[i,j] = get(dicts[i], key, missing) + end + end + return ordered_names, vals +end + +ind2sub(v, i) = Tuple(CartesianIndices(v)[i]) diff --git a/src/Inference/sample_init.jl b/src/Inference/sample_init.jl new file mode 100644 index 00000000..c45f03c9 --- /dev/null +++ b/src/Inference/sample_init.jl @@ -0,0 +1,16 @@ +function AbstractMCMC.sample_init!( + rng::AbstractRNG, + ℓ::ModelType, + spl::SamplerType, + N::Integer; + debug::Bool=false, + kwargs... +) where {ModelType<:AbstractPFModel, SamplerType<:SMCSampler} + + T = Trace{typeof(spl.vi),SMCTaskInfo{Float64}} + particles = T[ Trace(spl.vi, ℓ.task, SMCTaskInfo(), spl.uf.copy) for _ =1:N] + spl.pc = ParticleContainer{typeof(particles[1])}(particles,zeros(N),0.0,0) + + sample!(spl.pc, spl.alg, spl.uf, nothing) + +end diff --git a/src/Inference/step.jl b/src/Inference/step.jl new file mode 100644 index 00000000..93991a6e --- /dev/null +++ b/src/Inference/step.jl @@ -0,0 +1,155 @@ + +#SMC step +function AbstractMCMC.step!( + ::AbstractRNG, + model::AbstractPFModel, + spl::SPL, + ::Integer; + iteration::Integer, + kwargs... + ) where SPL <: SMCSampler + + particle = spl.pc[iteration] + + params = spl.uf.tonamedtuple(particle.vi) + return PFTransition(params, particle.taskinfo.logp, spl.pc.logE, weights(spl.pc)[iteration]) +end + + +# PG step +function AbstractMCMC.step!( + ::AbstractRNG, + model::AbstractPFModel, + spl::SPL, + ::Integer; + kwargs... + ) where SPL <: PGSampler + + n = spl.alg.n + + T = Trace{typeof(spl.vi), PGTaskInfo{Float64}} + + if spl.ref_traj !== nothing + particles = T[ Trace(spl.vi, model.task, PGTaskInfo(), spl.uf.copy) for _ =1:n-1] + pc = ParticleContainer{typeof(particles[1])}(particles,zeros(n-1),0.0,0) + # Reset Task + spl.ref_traj = forkr(spl.ref_traj, spl.uf.copy) + push!(pc, spl.ref_traj) + else + particles = T[ Trace(spl.vi, model.task, PGTaskInfo(), spl.uf.copy) for _ =1:n] + pc = ParticleContainer{typeof(particles[1])}(particles,zeros(n),0.0,0) + end + + sample!(pc, spl.alg, spl.uf, spl.ref_traj) + + indx = AdvancedPS.randcat(weights(pc)) + particle = spl.ref_traj = pc[indx] + params = spl.uf.tonamedtuple(particle.vi) + return PFTransition(params, particle.taskinfo.logp, pc.logE, weights(pc)[indx]) +end + +# +# # The resampler threshold is only imprtant for the first step! +# function samplePGAS!(pc::ParticleContainer,resampler::Function = resample_systematic ,ref_particle::Union{Particle,Nothing}=nothing ,joint_logp::Union{Tuple{Vector{Symbol},Function},Nothing}=nothing, resampler_threshold =0.5) +# if ref_particle === nothing +# # We do not have a reference trajectory yet, therefore, perform normal SMC! +# sampleSMC!(pc,resampler,resampler_threshold) +# +# else +# # Before starting, we need to copute the ancestor weights. +# # Note that there is a inconsistency with the original paper +# # http://jmlr.org/papers/volume15/lindsten14a/lindsten14a.pdf +# # Lindsten samples already for x1. This is not possible in +# # in this situation because we do not have access to the information +# # which variables belong to x1 and which to x0! +# +# # The procedure works as follows. The ancestor weights are only dependent on +# # the states x_{0:t-1}. We make us of this by computing the ancestor indices for +# # the next state. +# ancestor_index = length(pc) +# ancestor_particle::Union{typeof(pc[1]), Nothing} = nothing +# n = length(pc) +# +# num_totla_consume = typemax(Int64) +# first_round = true +# while consume(pc) != Val{:done} +# # compute weights +# Ws = weights(pc) +# # We need them for ancestor sampling... +# logws = copy(pc.logWs) +# logpseq = [pc[i].taskinfo.logpseq for i in 1:n ] +# # check that weights are not NaN +# @assert !any(isnan, Ws) +# # sample ancestor indices +# # Ancestor trajectory is not sampled +# nresamples = n-1 +# indx = resampler(Ws, nresamples) +# +# # Now the ancestor sampling starts. We do not use ancestor sampling in the +# # first step. This is due to the reason before. In addition, we do not need +# # to compute the ancestor index for the last step, because we are always +# #computing the ancestor index one step ahead. +# if ancestor_particle === nothing +# push!(indx,n) +# resample!(pc, indx,ref_particle) +# # In this case, we do have access to the joint_logp ! Therefore: +# if joint_logp !== nothing +# # We need to create a dictionary with symbol value paris, which we pass +# # to the joint_logp function. +# @assert isa(joint_logp[1], Vector{Symbols}) "[AdvancedPS] the first argument of the joint_logp tuble must be a vector of Symbols!" +# @assert isa(joint_logp[2], Function) "[AdvancedPS] the second argument of the joint_logp tuble must be a function of the for"* +# " f(num_produce, args...), which returns a value of Float64!" +# +# # We are executing the joint_logp function as a set of tasks. The idea behind +# # this is that we might extend the task to include parallelization. +# tasks = [] +# for i = 1:n +# args = pc.manipulators["get_AS_joint_logp_args"](joint_logp[1],pc[i],ref_particle) +# task = CTask( () -> begin result=joint_logp[2](pc.n_consume,args...); produce(result); end ) +# push(tasks,task) +# schedule(task) +# yield() +# end +# for (i,t) in enumerate(tasks) +# logws[i] += consume(t) -logpseq[i] # The ancestor weights w_ancstor = w_i *p(x_{0:t-1},x_{t:T})/p(x_{0:t-1}) +# end +# +# ancestor_index = randcat(softmax!(logws)) +# # We are one step behind.... +# selected_path = pc[ancestor_index] +# new_vi = pc.manipulators["merge_traj"](pc.manipulators["copy"](selected_path.vi),ref_particle.vi,pc.n_consume +1) +# ancestor_particle = Trace{typeof(new_vi),typeof(selected_path.taskinfo)}(new_vi,copy(selected_path.task),copy(selected_path.taskinfo)) +# else +# if pc.n_consume <= num_totla_consume-1 #We do not need to sample the last one... +# # The idea is rather simple, we extend the vs and let them run trough... +# pc_ancestor = ParticleContainer{typeof(pc[1].vi),typeof(pc[1].taskinfo)}() +# pc_ancestor.n_consume = pc.n_consume +# for i = 1:n +# new_vi = pc.manipulators["merge_traj"](pc.manipulators["copy"](pc[i].vi),ref_particle.vi) +# new_particle = Trace{typeof(new_vi),typeof(pc[i].taskinfo)}(new_vi,copy(pc[i].task),copy(pc[i].taskinfo)) +# push!(pc_ancestor,new_particle) +# end +# +# while consume(pc_ancestor) != Val{:done} end # No resampling, we just want to get log p(x_{0:t-1},x'_{t,T}) +# for i in 1:n +# logws[i] += pc_ancestor[i].taskinfo.logpseq -logpseq[i] # The ancestor weights w_ancstor = w_i *p(x_{0:t-1},x_{t:T})/p(x_{0:t-1}) +# end +# +# +# ancestor_index = randcat(softmax!(logws)) +# # We are one step behind.... +# selected_path = pc[ancestor_index] +# new_vi = pc.manipulators["merge_traj"](pc.manipulators["copy"](selected_path.vi),ref_particle.vi,pc.n_consume +1) +# ancestor_particle = Trace{typeof(new_vi),typeof(selected_path.taskinfo)}(new_vi,copy(selected_path.task),copy(selected_path.taskinfo)) +# num_total_consume = pc_ancestor.n_consume +# end +# end +# end +# # Reactivate the tasks, this is only important for synchorinty. +# for i =1:n +# consume(pc[i].task) +# end +# end +# end +# +# end diff --git a/test_with_turing.jl b/test_with_turing.jl new file mode 100644 index 00000000..e5d8be57 --- /dev/null +++ b/test_with_turing.jl @@ -0,0 +1,64 @@ +# We need to have the ParticleGibbsExtension branch!!! +using Turing +using Distributions +using CuArrays +## It works even for GPU's! +cache = Dict() + +arr = rand(10,10) +arr2 = rand(10,1) +Turing.@model gdemo(x1,x2,x3,x4) = begin + s ~ InverseGamma(2, 3) + m ~ Normal(0, sqrt(s)) + + + arr3 = arr*arr2 + + x1 ~ Normal(m, abs(s)) + h1 ~ Normal(m,abs(s)) + x2 ~ Normal(m, abs(h1)) + h2 ~ Normal(m,abs(h1)) + x3 ~ Normal(m, abs(h2)) + h3 ~ Normal(m,abs(h2)) + x4 ~ Normal(m, abs(h3)) + + k~ Exponential(0.1) +end + +# arrcu = cu(rand(10000,10000)) +# arr2cu = cu(rand(10000,1)) +# Turing.@model gdemocu(x, y) = begin +# s ~ InverseGamma(2, 3) +# m ~ Normal(0, sqrt(s)) +# +# +# arr3cu = arrcu*arr2cu +# +# x ~ Normal(m, sqrt(s)) +# h ~ Normal(m,abs(x)) +# +# y ~ Normal(m, sqrt(s)) +# +# k~ Exponential(0.1) +# +# end + +#@elapsed arr*arr2 +#@elapsed arrcu*arr2cu +#time2 = @elapsed c1 = Turing.sample(gdemocu(1.5, 2), Turing.PG(100), 100) + +#time = @elapsed c1 = Turing.sample(gdemo(1.5, 2), Turing.PG(10), 10) + +#c2 = Turing.sample(gdemo(1.5, 2), Turing.SMC(), 10) + + + +@elapsed c1 = Turing.sample(gdemo([1.5,2,2,2]...), Turing.PGAS(100),1000) +@elapsed c2 = Turing.sample(gdemo([1.5,2,2,2]...), Turing.PG(100),1000) +@elapsed c3 = Turing.sample(gdemo([1.5,2,2,2]...), Turing.SMC(),1000) +@elapsed c4 = Turing.sample(gdemo([1.5,2,2,2]...), Turing.MH(),100000) + +c2 +c1 +c3 +c4