feat: clean up importance sampling code

Author: sritchie

src/gen/inference/importance.cljc

@@ -1,39 +1,36 @@
 (ns gen.inference.importance
-  (:require [clojure.math :as math]
-            [gen.distribution.kixi :as dist]
+  (:require [gen.distribution.kixi :as dist]
             [gen.generative-function :as gf]))
 
 ;; This implementation comes from `fastmath.core`, ported here for cljc
 ;; purposes.
 
-(defn- logsumexp
-  "log(exp(x1)+...+exp(xn))."
-  ^double [xs]
-  (loop [[^double x & rst] xs
-         r 0.0
-         alpha ##-Inf]
-    (if (<= x alpha)
-      (let [nr (+ r (math/exp (- x alpha)))]
-        (if-not (seq rst)
-          (+ (math/log nr) alpha)
-          (recur rst nr alpha)))
-      (let [nr (inc (* r (math/exp (- alpha x))))]
-        (if-not (seq rst)
-          (+ (math/log nr) x)
-          (recur rst nr (double x)))))))
+(defn ^:nodoc logsumexp
+  "log(exp(x1)+...+exp(xn)). This version can handle infinities."
+  ([x] x)
+  ([x y]
+   (if (<= y x)
+     (+ x (Math/log (inc (Math/exp (- y x)))))
+     (+ y (Math/log (inc (Math/exp (- x y)))))))
+  ([x y & more]
+   (reduce logsumexp (logsumexp x y) more)))
 
-(defn resampling [gf args observations n-samples]
-  ;; https://github.com/probcomp/Gen.jl/blob/master/src/inference/importance.jl#L77...L95
-  (let [result (gf/generate gf args observations)
-        model-trace (volatile! (:trace result))
-        log-total-weight (volatile! (:weight result))]
-    (dotimes [_ (dec n-samples)]
-      (let [candidate (gf/generate gf args observations)
-            candidate-model-trace (:trace candidate)
-            log-weight (:weight candidate)]
-        (vswap! log-total-weight #(logsumexp [log-weight %]))
-        (when (dist/bernoulli (math/exp (- log-weight @log-total-weight)))
-          (vreset! model-trace candidate-model-trace))))
-    (let [log-ml-estimate (- @log-total-weight (math/log n-samples))]
-      {:trace @model-trace
-       :weight log-ml-estimate})))
+(defn resampling
+  "https://github.com/probcomp/Gen.jl/blob/master/src/inference/importance.jl#L77...L95"
+  [gf args observations n-samples]
+  (let [result (gf/generate gf args observations)]
+    (loop [i                (dec n-samples)
+           model-trace      (:trace result)
+           log-total-weight (:weight result)]
+      (if (zero? i)
+        (let [log-ml-estimate (- log-total-weight (Math/log n-samples))]
+          {:trace  model-trace
+           :weight log-ml-estimate})
+        (let [candidate       (gf/generate gf args observations)
+              candidate-trace (:trace candidate)
+              log-weight      (:weight candidate)
+              new-weight      (logsumexp log-total-weight log-weight)
+              new-trace       (if (dist/bernoulli (Math/exp (- log-weight new-weight)))
+                                candidate-trace
+                                model-trace)]
+          (recur (unchecked-dec-int i) new-trace new-weight))))))