mcmc.rkt

#lang typed/racket/base

(provide (all-defined-out))

(require racket/flonum)
(require racket/fixnum)
(require racket/require)
(require racket/unsafe/ops)

;; 19.402 seconds with assignments for 50,000,000 samples
;; 18.868 seconds without assignments
;; 16.788 seconds with fl/fx specific operations and f64vector
;; essentially ditto with unsafe variants

(: compiled-mcmc : (-> Flonum) (-> Nonnegative-Integer) Fixnum -> (Vector Fixnum Fixnum))
(define (compiled-mcmc random-real random-bit N)
  
  (: cpt_alarm_*_0_* FlVector)
  (define cpt_alarm_*_0_* (flvector 0.999d0 0.001d0
                                     0.060d0 0.940d0))
  (: cpt_alarm_0_*_* FlVector)
  (define cpt_alarm_0_*_* (flvector 0.999d0 0.001d0
                                     0.710d0 0.290d0))
  (: cpt_alarm_*_*_0 FlVector)
  (define cpt_alarm_*_*_0 (flvector 0.999d0 0.710d0
                                     0.060d0 0.050d0))
  (: cpt_johncalls_*_1 FlVector)
  (define cpt_johncalls_*_1 (flvector 0.050d0 0.900d0))
  (: cpt_marycalls_*_1 FlVector)
  (define cpt_marycalls_*_1 (flvector 0.010d0 0.700d0))
  
  (: random-biased-bit : Float -> Byte)
  (define (random-biased-bit p)
    (if (fl> (random-real) p) 1 0))
  
  (let: loop : (Vector Fixnum Fixnum)
    ((countTrue : Fixnum 0)
     (burglary : Nonnegative-Integer (random-bit))
     (earthquake : Nonnegative-Integer (random-bit))
     (alarm : Nonnegative-Integer (random-bit))
     (j : Fixnum 1))
    (if (fx>= j N)
        (vector (unsafe-fx- (unsafe-fx* 3 N) countTrue) countTrue)
        (let* (;; sample burglary
               (countTrue (unsafe-fx+ burglary countTrue))
               (p_burglary_0 (fl* 0.999d0 (unsafe-flvector-ref
                                           cpt_alarm_0_*_*
                                           (unsafe-fx+ (unsafe-fx* 2 earthquake) alarm))))
               (burglary (random-biased-bit p_burglary_0))
               
               ;; sample earthquake
               (countTrue (unsafe-fx+ burglary countTrue))
               (p_earthquake_0 (fl* 0.998d0 (unsafe-flvector-ref 
                                             cpt_alarm_*_0_* 
                                             (unsafe-fx+ (unsafe-fx* 2  burglary) earthquake))))
               (earthquake (random-biased-bit p_earthquake_0))
               
               ;; sample alarm
               (countTrue (unsafe-fx+ burglary countTrue))
               (p_alarm_0 (fl* (unsafe-flvector-ref
                                cpt_alarm_*_*_0 (unsafe-fx+ (unsafe-fx* 2  burglary) earthquake))
                               (fl* (unsafe-flvector-ref cpt_johncalls_*_1 alarm)
                                    (unsafe-flvector-ref cpt_marycalls_*_1 alarm))))
               (alarm (random-biased-bit p_alarm_0))
               (j (unsafe-fx+ 1 j)))
          (loop countTrue burglary earthquake alarm j)))))



(: normalize : (Vectorof Fixnum) -> FlVector)
(define (normalize counts)
  (: sum : Integer)
  (define sum (for/sum ([c (in-vector counts)]) c))
  (define sumfl (exact->inexact sum))
  (for/flvector #:length (vector-length counts) ((x (in-vector counts)))
    (fl/ (fx->fl x) sumfl)))

(module+ main
  (require racket/cmdline)
  (define r (current-pseudo-random-generator))
  (define iters (command-line #:args (#{iterations : String}) 
                          (cast (string->number iterations) Fixnum)))
   (normalize 
    (compiled-mcmc (λ () (random r)) (λ () (random 2 r)) iters)))