;;;; srfi-27-distributions.scm
;;;; Kon Lovett, May '06

; Chicken Generic Arithmetic! (could use fp routines)

(module srfi-27-distributions

  (;export
    make-random-exponentials
    make-random-normals
    make-random-triangles
    make-random-poissons
    make-random-bernoullis
    make-random-binomials
    make-random-geometrics
    make-random-lognormals
    make-random-cauchys
    make-random-gammas
    make-random-erlangs
    make-random-paretos
    make-random-levys
    make-random-weibulls)

  (import
    scheme
    chicken
    (only type-errors
      error-argument-type)
    (only type-checks
      check-procedure
      check-cardinal-integer
      check-real
      check-open-interval
      check-closed-interval)
    (only srfi-27-uniform-random
      make-uniform-random-reals))

  (require-library
    type-errors type-checks
    srfi-27-uniform-random)

;;; Chicken Generic Arithmetic Argument Checks

(define (check-nonzero-real loc obj #!optional argnam)
  (unless (and (real? obj) (not (zero? obj)))
    (error-argument-type loc obj "nonzero-real" argnam))
  obj )

(define (check-nonnegative-real loc obj #!optional argnam)
  (unless (and (real? obj) (not (negative? obj)))
    (error-argument-type loc obj "nonnegative-real" argnam))
  obj )

(define (check-positive-real loc obj #!optional argnam)
  (unless (and (real? obj) (positive? obj))
    (error-argument-type loc obj "positive-real" argnam))
  obj )

(define (check-real-open-interval loc obj mn mx #!optional argnam)
  (check-real loc obj argnam)
  (check-open-interval loc obj mn mx argnam)
  obj )

(define (check-real-closed-interval loc obj mn mx #!optional argnam)
  (check-real loc obj argnam)
  (check-closed-interval loc obj mn mx argnam)
  obj )

#;
(define (check-real-precision loc obj #!optional argnam)
  (check-real-open-interval loc obj 0 1 argnam)
  obj )

(define (check-real-unit loc obj #!optional argnam)
  (check-real-closed-interval loc obj 0 1 argnam)
  obj )

;;;

(define-constant PI     3.1415926535897932384626433832795028841972)
(define-constant FP1/3  0.3333333333333333333333333333333333333333)

;;;

; (in case special processing needed near limits TBD)
(define (*reciprocal n) (/ 1.0 n))
(define (*-reciprocal n) (/ -1.0 n))

;;; Normal distribution

;; Knuth's "The Art of Computer Programming", Vol. II, 2nd ed.,
;; Algorithm P of Section 3.4.1.C.

(define (*make-random-normals mu sigma randoms)
  (let ((next #f))
    (lambda ()
      (if next
        (let ((result next))
          (set! next #f)
          (+ mu (* sigma result)))
        (let loop ()
          (let* ((v1 (- (* 2.0 (randoms)) 1.0))
                 (v2 (- (* 2.0 (randoms)) 1.0))
                 (s (+ (* v1 v1) (* v2 v2))) )
            (if (<= 1.0 s)
              (loop)
              (let ((scale (sqrt (/ (* -2.0 (log s)) s))))
                (set! next (* scale v2))
                (+ mu (* sigma scale v1))))))))) )

(define (make-random-normals
          #!key (mu 0.0) (sigma 1.0) (randoms (make-uniform-random-reals)))
  (check-real 'make-random-normals mu 'mu)
  (check-nonzero-real 'make-random-normals sigma 'sigma)
  (check-procedure 'make-random-normals randoms 'randoms)
  (values
    (*make-random-normals mu sigma randoms)
    (lambda () (values mu sigma randoms))) )

;;; Exponential distribution

;; Knuth's "The Art of Computer Programming", Vol. II, 2nd ed.,
;; Section 3.4.1.D.

(define (*make-random-exponentials mu randoms)
  (if (= 1.0 mu)
    (lambda () (- (log (randoms))))
    (lambda () (* mu (- (log (randoms)))))) )

(define (make-random-exponentials
          #!key (mu 1.0) (randoms (make-uniform-random-reals)))
  (check-real-unit 'make-random-exponentials mu 'mu)
  (check-procedure 'make-random-exponentials randoms 'randoms)
  (values
    (*make-random-exponentials mu randoms)
    (lambda () (values mu randoms))) )

;;; Triangle distribution

;; s - smallest, m - most probable, l - largest

(define (*make-random-triangles s m l randoms)
  (let ((d1 (- m s))
        (d2 (- l s))
        (d3 (sqrt (- l m))))
    (let ((q1 (/ d1 d2))
          (p1 (sqrt (* d1 d2))))
      (lambda ()
        (let ((u (randoms)))
          (if (<= u q1)
            (+ s (* p1 (sqrt u)))
            (- l (* d3 (sqrt (- (* d2 u) d1))))))))) )

(define (make-random-triangles
          #!key (s 0.0) (m 0.5) (l 1.0) (randoms (make-uniform-random-reals)))
  (check-real 'make-random-triangles s 's)
  (check-real 'make-random-triangles m 'm)
  (check-real 'make-random-triangles l 'l)
  (check-real-open-interval 'make-random-triangles l s +inf.0 'l)
  (check-real-closed-interval 'make-random-triangles m s l 'm)
  (check-procedure 'make-random-triangles randoms 'randoms)
  (values
    (*make-random-triangles s m l randoms)
    (lambda () (values s m l randoms))) )

;;; Poisson distribution

(define (*make-random-poissons mu randoms)
  (let ((emu (exp (- mu))))
    (lambda ()
      ; FIXME O(mu) but O(log(mu)) desired for >> mu
      (do ((m 0 (fx+ 1 m))
           (prod (randoms) (* prod (randoms))))
          ((<= prod emu) m)))) )

(define (make-random-poissons
          #!key (mu 1.0) (randoms (make-uniform-random-reals)))
  (check-nonnegative-real 'make-random-poissons mu 'mu)
  (check-procedure 'make-random-poissons randoms 'randoms)
  (values
    (*make-random-poissons mu randoms)
    (lambda () (values mu randoms))) )

;;; Bernoulli distribution

(define (*make-random-bernoullis p randoms)
  (cond
    ((= 0.0 p) (lambda () #f))
    ((= 1.0 p) (lambda () #t))
    (else      (lambda () (<= (randoms) p)))) )

(define (make-random-bernoullis
          #!key (p 0.5) (randoms (make-uniform-random-reals)))
  (check-real-unit 'make-random-bernoullis p 'p)
  (check-procedure 'make-random-bernoullis randoms 'randoms)
  (values
    (*make-random-bernoullis p randoms)
    (lambda () (values p randoms))) )

;;; Binomial distribution

(define (*make-random-binomials t p randoms)
  (let ((bernoullis (*make-random-bernoullis p randoms)))
    ;FIXME O(t) but O(log(t)) desired for >> t
    (if (fixnum? t)
      (lambda ()
        (do ((i 0 (fx+ 1 i))
             (n 0 (if (bernoullis) (fx+ 1 n) n)))
            ((fx<= t i) n)))
      (lambda ()
        (do ((i 0 (add1 i))
             (n 0 (if (bernoullis) (add1 n) n)))
            ((<= t i) n))))) )

(define (make-random-binomials
          #!key (t 1) (p 0.5) (randoms (make-uniform-random-reals)))
  (check-cardinal-integer 'make-random-binomials t 't)
  (check-real-unit 'make-random-binomials p 'p)
  (check-procedure 'make-random-binomials randoms 'randoms)
  (values
    (*make-random-binomials t p randoms)
    (lambda () (values t p randoms))) )

;;; Geometric distribution

(define (*make-random-geometrics p randoms)
  (let ((log-p (log p)))
    (lambda ()
      (+ 1 (inexact->exact (floor (/ (log (- 1.0 (randoms))) log-p)))))) )

(define (make-random-geometrics
          #!key (p 0.5) (randoms (make-uniform-random-reals)))
  (check-real-unit 'make-random-geometrics p 'p)
  (check-procedure 'make-random-geometrics randoms 'randoms)
  (values
    (*make-random-geometrics p randoms)
    (lambda () (values p randoms))) )

;;; Lognormal distribution

(define (*make-random-lognormals mu sigma randoms)
  (let ((normals (*make-random-normals 0.0 1.0 randoms))
        (nmu (log (* mu (/ mu (sqrt (+ (* sigma sigma) (* mu mu)))))))
        (nsigma (sqrt (log (+ 1.0 (* sigma (/ sigma mu mu)))))) )
    (lambda ()
      (exp (+ nmu (* (normals) nsigma))))) )

(define (make-random-lognormals
          #!key (mu 1.0) (sigma 1.0) (randoms (make-uniform-random-reals)))
  (check-nonzero-real 'make-random-lognormals mu 'mu)
  (check-nonnegative-real 'make-random-lognormals sigma 'sigma)
  (check-procedure 'make-random-lognormals randoms 'randoms)
  (values
    (*make-random-lognormals mu sigma randoms)
    (lambda () (values mu sigma randoms))) )

;;; Cauchy distribution

(define (*make-random-cauchys median sigma randoms)
  (lambda () (+ median (* sigma (tan (* PI (- (randoms) 0.5)))))) )

(define (make-random-cauchys
          #!key (median 0.0) (sigma 1.0) (randoms (make-uniform-random-reals)))
  (check-real 'make-random-cauchys median 'median)
  (check-positive-real 'make-random-cauchys sigma 'sigma)
  (check-procedure 'make-random-cauchys randoms 'randoms)
  (values
    (*make-random-cauchys median sigma randoms)
    (lambda () (values median sigma randoms))) )

;;; Gamma distribution

;; "A Simple Method for Generating Gamma Variables", George Marsaglia & Wai Wan Tsang,
;; ACM Transactions on Mathematical Software, Vol. 26, No. 3, September 2000, Pages 363 372.

(define (*make-random-gammas alpha theta randoms)
  (if (= 1.0 alpha)
    ; then special case
    (lambda () (* theta (- (log (randoms)))) )
    ; else general case
    (let ((norms (*make-random-normals 0.0 1.0 randoms))
          (unis
            (if (< alpha 1.0)
              (let ((inv-alpha (*reciprocal alpha)))
                (lambda () (expt (randoms) inv-alpha) ) )
              randoms)))
      (let* ((d (- (if (< alpha 1.0) (+ 1.0 alpha) alpha) FP1/3))
             (c (*reciprocal (sqrt (* 9.0 d)))))
        (lambda ()
          (*
            theta
            (let loop ()
              (let* ((x (norms))
                     (v (+ 1.0 (* c x))))
                (if (and
                      (< 0.0 v)
                      (let ((v (* v v v))
                            (u (unis))
                            (x^2 (* x x)))
                        (or
                          (< u (- 1.0 (* 0.0331 x^2 x^2)))
                          (< (log u) (+ (* 0.5 x^2) (* d (- 1.0 (+ v (log v)))))))))
                   (* d v)
                   (loop) ) ) ) ) ) ) ) ) )

(define (make-random-gammas
          #!key (alpha 1.0) (theta 1.0) (randoms (make-uniform-random-reals)))
  (check-positive-real 'make-random-gammas alpha 'alpha)
  (check-positive-real 'make-random-gammas theta 'theta)
  (check-procedure 'make-random-gammas randoms 'randoms)
  (values
    (*make-random-gammas alpha theta randoms)
    (lambda () (values alpha theta randoms))) )

;;; Erlang distribution

(define (*make-random-erlangs alpha theta randoms)
  (*make-random-gammas (exact->inexact alpha) (exact->inexact theta) randoms) )

(define (make-random-erlangs
          #!key (alpha 1) (theta 1.0) (randoms (make-uniform-random-reals)))
  (check-positive-real 'make-random-erlangs alpha 'alpha)
  (check-positive-real 'make-random-erlangs theta 'theta)
  (check-procedure 'make-random-erlangs randoms 'randoms)
  (values
    (*make-random-erlangs alpha theta randoms)
    (lambda () (values alpha theta randoms))) )

;;; Pareto distribution

(define (*make-random-paretos alpha xmin randoms)
  (let ((gammas (*make-random-gammas alpha (*reciprocal xmin) randoms)))
    (*make-random-exponentials 1.0 (lambda () (*reciprocal (+ xmin (gammas)))))) )

(define (make-random-paretos
          #!key (alpha 1.0) (xmin 1.0) (randoms (make-uniform-random-reals)))
  (check-positive-real 'make-random-paretos alpha 'alpha)
  (check-positive-real 'make-random-paretos xmin 'xmin)
  (check-procedure 'make-random-paretos randoms 'randoms)
  (values
    (*make-random-paretos alpha xmin randoms)
    (lambda () (values alpha xmin randoms))) )

;;; Levy distribution

;; See Stable Distributions - John P. Nolan, Formula 1.12

(define (*make-random-levys gamma delta randoms)
  (if (and (= 1.0 gamma) (= 0.0 delta))
    (lambda () (let ((r (randoms))) (*reciprocal (* r r))))
    (lambda () (let ((r (randoms))) (+ delta (* gamma (*reciprocal (* r r))))))) )

(define (make-random-levys
          #!key (gamma 1.0) (delta 0.0) (randoms (make-uniform-random-reals)))
  (check-nonnegative-real 'make-random-levys delta 'delta)
  (check-positive-real 'make-random-levys gamma 'gamma)
  (check-procedure 'make-random-levys randoms 'randoms)
  (values
    (*make-random-levys gamma delta randoms)
    (lambda () (values gamma delta randoms))) )

;;; Weibull distribution

(define (*make-random-weibulls shape scale randoms)
  (let ((invscale (*-reciprocal scale))
        (invshape (*reciprocal shape)) )
    (lambda () (expt (* invscale (log (- 1.0 (randoms)))) invshape)) ) )

(define (make-random-weibulls
          #!key (shape 1.0) (scale 1.0) (randoms (make-uniform-random-reals)))
  (check-positive-real 'make-random-weibulls shape 'shape)
  (check-positive-real 'make-random-weibulls scale 'scale)
  (check-procedure 'make-random-weibulls randoms 'randoms)
  (values
    (*make-random-weibulls shape scale randoms)
    (lambda () (values shape scale randoms))) )

) ;module srfi-27-distributions