96 lines
3.9 KiB
Racket
96 lines
3.9 KiB
Racket
#lang typed/racket/base
|
|
|
|
(require "../../../flonum.rkt"
|
|
"../../../base.rkt"
|
|
"../../unsafe.rkt"
|
|
"../../distributions/impl/normal-cdf.rkt"
|
|
"../stirling-error.rkt"
|
|
"gamma-utils.rkt")
|
|
|
|
(provide flgamma-temme fllog-gamma-temme)
|
|
|
|
;; ===================================================================================================
|
|
;; Temme's series for the incomplete gamma functions (used when k ~ x and k is not small)
|
|
|
|
(define temme-iters (make-parameter 32))
|
|
|
|
(define num-fs 100)
|
|
(define fs
|
|
(let ()
|
|
(define: start-fs : (Vectorof Real) (vector 1 -1/3 1/12 -2/135))
|
|
(define: fs : (Vectorof Real) (make-vector num-fs 0))
|
|
(vector-copy! fs 0 start-fs)
|
|
;; DP algorithm to compute f coefficients
|
|
(for ([m (in-range 4 num-fs)])
|
|
(vector-set!
|
|
fs m
|
|
(* (- (/ (+ m 1) (+ m 2)))
|
|
(+ (* (/ (- m 1) (* 3 m)) (vector-ref fs (- m 1)))
|
|
(for/fold: ([sum : Real 0]) ([j (in-range 3 m)])
|
|
(+ sum (/ (* (vector-ref fs (- j 1)) (vector-ref fs (+ m 1 (- j))))
|
|
(+ m 2 (- j)))))))))
|
|
(vector->flvector fs)))
|
|
|
|
(: R-sum (Float Flonum -> Flonum))
|
|
(define (R-sum k n)
|
|
(define num-fs (temme-iters))
|
|
;; This originally filled a vector of bs, because imperative programmers don't know how to do
|
|
;; anything else besides bang an array full of values (sheesh)
|
|
(define-values (sum b2 b1)
|
|
(for/fold: ([sum : Flonum 0.0]
|
|
[b2 : Flonum (unsafe-flvector-ref fs (- num-fs 1))]
|
|
[b1 : Flonum (unsafe-flvector-ref fs (- num-fs 2))]
|
|
) ([m (in-range (- num-fs 3) 0 -1)])
|
|
(define c (unsafe-flvector-ref fs m))
|
|
(define b0 (fl+ c (fl/ (fl* (fl+ (fl m) 1.0) b2) k)))
|
|
(values (fl+ (fl* n sum) b0) b1 b0)))
|
|
sum)
|
|
|
|
(: R-log (Float Float -> (Values Float Float)))
|
|
;; Log-space version of `R' above
|
|
(define (R-log k n)
|
|
(define sum (R-sum k n))
|
|
(values
|
|
(fl- (fl- (fl- (fl+ (fllog (abs sum)) (fl* (fl* (fl* -0.5 k) n) n))
|
|
(fl* 0.5 (fllog k)))
|
|
(fl* 0.5 (fllog (fl* 2.0 pi))))
|
|
(flstirling k))
|
|
(flsgn sum)))
|
|
|
|
(: flgamma-temme (Float Float Any -> Float))
|
|
;; Computes a regularized incomplete gamma using Temme's series
|
|
(define (flgamma-temme k x upper?)
|
|
(define n
|
|
(let ([l (fl/ x k)])
|
|
(fl* (flsgn (fl- l 1.0)) (flsqrt (fl* 2.0 (fl- (fl- l 1.0) (fllog l)))))))
|
|
(define z (let ([z (fl* n (flsqrt k))])
|
|
(if upper? (- z) z)))
|
|
(define r (let ([r (fl* (R-sum k n) (flgamma-series-const k x))])
|
|
(if upper? (- r) r)))
|
|
(cond [(z . fl<= . 0.0) (fl- (standard-flnormal-cdf z) r)]
|
|
[else (fl- 1.0 (fl+ (standard-flnormal-cdf (- z)) r))]))
|
|
|
|
(: fllog-gamma-temme (Float Float Any -> Float))
|
|
(define (fllog-gamma-temme k x upper?)
|
|
(define n (let ([l (fl/ x k)])
|
|
(fl* (flsgn (fl- l 1.0)) (flsqrt (fl* 2.0 (fl- (fl- l 1.0) (fllog l)))))))
|
|
(define z (let ([z (fl* n (flsqrt k))])
|
|
(if upper? (- z) z)))
|
|
(define-values (log-r r-sgn) (let-values ([(log-r r-sgn) (R-log k n)])
|
|
(if upper? (values log-r (- r-sgn)) (values log-r r-sgn))))
|
|
(cond [(z . fl<= . 0.0)
|
|
(define norm-log-p (standard-flnormal-log-cdf z))
|
|
(define log-p (if (r-sgn . fl< . 0.0)
|
|
(lg+ norm-log-p log-r)
|
|
(lg- norm-log-p log-r)))
|
|
;; When norm-log-p ~ log-r, the above log-space arithmetic can go bad
|
|
;; Fortunately, this means we don't need any correctional terms - a normal approximation is
|
|
;; good enough
|
|
(if (rational? log-p) log-p norm-log-p)]
|
|
[else
|
|
(define norm-log-p (standard-flnormal-log-cdf (- z)))
|
|
(define log-p (if (r-sgn . fl< . 0.0)
|
|
(lg1- (lg- norm-log-p log-r))
|
|
(lg1- (lg+ norm-log-p log-r))))
|
|
(if (rational? log-p) log-p (lg1- norm-log-p))]))
|