Add getparams and get_params_varinfo functions

src/model.jl

@@ -172,15 +172,74 @@ 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, m.transformed)
+    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, m.transformed)
+    end
+end
+
+# TODO: add this function to `ADgradient` with `ReverseDiff` when compiled
+# see https://github.com/TuringLang/Turing.jl/pull/2097  
+"""
+    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)
+    if !m.transformed
+        return vcat([isa(vi[p], Real) ? vi[p] : vec(vi[p]) for p in m.parameters]...)
+    else # otherwise, we need to talk through the graph and transform the values to unconstrained space
+        vs = 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])
+                if linked_val isa AbstractArray
+                    push!(vs, vec(linked_val))
+                else
+                    push!(vs, linked_val)
+                end
+            end
+        end
+        return vcat(vs...)
     end
-    return SimpleVarInfo(d, vi.logp, vi.transformation)
 end
 
 """
@@ -359,7 +418,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
@@ -444,4 +503,4 @@ function AbstractPPL.evaluate!!(
         end
     end
     return vi, logp
-end
+end

test/inference.jl

@@ -1,3 +1,17 @@
+# 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
+
 # ReverseDiff
 
 @testset "trans-dim bijectors tape compilation" begin