f2dc2027f6
in the original GitHub fork: https://github.com/ntoronto/racket Some things about this are known to be broken (most egregious is that the array tests DO NOT RUN because of a problem in typed/rackunit), about half has no coverage in the tests, and half has no documentation. Fixes and docs are coming. This is committed now to allow others to find errors and inconsistency in the things that appear to be working, and to give the author a (rather incomplete) sense of closure.
209 lines
9.5 KiB
Racket
209 lines
9.5 KiB
Racket
#lang typed/racket/base
|
||
|
||
(require racket/fixnum
|
||
"../../../flonum.rkt"
|
||
"../../functions/beta.rkt"
|
||
"../../functions/incomplete-beta.rkt"
|
||
"beta-pdf.rkt"
|
||
"beta-utils.rkt"
|
||
"normal-inv-cdf.rkt")
|
||
|
||
(provide flbeta-inv-cdf)
|
||
|
||
;; =================================================================================================
|
||
;; Initial approximation
|
||
|
||
(define 1-eps (fl- 1.0 (fl* 0.5 epsilon.0)))
|
||
|
||
(: bound-estimate (Flonum -> Flonum))
|
||
(define (bound-estimate x)
|
||
(flmax +min.0 (flmin 1-eps x)))
|
||
|
||
(: log-const-numer (Flonum Flonum Flonum -> Flonum))
|
||
(define (log-const-numer a b x)
|
||
(fl+ (fl* a (fllog x)) (fl* b (fllog1p (- x)))))
|
||
|
||
(: recip-log-const-numer-diff (Flonum Flonum Flonum -> Flonum))
|
||
(define (recip-log-const-numer-diff a b x)
|
||
(fl/ (fl- (fl* x x) x)
|
||
(fl- (fl* (fl+ a b) x) a)))
|
||
|
||
(: flbeta-inv-log-cdf-appx-asym-0 (Flonum Flonum Flonum -> Flonum))
|
||
;; Approximates by inverting the multiplicative term in front of the hypergeometric series (see
|
||
;; flbeta-regularized-hypergeom in math/private/functions/incomplete-beta)
|
||
;; This is a *really good* approximation in the tails; for the middle, we'll interpolate between
|
||
;; two reasonable points near the tails, or use a normal approximation (whichever is better)
|
||
(define (flbeta-inv-log-cdf-appx-asym-0 a b log-p)
|
||
(define y (fl+ (fl+ log-p (fllog a)) (fllog-beta a b)))
|
||
;; ^^^ the easy part; now we need to invert log(p) = a*log(x) + b*log(1-x)
|
||
(define x0 (flexp (fl/ log-p a))) ; initial guess
|
||
(define x (flmax +min.0 (flmin (fl* 0.9 (fl/ a (fl+ a b))) x0))) ; bound the guess
|
||
;(printf "x = ~v~n" x)
|
||
(define fx (log-const-numer a b x))
|
||
(let loop ([x x] [fx fx] [fac 1.0] [i 0])
|
||
;; Newton says this is the change:
|
||
(define dx (fl* (fl- fx y) (recip-log-const-numer-diff a b x)))
|
||
;; Halve it until we get something that doesn't obviously overshoot, and adjust `fac'
|
||
(define-values (new-x new-dx new-fac)
|
||
(let: dx-loop : (Values Flonum Flonum Flonum)
|
||
([dx : Flonum (fl* fac (flmax -0.1 (flmin 0.1 dx)))]
|
||
[fac : Flonum fac]
|
||
[j : Nonnegative-Fixnum 0])
|
||
(define new-x (- x dx))
|
||
;(printf "dx: new-x = ~v dx = ~v~n" new-x dx)
|
||
(cond [(and (new-x . fl> . 0.0) (new-x . fl<= . 1.0)) (values new-x dx fac)]
|
||
[(not (rational? new-x)) (values x dx fac)]
|
||
[(j . fx< . 1000) (dx-loop (fl* 0.5 dx) (fl* 0.5 fac) (fx+ j 1))]
|
||
[else (values x dx fac)])))
|
||
;(printf "x = ~v dx = ~v fac = ~v~n" new-x (- new-x x) new-fac)
|
||
(define new-fx (log-const-numer a b new-x))
|
||
(cond [(or ((flabs (fl- fx new-fx)) . fl<= . (flabs (fl* (fl* 1000.0 epsilon.0) new-fx)))
|
||
(fl<= (flabs new-dx) (fl* (fl* 0.5 epsilon.0) new-x))
|
||
(not (rational? new-x)))
|
||
new-x]
|
||
[(i . fx< . 1000)
|
||
(loop new-x new-fx (flmin 1.0 (fl* new-fac 2.0)) (fx+ i 1))]
|
||
[else
|
||
new-x])))
|
||
|
||
(: flbeta-inv-log-cdf-appx-asym (Flonum Flonum Flonum Flonum -> Flonum))
|
||
(define (flbeta-inv-log-cdf-appx-asym a b log-p log-1-p)
|
||
(define σ (flbeta-stddev a b))
|
||
(define log-p0
|
||
(fl- (fl- (log-const-numer a b (bound-estimate (fl- (fl/ a (fl+ a b)) (fl* 0.85 σ))))
|
||
(fllog a))
|
||
(fllog-beta a b)))
|
||
(define log-1-p1
|
||
(fl- (fl- (log-const-numer b a (bound-estimate (fl- (fl/ b (fl+ a b)) (fl* 0.85 σ))))
|
||
(fllog b))
|
||
(fllog-beta a b)))
|
||
(cond [(log-p . fl< . log-p0)
|
||
(flbeta-inv-log-cdf-appx-asym-0 a b log-p)]
|
||
[(log-1-p . fl< . log-1-p1)
|
||
(fl- 1.0 (flbeta-inv-log-cdf-appx-asym-0 b a log-1-p))]
|
||
[else
|
||
(define x0 (flbeta-inv-log-cdf-appx-asym-0 a b log-p0))
|
||
(define x1 (fl- 1.0 (flbeta-inv-log-cdf-appx-asym-0 b a log-1-p1)))
|
||
(define s (fl/ (fl- x1 x0) (fl- (lg1- log-1-p1) log-p0)))
|
||
(define c (fl- x0 (fl* s log-p0)))
|
||
(fl+ (fl* s log-p) c)]))
|
||
|
||
(: flbeta-inv-log-cdf-appx-normal (Flonum Flonum Flonum Flonum -> Flonum))
|
||
(define (flbeta-inv-log-cdf-appx-normal a b log-p log-1-p)
|
||
(define m (flbeta-appx-median a b))
|
||
(define σ (flbeta-stddev a b))
|
||
(fl+ m (fl* σ (standard-flnormal-inv-log-cdf log-p))))
|
||
|
||
(: flbeta-inv-log-cdf-appx (Flonum Flonum Flonum Flonum -> (Values Flonum Flonum)))
|
||
(define (flbeta-inv-log-cdf-appx a b log-p log-1-p)
|
||
(cond [(or (a . fl< . 10.0) (b . fl< . 10.0))
|
||
(define x (bound-estimate (flbeta-inv-log-cdf-appx-asym a b log-p log-1-p)))
|
||
(values x (fllog-beta-inc a b x #f #t))]
|
||
[else
|
||
(define x0 (bound-estimate (flbeta-inv-log-cdf-appx-asym a b log-p log-1-p)))
|
||
(define x1 (bound-estimate (flbeta-inv-log-cdf-appx-normal a b log-p log-1-p)))
|
||
(define real-log-p0 (fllog-beta-inc a b x0 #f #t))
|
||
(define real-log-p1 (fllog-beta-inc a b x1 #f #t))
|
||
(if ((flabs (fl- log-p real-log-p0)) . fl< . (flabs (fl- log-p real-log-p1)))
|
||
(values x0 real-log-p0)
|
||
(values x1 real-log-p1))]))
|
||
|
||
;; =================================================================================================
|
||
;; Newton's method
|
||
|
||
(: newton-lower-log-iter (Flonum Flonum Flonum Flonum Flonum Flonum -> (Values Flonum Flonum)))
|
||
(define (newton-lower-log-iter a b log-p x real-log-p fac)
|
||
(define pdf-log-p (flbeta-log-pdf a b x))
|
||
(cond [(fl<= (flabs (fl- log-p real-log-p))
|
||
(flabs (fl* (fl* 100.0 epsilon.0) log-p)))
|
||
(values 0.0 fac)]
|
||
[else
|
||
(define new-dx (fl* (flexp (fl- real-log-p pdf-log-p)) (fl- log-p real-log-p)))
|
||
;; Limit dx, then halve until the x+dx is in bounds
|
||
(let loop ([new-dx (fl* fac (flmin 0.1 (flmax -0.1 new-dx)))]
|
||
[fac fac]
|
||
[#{j : Nonnegative-Fixnum} 0])
|
||
(define new-x (fl+ x new-dx))
|
||
(cond [(and (new-x . fl>= . 0.0) (new-x . fl<= . 1.0))
|
||
(values new-dx fac)]
|
||
[(not (rational? new-x))
|
||
(values 0.0 fac)]
|
||
[(j . fx< . 1000)
|
||
(loop (fl* 0.5 new-dx) (fl* 0.5 fac) (fx+ j 1))]
|
||
[else
|
||
(values 0.0 fac)]))]))
|
||
|
||
(: flbeta-inv-log-cdf-newton (Flonum Flonum Flonum Flonum Flonum -> Flonum))
|
||
(define (flbeta-inv-log-cdf-newton a b log-p x real-log-p)
|
||
(let loop ([dx 0.0] [x x] [real-log-p real-log-p] [fac 1.0] [#{c : Nonnegative-Fixnum} 1])
|
||
(define-values (new-dx new-fac) (newton-lower-log-iter a b log-p x real-log-p fac))
|
||
(define new-x (fl+ x new-dx))
|
||
(define new-real-log-p (fllog-beta-inc a b new-x #f #t))
|
||
;(printf "~v ~v ~v~n" new-x new-dx new-fac)
|
||
(cond [(or (fl<= (flabs (fl- real-log-p new-real-log-p))
|
||
(flabs (fl* (fl* 1000.0 epsilon.0) new-real-log-p)))
|
||
(fl<= (flabs new-dx) (fl* (fl* 0.5 epsilon.0) new-x))
|
||
(not (rational? new-x)))
|
||
new-x]
|
||
[(c . fx< . 1000)
|
||
(let ([new-fac (if (fl= (flsgn dx) (flsgn new-dx))
|
||
(flmin 1.0 (fl* new-fac 2.0))
|
||
(fl* (flmin new-fac 1.0) 0.5))])
|
||
(loop new-dx new-x (fllog-beta-inc a b new-x #f #t) new-fac (fx+ c 1)))]
|
||
[else
|
||
new-x])))
|
||
|
||
;; =================================================================================================
|
||
|
||
(: in-bounds? (Flonum Flonum Flonum -> Boolean))
|
||
(define (in-bounds? a b log-p)
|
||
(and (log-p . fl> . -inf.0) (log-p . fl< . 0.0)
|
||
(a . fl> . 0.0) (a . fl< . +inf.0)
|
||
(b . fl> . 0.0) (b . fl< . +inf.0)))
|
||
|
||
(: flbeta-inv-log-cdf-limits (Flonum Flonum Flonum -> Flonum))
|
||
(define (flbeta-inv-log-cdf-limits a b log-p)
|
||
(cond [(not (and (log-p . fl<= . 0.0) (a . fl>= . 0.0) (b . fl>= . 0.0))) +nan.0]
|
||
[(fl= log-p -inf.0) 0.0]
|
||
[(and (fl= a 0.0) (fl= b 0.0)) (if (log-p . fl< . (fllog 0.5)) 0.0 1.0)]
|
||
[(and (fl= a +inf.0) (fl= b +inf.0)) 0.5]
|
||
[(fl= a +inf.0) 1.0]
|
||
[(fl= b +inf.0) 0.0]
|
||
[(fl= a 0.0) 0.0]
|
||
[(fl= b 0.0) 1.0]
|
||
[(fl= log-p 0.0) 1.0]
|
||
[else +nan.0]))
|
||
|
||
(: flbeta-inv-log-cdf* (Flonum Flonum Flonum Flonum -> Flonum))
|
||
(define (flbeta-inv-log-cdf* a b log-p log-1-p)
|
||
(cond [(not (in-bounds? a b log-p))
|
||
(flbeta-inv-log-cdf-limits a b log-p)]
|
||
[else
|
||
(define mid-log-p
|
||
(if (a . fl> . b)
|
||
(lg1- (fllog-beta-inc b a (fl/ b (fl+ a b)) #f #t))
|
||
(fllog-beta-inc a b (fl/ a (fl+ a b)) #f #t)))
|
||
(let-values ([(a b log-p log-1-p 1-?) (if (log-p . fl< . mid-log-p)
|
||
(values a b log-p log-1-p #f)
|
||
(values b a log-1-p log-p #t))])
|
||
(define-values (x0 real-log-p) (flbeta-inv-log-cdf-appx a b log-p log-1-p))
|
||
(define x (flbeta-inv-log-cdf-newton a b log-p x0 real-log-p))
|
||
(if 1-? (fl- 1.0 x) x))]))
|
||
|
||
(: flbeta-inv-cdf (Flonum Flonum Flonum Any Any -> Flonum))
|
||
(define (flbeta-inv-cdf a b p log? 1-p?)
|
||
(cond [log?
|
||
(cond [(p . fl<= . 0.0)
|
||
(let-values ([(log-p log-1-p) (cond [1-p? (values (lg1- p) p)]
|
||
[else (values p (lg1- p))])])
|
||
(flbeta-inv-log-cdf* a b log-p log-1-p))]
|
||
[else
|
||
+nan.0])]
|
||
[else
|
||
(cond [(and (p . fl>= . 0.0) (p . fl<= . 1.0))
|
||
(let-values ([(log-p log-1-p) (cond [1-p? (values (fllog1p (- p)) (fllog p))]
|
||
[else (values (fllog p) (fllog1p (- p)))])])
|
||
(flbeta-inv-log-cdf* a b log-p log-1-p))]
|
||
[else
|
||
+nan.0])]))
|