racket/collects/math/private/distributions/binomial-dist.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

81 lines
3.0 KiB
Racket

#lang typed/racket/base
(require racket/performance-hint
racket/promise
"../../flonum.rkt"
"../unsafe.rkt"
"../functions/incomplete-beta.rkt"
(prefix-in impl: "impl/binomial-pdf.rkt")
"impl/binomial-random.rkt"
"normal-dist.rkt"
"dist-struct.rkt"
"utils.rkt")
(provide flbinomial-pdf
flbinomial-cdf
flbinomial-inv-cdf
flbinomial-sample
Binomial-Dist binomial-dist binomial-dist-count binomial-dist-prob)
(: flbinomial-pdf (Flonum Flonum Flonum Any -> Flonum))
(define (flbinomial-pdf n q k log?)
(cond [log? (impl:flbinomial-log-pdf n q k)]
[else (impl:flbinomial-pdf n q k)]))
(: flbinomial-cdf (Flonum Flonum Flonum Any Any -> Flonum))
(define (flbinomial-cdf n q k log? 1-p?)
(cond [(or (not (flinteger? n)) (n . fl< . 0.0)
(q . fl< . 0.0) (q . fl> . 1.0))
+nan.0]
[else
(let ([k (flfloor k)])
(cond [log? (fllog-beta-inc (fl+ k 1.0) (fl- n k) q (not 1-p?) #t)]
[else (flbeta-inc (fl+ k 1.0) (fl- n k) q (not 1-p?) #t)]))]))
(: flbinomial-inv-cdf (Flonum Flonum Flonum Any Any -> Flonum))
(define (flbinomial-inv-cdf n q p log? 1-p?)
(cond [(or (not (flinteger? n)) (n . fl< . 0.0)
(q . fl< . 0.0) (q . fl> . 1.0)
(not (flprobability? p log?)))
+nan.0]
[(flprobability-one? p log? 1-p?) n]
[(flprobability-zero? p log? 1-p?) 0.0]
[1-p?
(define z (flnormal-inv-cdf (fl* n q) (flsqrt (* n q (fl- 1.0 q))) p log? #t))
(flfind-least-integer
(λ: ([k : Flonum]) ((flbinomial-cdf n q k log? 1-p?) . fl< . p))
0.0 n
(flmax 0.0 (flmin n z)))]
[else
(define z (flnormal-inv-cdf (fl* n q) (flsqrt (* n q (fl- 1.0 q))) p log? #f))
(flfind-least-integer
(λ: ([k : Flonum]) ((flbinomial-cdf n q k log? 1-p?) . fl>= . p))
0.0 n
(flmax 0.0 (flmin n z)))]))
(define-real-dist: binomial-dist Binomial-Dist
binomial-dist-struct ([count : Flonum] [prob : Flonum]))
(begin-encourage-inline
(: binomial-dist (case-> (-> Binomial-Dist)
(Real -> Binomial-Dist)
(Real Real -> Binomial-Dist)))
(define (binomial-dist [n 1.0] [q 0.5])
(let ([n (fl n)] [q (fl q)])
(define pdf (opt-lambda: ([k : Real] [log? : Any #f])
(flbinomial-pdf n q (fl k) log?)))
(define cdf (opt-lambda: ([k : Real] [log? : Any #f] [1-p? : Any #f])
(flbinomial-cdf n q (fl k) log? 1-p?)))
(define inv-cdf (opt-lambda: ([p : Real] [log? : Any #f] [1-p? : Any #f])
(flbinomial-inv-cdf n q (fl p) log? 1-p?)))
(define sample (case-lambda:
[() (unsafe-flvector-ref (flbinomial-sample n q 1) 0)]
[([m : Integer]) (flvector->list (flbinomial-sample n q m))]))
(binomial-dist-struct
pdf sample cdf inv-cdf
0.0 +inf.0 (delay (flfloor (fl* n q)))
n q)))
)