feat(inference): add nuts sampler

Author: chentoast

src/inference/hmc.jl

@@ -1,15 +1,3 @@
-function sample_momenta(n::Int)
-    Float64[random(normal, 0, 1) for _=1:n]
-end
-
-function assess_momenta(momenta)
-    logprob = 0.
-    for val in momenta
-        logprob += logpdf(normal, val, 0, 1)
-    end
-    logprob
-end
-
 """
     (new_trace, accepted) = hmc(
         trace, selection::Selection; L=10, eps=0.1,

src/inference/hmc_common.jl

@@ -0,0 +1,11 @@
+function sample_momenta(n::Int)
+  Float64[random(normal, 0, 1) for _=1:n]
+end
+
+function assess_momenta(momenta)
+  logprob = 0.
+  for val in momenta
+      logprob += logpdf(normal, val, 0, 1)
+  end
+  logprob
+end

src/inference/inference.jl

@@ -14,10 +14,13 @@ export logsumexp
 
 include("trace_translators.jl")
 
+include("hmc_common.jl")
+
 # mcmc
 include("kernel_dsl.jl")
 include("mh.jl")
 include("hmc.jl")
+include("nuts.jl")
 include("mala.jl")
 include("elliptical_slice.jl")
 

src/inference/nuts.jl

@@ -0,0 +1,158 @@
+using LinearAlgebra: dot
+
+Tree = @NamedTuple begin
+  val_left
+  momenta_left
+  val_right
+  momenta_right
+  val_sample
+  n :: Int
+  weight :: Float64
+  stop :: Bool
+end
+
+Stats = @NamedTuple begin
+  depth
+  n
+  accept
+end
+
+function u_turn(values_left, values_right, momenta_left, momenta_right)
+  return (dot(values_left - values_right, momenta_right) >= 0) &&
+  (dot(values_right - values_left, momenta_left) >= 0)
+end
+
+function leapfrog(values, momenta, eps, integrator_state)
+  values_trie, selection, retval_grad, trace = integrator_state
+
+  values_trie = from_array(values_trie, values)
+  (trace, _, _) = update(trace, values_trie)
+  (_, _, gradient_trie) = choice_gradients(trace, selection, retval_grad)
+  gradient = to_array(gradient_trie, Float64)
+
+  # half step on momenta
+  momenta += (eps / 2) * gradient
+
+  # full step on positions
+  values += eps * momenta
+
+  # get new gradient
+  values_trie = from_array(values_trie, values)
+  (trace, _, _) = update(trace, values_trie)
+  (_, _, gradient_trie) = choice_gradients(trace, selection, retval_grad)
+  gradient = to_array(gradient_trie, Float64)
+
+  # half step on momenta
+  momenta += (eps / 2) * gradient
+  return values, momenta, get_score(trace)
+end
+
+function build_root(val, momenta, eps, direction, integrator_state)
+  val, momenta, lp = leapfrog(val, momenta, direction * eps, integrator_state)
+  weight = lp + assess_momenta(momenta)
+
+  return Tree((val, momenta, val, momenta, val, 1, weight, false))
+end
+
+function merge_trees(tree_left, tree_right)
+  # multinomial sampling
+  if log(rand()) < tree_right.weight - tree_left.weight
+    sample = tree_right.val_sample
+  else
+    sample = tree_left.val_sample
+  end
+
+  weight = logsumexp(tree_left.weight, tree_right.weight)
+  n = tree_left.n + tree_right.n
+  stop = tree_left.stop || tree_right.stop || u_turn(
+    tree_left.val_left, tree_right.val_right, tree_left.momenta_left, tree_right.momenta_right
+  )
+
+  return Tree((tree_left.val_left, tree_left.momenta_left, tree_right.val_right,
+    tree_right.momenta_right, sample, n, weight, stop))
+end
+
+function build_tree(val, momenta, depth, eps, direction, integrator_state)
+  if depth == 0
+    return build_root(val, momenta, eps, direction, integrator_state)
+  end
+
+  tree = build_tree(val, momenta, depth - 1, eps, direction, integrator_state)
+
+  if tree.stop
+    return tree
+  end
+
+  if direction == 1
+    other_tree = build_tree(tree.val_right, tree.momenta_right, depth - 1, eps, direction, integrator_state)
+  else
+    other_tree = build_tree(tree.val_left, tree.momenta_left, depth - 1, eps, direction, integrator_state)
+  end
+
+  if direction == 1
+    return merge_trees(tree, other_tree)
+  else
+    return merge_trees(other_tree, tree)
+  end
+end
+
+function nuts(
+  trace::Trace, selection::Selection; eps=0.1, max_treedepth=15,
+  check=false, observations=EmptyChoiceMap())
+    prev_model_score = get_score(trace)
+    retval_grad = accepts_output_grad(get_gen_fn(trace)) ? zero(get_retval(trace)) : nothing
+
+    # values needed for a leapfrog step
+    (_, values_trie, _) = choice_gradients(trace, selection, retval_grad)
+    values = to_array(values_trie, Float64)
+    integrator_state = (values_trie, selection, retval_grad, trace)
+
+    momenta = sample_momenta(length(values))
+    prev_momenta_score = assess_momenta(momenta)
+
+    tree = Tree((values, momenta, values, momenta, values, 1, -Inf, false))
+
+    direction = 0
+    depth = 0
+    stop = false
+    while depth < max_treedepth
+      direction = rand([-1, 1])
+
+      if direction == 1 # going right
+        other_tree = build_tree(tree.val_right, tree.momenta_right, depth, eps, direction, integrator_state)
+        tree = merge_trees(tree, other_tree)
+      else # going left
+        other_tree = build_tree(tree.val_left, tree.momenta_left, depth, eps, direction, integrator_state)
+        tree = merge_trees(other_tree, tree)
+      end
+
+      stop = stop || tree.stop
+      if stop
+        break
+      end
+      depth += 1
+    end
+
+    (new_trace, _, _) = update(trace, from_array(values_trie, tree.val_sample))
+    check && check_observations(get_choices(new_trace), observations)
+
+    # assess new model score (negative potential energy)
+    new_model_score = get_score(new_trace)
+
+    # assess new momenta score (negative kinetic energy)
+    if direction == 1
+      new_momenta_score = assess_momenta(-tree.momenta_right)
+    else
+      new_momenta_score = assess_momenta(-tree.momenta_left)
+    end
+
+    # accept or reject
+    alpha = new_model_score - prev_model_score + new_momenta_score - prev_momenta_score
+    if log(rand()) < alpha
+      return (new_trace, Stats((depth, tree.n, true)))
+    else
+      return (trace, Stats((depth, tree.n, false)))
+    end
+end
+
+export nuts