144 lines
5.2 KiB
Racket
144 lines
5.2 KiB
Racket
#lang typed/racket/base
|
|
|
|
(require racket/performance-hint
|
|
racket/promise
|
|
"../../flonum.rkt")
|
|
|
|
(provide
|
|
PDF Sample CDF Inverse-CDF
|
|
(struct-out distribution)
|
|
(struct-out ordered-dist)
|
|
Real-Dist
|
|
pdf sample cdf inv-cdf real-dist-prob)
|
|
|
|
(define-type (PDF In)
|
|
(case-> (In -> Flonum)
|
|
(In Any -> Flonum)))
|
|
|
|
(define-type (Sample Out)
|
|
(case-> (-> Out)
|
|
(Integer -> (Listof Out))))
|
|
|
|
(define-type (CDF In)
|
|
(case-> (In -> Flonum)
|
|
(In Any -> Flonum)
|
|
(In Any Any -> Flonum)))
|
|
|
|
(define-type (Inverse-CDF Out)
|
|
(case-> (Real -> Out)
|
|
(Real Any -> Out)
|
|
(Real Any Any -> Out)))
|
|
|
|
(struct: (In Out) distribution ([pdf : (PDF In)]
|
|
[sample : (Sample Out)])
|
|
#:transparent)
|
|
|
|
(struct: (In Out) ordered-dist distribution
|
|
([cdf : (CDF In)]
|
|
[inv-cdf : (Inverse-CDF Out)]
|
|
[min : Out]
|
|
[max : Out]
|
|
[median : (Promise Out)])
|
|
#:transparent)
|
|
|
|
(define-type Real-Dist (ordered-dist Real Flonum))
|
|
|
|
;; =================================================================================================
|
|
|
|
(begin-encourage-inline
|
|
|
|
(: pdf (All (In Out) (case-> ((distribution In Out) In -> Flonum)
|
|
((distribution In Out) In Any -> Flonum))))
|
|
(define (pdf d v [log? #f])
|
|
((distribution-pdf d) v log?))
|
|
|
|
(: sample (All (In Out) (case-> ((distribution In Out) -> Out)
|
|
((distribution In Out) Integer -> (Listof Out)))))
|
|
(define sample
|
|
(case-lambda
|
|
[(d) ((distribution-sample d))]
|
|
[(d n) ((distribution-sample d) n)]))
|
|
|
|
(: cdf (All (In Out) (case-> ((ordered-dist In Out) In -> Flonum)
|
|
((ordered-dist In Out) In Any -> Flonum)
|
|
((ordered-dist In Out) In Any Any -> Flonum))))
|
|
(define (cdf d v [log? #f] [1-p? #f])
|
|
((ordered-dist-cdf d) v log? 1-p?))
|
|
|
|
(: inv-cdf (All (In Out) (case-> ((ordered-dist In Out) Real -> Out)
|
|
((ordered-dist In Out) Real Any -> Out)
|
|
((ordered-dist In Out) Real Any Any -> Out))))
|
|
(define (inv-cdf d p [log? #f] [1-p? #f])
|
|
((ordered-dist-inv-cdf d) p log? 1-p?))
|
|
|
|
)
|
|
|
|
(: real-dist-prob* (Real-Dist Flonum Flonum Any -> Flonum))
|
|
;; Assumes a <= b
|
|
(define (real-dist-prob* d a b 1-p?)
|
|
(define c (force (ordered-dist-median d)))
|
|
(define cdf (ordered-dist-cdf d))
|
|
(define p
|
|
(cond [(a . fl= . b) (if 1-p? 1.0 0.0)]
|
|
[1-p? (fl+ (cdf a #f #f) (cdf b #f #t))]
|
|
[(b . fl<= . c)
|
|
;; Both less than the median; use lower tail only
|
|
(fl- (cdf b #f #f) (cdf a #f #f))]
|
|
[(a . fl>= . c)
|
|
;; Both greater than the median; use upper tail only
|
|
(fl- (cdf a #f #t) (cdf b #f #t))]
|
|
[else
|
|
;; Median between a and b; use lower for (a,c] and upper for (c,b]
|
|
(define P_x<=a (cdf a #f #f))
|
|
(define P_x>b (cdf b #f #t))
|
|
(fl+ (fl- 0.5 P_x<=a) (fl- 0.5 P_x>b))]))
|
|
(max 0.0 (min 1.0 p)))
|
|
|
|
(: real-dist-log-prob* (Real-Dist Flonum Flonum Any -> Flonum))
|
|
;; Assumes a <= b
|
|
(define (real-dist-log-prob* d a b 1-p?)
|
|
(define c (force (ordered-dist-median d)))
|
|
(define cdf (ordered-dist-cdf d))
|
|
(define log-p
|
|
(cond [(a . fl= . b) (if 1-p? 0.0 -inf.0)]
|
|
[1-p? (lg+ (cdf a #t #f) (cdf b #t #t))]
|
|
[(b . fl<= . c)
|
|
;; Both less than the median; use lower tail only
|
|
(define log-P_x<=a (cdf a #t #f))
|
|
(define log-P_x<=b (cdf b #t #f))
|
|
(cond [(log-P_x<=b . fl< . log-P_x<=a) -inf.0]
|
|
[else (lg- log-P_x<=b log-P_x<=a)])]
|
|
[(a . fl>= . c)
|
|
;; Both greater than the median; use upper tail only
|
|
(define log-P_x>a (cdf a #t #t))
|
|
(define log-P_x>b (cdf b #t #t))
|
|
(cond [(log-P_x>a . fl< . log-P_x>b) -inf.0]
|
|
[else (lg- log-P_x>a log-P_x>b)])]
|
|
[else
|
|
;; Median between a and b; try 1-upper first
|
|
(define log-P_x<=a (cdf a #t #f))
|
|
(define log-P_x>b (cdf b #t #t))
|
|
(define log-p (lg1- (lg+ log-P_x<=a log-P_x>b)))
|
|
(cond [(log-p . fl> . (log 0.1)) log-p]
|
|
[else
|
|
;; Subtracting from 1.0 (in log space) lost bits; split and add instead
|
|
(define log-P_a<x<=c
|
|
(cond [((fllog 0.5) . fl< . log-P_x<=a) -inf.0]
|
|
[else (lg- (fllog 0.5) log-P_x<=a)]))
|
|
(define log-P_c<x<=b
|
|
(cond [((fllog 0.5) . fl< . log-P_x>b) -inf.0]
|
|
[else (lg- (fllog 0.5) log-P_x>b)]))
|
|
(lg+ log-P_a<x<=c log-P_c<x<=b)])]))
|
|
(min 0.0 log-p))
|
|
|
|
(: real-dist-prob (case-> (Real-Dist Real Real -> Flonum)
|
|
(Real-Dist Real Real Any -> Flonum)
|
|
(Real-Dist Real Real Any Any -> Flonum)))
|
|
(define (real-dist-prob d a b [log? #f] [1-p? #f])
|
|
(let ([a (fl a)] [b (fl b)])
|
|
(let ([a (flmin a b)] [b (flmax a b)])
|
|
(cond [log? (define p (real-dist-prob* d a b 1-p?))
|
|
(cond [(and (p . fl> . +max-subnormal.0) (p . fl< . 0.9)) (fllog p)]
|
|
[else (real-dist-log-prob* d a b 1-p?)])]
|
|
[else (real-dist-prob* d a b 1-p?)]))))
|