racket/collects/math/private/distributions/impl/binomial-random.rkt
Neil Toronto 6009eed8d2 Moved flvector functions into math/flonum
Sped up normal distribution sampling procedure (2x for large samples)
2012-11-29 15:45:17 -07:00

103 lines
4.1 KiB
Racket
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#lang typed/racket/base
#|
Wolfgang Hormann. The Generation of Binomial Random Variates.
|#
(require "../../../base.rkt"
"../../../flonum.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)]))