Add getparams and get_params_varinfo functions

src/model.jl

@@ -172,15 +172,59 @@ function initialize_var_store(data, vars, array_sizes)
     return var_store
 end
 
+"""
+    get_params_varinfo(m::BUGSModel[, vi::SimpleVarInfo])
+
+Returns a `SimpleVarInfo` object containing only the parameter values of the model.
+If `vi` is provided, it will be used; otherwise, `m.varinfo` will be used.
+"""
 function get_params_varinfo(m::BUGSModel)
     return get_params_varinfo(m, m.varinfo)
 end
 function get_params_varinfo(m::BUGSModel, vi::SimpleVarInfo)
-    d = Dict{VarName,Any}()
-    for param in m.parameters
-        d[param] = vi[param]
+    if !m.transformed
+        d = Dict{VarName,Any}()
+        for param in m.parameters
+            d[param] = vi[param]
+        end
+        return SimpleVarInfo(d, vi.logp, DynamicPPL.NoTransformation())
+    else
+        d = Dict{VarName,Any}()
+        g = m.g
+        for vn in m.sorted_nodes
+            ni = g[vn]
+            @unpack node_type, node_function_expr, node_args = ni
+            args = Dict(getsym(arg) => vi[arg] for arg in node_args)
+            expr = node_function_expr.args[2]
+            if vn in m.parameters
+                dist = _eval(expr, args)
+                linked_val = DynamicPPL.link(dist, vi[vn])
+                d[vn] = linked_val
+            end
+        end
+        return SimpleVarInfo(d, vi.logp, DynamicPPL.DynamicTransformation())
     end
-    return SimpleVarInfo(d, vi.logp, vi.transformation)
+end
+
+# TODO: add this function to `ADgradient` with `ReverseDiff` when compiled
+# see https://github.com/TuringLang/Turing.jl/pull/2097
+# TODO: a static allocated version is possible because we know the size, but not worth the complexity now
+"""
+    getparams(m::BUGSModel[, vi::SimpleVarInfo])
+
+Return the values of the parameters in the model as a vector, the values are flattened 
+in the order of `m.parameters` (also the topological order).
+"""
+function getparams(m::BUGSModel)
+    return getparams(m, m.varinfo)
+end
+function getparams(m::BUGSModel, vi::SimpleVarInfo)
+    params_vi = get_params_varinfo(m, vi)
+    return vcat(
+        [
+            isa(params_vi[p], Real) ? params_vi[p] : vec(params_vi[p]) for p in m.parameters
+        ]...,
+    )
 end
 
 """
@@ -359,7 +403,7 @@ function AbstractPPL.evaluate!!(
         @unpack node_type, node_function_expr, node_args = ni
         args = Dict(getsym(arg) => vi[arg] for arg in node_args)
         expr = node_function_expr.args[2]
-        if node_type == JuliaBUGS.Logical
+        if node_type == JuliaBUGS.Logical # be conservative -- always propagate values of logical nodes
             value = _eval(expr, args)
             vi = setindex!!(vi, value, vn)
         else

test/inference.jl

@@ -1,14 +1,15 @@
 # ReverseDiff
 
 @testset "trans-dim bijectors tape compilation" begin
-    # `birats` contains Dirichlet distribution 
+    # `birats` contains Wishart distribution 
     model_def = JuliaBUGS.BUGSExamples.birats.model_def
     data = JuliaBUGS.BUGSExamples.birats.data
     inits = JuliaBUGS.BUGSExamples.birats.inits[1]
     model = compile(model_def, data, inits)
     ad_model = ADgradient(:ReverseDiff, model; compile=Val(false))
-    D = LogDensityProblems.dimension(model)
-    initial_θ = rand(D)
+    # random initialization sometimes fails because some parameters are supposed to be from
+    # PD matrix
+    initial_θ = JuliaBUGS.getparams(model)
     LogDensityProblems.logdensity_and_gradient(ad_model, initial_θ)
 end
 

test/logp_tests/blockers.jl

@@ -49,8 +49,6 @@ dppl_logp =
 bugs_logp = JuliaBUGS.evaluate!!(JuliaBUGS.settrans(bugs_model, true), DefaultContext())[2]
 dppl_logp =
     DynamicPPL.evaluate!!(
-        dppl_model,
-        DynamicPPL.link!!(get_params_varinfo(bugs_model), dppl_model),
-        DynamicPPL.DefaultContext(),
+        dppl_model, get_params_varinfo(bugs_model), DynamicPPL.DefaultContext()
     )[2].logp
 @test bugs_logp ≈ dppl_logp rtol = 1E-6

test/logp_tests/bones.jl

@@ -51,8 +51,6 @@ dppl_logp =
 bugs_logp = JuliaBUGS.evaluate!!(JuliaBUGS.settrans(bugs_model, true), DefaultContext())[2]
 dppl_logp =
     DynamicPPL.evaluate!!(
-        dppl_model,
-        DynamicPPL.link!!(get_params_varinfo(bugs_model), dppl_model),
-        DynamicPPL.DefaultContext(),
+        dppl_model, get_params_varinfo(bugs_model), DynamicPPL.DefaultContext()
     )[2].logp
 @test bugs_logp ≈ dppl_logp rtol = 1E-6

test/logp_tests/dogs.jl

@@ -51,8 +51,6 @@ dppl_logp =
 bugs_logp = JuliaBUGS.evaluate!!(JuliaBUGS.settrans(bugs_model, true), DefaultContext())[2]
 dppl_logp =
     DynamicPPL.evaluate!!(
-        dppl_model,
-        DynamicPPL.link!!(get_params_varinfo(bugs_model), dppl_model),
-        DynamicPPL.DefaultContext(),
+        dppl_model, get_params_varinfo(bugs_model), DynamicPPL.DefaultContext()
     )[2].logp
 @test bugs_logp ≈ dppl_logp rtol = 1E-6

test/logp_tests/rats.jl

@@ -69,9 +69,7 @@ dppl_logp =
 
 dppl_logp =
     DynamicPPL.evaluate!!(
-        dppl_model,
-        link!!(get_params_varinfo(bugs_model), dppl_model),
-        DynamicPPL.DefaultContext(),
+        dppl_model, get_params_varinfo(bugs_model), DynamicPPL.DefaultContext()
     )[2].logp
 bugs_logp = JuliaBUGS.evaluate!!(JuliaBUGS.settrans(bugs_model, true), DefaultContext())[2]
 @test bugs_logp ≈ dppl_logp rtol = 1E-6

test/runtests.jl

@@ -100,6 +100,10 @@ end
     model = compile(model_def, data, inits)
 end
 
+@testset "Utils" begin
+    include("utils.jl")
+end
+
 @testset "Log Probability Test" begin
     include("run_logp_tests.jl")
     @testset "Single stochastic variable test" begin

test/utils.jl

@@ -0,0 +1,13 @@
+# Tests for `getparams`, using `Rats` model
+@testset "`getparams` with Rats" begin
+    m = :rats
+    data = JuliaBUGS.BUGSExamples.VOLUME_I[m].data
+    inits = JuliaBUGS.BUGSExamples.VOLUME_I[m].inits[1]
+    model = JuliaBUGS.compile(JuliaBUGS.BUGSExamples.VOLUME_I[m].model_def, data, inits)
+    model_notran = JuliaBUGS.settrans(model, false)
+
+    @test LogDensityProblems.logdensity(model, JuliaBUGS.getparams(model)) ==
+        JuliaBUGS.evaluate!!(model)[2]
+    @test LogDensityProblems.logdensity(model_notran, JuliaBUGS.getparams(model_notran)) ==
+        JuliaBUGS.evaluate!!(model_notran)[2]
+end