;;                   NESTED SAMPLING MAIN PROGRAM
;; (GNU General Public License software, (C) Sivia and Skilling 2006)
;; Translated to Lush by Issac Trotts in 2007

(de uniform () (rand)) ; Uniform inside (0,1)

;; Logarithmic addition: return log(a+b) given x=log(a) and y=log(b)
(de plus (x y) 
	(if (> x y) 
		(+ x (log (+ 1 (exp (- y x)))))
		(+ y (log (+ 1 (exp (- x y)))))))

#? (remove! <item> <ls>)
;; Remove all items with a given value from a list.
(dmd remove! (_item _ls)
  ;; The underscores  are a hacky attempt to avoid variable capture
  ;; since Lush doesn't have gensym or hygienic macros.
  `(setq ,_ls (filter (lambda (x) (not (= x ,_item))) ,_ls)))

;; Return a uniform randomly chosen item from a list.
(de random-item (ls)
  (let ((n (length ls)))
    (nth (int (* (rand) n)) ls)))

#? (min2 <objs> <f>)
;; Given a list of objects and a function <f> taking objects to
;; numbers, returns the least value of <f> over the list, and 
;; the first object having this value.
;; 
(de min2 (objs f)
  (let* ((obest (car objs))
         (fbest (f obest)))
    (each ((o (cdr objs)))
      (let ((fo (f o)))
        (when (< fo fbest)
          (setq obest o)
          (setq fbest fo))))
    (list fbest obest)))

(dmd assert (stuff)
  (let ((msg (sprintf "Assertion failed: %l" stuff)))
    `(if (not ,stuff)
         (error ,msg))))

#? (push! <o> <ls>)
;; Pushes an object <o> onto the front of a list <ls>
(dmd push! (o ls)
  `(setq ,ls (cons ,o ,ls)))

;; Returns a list with n items, each of which is the result of function
;; f applied to no arguments.
;;
(de list-from-func (n f)
  (let ((ls '()))
    (for (i 1 n)
         (push! (f) ls))
    ls))

#? (nested_sampling <n> <max_iter> <sample_from_prior> <explore!>)
;;.VP
;; ((-real-) n)                      ; the number of objects to evolve
;; ((-real-) max_iter)               ; maximum number of iterations
;; ((-function-) sample_from_prior)  ; takes no args, returns obj from prior
;; ((-function-) explore!)           ; evolves one obj within constraint
;;
;; This is an implementation of John Skilling's Nested Sampling algorithm
;; for computing the normalizing constant of a probability distribution
;; (usually the posterior in Bayesian inference).
;; 
;; The objects generated from the prior are expected to have slots
;; named logL and logWt.  They must also have a method called "copy".
;;
;; The function <explore!> is expected to modify its argument in place.
;;
;; The return value is an association list with the following entries:
;;      samples
;;      num-iterations
;;      logZ
;;      logZ-sdev
;;      info-nats
;;      info-sdev
;;
;; More information is available here:
;;
;; http://www.inference.phy.cam.ac.uk/bayesys/
;;
(de nested_sampling (n max_iter sample_from_prior explore!)
  (let ((Samples '())       ; Objects stored for posterior results
        (H        0.0)      ; Information, initially 0
        (logZ    -1e300)    ; ln(Evidence), initially ln(0)

        ;; Initialize Objs to an array of n objects from the prior
        (Objs (list-from-func n sample_from_prior))

        ;; ln(width in prior mass)
        ;; Outermost interval of prior mass
        (logwidth (log (- 1.0 (exp (- (/ n)))))))

    (each ((o Objs))
      (if (not o) 
          (error (concat "Found a null object.  Did you remember to return "
                         "an object from the prior sampler?"))))
      
    ;; NESTED SAMPLING LOOP ___________________________________________
    (for (nest 1 max_iter)
         (let* ((min2result (min2 Objs (lambda (o) :o:logL)))
                (worst-logL (car min2result))
                (worst-obj  (cadr min2result))
                (worst-logWt (+ logwidth worst-logL))

                ;; Update Evidence Z and Information H
                (logZnew (plus logZ worst-logWt)))
                
           (setq H (- (+ (* (exp (- worst-logWt logZnew)) worst-logL) 
                         (* (exp (- logZ        logZnew)) (+ H logZ)))
                      logZnew))
           
           (setq :worst-obj:logWt worst-logWt)
           (setq logZ logZnew)
           
           ;; Posterior Samples (optional)
           (push! worst-obj Samples)
           
           (let ((logL* worst-logL)) ; new likelihood constraint
             (if (> n 1) 
                 (progn ; Kill worst obj in favor of copy of different survivor
                   (remove! worst-obj Objs)
                   (let ((my-copy (==> (random-item Objs) copy)))
                     (if (not my-copy)
                         (error (concat "Copied object is null.  "
"Did you forget to return an object from the copy method?")))
                     (explore! my-copy logL*) ; Evolve copy within constraint
                     (push! my-copy Objs)))
               
               (explore! (car Objs) logL*)) ; Evolve the only object

             ;; Shrink interval
             (setq logwidth (- logwidth (/ 1.0 n))))))
    
    ;; Exit with evidence Z, information H, and optional posterior Samples
    (let ((sdev_H (/ H (log 2.)))
          (sdev_logZ  (sqrt (/ H n))))
      `((samples . ,Samples)
        (num-iterations . ,max_iter)
        (logZ . ,logZ)
        (logZ-sdev . ,sdev_logZ)
        (info-nats . ,H)
        (info-sdev . ,sdev_H)))))