racket/collects/tests/mzscheme/em-imp.ss
2005-05-27 18:56:37 +00:00

468 lines
15 KiB
Scheme

;;; -*- scheme -*-
;;; Fortran-style implementation of an EM clustering algorithm.
;;;
;;; Written by Jeffrey Mark Siskind (qobi@cs.toronto.edu)
;;; R4RS-ified by by Lars Thomas Hansen (lth@cs.uoregon.edu)
;;; Random number generator by Ozan Yigit.
;;;
;;; To run: (run-benchmark)
;;; You must provide your own timer function.
;;;
;;; Some benchmark times:
;;;
;;; Chez Scheme 4.1 for SunOS running on Sparc 10/51 (1MB,96MB,50MHz), Solaris:
;;; Optimize-level 2: 112s run (CPU), 2.8s gc, 326 MB allocated, 1181 GCs
;;; Optimize-level 3: 79s run (CPU), 2.8s gc, 326 MB allocated, 1163 GCs
(define make-model vector)
(define (model-pi model) (vector-ref model 0))
(define (set-model-pi! model x) (vector-set! model 0 x))
(define (model-mu model) (vector-ref model 1))
(define (model-sigma model) (vector-ref model 2))
(define (model-log-pi model) (vector-ref model 3))
(define (set-model-log-pi! model x) (vector-set! model 3 x))
(define (model-sigma-inverse model) (vector-ref model 4))
(define (model-log-determinant-sigma model) (vector-ref model 5))
(define (set-model-log-sigma-determinant! model x) (vector-set! model 5 x))
;---------------------------------------------------------------------------
; Minimal Standard Random Number Generator
; Park & Miller, CACM 31(10), Oct 1988, 32 bit integer version.
; better constants, as proposed by Park.
; By Ozan Yigit
(define *seed* 1)
(define (srand seed)
(set! *seed* seed)
*seed*)
(define (rand)
(let ((A 48271)
(M 2147483647)
(Q 44488)
(R 3399))
(let* ((hi (quotient *seed* Q))
(lo (modulo *seed* Q))
(test (- (* A lo) (* R hi))))
(if (> test 0)
(set! *seed* test)
(set! *seed* (+ test M)))))
*seed*)
;---------------------------------------------------------------------------
(define (panic s) (error 'panic s))
(define *rand-max* 2147483648)
(define log-math-precision 35.0)
(define minus-infinity (- *rand-max*))
(define first car)
(define second cadr)
(define rest cdr)
(define (reduce f l i)
(cond ((null? l) i)
((null? (rest l)) (first l))
(else (let loop ((l (rest l)) (c (first l)))
(if (null? l) c (loop (rest l) (f c (first l))))))))
(define (every-n p n)
(let loop ((i 0)) (or (>= i n) (and (p i) (loop (+ i 1))))))
(define (sum f n)
(let loop ((n (- n 1)) (c 0.0))
(if (negative? n) c (loop (- n 1) (+ c (f n))))))
(define (add-exp e1 e2)
(let* ((e-max (max e1 e2))
(e-min (min e1 e2))
(factor (floor e-min)))
(if (= e-max minus-infinity)
minus-infinity
(if (> (- e-max factor) log-math-precision)
e-max
(+ (log (+ (exp (- e-max factor)) (exp (- e-min factor))))
factor)))))
(define (map-n f n)
(let loop ((i 0) (c '()))
(if (< i n) (loop (+ i 1) (cons (f i) c)) (reverse c))))
(define (map-n-vector f n)
(let ((v (make-vector n)))
(let loop ((i 0))
(if (< i n)
(begin (vector-set! v i (f i))
(loop (+ i 1)))))
v))
(define (remove-if-not p l)
(let loop ((l l) (c '()))
(cond ((null? l) (reverse c))
((p (first l)) (loop (rest l) (cons (first l) c)))
(else (loop (rest l) c)))))
(define (positionv x l)
(let loop ((l l) (i 0))
(cond ((null? l) #f)
((eqv? x (first l)) i)
(else (loop (rest l) (+ i 1))))))
(define (make-matrix m n)
(map-n-vector (lambda (i) (make-vector n)) m))
(define (make-matrix-initial m n initial)
(map-n-vector (lambda (i) (make-vector n initial)) m))
(define (matrix-rows a) (vector-length a))
(define (matrix-columns a) (vector-length (vector-ref a 0)))
(define (matrix-ref a i j) (vector-ref (vector-ref a i) j))
(define (matrix-set! a i j x) (vector-set! (vector-ref a i) j x))
(define (matrix-row-ref a i) (vector-ref a i))
(define (matrix-row-set! a i v) (vector-set! a i v))
(define (determinant a)
(if (not (= (matrix-rows a) (matrix-columns a)))
(panic "Can only find determinant of a square matrix"))
(call-with-current-continuation
(lambda (return)
(let* ((n (matrix-rows a))
(b (make-matrix n n))
(d 1.0))
(do ((i 0 (+ i 1))) ((= i n))
(do ((j 0 (+ j 1))) ((= j n)) (matrix-set! b i j (matrix-ref a i j))))
(do ((i 0 (+ i 1))) ((= i n))
;; partial pivoting reduces rounding errors
(let ((greatest (abs (matrix-ref b i i)))
(index i))
(do ((j (+ i 1) (+ j 1))) ((= j n))
(let ((x (abs (matrix-ref b j i))))
(if (> x greatest) (begin (set! index j) (set! greatest x)))))
(if (= greatest 0.0) (return 0.0))
(if (not (= index i))
(let ((v (matrix-row-ref b i)))
(matrix-row-set! b i (matrix-row-ref b index))
(matrix-row-set! b index v)
(set! d (- d))))
(let ((c (matrix-ref b i i)))
(set! d (* d c))
(do ((j i (+ j 1))) ((= j n))
(matrix-set! b i j (/ (matrix-ref b i j) c)))
(do ((j (+ i 1) (+ j 1))) ((= j n))
(let ((e (matrix-ref b j i)))
(do ((k (+ i 1) (+ k 1))) ((= k n))
(matrix-set!
b j k (- (matrix-ref b j k) (* e (matrix-ref b i k))))))))))
d))))
(define (invert-matrix! a b)
(if (not (= (matrix-rows a) (matrix-columns a)))
(panic "Can only invert a square matrix"))
(let* ((n (matrix-rows a))
(c (make-matrix n n)))
(do ((i 0 (+ i 1))) ((= i n))
(do ((j 0 (+ j 1))) ((= j n))
(matrix-set! b i j 0.0)
(matrix-set! c i j (matrix-ref a i j))))
(do ((i 0 (+ i 1))) ((= i n)) (matrix-set! b i i 1.0))
(do ((i 0 (+ i 1))) ((= i n))
(if (zero? (matrix-ref c i i))
(call-with-current-continuation
(lambda (return)
(do ((j 0 (+ j 1))) ((= j n))
(if (and (> j i) (not (zero? (matrix-ref c j i))))
(begin
(let ((e (vector-ref c i)))
(vector-set! c i (vector-ref c j))
(vector-set! c j e))
(let ((e (vector-ref b i)))
(vector-set! b i (vector-ref b j))
(vector-set! b j e))
(return #f))))
(panic "Matrix is singular"))))
(let ((d (/ (matrix-ref c i i))))
(do ((j 0 (+ j 1))) ((= j n))
(matrix-set! c i j (* d (matrix-ref c i j)))
(matrix-set! b i j (* d (matrix-ref b i j))))
(do ((k 0 (+ k 1))) ((= k n))
(let ((d (- (matrix-ref c k i))))
(if (not (= k i))
(do ((j 0 (+ j 1))) ((= j n))
(matrix-set!
c k j (+ (matrix-ref c k j) (* d (matrix-ref c i j))))
(matrix-set!
b k j (+ (matrix-ref b k j) (* d (matrix-ref b i j))))))))))))
(define (jacobi! a)
(if (not (and (= (matrix-rows a) (matrix-columns a))
(every-n (lambda (i)
(every-n (lambda (j)
(= (matrix-ref a i j) (matrix-ref a j i)))
(matrix-rows a)))
(matrix-rows a))))
(panic "Can only compute eigenvalues/eigenvectors of a symmetric matrix"))
(let* ((n (matrix-rows a))
(d (make-vector n))
(v (make-matrix-initial n n 0.0))
(b (make-vector n))
(z (make-vector n 0.0)))
(do ((ip 0 (+ ip 1))) ((= ip n))
(matrix-set! v ip ip 1.0)
(vector-set! b ip (matrix-ref a ip ip))
(vector-set! d ip (matrix-ref a ip ip)))
(let loop ((i 0))
(if (> i 50) (panic "Too many iterations in JACOBI!"))
(let ((sm (sum (lambda (ip)
(sum (lambda (ir)
(let ((iq (+ ip ir 1)))
(abs (matrix-ref a ip iq))))
(- n ip 1)))
(- n 1))))
(if (not (zero? sm))
(begin
(let ((tresh (if (< i 3) (/ (* 0.2 sm) (* n n)) 0.0)))
(do ((ip 0 (+ ip 1))) ((= ip (- n 1)))
(do ((ir 0 (+ ir 1))) ((= ir (- n ip 1)))
(let* ((iq (+ ip ir 1))
(g (* 100.0 (abs (matrix-ref a ip iq)))))
(cond
((and (> i 3)
(= (+ (abs (vector-ref d ip)) g)
(abs (vector-ref d ip)))
(= (+ (abs (vector-ref d iq)) g)
(abs (vector-ref d iq))))
(matrix-set! a ip iq 0.0))
((> (abs (matrix-ref a ip iq)) tresh)
(let* ((h (- (vector-ref d iq) (vector-ref d ip)))
(t (if (= (+ (abs h) g) (abs h))
(/ (matrix-ref a ip iq) h)
(let ((theta (/ (* 0.5 h)
(matrix-ref a ip iq))))
(if (negative? theta)
(- (/ (+ (abs theta)
(sqrt (+ (* theta theta) 1.0)))))
(/ (+ (abs theta)
(sqrt (+ (* theta theta) 1.0))))))))
(c (/ (sqrt (+ (* t t) 1.0))))
(s (* t c))
(tau (/ s (+ c 1.0)))
(h (* t (matrix-ref a ip iq))))
(define (rotate a i j k l)
(let ((g (matrix-ref a i j))
(h (matrix-ref a k l)))
(matrix-set! a i j (- g (* s (+ h (* g tau)))))
(matrix-set! a k l (+ h (* s (- g (* h tau)))))))
(vector-set! z ip (- (vector-ref z ip) h))
(vector-set! z iq (+ (vector-ref z iq) h))
(vector-set! d ip (- (vector-ref d ip) h))
(vector-set! d iq (+ (vector-ref d iq) h))
(matrix-set! a ip iq 0.0)
(do ((j 0 (+ j 1))) ((= j n))
(cond ((< j ip) (rotate a j ip j iq))
((< ip j iq) (rotate a ip j j iq))
((< iq j) (rotate a ip j iq j)))
(rotate v j ip j iq)))))))))
(do ((ip 0 (+ ip 1))) ((= ip n))
(vector-set! b ip (+ (vector-ref b ip) (vector-ref z ip)))
(vector-set! d ip (vector-ref b ip))
(vector-set! z ip 0.0))
(loop (+ i 1))))))
(do ((i 0 (+ i 1))) ((= i (- n 1)))
(let ((k i)
(p (vector-ref d i)))
(do ((l 0 (+ l 1))) ((= l (- n i 1)))
(let* ((j (+ i l 1)))
(if (>= (vector-ref d j) p)
(begin (set! k j) (set! p (vector-ref d j))))))
(if (not (= k i))
(begin (vector-set! d k (vector-ref d i))
(vector-set! d i p)
(do ((j 0 (+ j 1))) ((= j n))
(let ((p (matrix-ref v j i)))
(matrix-set! v j i (matrix-ref v j k))
(matrix-set! v j k p)))))))
(list d v)))
(define (clip-eigenvalues! a v)
(let* ((j (jacobi! a))
(l (first j))
(e (second j)))
(do ((k1 0 (+ k1 1))) ((= k1 (vector-length a)))
(let ((a-k1 (vector-ref a k1))
(e-k1 (vector-ref e k1)))
(do ((k2 0 (+ k2 1))) ((= k2 (vector-length a-k1)))
(let ((e-k2 (vector-ref e k2))
(s 0.0))
(do ((k 0 (+ k 1))) ((= k (vector-length a)))
(set! s (+ s (* (max (vector-ref v k) (vector-ref l k))
(vector-ref e-k1 k)
(vector-ref e-k2 k)))))
(vector-set! a-k1 k2 s)))))))
;;; EM
(define (e-step! x z models)
(do ((i 0 (+ i 1))) ((= i (vector-length x)))
(let ((xi (vector-ref x i))
(zi (vector-ref z i)))
(do ((j 0 (+ j 1))) ((= j (vector-length models)))
;; Compute for each model.
(let* ((model (vector-ref models j))
(log-pi (model-log-pi model))
(mu (model-mu model))
(sigma-inverse (model-sigma-inverse model))
(log-determinant-sigma (model-log-determinant-sigma model))
(t 0.0))
;; Compute likelihoods (note: up to constant for all models).
(set! t 0.0)
(do ((k1 0 (+ k1 1))) ((= k1 (vector-length xi)))
(let ((sigma-inverse-k1 (vector-ref sigma-inverse k1)))
(do ((k2 0 (+ k2 1))) ((= k2 (vector-length xi)))
(set! t (+ t (* (- (vector-ref xi k1) (vector-ref mu k1))
(vector-ref sigma-inverse-k1 k2)
(- (vector-ref xi k2) (vector-ref mu k2))))))))
(vector-set! zi j (- log-pi (* 0.5 (+ log-determinant-sigma t))))))))
(let ((l 0.0))
(do ((i 0 (+ i 1))) ((= i (vector-length x)))
(let ((s minus-infinity)
(zi (vector-ref z i)))
;; Normalize ownerships to sum to one.
(do ((j 0 (+ j 1))) ((= j (vector-length models)))
(set! s (add-exp s (vector-ref zi j))))
(do ((j 0 (+ j 1))) ((= j (vector-length models)))
(vector-set! zi j (exp (- (vector-ref zi j) s))))
(set! l (+ l s))))
;; Return log likelihood.
l))
(define (m-step! x models z clip)
(let ((kk (vector-length (vector-ref x 0))))
;; For each model, optimize parameters.
(do ((j 0 (+ j 1))) ((= j (vector-length models)))
(let* ((model (vector-ref models j))
(mu (model-mu model))
(sigma (model-sigma model))
(s 0.0))
;; Optimize values.
(do ((k 0 (+ k 1))) ((= k kk))
(do ((i 0 (+ i 1))) ((= i (vector-length x)))
(set! s (+ s (vector-ref (vector-ref z i) j)))))
(do ((k 0 (+ k 1))) ((= k kk))
(let ((m 0.0))
(do ((i 0 (+ i 1))) ((= i (vector-length x)))
(set! m (+ m (* (vector-ref (vector-ref z i) j)
(vector-ref (vector-ref x i) k)))))
(vector-set! mu k (/ m s))))
(do ((k1 0 (+ k1 1))) ((= k1 kk))
(let ((sigma-k1 (vector-ref sigma k1))
(mu-k1 (vector-ref mu k1)))
(do ((k2 0 (+ k2 1))) ((= k2 kk))
(let ((mu-k2 (vector-ref mu k2))
(m 0.0))
(do ((i 0 (+ i 1))) ((= i (vector-length x)))
(set! m (+ m (* (vector-ref (vector-ref z i) j)
(- (vector-ref (vector-ref x i) k1) mu-k1)
(- (vector-ref (vector-ref x i) k2) mu-k2)))))
(vector-set! sigma-k1 k2 (/ m s))))))
(clip-eigenvalues! sigma clip)
(set-model-pi! model (/ s (vector-length x)))
(set-model-log-pi! model (log (/ s (vector-length x))))
(invert-matrix! sigma (model-sigma-inverse model))
(set-model-log-sigma-determinant! model (log (determinant sigma)))))))
(define (em! x z models clip em-kick-off-tolerance em-convergence-tolerance)
(let loop ((old-log-likelihood minus-infinity) (starting? #t))
(let ((log-likelihood (e-step! x z models)))
(cond
((or (and starting? (> log-likelihood old-log-likelihood))
(> log-likelihood (+ old-log-likelihood em-convergence-tolerance)))
(m-step! x models z clip)
(loop log-likelihood
(and starting?
(not (= (vector-length models) 1))
(or (= old-log-likelihood minus-infinity)
(< log-likelihood
(+ old-log-likelihood em-kick-off-tolerance))))))
(else old-log-likelihood)))))
(define (noise epsilon) (- (* 2.0 epsilon (/ (rand) *rand-max*)) epsilon))
(define (initial-z ii jj)
(map-n-vector
(lambda (i)
(let ((zi (map-n-vector (lambda (j) (+ (/ jj) (noise (/ jj)))) jj))
(s 0.0))
(do ((j 0 (+ j 1))) ((= j jj)) (set! s (+ s (vector-ref zi j))))
(do ((j 0 (+ j 1))) ((= j jj)) (vector-set! zi j (/ (vector-ref zi j) s)))
zi))
ii))
(define (ems x clip em-kick-off-tolerance em-convergence-tolerance
ems-convergence-tolerance)
(let loop ((jj 1)
(old-z #f)
(old-models #f)
(old-log-likelihood minus-infinity))
(let* ((kk (vector-length (vector-ref x 0)))
(z (initial-z (vector-length x) jj))
(models (map-n-vector
(lambda (j)
(make-model 0.0
(make-vector kk)
(make-matrix kk kk)
0.0
(make-matrix kk kk)
0.0))
jj)))
(m-step! x models z clip)
(let ((new-log-likelihood
(em!
x z models clip em-kick-off-tolerance em-convergence-tolerance)))
(if (> (- (/ old-log-likelihood new-log-likelihood) 1.0)
ems-convergence-tolerance)
(loop (+ jj 1) z models new-log-likelihood)
(list old-z old-models))))))
(define (em-clusterer x clip em-kick-off-tolerance em-convergence-tolerance
ems-convergence-tolerance)
(let* ((z-models (ems x clip em-kick-off-tolerance
em-convergence-tolerance
ems-convergence-tolerance))
(z (first z-models))
(models (second z-models)))
(e-step! x z models)
(let ((clusters
(map-n (lambda (i)
(let ((zi (vector->list (vector-ref z i))))
(list i (positionv (reduce max zi minus-infinity) zi))))
(vector-length z))))
(map-n (lambda (j)
(map (lambda (cluster) (vector-ref x (first cluster)))
(remove-if-not (lambda (cluster) (= (second cluster) j))
clusters)))
(vector-length (vector-ref z 0))))))
(define (go)
(em-clusterer
'#(#(1.0) #(2.0) #(3.0) #(11.0) #(12.0) #(13.0)) '#(1.0) 10.0 1.0 0.01))
(define (run-benchmark)
(srand 1)
(do ((i 0 (+ i 1))) ((= i 100))
(write (go))
(newline)))
; eof