initial API idea monte carlo

Author: dharasim

Project.toml

@@ -11,6 +11,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

src/SimpleProbabilisticPrograms.jl

@@ -24,6 +24,8 @@ export iid
 export UniformCategorical, Dirac # specific distributions
 export BetaGeometric, DirCat, symdircat
 export DictCond # simple conditional distribution
+export Conditioned, condition, mergetraces, logdensity
+export montecarlo, MonteCarloMethod, LikelihoodWeighting, likelihood_weighting
 
 using SpecialFunctions: digamma, logbeta
 using LogExpFunctions: logsumexp, logaddexp
@@ -32,7 +34,7 @@ using Distributions: Beta, Binomial, BetaBinomial, Geometric
 using Random: AbstractRNG, default_rng
 using MacroTools: @capture, splitdef, combinedef, postwalk, prewalk
 using Accessors: insert, PropertyLens, @set
-using StaticArrays: @SVector
+using StaticArrays: @SVector, sacollect
 
 import Base: show, rand
 import Distributions: logpdf, insupport
@@ -48,7 +50,7 @@ struct ProbProg{NAME, A, KW}
   kwargs::KW
 end
 
-show(io::IO, ::ProbProg{NAME}) where NAME = print(io, "ProbProg{$NAME}(...)")
+show(io::IO, ::ProbProg{NAME}) where NAME = print(io, "ProbProg{:$NAME}(...)")
 recover_trace(::ProbProg, trace) = trace
 
 function logpdf(model::ProbProg, x)
@@ -120,7 +122,7 @@ For example, `heads ~ Bernoulli(0.5)` is transformed into
 The lens is used as getter and setter by the interpreters of the program.
 """
 macro probprog(ex)
-  interpreter(i) = Symbol('i')
+  interpreter(i) = Symbol("interpreter")
   i = 0
 
   function rewrite_return_expr(ex)
@@ -418,6 +420,73 @@ end
 (cond::DictCond)(x) = cond.dists[x]
 DictCond(dists...) = DictCond(Dict(dists...))
 
+#################################
+### Conditioned distributions ###
+#################################
+
+struct Conditioned{D, C}
+  joint :: D # joint distribution
+  conds :: C # values conditioned on (as a trace)
+end
+
+condition(joint; on) = Conditioned(joint, on)
+function rand(rng::AbstractRNG, cond::Conditioned{<:ProbProg})
+   interpret(cond.joint, RandTrace(rng, cond.conds)).interpreter.trace
+end
+
+logdensity(dist, x) = logpdf(dist, x)
+
+function logdensity(cond::Conditioned, trace)
+  logdensity(cond.joint, mergetraces(cond.conds, trace))  
+end
+
+mergetraces(t1::NamedTuple, t2::NamedTuple) = (; t1..., t2...)
+
+#########################
+### Inference Methods ###
+#########################
+
+abstract type MonteCarloMethod end
+
+function montecarlo(f, dist, mcm::MonteCarloMethod, rng::AbstractRNG=default_rng())
+  montecarlo(dist, mcm, rng)(f)
+end
+
+function montecarlo(dist, mcm::MonteCarloMethod, rng::AbstractRNG=default_rng())
+  samples, logweights = weighted_samples(mcm, dist, rng)
+  function 𝔼(f)
+    # Dispatch on whether f is multivariate.
+    function calc_expectation(::Type{<:AbstractArray{T}}) where T
+      weighted = Matrix{T}(undef, length(y), length(samples))
+      for j in eachindex(samples)
+        weighted[:, j] = log.(samples[j]) .+ logweights[j]
+      end
+      exp(logsumexp(weighted, dims=2) .- logsumexp(logweights))
+    end
+    function calc_expectation(::Type{<:Any})
+      weighted = map(samples, logweights) do x, logweight
+        log(f(x)) + logweight
+      end
+      exp(logsumexp(weighted) - logsumexp(logweights))
+    end
+    calc_expectation(typeof(f(first(samples))))
+  end
+end
+
+struct LikelihoodWeighting <: MonteCarloMethod
+  num_samples :: Int
+end
+
+function weighted_samples(mcm::LikelihoodWeighting, dist, rng) 
+  likelihood_weighting(dist, mcm.num_samples, rng)
+end
+
+function likelihood_weighting(model::Conditioned{<:ProbProg}, num_samples, rng::AbstractRNG=default_rng())
+  traces = rand(rng, iid(model, num_samples))
+  logweights = logdensity.((model,), traces)
+  traces, logweights
+end
+
 ##############################################
 ### Interpreters of probabilistic programs ###
 ##############################################
@@ -495,22 +564,25 @@ function sample(i::EvalTrace, dist, lens)
 end
 
 """
-    RandTrace([rng])
+    RandTrace([rng], [conds])
 
 Interpreter that runs a probabilistic program and traces the random choices
 made in the sample statements in a named tuple.
 """
-struct RandTrace{R,T} <: Interpreter
+struct RandTrace{R<:AbstractRNG,T<:NamedTuple} <: Interpreter
   rng::R
   trace::T
 end
 
 RandTrace(rng::AbstractRNG) = RandTrace(rng, (;))
-RandTrace() = RandTrace(default_rng())
 
-function sample(i::RandTrace, dist, lens)
-  x = rand(i.rng, dist)
-  return RandTrace(i.rng, insert(i.trace, lens, x)), x
+function sample(i::RandTrace, dist, lens::PropertyLens{name}) where name
+  if hasproperty(i.trace, name)
+    i, lens(i.trace)
+  else
+    x = rand(i.rng, dist)
+    RandTrace(i.rng, insert(i.trace, lens, x)), x
+  end
 end
 
 """

test/runtests.jl

@@ -1,7 +1,9 @@
 using SimpleProbabilisticPrograms
 using Test
-using Random
+
+using Random: MersenneTwister
 using Distributions: Beta, Bernoulli
+using Statistics: mean
 
 @testset "basic tests" begin
   @probprog function beta_bernoulli_model(a, b, n)
@@ -10,8 +12,8 @@ using Distributions: Beta, Bernoulli
     return (; bias, coins)
   end
   model = beta_bernoulli_model(3, 4, 10)
-  trace = rand(Random.MersenneTwister(42), model)
-  trace_static = rand(Random.MersenneTwister(42), beta_bernoulli_model(3, 4, Val(10)))
+  trace = rand(MersenneTwister(42), model)
+  trace_static = rand(MersenneTwister(42), beta_bernoulli_model(3, 4, Val(10)))
   @test trace_static == trace
   @test -Inf < logpdf(model, trace) < 0
   @test insupport(model, trace)
@@ -109,12 +111,43 @@ end
   @test rand(cond('b')) in 10:15
 end
 
-using Distributions: Beta, Bernoulli
-@probprog function beta_bernoulli_model(a, b, n)
-  bias ~ Beta(a, b)
-  coins ~ iid(Bernoulli(bias), n)
-  return (; bias, coins)
+@testset "Monte Carlo" begin
+  @probprog function bbm(a, b, n) # beta bernoulli model
+    bias ~ Beta(a, b)
+    coins ~ iid(Bernoulli(bias), n)
+    return (; bias, coins)
+  end
+
+  N = 1000
+  model = bbm(1, 1, N)
+  data = rand(iid(Bernoulli(0.4), N))
+  E = montecarlo(condition(model, on=(; coins=data)), LikelihoodWeighting(10_000)) 
+  @test E(trace -> trace.bias) ≈ mean(data) atol=0.01
+
+  # TODO: Test montecarlo with multivariate function
 end
-@time model = beta_bernoulli_model(3, 4, 1000)
-@time trace = rand(model)
-@time logpdf(model, trace)
+
+
+# t1 = (a=1, b=(c=2, d=3, e=4))
+# t2 = (b=(c=5, d=6), f=7)
+# @test mergetraces(t1, t2) == (a=1, b=(c=5, d=6, e=4), f=7)
+
+# T1 = typeof(t1)
+# T2 = typeof(t2)
+# fieldnames(T1)
+# fieldtypes(T1)
+# fieldnames(T2)
+# fieldtypes(T2)
+
+# all_names = union(fieldnames(T1), fieldnames(T2))
+# map(all_names) do n
+#   if n in fieldnames(T1)
+#     if n in fieldnames(T2)
+#       if 
+#     else
+#       :($n = t1.$n)
+#     end
+#   else
+#     :($n = t2.$n)
+#   end
+# end