racket/collects/math/private/distributions/gamma-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

74 lines
2.5 KiB
Racket

#lang typed/racket/base
(require racket/performance-hint
racket/promise
"../../flonum.rkt"
"../unsafe.rkt"
"../functions/incomplete-gamma.rkt"
"impl/gamma-pdf.rkt"
"impl/gamma-inv-cdf.rkt"
"impl/gamma-random.rkt"
"dist-struct.rkt"
"utils.rkt")
(provide flgamma-pdf
flgamma-cdf
flgamma-inv-cdf
flgamma-sample
Gamma-Dist gamma-dist gamma-dist-shape gamma-dist-scale)
(: flgamma-pdf (Float Float Float Any -> Float))
(define (flgamma-pdf k s x log?)
(cond [(k . fl< . 0.0) +nan.0]
[else
((make-one-sided-scale-flpdf
(λ: ([x : Float] [log? : Any])
(cond [log? (standard-flgamma-log-pdf k x)]
[else (standard-flgamma-pdf k x)])))
s x log?)]))
(: flgamma-cdf (Float Float Float Any Any -> Float))
(define (flgamma-cdf k s x log? 1-p?)
(cond [(k . fl< . 0.0) +nan.0]
[else
((make-one-sided-scale-flcdf
(λ: ([x : Float] [log? : Any] [1-p? : Any])
(cond [log? (fllog-gamma-inc k x 1-p? #t)]
[else (flgamma-inc k x 1-p? #t)])))
s x log? 1-p?)]))
(: flgamma-inv-cdf (Float Float Float Any Any -> Float))
(define (flgamma-inv-cdf k s p log? 1-p?)
(cond [(k . fl< . 0.0) +nan.0]
[else
((make-one-sided-scale-flinv-cdf
(λ: ([p : Float] [log? : Any] [1-p? : Any])
(standard-flgamma-inv-cdf k p log? 1-p?)))
s p log? 1-p?)]))
(define-real-dist: gamma-dist Gamma-Dist
gamma-dist-struct ([shape : Flonum] [scale : Flonum]))
(begin-encourage-inline
(: gamma-dist (case-> (-> Gamma-Dist)
(Real -> Gamma-Dist)
(Real Real -> Gamma-Dist)))
(define (gamma-dist [k 1.0] [s 1.0])
(let ([k (fl k)] [s (fl s)])
(define pdf (opt-lambda: ([x : Real] [log? : Any #f])
(flgamma-pdf k s (fl x) log?)))
(define cdf (opt-lambda: ([x : Real] [log? : Any #f] [1-p? : Any #f])
(flgamma-cdf k s (fl x) log? 1-p?)))
(define inv-cdf (opt-lambda: ([p : Real] [log? : Any #f] [1-p? : Any #f])
(flgamma-inv-cdf k s (fl p) log? 1-p?)))
(define sample (case-lambda:
[() (unsafe-flvector-ref (flgamma-sample k s 1) 0)]
[([n : Integer]) (flvector->list (flgamma-sample k s n))]))
(gamma-dist-struct
pdf sample cdf inv-cdf
0.0 +inf.0 (delay (flgamma-inv-cdf k s 0.5 #f #f))
k s)))
)