Stochastic search trees

gistfile1.scm

;; Stochastic search trees (from Kiselyov and Shan's HANSEI)

(import (delimcc-simple-ikarus)
        (printing)
        (_srfi :1))

;; the core: stochastic search trees
(define (pv-unit v) (list (list 1.0 `(V ,v))))

(define (dist pvs)
  (shift k (map (lambda (pv) (list (first pv) 
                                   ;; choice branch
                                   ;; label with value
                                   ;; the thunk: rest of the computation after this choice
                                   `(C ,(second pv) ,(lambda () (k (second pv)))))) pvs)))

(define (reify0 thunk)
  (reset (pv-unit (thunk))))

(define (fail) (dist '()))

;; Convenience functions

(define (flip p) (dist `((,p #t) (,(- 1.0 p) #f))))

(define (value? pv-e) (or (null? pv-e) (eq? 'V (car (cadr pv-e)))))
(define (thunk? pv-e) (and (eq? 'C (car (cadr pv-e)))
                           (procedure? (pv-e->val pv-e))))
(define (branch? pv-e) (and (eq? 'C (car (cadr pv-e)))
                           (list? (pv-e->val pv-e))))

(define (pv-e->val pv-e) (caddr (cadr pv-e)))
(define (pv-e->label pv-e) (cadr (cadr pv-e)))

(define (pv-e->prob pv-e) (car pv-e))
(define (PV p v) (list p v))
(define (V v) `(V ,v))
(define (C label c) `(C ,label ,c))

;; Inefficient schemes for unfolding the tree to variable depths
(define (extend-once pv)
  (define (transform pv-e)
    (cond [(value? pv-e) pv-e]
          [(thunk? pv-e) (PV (pv-e->prob pv-e) (C (pv-e->label pv-e)
                                                  ((pv-e->val pv-e))))]
          [(branch? pv-e) (cond [(null? (pv-e->val pv-e)) '()]
                                [else (PV (pv-e->prob pv-e) 
                                          (C (pv-e->label pv-e) (map transform (pv-e->val pv-e))))])]))
  (map transform pv))

(define (iterate n f x)
  (define (loop n res)
    (cond [(= 0 n) res]
          [else (loop (- n 1) (f res))]))
  (loop n x))

(define (iterate-fix f x)
  (define (loop res)
    (let* ([next (f res)])
      (cond [(equal? next res) res]
            [else (loop next)])))
  (loop x))

(define (explore pv . depth)
  (cond [(null? depth) (iterate-fix extend-once pv)]
        [else (iterate (car depth) extend-once pv)]))

;; Demo (from the website and paper)

;; P(rain | grass wet)
(define (grass-model)
  (let* ([rain (flip 0.3)]
         [sprinkler (flip 0.5)]
         [grass_is_wet (or (and (flip 0.9) rain)
                           (and (flip 0.9) sprinkler))])
         (if grass_is_wet rain (fail))))

(pretty-print (explore (reify0 grass-model)))