<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 racket/class
         "interfaces.rkt"
         "context.rkt"
         "../dist.rkt")
(provide (all-defined-out))

;; mem and sample wrappers

(define (mem f) (send (current-stochastic-ctx) mem f))
(define (observe-at dist val) (send (current-stochastic-ctx) observe-at dist val))
(define (fail [reason #f]) (send (current-stochastic-ctx) fail reason))
(define (sample* dist) (send (current-stochastic-ctx) sample dist))

(define (sample dist)
  (define obs (observing?))
  (cond [(observation? obs)
         (with-continuation-mark
             obs-mark 'ok ;; overwrite this frame's 'unknown mark
           (let ()
             (define cc (get-observe-context))
             (define value (observe-context-adjust-value cc obs))
             (when (verbose?)
               (eprintf "** CC = ~e\n" cc)
               (eprintf "** sample -> condition: ~e @ ~e\n" dist value))
             (observe-at dist value)
             value))]
        [else (sample* dist)]))

;; == Finite distributions ==

;; flip : Prob -> (U #t #f)
(define (flip [prob 1/2])
  (positive?
   (sample (bernoulli-dist prob))))

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

;; categorical : (Vectorof Prob) -> Nat
(define (categorical weights)
  (sample (categorical-dist weights)))

;; discrete : Nat -> Nat
(define (discrete dict)
  (sample (make-discrete-dist dict)))

;; discrete* : (Listof/Vectorof A) (Listof/Vectorof Prob) -> A
(define (discrete* vals0 [weights0 #f])
  (let ([vals (if (list? vals0) (list->vector vals0) vals0)])
    (cond [(eq? weights0 #f)
           (sample (make-discrete-dist* vals))]
          [else
           (let ([weights (if (list? weights0) (list->vector weights0) weights0)])
             (unless (= (vector-length vals) (vector-length weights))
               (error 'discrete*
                      "values and weights have unequal lengths\n  values: ~e\n  weights: ~e"
                      vals0 weights0))
             (sample (make-discrete-dist* vals weights)))])))

;; discrete-uniform : Nat -> Nat
(define (discrete-uniform n)
  (sample (categorical-dist (make-vector n (/ n)))))

;; == Countable distributions ==

;; binomial : Nat Prob -> Integer
(define (binomial n p)
  (sample (binomial-dist n p)))

;; geometric : Prob -> Integer
(define (geometric [p 1/2])
  (sample (geometric-dist p)))

;; poisson : Real -> Integer
(define (poisson mean)
  (sample (poisson-dist mean)))

;; == Continuous distributions ==

;; beta : PositiveReal PositiveReal -> Real in [0,1]
(define (beta a b)
  (sample (beta-dist a b)))

(define (cauchy [mode 0] [scale 1])
  (sample (cauchy-dist mode scale)))

;; exponential : PositiveReal -> PositiveReal
;; NOTE: mean aka scale = 1/rate
(define (exponential [mean 1])
  (sample (exponential-dist mean)))

;; gamma : PositiveReal PositiveReal -> Real
;; NOTE: scale = 1/rate
(define (gamma [shape 1] [scale 1])
  (sample (gamma-dist shape scale)))

(define (inverse-gamma [shape 1] [scale 1])
  (sample (inverse-gamma shape scale)))

;; logistic : Real Real -> Real
(define (logistic [mean 0] [scale 1])
  (sample (logistic-dist mean scale)))

;; normal : Real PositiveReal -> Real
;; NOTE: stddev = (sqrt variance)
(define (normal [mean 0] [stddev 1])
  (sample (normal-dist mean stddev)))

;; pareto : Real Real -> Real
(define (pareto shape scale)
  (sample (pareto-dist shape scale)))

;; uniform : Real Real -> Real
(define uniform
  (case-lambda
    [() (uniform 0 1)]
    [(max) (uniform 0 max)]
    [(min max)
     (sample (uniform-dist min max))]))

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

(define (dirichlet alpha)
  (sample (dirichlet-dist alpha)))<commit_msg>fix inverse-gamma convenience function<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 racket/class
         "interfaces.rkt"
         "context.rkt"
         "../dist.rkt")
(provide (all-defined-out))

;; mem and sample wrappers

(define (mem f) (send (current-stochastic-ctx) mem f))
(define (observe-at dist val) (send (current-stochastic-ctx) observe-at dist val))
(define (fail [reason #f]) (send (current-stochastic-ctx) fail reason))
(define (sample* dist) (send (current-stochastic-ctx) sample dist))

(define (sample dist)
  (define obs (observing?))
  (cond [(observation? obs)
         (with-continuation-mark
             obs-mark 'ok ;; overwrite this frame's 'unknown mark
           (let ()
             (define cc (get-observe-context))
             (define value (observe-context-adjust-value cc obs))
             (when (verbose?)
               (eprintf "** CC = ~e\n" cc)
               (eprintf "** sample -> condition: ~e @ ~e\n" dist value))
             (observe-at dist value)
             value))]
        [else (sample* dist)]))

;; == Finite distributions ==

;; flip : Prob -> (U #t #f)
(define (flip [prob 1/2])
  (positive?
   (sample (bernoulli-dist prob))))

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

;; categorical : (Vectorof Prob) -> Nat
(define (categorical weights)
  (sample (categorical-dist weights)))

;; discrete : Nat -> Nat
(define (discrete dict)
  (sample (make-discrete-dist dict)))

;; discrete* : (Listof/Vectorof A) (Listof/Vectorof Prob) -> A
(define (discrete* vals0 [weights0 #f])
  (let ([vals (if (list? vals0) (list->vector vals0) vals0)])
    (cond [(eq? weights0 #f)
           (sample (make-discrete-dist* vals))]
          [else
           (let ([weights (if (list? weights0) (list->vector weights0) weights0)])
             (unless (= (vector-length vals) (vector-length weights))
               (error 'discrete*
                      "values and weights have unequal lengths\n  values: ~e\n  weights: ~e"
                      vals0 weights0))
             (sample (make-discrete-dist* vals weights)))])))

;; discrete-uniform : Nat -> Nat
(define (discrete-uniform n)
  (sample (categorical-dist (make-vector n (/ n)))))

;; == Countable distributions ==

;; binomial : Nat Prob -> Integer
(define (binomial n p)
  (sample (binomial-dist n p)))

;; geometric : Prob -> Integer
(define (geometric [p 1/2])
  (sample (geometric-dist p)))

;; poisson : Real -> Integer
(define (poisson mean)
  (sample (poisson-dist mean)))

;; == Continuous distributions ==

;; beta : PositiveReal PositiveReal -> Real in [0,1]
(define (beta a b)
  (sample (beta-dist a b)))

(define (cauchy [mode 0] [scale 1])
  (sample (cauchy-dist mode scale)))

;; exponential : PositiveReal -> PositiveReal
;; NOTE: mean aka scale = 1/rate
(define (exponential [mean 1])
  (sample (exponential-dist mean)))

;; gamma : PositiveReal PositiveReal -> Real
;; NOTE: scale = 1/rate
(define (gamma [shape 1] [scale 1])
  (sample (gamma-dist shape scale)))

(define (inverse-gamma [shape 1] [scale 1])
  (sample (inverse-gamma-dist shape scale)))

;; logistic : Real Real -> Real
(define (logistic [mean 0] [scale 1])
  (sample (logistic-dist mean scale)))

;; normal : Real PositiveReal -> Real
;; NOTE: stddev = (sqrt variance)
(define (normal [mean 0] [stddev 1])
  (sample (normal-dist mean stddev)))

;; pareto : Real Real -> Real
(define (pareto shape scale)
  (sample (pareto-dist shape scale)))

;; uniform : Real Real -> Real
(define uniform
  (case-lambda
    [() (uniform 0 1)]
    [(max) (uniform 0 max)]
    [(min max)
     (sample (uniform-dist min max))]))

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

(define (dirichlet alpha)
  (sample (dirichlet-dist alpha)))
{'fraction_non_alphanumeric': 0.16770186335403728, 'fraction_numerical': 0.010869565217391304, 'mean_word_length': 3.09375, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 43, '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': '5801083', 'n_tokens_mistral': 2612, 'n_tokens_neox': 2445, 'n_words': 794}