<commit_before>;; Copyright (c) 2014 <NAME>
;; Released under the terms of the 2-clause BSD license.
;; See the file COPYRIGHT for details.

#lang racket/base
(require (rename-in racket/match [match-define defmatch])
         racket/class
         "../dist.rkt"
         "prob-lazy-tree.rkt"
         "pairingheap.rkt"
         "interfaces.rkt")
(provide enumerate*
         enum-importance-sampler*)

;; == Overview ==
;;
;; Enumeration using lazy tree of probability-labeled possibilities,
;; based on delimited continuations (only for discrete ERPs).
;;
;; For enumeration to work:
;;
;;  - All paths must either terminate or evaluate infinitely many ERPS.
;;    (and every ERP must generate at least two values with nonzero prob.)
;;
;;  - The predicate must accept a nonempty set of paths.
;;
;; Without knowing structure of paths and conditioning predicate, can't
;; make smart distinctions between incomplete paths.
;;
;; Would be nice if we could tell whether a path was viable or not wrt
;; condition. Seems like it would require drastic changes to model of
;; computation (eg, like symbolic execution) to support non-trivial
;; conditions. Would that be a profitable place to spend effort?

;; ============================================================

;; enumerate* : (-> A) etc -> (Listof (List A Prob))
(define (enumerate* thunk limit normalize?)
  (vprintf "Enumerating\n")
  (define tree (reify-tree thunk))
  (define table (explore tree limit))
  (when (and normalize? (zero? (hash-count table)))
    (error 'enumerate "condition accepted no paths"))
  (make-discrete-dist table #:normalize? normalize?))

;; ----------------------------------------

;; https://hips.seas.harvard.edu/blog/2013/01/09/computing-log-sum-exp/

;; logspace+ : Real Real -> Real
;; Given (log A) and (log B), return (log (+ A B)).
(define logspace+*
  (case-lambda
    [() -inf.0]
    [(a) a]
    [(a b)
     (define m (max a b))
     (if (= m -inf.0)
         -inf.0
         (+ m (log (+ (exp (- a m)) (exp (- b m))))))]
    [as
     (define m (apply max as))
     (if (= m -inf.0)
         -inf.0
         (+ m (log (for/sum ([a (in-list as)]) (exp (- a m))))))]))

(define (logspace+ a b)
  (define r (logspace+* a b))
  (when (eqv? r +nan.0)
    (error 'logspace+ "~s + ~s => ~s" a b r))
  r)

(define (logspace-compl a)
  (log (- 1.0 (exp a))))

;; An ExpST is (expst Prob Prob Hash[A => Prob] Nat)
;; where Prob is Real or Real+ depending on logspace?,
;; and dim is density dimension of every entry in table.
(struct expst (explored accepted table table-ddim) #:transparent)

(define (expst->logspace st)
  (defmatch (expst explored accepted table table-ddim) st)
  (expst (log (exact->inexact explored)) (log (exact->inexact accepted))
         (table->logspace table) table-ddim))

(define (table->logspace table)
  (for/hash ([(k v) (in-hash table)])
    (values k (log (exact->inexact v)))))

;; A (HeapEntry A) is (list* Prob Nat (-> (EnumTree A))).
(define (entry->=? x y) (>= (car x) (car y)))

(define (heap->logspace h)
  (for/fold ([new-h (heap entry->=?)])
            ([entry (in-list (heap->list h))])
    (heap-insert new-h (cons (log (car entry)) (cdr entry)))))

;; explore : (EnumTree A) (U #f Real>0) -> Hash[A => Real+]
;; * Starts with logspace? = #f, but switches to logspace if an inexact
;;   density is ever seen.
;; * Limit means: explore until potential accepted dist is < limit unaccounted for,
;;   ie, unexplored < limit * accepted. Thus, limit is nonlocal; use BFS.
;;   Limit is always realspace, not logspace. If no limit, prob-{explored,accepted}
;;   can be meaningless. Note: Cannot observe continuous dist if limit given.
(define (explore tree limit)
  (define initial-st (expst 0 0 '#hash() +inf.0))
  (define initial-heap
    ;; lib provides "min-heap", but want max prob, so use >= comparison
    (heap entry->=?))
  (define initial-ddim 0)
  (define initial-p 1)
  (explore* limit #f initial-heap initial-st tree initial-ddim initial-p))

;; explore* : (EnumTree A) (U #f Real>0) ExpSt Heap Nat -> Hash[A => Real+]
;; - st   : ExpST, rep of Prob depends on logspace?
;; - h    : heap[(cons Prob (-> (EnumTree A)))]
;; - ddim : Nat, density dimension
(define (explore* limit logspace? h st tree tree-ddim tree-p)
  ;; type Prob = Real
  ;; The representation of probabilities in this code is determined by
  ;; the logspace? argument.  If logspace? is true, then values
  ;; represent log of prob; if false, values represent prob directly.
  (define (->p a) ;; inexact to avoid error on exact 0
    (when (negative? a) (error '->p "negative: ~e" a))
    (if logspace? (log (exact->inexact a)) a))
  (define (p->real a)
    (if logspace? (exp a) a))
  (define (p+ a b)
    (if logspace? (logspace+ a b) (+ a b)))
  (define (pcompl a)
    (if logspace? (logspace-compl a) (- 1 a)))
  (define (p* a b)
    (if logspace? (+ a b) (* a b)))
  (define pzero (->p 0))

  (define (switch-to-logspace h st tree tree-ddim tree-p)
    (explore* limit #t (heap->logspace h) (expst->logspace st) tree tree-ddim (log tree-p)))<commit_msg>fix enumerate again wrt density dimension<commit_after>;; Copyright (c) 2014 <NAME>
;; Released under the terms of the 2-clause BSD license.
;; See the file COPYRIGHT for details.

#lang racket/base
(require (rename-in racket/match [match-define defmatch])
         racket/class
         "../dist.rkt"
         "prob-lazy-tree.rkt"
         "pairingheap.rkt"
         "interfaces.rkt")
(provide enumerate*
         enum-importance-sampler*)

;; == Overview ==
;;
;; Enumeration using lazy tree of probability-labeled possibilities,
;; based on delimited continuations (only for discrete ERPs).
;;
;; For enumeration to work:
;;
;;  - All paths must either terminate or evaluate infinitely many ERPS.
;;    (and every ERP must generate at least two values with nonzero prob.)
;;
;;  - The predicate must accept a nonempty set of paths.
;;
;; Without knowing structure of paths and conditioning predicate, can't
;; make smart distinctions between incomplete paths.
;;
;; Would be nice if we could tell whether a path was viable or not wrt
;; condition. Seems like it would require drastic changes to model of
;; computation (eg, like symbolic execution) to support non-trivial
;; conditions. Would that be a profitable place to spend effort?

;; ============================================================

;; enumerate* : (-> A) etc -> (Listof (List A Prob))
(define (enumerate* thunk limit normalize?)
  (vprintf "Enumerating\n")
  (define tree (reify-tree thunk))
  (define table (explore tree limit))
  (when (and normalize? (zero? (hash-count table)))
    (error 'enumerate "condition accepted no paths"))
  (make-discrete-dist table #:normalize? normalize?))

;; ----------------------------------------

;; https://hips.seas.harvard.edu/blog/2013/01/09/computing-log-sum-exp/

;; logspace+ : Real Real -> Real
;; Given (log A) and (log B), return (log (+ A B)).
(define logspace+*
  (case-lambda
    [() -inf.0]
    [(a) a]
    [(a b)
     (define m (max a b))
     (if (= m -inf.0)
         -inf.0
         (+ m (log (+ (exp (- a m)) (exp (- b m))))))]
    [as
     (define m (apply max as))
     (if (= m -inf.0)
         -inf.0
         (+ m (log (for/sum ([a (in-list as)]) (exp (- a m))))))]))

(define (logspace+ a b)
  (define r (logspace+* a b))
  (when (eqv? r +nan.0)
    (error 'logspace+ "~s + ~s => ~s" a b r))
  r)

(define (logspace-compl a)
  (log (- 1.0 (exp a))))

;; An ExpST is (expst Prob Prob Hash[A => Prob] Nat)
;; where Prob is Real or Real+ depending on logspace?,
;; and dim is density dimension of every entry in table.
(struct expst (explored accepted table table-ddim) #:transparent)

(define (expst->logspace st)
  (defmatch (expst explored accepted table table-ddim) st)
  (expst (log (exact->inexact explored)) (log (exact->inexact accepted))
         (table->logspace table) table-ddim))

(define (table->logspace table)
  (for/hash ([(k v) (in-hash table)])
    (values k (log (exact->inexact v)))))

;; A (HeapEntry A) is (list* Prob Nat (-> (EnumTree A))).
(define (entry->=? x y) (>= (car x) (car y)))

(define (heap->logspace h)
  (for/fold ([new-h (heap entry->=?)])
            ([entry (in-list (heap->list h))])
    (heap-insert new-h (cons (log (car entry)) (cdr entry)))))

;; explore : (EnumTree A) (U #f Real>0) -> Hash[A => Real+]
;; * Starts with logspace? = #f, but switches to logspace if an inexact
;;   density is ever seen.
;; * Limit means: explore until potential accepted dist is < limit unaccounted for,
;;   ie, unexplored < limit * accepted. Thus, limit is nonlocal; use BFS.
;;   Limit is always realspace, not logspace. If no limit, prob-{explored,accepted}
;;   can be meaningless. Note: Cannot observe continuous dist if limit given.
(define (explore tree limit)
  (define initial-st (expst 0 0 '#hash() (if limit 0 +inf.0)))
  (define initial-heap
    ;; lib provides "min-heap", but want max prob, so use >= comparison
    (heap entry->=?))
  (define initial-ddim 0)
  (define initial-p 1)
  (explore* limit #f initial-heap initial-st tree initial-ddim initial-p))

;; explore* : (EnumTree A) (U #f Real>0) ExpSt Heap Nat -> Hash[A => Real+]
;; - st   : ExpST, rep of Prob depends on logspace?
;; - h    : heap[(cons Prob (-> (EnumTree A)))]
;; - ddim : Nat, density dimension
(define (explore* limit logspace? h st tree tree-ddim tree-p)
  ;; type Prob = Real
  ;; The representation of probabilities in this code is determined by
  ;; the logspace? argument.  If logspace? is true, then values
  ;; represent log of prob; if false, values represent prob directly.
  (define (->p a) ;; inexact to avoid error on exact 0
    (when (negative? a) (error '->p "negative: ~e" a))
    (if logspace? (log (exact->inexact a)) a))
  (define (p->real a)
    (if logspace? (exp a) a))
  (define (p+ a b)
    (if logspace? (logspace+ a b) (+ a b)))
  (define (pcompl a)
    (if logspace? (logspace-compl a) (- 1 a)))
  (define (p* a b)
    (if logspace? (+ a b) (* a b)))
  (define pzero (->p 0))

  (define (switch-to-logspace h st tree tree-ddim tree-p)
    (explore* limit #t (heap->logspace h) (expst->logspace st) tree tree-ddim (log tree-p)))
{'fraction_non_alphanumeric': 0.1723832923832924, 'fraction_numerical': 0.006191646191646191, 'mean_word_length': 3.2032218091697646, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 71, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2946899', 'n_tokens_mistral': 3487, 'n_tokens_neox': 3311, 'n_words': 1275}