working metropolis hastings (with hacks)

Author: sbrantq

src/ProbProg.jl

@@ -14,10 +14,14 @@ mutable struct ProbProgTrace
     choices::Dict{Symbol,Any}
     retval::Any
     weight::Any
+    fn::Union{Nothing,Function}
+    args::Union{Nothing,Tuple}
 
-    function ProbProgTrace()
-        return new(Dict{Symbol,Any}(), nothing, nothing)
+    function ProbProgTrace(fn::Function, args::Tuple)
+        return new(Dict{Symbol,Any}(), nothing, nothing, fn, args)
     end
+
+    ProbProgTrace() = new(Dict{Symbol,Any}(), nothing, nothing, nothing, ())
 end
 
 function addSampleToTraceLowered(
@@ -292,7 +296,7 @@ function call_internal(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
 end
 
 function generate(f::Function, args::Vararg{Any,Nargs}; constraints=nothing) where {Nargs}
-    trace = ProbProgTrace()
+    trace = ProbProgTrace(f, (args...,))
 
     weight, res = @jit sync = true optimize = :probprog generate_internal(
         f, args...; trace, constraints
@@ -416,7 +420,7 @@ function generate_internal(
 end
 
 function simulate(f::Function, args::Vararg{Any,Nargs}) where {Nargs}
-    trace = ProbProgTrace()
+    trace = ProbProgTrace(f, (args...,))
 
     res = @jit optimize = :probprog sync = true simulate_internal(f, args...; trace)
 
@@ -571,4 +575,38 @@ function Base.show(io::IO, trace::ProbProgTrace)
     end
 end
 
+struct Selection
+    symbols::Vector{Symbol}
+end
+
+select(symbol::Symbol) = Selection([symbol])
+
+choicemap() = Dict{Symbol,Any}()
+get_choices(trace::ProbProgTrace) = trace.choices
+
+function metropolis_hastings(trace::ProbProgTrace, sel::Selection)
+    if trace.fn === nothing
+        error("MH requires a trace with fn and args recorded")
+    end
+
+    constraints = Dict{Symbol,Any}()
+    for (sym, val) in trace.choices
+        sym in sel.symbols && continue
+        constraints[sym] = [val]
+    end
+
+    new_trace, _ = generate(trace.fn, trace.args...; constraints)
+    rng_state = new_trace.retval[1] # TODO: this is a temporary hack
+
+    log_alpha = new_trace.weight - trace.weight
+
+    if log(rand()) < log_alpha
+        new_trace.args = (rng_state, new_trace.args[2:end]...)
+        return (new_trace, true)
+    else
+        trace.args = (rng_state, trace.args[2:end]...)
+        return (trace, false)
+    end
+end
+
 end

test/probprog/linear_regression.jl

@@ -0,0 +1,77 @@
+using Reactant, Test, Random
+using Reactant: ProbProg
+
+# Reference: https://www.gen.dev/docs/stable/getting_started/linear_regression/
+
+normal(rng, μ, σ, shape) = μ .+ σ .* randn(rng, shape)
+normal_logpdf(x, μ, σ, _) = -sum(log.(σ)) - sum((μ .- x) .^ 2) / (2 * σ^2)
+
+function my_model(seed, xs)
+    rng = Random.default_rng()
+    Random.seed!(rng, seed)
+
+    slope = ProbProg.sample(
+        normal, rng, 0.0, 2.0, (1,); symbol=:slope, logpdf=normal_logpdf
+    )
+    intercept = ProbProg.sample(
+        normal, rng, 0.0, 10.0, (1,); symbol=:intercept, logpdf=normal_logpdf
+    )
+
+    ys = ProbProg.sample(
+        normal,
+        rng,
+        slope .* xs .+ intercept,
+        1.0,
+        (length(xs),);
+        symbol=:ys,
+        logpdf=normal_logpdf,
+    )
+
+    return rng.seed, ys
+end
+
+function my_inference_program(xs, ys, num_iters)
+    xs_r = Reactant.to_rarray(xs)
+
+    constraints = ProbProg.choicemap()
+    constraints[:ys] = [ys]
+
+    seed = Reactant.to_rarray(UInt64[1, 4])
+
+    trace, _ = ProbProg.generate(my_model, seed, xs_r; constraints)
+    trace.args = (trace.retval[1], trace.args[2:end]...) # TODO: this is a temporary hack
+
+    for i in 1:num_iters
+        trace, _ = ProbProg.metropolis_hastings(trace, ProbProg.select(:slope))
+        trace, _ = ProbProg.metropolis_hastings(trace, ProbProg.select(:intercept))
+        choices = ProbProg.get_choices(trace)
+        @show i, choices[:slope], choices[:intercept]
+    end
+
+    choices = ProbProg.get_choices(trace)
+    return (choices[:slope], choices[:intercept])
+end
+
+@testset "linear_regression" begin
+    @testset "simulate" begin
+        seed = Reactant.to_rarray(UInt64[1, 4])
+
+        xs = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+        xs_r = Reactant.to_rarray(xs)
+
+        trace = ProbProg.simulate(my_model, seed, xs_r)
+
+        @test haskey(trace.choices, :slope)
+        @test haskey(trace.choices, :intercept)
+        @test haskey(trace.choices, :ys)
+    end
+
+    @testset "inference" begin
+        xs = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+        ys = [8.23, 5.87, 3.99, 2.59, 0.23, -0.66, -3.53, -6.91, -7.24, -9.90]
+
+        slope, intercept = my_inference_program(xs, ys, 1000)
+
+        @show slope, intercept
+    end
+end