0936d8c20b
Added docs for math/distributions (about 75% finished) Started docs for math/array (very incomplete)
104 lines
4.1 KiB
Racket
104 lines
4.1 KiB
Racket
#lang typed/racket/base
|
||
|
||
#|
|
||
Wolfgang Hormann. The Generation of Binomial Random Variates.
|
||
|#
|
||
|
||
(require "../../../base.rkt"
|
||
"../../../flonum.rkt"
|
||
"../../../vector.rkt"
|
||
"../../unsafe.rkt"
|
||
"normal-random.rkt")
|
||
|
||
(provide flbinomial-sample)
|
||
|
||
(: flbinomial-sample-small (Flonum Flonum Natural -> FlVector))
|
||
;; For n*min(p,1-p) <= 30
|
||
(define (flbinomial-sample-small n p m)
|
||
(let-values ([(p q s?) (cond [(p . fl< . 0.5) (values p (fl- 1.0 p) #f)]
|
||
[else (values (fl- 1.0 p) p #t)])])
|
||
(define q^n (flexpt q n))
|
||
(define r (fl/ p q))
|
||
(define g (fl* r (fl+ n 1.0)))
|
||
(build-flvector
|
||
m (λ (_)
|
||
(define k
|
||
(let: reject : Flonum ()
|
||
(let loop ([k 0.0] [f q^n] [u (random)])
|
||
(cond [(u . fl< . f) k]
|
||
[(k . fl> . 110.0) (reject)]
|
||
[else (let ([k (fl+ k 1.0)])
|
||
(loop k (fl* f (fl- (fl/ g k) r)) (fl- u f)))]))))
|
||
(if s? (fl- n k) k)))))
|
||
|
||
(: flbinomial-sample-hormann (Flonum Flonum Natural -> FlVector))
|
||
;; For n*min(p,1-p) >= 10
|
||
(define (flbinomial-sample-hormann n p j)
|
||
(let-values ([(p q s?) (cond [(p . fl< . 0.5) (values p (fl- 1.0 p) #f)]
|
||
[else (values (fl- 1.0 p) p #t)])])
|
||
(define σ (flsqrt (* n p q)))
|
||
(define m (flfloor (fl* (fl+ n 1.0) p)))
|
||
|
||
(define b (fl+ 1.15 (fl* 2.53 σ)))
|
||
(define a (+ -0.0873 (fl* 0.0248 b) (fl* 0.01 p)))
|
||
(define c (fl+ 0.5 (fl* n p)))
|
||
(define α (fl* σ (fl+ 2.83 (fl/ 5.1 b))))
|
||
(define vr (fl- 0.92 (fl/ 4.2 b)))
|
||
(build-flvector
|
||
j (λ (_)
|
||
(define k
|
||
(let: loop : Flonum ()
|
||
(define v (random))
|
||
(define u (fl- (random) 0.5))
|
||
(define us (fl- 0.5 (flabs u)))
|
||
(define k (flfloor (fl+ c (fl* u (fl+ b (fl/ (fl* 2.0 a) us))))))
|
||
(cond [(or (k . fl< . 0.0) (k . fl> . n)) (loop)]
|
||
[(and (us . fl>= . 0.07) (v . fl<= . vr)) k]
|
||
[else
|
||
(let ([v (fl* v (fl/ α (fl+ b (fl/ a (fl* us us)))))])
|
||
(define h (+ (fllog-factorial m)
|
||
(fllog-factorial (fl- n m))
|
||
(- (fllog-factorial k))
|
||
(- (fllog-factorial (fl- n k)))
|
||
(fl* (fl- k m) (fllog (fl/ p q)))))
|
||
(cond [((fllog v) . fl<= . h) k]
|
||
[else (loop)]))])))
|
||
(if s? (fl- n k) k)))))
|
||
|
||
(: flbinomial-sample-normal (Flonum Flonum Natural -> FlVector))
|
||
(define (flbinomial-sample-normal n p m)
|
||
(define q (fl- 1.0 p))
|
||
(define μ (fl- (fl* (fl+ n 1.0) p) 0.5))
|
||
(define σ (flsqrt (* (+ 1.0 n) p q)))
|
||
(define γ (fl/ (fl- q p) σ))
|
||
(build-flvector
|
||
m (λ (_)
|
||
(let loop ()
|
||
(define z (unsafe-flvector-ref (flnormal-sample 0.0 1.0 1) 0))
|
||
(define k (flround (fl+ μ (fl* σ (fl+ z (fl/ (fl* γ (fl- (fl* z z) 1.0)) 6.0))))))
|
||
(if (and (k . fl>= . 0.0) (k . fl<= . n)) k (loop))))))
|
||
|
||
(: flbinomial-normal-appx-error-bound (Flonum Flonum -> Flonum))
|
||
;; Returns a bound on the integrated difference between the normal and binomial cdfs
|
||
;; See the Berry-Esséen theorem
|
||
(define (flbinomial-normal-appx-error-bound n p)
|
||
(define q (fl- 1.0 p))
|
||
(fl/ (fl* 0.4784 (fl+ (fl* p p) (fl* q q))) (flsqrt (* n p q))))
|
||
|
||
(: flbinomial-sample (Flonum Flonum Integer -> FlVector))
|
||
(define (flbinomial-sample n p m)
|
||
(cond [(m . < . 0) (raise-argument-error 'flbinomial-sample "Natural" 2 n p m)]
|
||
[(or (not (integer? n)) (n . fl< . 0.0) (p . fl< . 0.0) (p . fl> . 1.0))
|
||
(build-flvector m (λ (_) +nan.0))]
|
||
[(or (fl= n 0.0) (fl= p 0.0))
|
||
(build-flvector m (λ (_) 0.0))]
|
||
[(fl= p 1.0)
|
||
(build-flvector m (λ (_) n))]
|
||
[(and (n . fl> . 1e8)
|
||
((flbinomial-normal-appx-error-bound n p) . fl< . (flexp -10.0)))
|
||
(flbinomial-sample-normal n p m)]
|
||
[((fl* n (flmin p (fl- 1.0 p))) . fl>= . 10.0)
|
||
(flbinomial-sample-hormann n p m)]
|
||
[else
|
||
(flbinomial-sample-small n p m)]))
|