feat!: port to genjl style interfaces (#60)

examples/editable.clj

@@ -5,12 +5,11 @@
   `intro-to-modeling`, and also changed the `pmap` call in `prepeatedly` into
   `map`, since we don't have `pmap` available in the browser."
   {:nextjournal.clerk/toc true}
-  (:require [gen.choice-map]
-            [gen.dynamic :as dynamic :refer [gen]]
-            [gen.dynamic.choice-map :refer [choice-map]]
+  (:require [gen.choicemap :refer [choicemap]]
             [gen.clerk.callout :as callout]
             [gen.clerk.viewer :as viewer]
             [gen.distribution.kixi :as dist]
+            [gen.dynamic :as dynamic :refer [gen]]
             [gen.generative-function :as gf]
             [gen.inference.importance :as importance]
             [gen.trace :as trace]
@@ -586,7 +585,7 @@ math/PI
   ;; them to be inferred.
   (let [observations (reduce (fn [observations [i y]]
                                (assoc observations [:y i] y))
-                             (choice-map {})
+                             (choicemap {})
                              (map-indexed vector ys))]
     (:trace (importance/resampling model [xs] observations amount-of-computation))))
 
@@ -652,7 +651,7 @@ math/PI
 
 ;; For example:
 
-(def predicting-constraints (choice-map {:slope 0 :intercept 0}))
+(def predicting-constraints (choicemap {:slope 0 :intercept 0}))
 (def predicting-trace (:trace (gf/generate line-model [xs] predicting-constraints)))
 
 (def predict-opts
@@ -684,7 +683,7 @@ math/PI
   ;; constraints.
   (let [constraints (reduce (fn [cm param-addr]
                               (assoc cm param-addr (get trace param-addr)))
-                            (choice-map {})
+                            (choicemap {})
                             param-addrs)
 
         ;; Run the model with new x coordinates, and with parameters
@@ -944,7 +943,7 @@ math/PI
 ;; collisions for complex models.
 
 ;; Hierarchical traces are represented using nested choice maps
-;; (`gen.dynamic.choice-map/ChoiceMap`). Hierarchical addresses can be accessed
+;; (`gen.dynamic.choicemap/ChoiceMap`). Hierarchical addresses can be accessed
 ;; using `clojure.core` functions like `clojure.core/get-in`.
 
 (get-in bar-with-key-trace [:z :y])

examples/intro_to_modeling.clj

@@ -1,12 +1,11 @@
 ^{:nextjournal.clerk/visibility {:code :hide :result :hide}}
 (ns intro-to-modeling
   {:nextjournal.clerk/toc true}
-  (:require [gen.choice-map]
-            [gen.dynamic :as dynamic :refer [gen]]
-            [gen.dynamic.choice-map :refer [choice-map]]
+  (:require [gen.choicemap :as choicemap :refer [choicemap]]
             [gen.clerk.callout :as callout]
             [gen.clerk.viewer :as viewer]
             [gen.distribution.kixi :as dist]
+            [gen.dynamic :as dynamic :refer [gen]]
             [gen.generative-function :as gf]
             [gen.inference.importance :as importance]
             [gen.trace :as trace]
@@ -51,7 +50,7 @@
 ;; - **Modeling**: You first need to frame the problem — and any assumptions you
 ;;   bring to the table — as a probabilistic model. A huge variety of problems
 ;;   can be viewed from a modeling & inference lens, if you set them up
-;;   properly.  **This notebook is about how to think of problems in this light,
+;;   properly. **This notebook is about how to think of problems in this light,
 ;;   and how to use Gen** **to formally specify your assumptions and the tasks
 ;;   you wish to solve.**
 
@@ -61,7 +60,7 @@
 ;;   proposal distributions. With enough computation, the algorithm can in
 ;;   theory solve any modeling and inference problem, but in practice, for most
 ;;   problems of interest, it is too slow to achieve accurate results in a
-;;   reasonable amount of time.  **Future tutorials introduce some of Gen's**
+;;   reasonable amount of time. **Future tutorials introduce some of Gen's**
 ;;   **programmable inference features**, which let you tailor the inference
 ;;   algorithm for use with more complex models (Gen will still automate the
 ;;   math!).
@@ -222,32 +221,27 @@
 ^{::clerk/visibility {:result :hide}}
 (def line-model
   (gen [xs]
-
-    ;; We begin by sampling a slope and intercept for the line.  Before we have
+    ;; We begin by sampling a slope and intercept for the line. Before we have
     ;; seen the data, we don't know the values of these parameters, so we treat
-    ;; them as random choices. The distributions they are drawn from represent our
-    ;; prior beliefs about the parameters: in this case, that neither the slope
-    ;; nor the intercept will be more than a couple points away from 0.
-
+    ;; them as random choices. The distributions they are drawn from represent
+    ;; our prior beliefs about the parameters: in this case, that neither the
+    ;; slope nor the intercept will be more than a couple points away from 0.
     (let [slope (dynamic/trace! :slope dist/normal 0 1)
           intercept (dynamic/trace! :intercept dist/normal 0 2)
 
           ;; We define a function to compute y for a given x.
-
           y (fn [x]
               (+ (* slope x)
                  intercept))]
 
       ;; Given the slope and intercept, we can sample y coordinates for each of
       ;; the x coordinates in our input vector.
-
       (doseq [[i x] (map vector (range) xs)]
         (dynamic/trace! [:y i] dist/normal (y x) 0.1))
 
-      ;; Most of the time, we don't care about the return
-      ;; value of a model, only the random choices it makes.
-      ;; It can sometimems be useful to return something
-      ;; meaningful, however; here, we return the function `y`.
+      ;; Most of the time, we don't care about the return value of a model, only
+      ;; the random choices it makes. It can sometimems be useful to return
+      ;; something meaningful, however; here, we return the function `y`.
       y)))
 
 ;; The generative function takes as an argument a vector of x-coordinates. We
@@ -304,7 +298,6 @@
 ;; (require '[gen.trace :as trace])
 ;; ```
 
-
 (trace/get-args trace)
 
 ;; The trace also contains the value of the random choices, stored in a map from
@@ -326,15 +319,6 @@
 
 (:slope (trace/get-choices trace))
 
-;; We can also read the value of a random choice directly from the trace,
-;; without having to use `gen.trace/get-choices` first:
-
-(get trace :slope)
-
-(trace :slope)
-
-(:slope trace)
-
 ;; The return value is also recorded in the trace, and is accessible with the
 ;; `trace/get-retval` API method:
 
@@ -352,8 +336,9 @@
   [trace & {:keys [clip x-domain y-domain] :or {clip false}}]
   (let [[xs] (trace/get-args trace) ; Pull out the xs from the trace.
         y (trace/get-retval trace) ; Pull out the return value, useful for plotting.
+        choices (trace/get-choices trace)
         ys (for [i (range (count xs))]
-             (trace [:y i]))
+             (choicemap/get-value choices [:y i]))
         data (mapv (fn [x y]
                      {:x x :y y})
                    xs
@@ -472,8 +457,9 @@ math/PI
         min-x (apply min xs)
         max-x (apply max xs)
         y (trace/get-retval trace) ; Pull out the return value, useful for plotting.
+        choices (trace/get-choices trace)
         ys (for [i (range (count xs))]
-             (get trace [:y i]))
+             (choicemap/get-value choices [:y i]))
         data (mapv (fn [x y]
                      {:x x :y y})
                    xs
@@ -585,7 +571,7 @@ math/PI
 
 ;; ```clojure
 ;; (require '[gen.inference.importance :as importance]
-;;          '[gen.dynamic.choice-map :refer [choice-map]])
+;;          '[gen.choicemap :as choicemap :refer [choicemap]])
 ;; ```
 
 (defn do-inference
@@ -595,7 +581,7 @@ math/PI
   ;; them to be inferred.
   (let [observations (reduce (fn [observations [i y]]
                                (assoc observations [:y i] y))
-                             (choice-map {})
+                             (choicemap {})
                              (map-indexed vector ys))]
     (:trace (importance/resampling model [xs] observations amount-of-computation))))
 
@@ -661,7 +647,7 @@ math/PI
 
 ;; For example:
 
-(def predicting-constraints (choice-map {:slope 0 :intercept 0}))
+(def predicting-constraints {:slope 0 :intercept 0})
 (def predicting-trace (:trace (gf/generate line-model [xs] predicting-constraints)))
 
 (def predict-opts
@@ -691,17 +677,19 @@ math/PI
   [model trace new-xs param-addrs]
   ;; Copy parameter values from the inferred trace (`trace`) into a fresh set of
   ;; constraints.
-  (let [constraints (reduce (fn [cm param-addr]
-                              (assoc cm param-addr (get trace param-addr)))
-                            (choice-map {})
+  (let [prev-choices (trace/get-choices trace)
+        constraints (reduce (fn [cm param-addr]
+                              (assoc cm param-addr
+                                     (choicemap/get-value prev-choices param-addr)))
+                            (choicemap)
                             param-addrs)
-
         ;; Run the model with new x coordinates, and with parameters
         ;; fixed to be the inferred values.
-        {new-trace :trace} (gf/generate model [new-xs] constraints)]
+        {new-trace :trace} (gf/generate model [new-xs] constraints)
+        choices (trace/get-choices new-trace)]
 
     ;; Pull out the y-values and return them.
-    (mapv #(get new-trace [:y %])
+    (mapv #(choicemap/get-value choices [:y %])
           (range (count new-xs)))))
 
 ;; To illustrate, we call the function above given the previous trace (which
@@ -953,8 +941,8 @@ math/PI
 ;; collisions for complex models.
 
 ;; Hierarchical traces are represented using nested choice maps
-;; (`gen.dynamic.choice-map/ChoiceMap`). Hierarchical addresses can be accessed
-;; using `clojure.core` functions like `clojure.core/get-in`.
+;; (`gen.choicemap/IChoiceMap`). Hierarchical addresses can be accessed using
+;; `clojure.core` functions like `clojure.core/get-in`.
 
 (get-in bar-with-key-trace [:z :y])
 
@@ -1209,15 +1197,16 @@ math/PI
 
 (defn render-changepoint-model-trace
   [trace]
-  (let [[xs] (trace/get-args trace)
-        ys (for [i (range (count xs))]
-             (trace [:y i]))
-        node (trace/get-retval trace)
+  (let [[xs]       (trace/get-args trace)
+        choices    (trace/get-choices trace)
+        ys         (for [i (range (count xs))]
+                     (choicemap/get-value choices [:y i]))
+        node       (trace/get-retval trace)
         node-layer (node-vl-spec node)
         data-layer (scatter-spec xs ys :color "grey" :fillOpacity 0.3 :strokeOpacity 1.0)]
-    (clerk/vl {:schema "https://vega.github.io/schema/vega-lite/v5.json"
+    (clerk/vl {:schema     "https://vega.github.io/schema/vega-lite/v5.json"
                :embed/opts {:actions false}
-               :layer [node-layer data-layer]})))
+               :layer      [node-layer data-layer]})))
 
 {::clerk/visibility {:result :show}}
 

examples/introduction.clj

@@ -4,6 +4,7 @@
             [clojure.repl :as repl]
             [gen.distribution.commons-math :as dist]
             [gen.dynamic :as dynamic :refer [gen]]
+            [gen.choicemap :as choicemap]
             [gen.generative-function :as gf]
             [gen.trace :as trace]
             [nextjournal.clerk :as clerk]))
@@ -298,12 +299,14 @@
   (let [data (apply concat
                     (for [p ps]
                       (->> (repeatedly #(trace/get-choices (gf/simulate gen-f [p])))
-                           (filter (fn [trace]
-                                     (= observed-fp (get trace :fp))))
-                           (take 1000)
-                           (mapv (fn [trace]
+                           (filter (fn [choices]
+                                     (= observed-fp
+                                        (choicemap/get-value choices :fp))))
+                           (take 100)
+                           (mapv (fn [choices]
                                    {:p p
-                                    :if-test (get trace :if-test)})))))]
+                                    :if-test
+                                    (choicemap/get-value choices :if-test)})))))]
     (clerk/vl {:schema "https://vega.github.io/schema/vega-lite/v5.json"
                :embed/opts {:actions false}
                :data {:values data}
@@ -437,30 +440,9 @@
 ;; where `:a` is always true and `:c` is always false. We first construct a
 ;; choice map containing these constraints:
 
-(require '[gen.dynamic.choice-map :as dynamic.choice-map]
-         '[gen.choice-map :as choice-map])
-
 (def constraints
-  (dynamic.choice-map/choice-map
-   :a true
-   :c false))
-
-#_
-(choice-map/submaps
- (dynamic.choice-map/choice-map
-  :a true
-  :c false))
-
-;; The `gen.dynamic.choice-map/choice-map` constructor above took two elements
-;; of the form (address, value). This is equivalent to constructing an empty
-;; choice map and then populating it:
-
-(def choices
-  (assoc (dynamic.choice-map/choice-map)
-         :a true
-         :c false))
-
-(choice-map/submaps choices)
+  {:a true
+   :c false})
 
 ;; Then, we pass the constraints as the third argument to
 ;; `gen.generative-function/generate`, after the function itself and the
@@ -613,7 +595,7 @@
 (def samples
   (let [n-particles [1 10 100]]
     (zipmap n-particles
-            (mapv #(draw-samples % #gen/choice-map {:c false})
+            (mapv #(draw-samples % {:c false})
                   n-particles))))
 
 (clerk/vl {:schema "https://vega.github.io/schema/vega-lite/v5.json"
@@ -696,13 +678,13 @@
 
 ;; Consider the function `foo` from above. Let's obtain an initial trace:
 
-(def update-trace (:trace (gf/generate foo [0.3] #gen/choice-map {:a true :b true :c true})))
+(def update-trace (:trace (gf/generate foo [0.3] {:a true :b true :c true})))
 (trace/get-choices update-trace)
 
 ;; Now, we use the `update` function, to change the value of `:c` from `true` to
 ;; `false`:
 
-(def updated (trace/update update-trace #gen/choice-map {:c false}))
+(def updated (trace/update update-trace {:c false}))
 (trace/get-choices (:trace updated))
 
 ;; The `update` function returns the new trace, as well as a weight, which the
@@ -724,9 +706,9 @@
 
 ;; Doing an update can also cause some addresses to leave the choice map
 ;; altogether. For example, if we set `:a` to `false`, then choice at address
-;; `:b` is no longer include in the choice map.
+;; `:b` is no longer included in the choice map.
 
-(def update-a-true (trace/update update-trace #gen/choice-map {:a false}))
+(def update-a-true (trace/update update-trace {:a false}))
 (trace/get-choices (:trace update-a-true))
 
 ;; The *discard* choice map that is returned by `update` contains the valus for

src/data_readers.cljc

@@ -1,2 +1,2 @@
-{gen/choice gen.dynamic.choice-map/parse-choice
- gen/choice-map gen.dynamic.choice-map/parse-choice-map}
+{gen/choice gen.choicemap/parse-choice
+ gen/choicemap gen.choicemap/parse-choicemap}

src/gen/choicemap.cljc

@@ -670,7 +670,7 @@
 
 ;; ### ChoiceMap interactions
 
-(defn- equiv
+(defn ^:no-doc equiv
   "Returns true if `r` is a choicemap with equivalent submaps to `l`, false
   otherwise.