;;; PLT Scheme Science Collection
;;; histogram-2d.ss
;;; Copyright (c) 2004 M. Douglas Williams
;;;
;;; This library is free software; you can redistribute it and/or
;;; modify it under the terms of the GNU Lesser General Public
;;; License as published by the Free Software Foundation; either
;;; version 2.1 of the License, or (at your option) any later version.
;;;
;;; This library is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; Lesser General Public License for more details.
;;;
;;; You should have received a copy of the GNU Lesser General Public
;;; License along with this library; if not, write to the Free
;;; Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
;;; 02111-1307 USA.
;;;
;;; -------------------------------------------------------------------
;;;
;;; This module implements 2D histograms.  It is based on the
;;; Histograms in the GNU Scientific Library (GSL).
;;;
;;; Version  Date      Description
;;; 1.0.0    09/29/04  Marked as ready for Release 1.0.  (Doug
;;;                    Williams)

(module histogram-2d mzscheme
  
  ;; Data Definition
  
  ;; Histogram structure with 6 fields:
  ;;   0 - nx, number of bins in the x direction
  ;;   1 - ny, number of bins in the y direction
  ;;   2 - x-ranges
  ;;   3 - y-ranges
  ;;   4 - bins
  ;;   5 - ranges-uniform?
  
  (define-values (struct:histogram-2d
                  histogram-2d-constructor
                  histogram-2d?
                  histogram-2d-field-ref
                  set-histogram-2d-field!)
    (make-struct-type 'histogram-2d #f 6 0))
  
  ;; Contracts
  
  (require (lib "contract.ss"))
  
  (provide/contract
   (histogram-2d?
    (-> any/c boolean?))
   (make-histogram-2d
    (-> (integer-in 1 +inf.0) (integer-in 1 +inf.0) histogram-2d?))
   (make-histogram-2d-with-ranges-uniform
    (->r ((nx (integer-in 1 +inf.0))
          (xy (integer-in 1 +inf.0))
          (x-min real?)
          (x-max (>/c x-min))
          (y-min real?)
          (y-max (>/c y-min)))
         histogram-2d?))
   (histogram-2d-nx
    (-> histogram-2d? (integer-in 1 +inf.0)))
   (histogram-2d-ny
    (-> histogram-2d? (integer-in 1 +inf.0)))
   (histogram-2d-x-ranges
    (-> histogram-2d? (vectorof real?)))
   (histogram-2d-y-ranges
    (-> histogram-2d? (vectorof real?)))
   (set-histogram-2d-ranges!
    (-> histogram-2d? (vectorof real?) (vectorof real?) void?))
   (set-histogram-2d-ranges-uniform!
    (->r ((h histogram-2d?)
          (x-min real?)
          (x-max (>/c x-min))
          (y-min real?)
          (y-max (>/c y-min)))
         void?))
   (histogram-2d-bins
    (-> histogram-2d? (vectorof real?)))
   (histogram-2d-increment!
    (-> histogram-2d? real? real? void?))
   (histogram-2d-accumulate!
    (-> histogram-2d? real? real? (>=/c 0.0) void?))
   (histogram-2d-get
    (-> histogram-2d? natural-number/c natural-number/c (>=/c 0.0)))
   (histogram-2d-get-x-range
    (-> histogram-2d? natural-number/c (values real? real?)))
   (histogram-2d-get-y-range
    (-> histogram-2d? natural-number/c (values real? real?)))
   (histogram-2d-max
    (-> histogram-2d? (>=/c 0.0)))
   (histogram-2d-min
    (-> histogram-2d? (>=/c 0.0)))
   (histogram-2d-sum
    (-> histogram-2d? (>=/c 0.0)))
   (histogram-2d-x-mean
    (-> histogram-2d? real?))
   (histogram-2d-y-mean
    (-> histogram-2d? real?))
   (histogram-2d-x-sigma
    (-> histogram-2d? (>=/c 0.0)))
   (histogram-2d-y-sigma
    (-> histogram-2d? (>=/c 0.0)))
   (histogram-2d-covariance
    (-> histogram-2d? (>=/c 0.0))))

  ;; make-histogram-2d: integer x integer -> histogram
  (define (make-histogram-2d nx ny)
    (histogram-2d-constructor
     nx ny
     (make-vector (+ nx 1))
     (make-vector (+ ny 1))
     (make-vector (* nx ny))
     #f))
  
  (define (make-histogram-2d-with-ranges-uniform
           nx ny x-min x-max y-min y-max)
    (let ((h (make-histogram-2d nx ny)))
      (set-histogram-2d-ranges-uniform! h x-min x-max y-min y-max)
      h))
  
  (define histogram-2d-nx
    (make-struct-field-accessor histogram-2d-field-ref 0 'nx))
  
  (define histogram-2d-ny
    (make-struct-field-accessor histogram-2d-field-ref 1 'ny))
  
  (define histogram-2d-x-ranges
    (make-struct-field-accessor histogram-2d-field-ref 2 'x-ranges))
  
  (define histogram-2d-y-ranges
    (make-struct-field-accessor histogram-2d-field-ref 3 'y-ranges))
  
  (define (set-histogram-2d-ranges! h x-ranges y-ranges)
    (let ((nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (x-ranges (histogram-2d-x-ranges h))
          (y-ranges (histogram-2d-y-ranges h))
          (bins (histogram-2d-bins h)))
      ;; Set x ranges
      (do ((i 0 (+ i 1)))
          ((> i nx) (void))
        (vector-set! x-ranges i
                     (vector-ref x-ranges i)))
      ;; Set y ranges
      (do ((i 0 (+ i 1)))
          ((> i ny) (void))
        (vector-set! y-ranges i
                     (vector-ref y-ranges i)))
      ;; Re-initialize the bins
      (do ((i 0 (+ i 1)))
          ((= i (* nx ny)) (void))
        (vector-set! bins i 0))
      ;; Clear ranges-uniform?
      (set-histogram-2d-ranges-uniform?! h #f)
      (void)))
  
  (define (set-histogram-2d-ranges-uniform!
           h x-min x-max y-min y-max)
    (let* ((nx (histogram-2d-nx h))
           (ny (histogram-2d-ny h))
           (x-ranges (histogram-2d-x-ranges h))
           (y-ranges (histogram-2d-y-ranges h))
           (bins (histogram-2d-bins h))
           (x-bin-width (/ (- x-max x-min) nx))
           (y-bin-width (/ (- y-max y-min) ny)))
      ;; Set x ranges
      (do ((i 0 (+ i 1)))
          ((= i nx) (void))
        (vector-set! x-ranges i (+ x-min (* i x-bin-width))))
      (vector-set! x-ranges nx x-max)
      ;; Set y ranges
      (do ((i 0 (+ i 1)))
          ((= i ny) (void))
        (vector-set! y-ranges i (+ y-min (* i y-bin-width))))
      (vector-set! y-ranges ny y-max)
      ;; Re-initialize the bins
      (do ((i 0 (+ i 1)))
          ((= i (* nx ny)) (void))
        (vector-set! bins i 0))
      ;; Set uniform-ranges?
      (set-histogram-2d-ranges-uniform?! h #t)
      (void)))
  
  (define histogram-2d-bins
    (make-struct-field-accessor histogram-2d-field-ref 4 'bins))
  
  (define histogram-2d-ranges-uniform?
    (make-struct-field-accessor histogram-2d-field-ref 5 'ranges-uniform?))
  
  (define set-histogram-2d-ranges-uniform?!
    (make-struct-field-mutator set-histogram-2d-field! 5 'ranges-uniform?))
  
  (define (histogram-2d-increment! h x y)
    (histogram-2d-accumulate! h x y 1))
  
  (define (histogram-2d-accumulate! h x y weight)
    (let ((nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (x-ranges (histogram-2d-x-ranges h))
          (y-ranges (histogram-2d-y-ranges h))
          (bins (histogram-2d-bins h))
          (ranges-uniform? (histogram-2d-ranges-uniform? h)))
      (if ranges-uniform?
          ;; Compute bin
          (let ((i (inexact->exact
                    (floor (/ (- x (vector-ref x-ranges 0))
                              (/ (- (vector-ref x-ranges nx)
                                    (vector-ref x-ranges 0))
                                 nx)))))
                (j (inexact->exact
                    (floor (/ (- y (vector-ref y-ranges 0))
                              (/ (- (vector-ref y-ranges ny)
                                    (vector-ref y-ranges 0))
                                 ny))))))
            (if (and (<= 0 i (- nx 1))
                     (<= 0 j (- ny 1)))
                (let ((bin (+ (* i ny) j)))
                  (vector-set! bins bin
                               (+ (vector-ref bins bin) weight)))))
          ;; Search for bin
          (let/ec exit
            (if (or (< x (vector-ref x-ranges 0))
                    (< y (vector-ref y-ranges 0)))
                (exit))
            (do ((i 0 (+ i 1)))
                ((= i nx) (void))
              (if (< x (vector-ref x-ranges (+ i 1)))
                  (do ((j 0 (+ j i)))
                      ((= 1 ny) (exit))
                    (if (< y (vector-ref y-ranges (+ j 1)))
                        (begin
                          (let ((bin (+ (* i ny) j)))
                            (vector-set! bins bin
                                         (+ (vector-ref bins bin) weight)))
                          (exit))))))))))
  
  (define (histogram-2d-get h i j)
    (let* ((nx (histogram-2d-nx h))
           (ny (histogram-2d-ny h))
           (bins (histogram-2d-bins h))
           (bin (+ (* i ny) j)))
      (vector-ref bins bin)))
  
  (define (histogram-2d-get-x-range h i)
    (let ((nx (histogram-2d-nx h))
          (x-ranges (histogram-2d-x-ranges h)))
      (values (vector-ref x-ranges i)
              (vector-ref x-ranges (+ i 1)))))
  
  (define (histogram-2d-get-y-range h j)
    (let ((ny (histogram-2d-ny h))
          (y-ranges (histogram-2d-y-ranges h)))
      (values (vector-ref y-ranges j)
              (vector-ref y-ranges (+ j 1)))))
  
  ;; Histogram-2d Statistics
  
  (define (histogram-2d-max h)
    (let ((nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (bins (histogram-2d-bins h))
          (max-value -inf.0))
      (do ((i 0 (+ i 1)))
          ((= i nx) max-value)
        (do ((j 0 (+ j 1)))
            ((= j ny) (void))
          (let ((bin (+ (* i ny) j)))
            (set! max-value (max max-value (vector-ref bins bin))))))))
  
  (define (histogram-2d-min h)
    (let ((nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (bins (histogram-2d-bins h))
          (min-value +inf.0))
      (do ((i 0 (+ i 1)))
          ((= i nx) min-value)
        (do ((j 0 (+ j 1)))
            ((= j ny) (void))
          (let ((bin (+ (* i ny) j)))
            (set! min-value (min min-value (vector-ref bins bin))))))))
  
  (define (histogram-2d-sum h)
    (let ((n (* (histogram-2d-nx h) (histogram-2d-ny h)))
          (bins (histogram-2d-bins h))
          (sum 0))
      (do ((i 0 (+ i 1)))
          ((= i n) sum)
        (set! sum (+ sum (vector-ref bins i))))))
  
  (define (histogram-2d-x-mean h)
    (let ((nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (bins (histogram-2d-bins))
          (x-ranges (histogram-2d-x-ranges h))
          (wmean 0.0)
          (W 0.0))
      (do ((i 0 (+ i 1)))
          ((= i nx) wmean)
        (let ((xi (/ (+ (vector-ref x-ranges i)
                        (vector-ref x-ranges (+ i 1)))
                     2.0))
              (wi 0.0))
          (do ((j 0 (+ j 1)))
              ((= j ny) (void))
            (let ((wij (vector-ref bins (+ (* i ny) j))))
              (if (> wij 0)
                  (set! wi (+ wi wij)))))
          (if (> wi 0)
              (begin
                (set! W (+ W wi))
                (set! wmean (+ wmean (* (- xi wmean) (/ wi W))))))))))
  
  (define (histogram-2d-y-mean h)
    (let ((nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (bins (histogram-2d-bins h))
          (y-ranges (histogram-2d-x-ranges h))
          (wmean 0.0)
          (W 0.0))
      (do ((j 0 (+ j 1)))
          ((= j ny) wmean)
        (let ((yj (/ (+ (vector-ref y-ranges j)
                        (vector-ref y-ranges (+ j 1)))
                     2.0))
              (wj 0.0))
          (do ((i 0 (+ i 1)))
              ((= i nx) (void))
            (let ((wij (vector-ref bins (+ (* i ny) j))))
              (if (> wij 0)
                  (set! wj (+ wj wij)))))
          (if (> wj 0)
              (begin
                (set! W (+ W wj))
                (set! wmean (+ wmean (* (- yj wmean) (/ wj W))))))))))
  
  (define (histogram-2d-x-sigma h)
    (let ((x-mean (histogram-2d-x-mean h))
          (nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (bins (histogram-2d-bins h))
          (x-ranges (histogram-2d-x-ranges h))
          (wvariance 0.0)
          (W 0.0))
      (do ((i 0 (+ i 1)))
          ((= i nx) (sqrt wvariance))
        (let ((xi (/ (+ (vector-ref x-ranges i)
                        (vector-ref x-ranges (+ i 1)))
                     (- 2.0 x-mean)))
              (wi 0.0))
          (do ((j 0 (+ j 1)))
              ((= j ny) (void))
            (let ((wij (vector-ref bins (+ (* i ny) j))))
              (if (> wij 0)
                  (set! wi (+ wi wij)))))
          (if (> wi 0)
              (begin
                (set! W (+ W wi))
                (set! wvariance (+ wvariance (* (- (* xi xi) wvariance)
                                                (/ wi W))))))))))
    
  (define (histogram-2d-y-sigma h)
    (let ((y-mean (histogram-2d-y-mean h))
          (nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (bins (histogram-2d-bins h))
          (y-ranges (histogram-2d-y-ranges h))
          (wvariance 0.0)
          (W 0.0))
      (do ((j 0 (+ j 1)))
          ((= j ny) (sqrt wvariance))
        (let ((yj (/ (+ (vector-ref y-ranges j)
                        (vector-ref y-ranges (+ j 1)))
                     (- 2.0 y-mean)))
              (wj 0.0))
          (do ((i 0 (+ i 1)))
              ((= i nx) (void))
            (let ((wij (vector-ref bins (+ (* i ny) j))))
              (if (> wij 0)
                  (set! wj (+ wj wij)))))
          (if (> wj 0)
              (begin
                (set! W (+ W wj))
                (set! wvariance (+ wvariance (* (- (* yj yj) wvariance)
                                                (/ wj W))))))))))

  (define (histogram-2d-covariance h)
    (let ((x-mean (histogram-2d-x-mean h))
          (y-mean (histogram-2d-y-mean h))
          (nx (histogram-2d-nx h))
          (ny (histogram-2d-ny h))
          (bins (histogram-2d-bins h))
          (x-ranges (histogram-2d-x-ranges h))
          (y-ranges (histogram-2d-y-ranges h))
          (wcovariance 0.0)
          (W 0.0))
      (do ((i 0 (+ i 1)))
          ((= i nx) wcovariance)
        (let ((xi (/ (+ (vector-ref x-ranges i)
                        (vector-ref x-ranges (+ i 1)))
                     (- 2.0 x-mean))))
          (do ((j 0 (+ j 1)))
              ((= j ny) (void))
            (let ((yj (/ (+ (vector-ref y-ranges j)
                            (vector-ref y-ranges (+ j 1)))
                         (- 2.0 y-mean)))
                  (wij (vector-ref bins (+ (* i ny) j))))
              (if (> wij 0)
                  (begin
                    (set! W (+ W wij))
                    (set! wcovariance (+ wcovariance
                                         (* (- (* xi yj) wcovariance)
                                            (/ wij W))))))))))))
  
)