<commit_before>#lang racket/base
(require math/distributions
         racket/contract/base
         "prob-hooks.rkt"
         "util.rkt")
(provide rejection-sample
         (contract-out
          [mem (-> procedure? procedure?)])
         ERP)

;; Rejection sampling

(define (rejection-sample thunk pred [project values])
  (let ([v (thunk)])
    (if (pred v)
        (project v)
        (rejection-sample thunk pred project))))

;; mem and ERP wrappers

;; NOTE: enum-ERP knows about discrete dists by tag
;; FIXME: add discrete-dist? to math/distributions
(define (mem f) ((current-mem) f))
(define (ERP tag dist) ((current-ERP) tag dist))

;; ==  Finite distributions ==

(provide
 (contract-out
  [flip
   (->* [] [probability?] boolean?)]
  [d2
   (->* [] [probability?] (or/c 1 0))])
 discrete)

;; flip : Prob -> (U #t #f)
(define (flip [prob 1/2])
  (positive?
   (ERP `(flip ,prob)
        (make-dist bernoulli #:params (prob) #:enum 2))))

;; d2 : Prob -> (U 1 0)
(define (d2 [prob 1/2])
  (bernoulli prob))

(define (bernoulli [prob 1/2])
  (inexact->exact
   (ERP `(bernoulli ,prob)
        (make-dist bernoulli #:params (prob) #:enum 2))))

;; Used in #:params, syntactically required to be identifier.
(define ZERO 0.0)

;; discrete : Nat -> Nat
;; discrete : (listof (list A Prob))) -> A
(define discrete
  (case-lambda
    [(n/vals)
     (cond [(and (list? n/vals) (pair? n/vals))
            (inexact->exact
             (floor
              (ERP `(discrete ,n/vals)
                   (make-dist uniform #:params (ZERO n/vals) #:enum n/vals))))]
           [(exact-positive-integer? n/vals)
            (let ([n (length n/vals)])
              (list-ref n/vals
                        (inexact->exact
                         (floor
                          (ERP `(discrete ,n/vals)
                               (make-dist uniform #:params (ZERO n) #:enum n))))))]
           [else
            (raise-argument-error 'discrete
              "(or/c exact-positive-integer? (and/c list? pair?))" 0 n/vals)])]
    [(vals probs)
     (unless (and (list? vals) (pair? vals))
       (raise-argument-error 'discrete "(and/c list? pair?)" 0 vals probs))
     (unless (and (list? probs) (pair? probs) (andmap real? probs) (andmap positive? probs))
       (raise-argument-error 'discrete "(non-empty-listof (>/c 0))" 1 vals probs))
     (unless (= (length vals) (length probs))
       (error 'discrete
              "values and probability weights have different lengths\n  values: ~e\n  weights: ~e"
              vals probs))
     (define n (length vals))
     (lookup-discrete vals probs 
                      (ERP `(discrete ,vals ,probs)
                           (make-dist uniform #:params (ZERO n) #:enum n)))]))

(define (lookup-discrete vals probs p)
  (cond [(null? vals)
         (error 'discrete "internal error: out of values")]
        [(< p (car probs))
         (car vals)]
        [else (lookup-discrete (cdr vals) (cdr probs) (- p (car probs)))]))

;; == Countable distributions ==<commit_msg>fix discrete<commit_after>#lang racket/base
(require math/distributions
         racket/contract/base
         "prob-hooks.rkt"
         "util.rkt")
(provide rejection-sample
         (contract-out
          [mem (-> procedure? procedure?)])
         ERP)

;; Rejection sampling

(define (rejection-sample thunk pred [project values])
  (let ([v (thunk)])
    (if (pred v)
        (project v)
        (rejection-sample thunk pred project))))

;; mem and ERP wrappers

;; NOTE: enum-ERP knows about discrete dists by tag
;; FIXME: add discrete-dist? to math/distributions
(define (mem f) ((current-mem) f))
(define (ERP tag dist) ((current-ERP) tag dist))

;; ==  Finite distributions ==

(provide
 (contract-out
  [flip
   (->* [] [probability?] boolean?)]
  [d2
   (->* [] [probability?] (or/c 1 0))])
 discrete)

;; flip : Prob -> (U #t #f)
(define (flip [prob 1/2])
  (positive?
   (ERP `(flip ,prob)
        (make-dist bernoulli #:params (prob) #:enum 2))))

;; d2 : Prob -> (U 1 0)
(define (d2 [prob 1/2])
  (bernoulli prob))

(define (bernoulli [prob 1/2])
  (inexact->exact
   (ERP `(bernoulli ,prob)
        (make-dist bernoulli #:params (prob) #:enum 2))))

;; Used in #:params, syntactically required to be identifier.
(define ZERO 0.0)

;; discrete : Nat -> Nat
;; discrete : (listof (list A Prob))) -> A
(define discrete
  (case-lambda
    [(n/vals)
     (cond [(and (list? n/vals) (pair? n/vals))
            (let ([n (length n/vals)])
              (list-ref n/vals
                        (inexact->exact
                         (floor
                          (ERP `(discrete ,n/vals)
                               (make-dist uniform #:params (ZERO n) #:enum n))))))]
           [(exact-positive-integer? n/vals)
            (let ([n n/vals])
              (inexact->exact
               (floor
                (ERP `(discrete ,n)
                     (make-dist uniform #:params (ZERO n) #:enum n)))))]
           [else
            (raise-argument-error 'discrete
              "(or/c exact-positive-integer? (and/c list? pair?))" 0 n/vals)])]
    [(vals probs)
     (unless (and (list? vals) (pair? vals))
       (raise-argument-error 'discrete "(and/c list? pair?)" 0 vals probs))
     (unless (and (list? probs) (pair? probs) (andmap real? probs) (andmap positive? probs))
       (raise-argument-error 'discrete "(non-empty-listof (>/c 0))" 1 vals probs))
     (unless (= (length vals) (length probs))
       (error 'discrete
              "values and probability weights have different lengths\n  values: ~e\n  weights: ~e"
              vals probs))
     (define n (length vals))
     (lookup-discrete vals probs 
                      (ERP `(discrete ,vals ,probs)
                           (make-dist uniform #:params (ZERO n) #:enum n)))]))

(define (lookup-discrete vals probs p)
  (cond [(null? vals)
         (error 'discrete "internal error: out of values")]
        [(< p (car probs))
         (car vals)]
        [else (lookup-discrete (cdr vals) (cdr probs) (- p (car probs)))]))

;; == Countable distributions ==
{'fraction_non_alphanumeric': 0.16032921810699588, 'fraction_numerical': 0.006913580246913581, 'mean_word_length': 2.172845953002611, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 25, 'https://': 0, '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': '1849301', 'n_tokens_mistral': 2116, 'n_tokens_neox': 2025, 'n_words': 598}