;; apply.c     "LIGHTHOUSE" NESTED SAMPLING APPLICATION
;; (GNU General Public License software, (C) Sivia and Skilling 2006)
;; Translated to Lush by Issac Trotts in 2007
;;
;;              u=0                                 u=1
;;               -------------------------------------
;;          y=2 |:::::::::::::::::::::::::::::::::::::| v=1
;;              |::::::::::::::::::::::LIGHT::::::::::|
;;         north|::::::::::::::::::::::HOUSE::::::::::|
;;              |:::::::::::::::::::::::::::::::::::::|
;;              |:::::::::::::::::::::::::::::::::::::|
;;          y=0 |:::::::::::::::::::::::::::::::::::::| v=0
;; --*--------------*----*--------*-**--**--*-*-------------*--------
;;             x=-2          coastline -->east      x=2
;; Problem:
;;  Lighthouse at (x,y) emitted n flashes observed at D[.] on coast.
;; Inputs:
;;  Prior(u)    is uniform (=1) over (0,1), mapped to x = 4*u - 2; and
;;  Prior(v)    is uniform (=1) over (0,1), mapped to y = 2*v; so that
;;  Position    is 2-dimensional -2 < x < 2, 0 < y < 2 with flat prior
;;  Likelihood  is L(x,y) = PRODUCT[k] (y/pi) / ((D[k] - x)^2 + y^2)
;; Outputs:
;;  Evidence    is Z = INTEGRAL L(x,y) Prior(x,y) dxdy
;;  Posterior   is P(x,y) = L(x,y) / Z estimating lighthouse position
;;  Information is H = INTEGRAL P(x,y) log(P(x,y)/Prior(x,y)) dxdy

(load "mininest.lsh")

(defclass LighthouseObj object 
  u      ; Uniform-prior controlling parameter for x
  v      ; Uniform-prior controlling parameter for y
  x      ; Geographical easterly position of lighthouse
  y      ; Geographical northerly position of lighthouse
  logL   ; logLikelihood = ln Prob(data | position)
  logWt) ; log(Weight), adding to SUM(Wt) = Evidence Z

;; It would be handy to have a macro called defclass-with-copy
;; that would write this method automatically.
;;
(defmethod LighthouseObj copy () 
  (let ((ret (new LighthouseObj)))
    (setq :ret:u :this:u)
    (setq :ret:v :this:v)
    (setq :ret:x :this:x)
    (setq :ret:y :this:y)
    (setq :ret:logL :this:logL)
    (setq :ret:logWt :this:logWt)
    ret))

;; logLikelihood function
(de logLhood(x  ; Easterly position
             y) ; Northerly position
  (let ((D '(4.73  0.45 -1.73  1.09  2.19  0.12
           1.31  1.00  1.32  1.07  0.86 -0.49 -2.59  1.73  2.11
           1.61  4.98  1.71  2.23 -57.20  0.96  1.25 -1.56  2.45
           1.19  2.17 -10.66  1.91 -4.16  1.92  0.10  1.98 -2.51
           5.55 -0.47  1.91  0.95 -0.78 -0.84  1.72 -0.01  1.48
           2.70  1.21  4.41 -4.79  1.33  0.81  0.20  1.58  1.29
           16.19  2.75 -2.38 -1.79  6.50 -18.53  0.72  0.94  3.64
           1.94 -0.11  1.57  0.57)) ; data

        (logL 0.0))          ; logLikelihood accumulator

    (each ((d D))
      (+= logL (log (/ (/ y 3.1416) 
                       (+ (square (- d x)) (* y y))))))
    logL))

(de sample-from-prior ()
  (let ((obj (new LighthouseObj)))
    (setq :obj:u (uniform))
    (setq :obj:v (uniform))
    (setq :obj:x (- (* 4.0 :obj:u) 2.0)) ; map to x
    (setq :obj:y (* 2.0 :obj:v))          ;; map to y
    (setq :obj:logL (logLhood :obj:x :obj:y))
    obj))

(dmd *= (x amt)
    `(setq ,x (* ,x ,amt)))

(dmd /= (x amt)
    `(setq ,x (/ ,x ,amt)))

(dmd -= (x amt)
    `(setq ,x (- ,x ,amt)))

;; Evolve object within likelihood constraint
(de explore! (Obj    ; Object being evolved
              logL*) ; Likelihood constraint L > Lstar

  (let ((step 0.1) ; Initial guess suitable step-size in (0,1)
        (accept 0) ; # MCMC acceptances
        (reject 0) ; # MCMC rejections
        (Try (new LighthouseObj))) ; Trial object

    (for (m 1 20) ; MCMC counter (pre-judged # steps)
;; Update trial object
        (setq :Try:u (+ :Obj:u (* step (- (* 2. (uniform)) 1.)))) ; |move| < step
        (setq :Try:v (+ :Obj:v (* step (- (* 2. (uniform)) 1.)))) ; |move| < step
        (-= :Try:u (int :Try:u))      ; wraparound to stay within (0,1)
        (-= :Try:v (int :Try:v))      ; wraparound to stay within (0,1)
        (setq :Try:x (- (* 4.0 :Try:u) 2.0))  ; map to x
        (setq :Try:y (* 2.0 :Try:v))          ; map to y
        (setq :Try:logL (logLhood :Try:x :Try:y))  ; trial likelihood value

;; Accept if and only if within hard likelihood constraint
        (if (> :Try:logL logL*)
            (progn
                (setq Obj Try) 
                (incr accept))
            (incr reject))
;; Refine step-size to let acceptance ratio converge around 50%
        (if (> accept reject)    (*= step (exp (/ accept))))
        (if (< accept reject)    (/= step (exp (/ reject)))))))

(de square (x) (* x x))


(let* ((n 100) ; # Objects
       (max-iter 1000) ; # iterates
       (results   (nested-sampling n max-iter sample-from-prior explore!))
       (samples   (alist-get 'samples results))
       (nest      (alist-get 'num_iterations results))
       (logZ      (alist-get 'logZ results))
       (logZ-sdev (alist-get 'logZ-sdev results))
       (info      (alist-get 'info-nats results))
       (x 0.0) (xx 0.0)   ; 1st and 2nd moments of x
       (y 0.0) (yy 0.0))  ; 1st and 2nd moments of y
  (each ((s samples))
    (let ((w (exp (- :s:logWt logZ)))) ; Proportional weight
      (+= x (* w :s:x))
      (+= xx (* w (square :s:x)))
      (+= y  (* w :s:y))
      (+= yy (* w (square :s:y)))))
  
  (printf "Evidence: ln(Z) = %g +- %g\n" logZ logZ-sdev)
  (printf "Information: H = %g nats = %g bits\n"
          info (/ info (log 2.0)))
  (printf "mean(x) = %g, stddev(x) = %g\n" x (sqrt (- xx (* x x))))
  (printf "mean(y) = %g, stddev(y) = %g\n" y (sqrt (- yy (* y y)))))