1#lang typed/racket/base
2
3#|
4Wolfgang Hormann. The Generation of Binomial Random Variates.
5|#
6
7(require "../../../base.rkt"
8         "../../../flonum.rkt"
9         "../../unsafe.rkt"
10         "normal-random.rkt")
11
12(provide flbinomial-sample)
13
14(: flbinomial-sample-small (Flonum Flonum Natural -> FlVector))
15;; For n*min(p,1-p) <= 30
16(define (flbinomial-sample-small n p m)
17  (let-values ([(p q s?)  (cond [(p . fl< . 0.5)  (values p (fl- 1.0 p) #f)]
18                                [else             (values (fl- 1.0 p) p #t)])])
19    (define q^n (flexpt q n))
20    (define r (fl/ p q))
21    (define g (fl* r (fl+ n 1.0)))
22    (build-flvector
23     m (λ (_)
24         (define k
25           (let: reject : Flonum ()
26             (let loop ([k 0.0] [f q^n] [u  (random)])
27               (cond [(u . fl< . f)  k]
28                     [(k . fl> . 110.0)  (reject)]
29                     [else  (let ([k  (fl+ k 1.0)])
30                              (loop k (fl* f (fl- (fl/ g k) r)) (fl- u f)))]))))
31         (if s? (fl- n k) k)))))
32
33(: flbinomial-sample-hormann (Flonum Flonum Natural -> FlVector))
34;; For n*min(p,1-p) >= 10
35(define (flbinomial-sample-hormann n p j)
36  (let-values ([(p q s?)  (cond [(p . fl< . 0.5)  (values p (fl- 1.0 p) #f)]
37                                [else             (values (fl- 1.0 p) p #t)])])
38    (define σ (flsqrt (* n p q)))
39    (define m (flfloor (fl* (fl+ n 1.0) p)))
40
41    (define b (fl+ 1.15 (fl* 2.53 σ)))
42    (define a (+ -0.0873 (fl* 0.0248 b) (fl* 0.01 p)))
43    (define c (fl+ 0.5 (fl* n p)))
44    (define α (fl* σ (fl+ 2.83 (fl/ 5.1 b))))
45    (define vr (fl- 0.92 (fl/ 4.2 b)))
46    (build-flvector
47     j (λ (_)
48         (define k
49           (let: loop : Flonum ()
50             (define v (random))
51             (define u (fl- (random) 0.5))
52             (define us (fl- 0.5 (flabs u)))
53             (define k (flfloor (fl+ c (fl* u (fl+ b (fl/ (fl* 2.0 a) us))))))
54             (cond [(or (k . fl< . 0.0) (k . fl> . n))  (loop)]
55                   [(and (us . fl>= . 0.07) (v . fl<= . vr))  k]
56                   [else
57                    (let ([v  (fl* v (fl/ α (fl+ b (fl/ a (fl* us us)))))])
58                      (define h (+ (fllog-factorial m)
59                                   (fllog-factorial (fl- n m))
60                                   (- (fllog-factorial k))
61                                   (- (fllog-factorial (fl- n k)))
62                                   (fl* (fl- k m) (fllog (fl/ p q)))))
63                      (cond [((fllog v) . fl<= . h)  k]
64                            [else  (loop)]))])))
65         (if s? (fl- n k) k)))))
66
67(: flbinomial-sample-normal (Flonum Flonum Natural -> FlVector))
68(define (flbinomial-sample-normal n p m)
69  (define q (fl- 1.0 p))
70  (define μ (fl- (fl* (fl+ n 1.0) p) 0.5))
71  (define σ (flsqrt (* (+ 1.0 n) p q)))
72  (define γ (fl/ (fl- q p) σ))
73  (build-flvector
74   m (λ (_)
75       (let loop ()
76         (define z (unsafe-flvector-ref (flnormal-sample 0.0 1.0 1) 0))
77         (define k (flround (fl+ μ (fl* σ (fl+ z (fl/ (fl* γ (fl- (fl* z z) 1.0)) 6.0))))))
78         (if (and (k . fl>= . 0.0) (k . fl<= . n)) k (loop))))))
79
80(: flbinomial-normal-appx-error-bound (Flonum Flonum -> Flonum))
81;; Returns a bound on the integrated difference between the normal and binomial cdfs
82;; See the Berry-Esséen theorem
83(define (flbinomial-normal-appx-error-bound n p)
84  (define q (fl- 1.0 p))
85  (fl/ (fl* 0.4784 (fl+ (fl* p p) (fl* q q))) (flsqrt (* n p q))))
86
87(: flbinomial-sample (Flonum Flonum Integer -> FlVector))
88(define (flbinomial-sample n p m)
89  (cond [(m . < . 0)  (raise-argument-error 'flbinomial-sample "Natural" 2 n p m)]
90        [(or (not (integer? n)) (n . fl< . 0.0) (p . fl< . 0.0) (p . fl> . 1.0))
91         (build-flvector m (λ (_) +nan.0))]
92        [(or (fl= n 0.0) (fl= p 0.0))
93         (build-flvector m (λ (_) 0.0))]
94        [(fl= p 1.0)
95         (build-flvector m (λ (_) n))]
96        [(and (n . fl> . 1e8)
97              ((flbinomial-normal-appx-error-bound n p) . fl< . (flexp -10.0)))
98         (flbinomial-sample-normal n p m)]
99        [((fl* n (flmin p (fl- 1.0 p))) . fl>= . 10.0)
100         (flbinomial-sample-hormann n p m)]
101        [else
102         (flbinomial-sample-small n p m)]))
103