basis/math/statistics/statistics.factor

   1 ! Copyright (C) 2008 Doug Coleman, Michael Judge.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: assocs combinators generalizations kernel locals math
   4 math.functions math.order math.vectors sequences
   5 sequences.private sorting fry ;
   6 IN: math.statistics
   7
   8 : mean ( seq -- x )
   9     [ sum ] [ length ] bi / ;
  10
  11 : geometric-mean ( seq -- x )
  12     [ length ] [ product ] bi nth-root ;
  13
  14 : harmonic-mean ( seq -- x )
  15     [ recip ] map-sum recip ;
  16
  17 <PRIVATE
  18
  19 :: (kth-object) ( seq k quot: ( x y -- ? ) nth-quot exchange-quot -- elt )
  20     #! Wirth's method, Algorithm's + Data structues = Programs p. 84
  21     #! The algorithm modifiers seq, so we clone it
  22     k seq bounds-check 2drop
  23     seq clone :> seq
  24     0 :> i!
  25     0 :> j!
  26     0 :> l!
  27     0 :> x!
  28     seq length 1 - :> m!
  29     [ l m < ]
  30     [
  31         k seq nth x!
  32         l i!
  33         m j!
  34         [ i j <= ]
  35         [
  36             [ i seq nth-quot call x quot call ] [ i 1 + i! ] while
  37             [ x j seq nth-quot call quot call ] [ j 1 - j! ] while
  38             i j <= [
  39                 i j seq exchange-quot call
  40                 i 1 + i!
  41                 j 1 - j!
  42             ] when
  43         ] do while
  44
  45         j k < [ i l! ] when
  46         k i < [ j m! ] when
  47     ] while
  48     k seq nth ; inline
  49
  50 : kth-object-unsafe ( seq k quot: ( x y -- ? ) -- elt )
  51     [ nth-unsafe ] [ exchange-unsafe ] (kth-object) ; inline
  52
  53 PRIVATE>
  54
  55 : kth-object ( seq k quot: ( x y -- ? ) -- elt )
  56     [ nth ] [ exchange ] (kth-object) ; inline
  57
  58 : kth-smallest ( seq k -- elt ) [ < ] kth-object-unsafe ;
  59
  60 : kth-largest ( seq k -- elt ) [ > ] kth-object-unsafe ;
  61
  62 : count-relative ( seq k -- lt eq gt )
  63     [ 0 0 0 ] 2dip '[
  64         _ <=> {
  65             { +lt+ [ [ 1 + ] 2dip ] }
  66             { +gt+ [ 1 + ] }
  67             { +eq+ [ [ 1 + ] dip ] }
  68         } case
  69     ] each ;
  70
  71 : minmax-relative ( seq k -- lt eq gt lt-max gt-min )
  72     [ 0 0 0 -1/0. 1/0. ] 2dip '[
  73         dup _ <=> {
  74             { +lt+ [ [ 1 + ] 5 ndip '[ _ max ] dip ] }
  75             { +gt+ [ [ 1 + ] 3dip min ] }
  76             { +eq+ [ [ 1 + ] 4dip drop ] }
  77         } case
  78     ] each ;
  79
  80 : lower-median ( seq -- elt )
  81     [ ] [ ] [ length odd? ] tri
  82     [ midpoint@ ] [ midpoint@ 1 - ] if kth-smallest ;
  83
  84 : upper-median ( seq -- elt )
  85     dup midpoint@ kth-smallest ;
  86
  87 : medians ( seq -- lower upper )
  88     [ lower-median ] [ upper-median ] bi ;
  89
  90 : median ( seq -- x )
  91     [ ] [ length odd? ] bi [ lower-median ] [ medians + 2 / ] if ;
  92
  93 <PRIVATE
  94
  95 : (sequence>assoc) ( seq map-quot: ( x -- ..y ) insert-quot: ( ..y assoc -- ) assoc -- assoc )
  96     [ swap curry compose each ] keep ; inline
  97
  98 PRIVATE>
  99
 100 : sequence>assoc! ( assoc seq map-quot: ( x -- ..y ) insert-quot: ( ..y assoc -- ) -- assoc )
 101     4 nrot (sequence>assoc) ; inline
 102
 103 : sequence>assoc ( seq map-quot: ( x -- ..y ) insert-quot: ( ..y assoc -- ) exemplar -- assoc )
 104     clone (sequence>assoc) ; inline
 105
 106 : sequence>hashtable ( seq map-quot: ( x -- ..y ) insert-quot: ( ..y assoc -- ) -- hashtable )
 107     H{ } sequence>assoc ; inline
 108
 109 : histogram! ( hashtable seq -- hashtable )
 110     [ ] [ inc-at ] sequence>assoc! ;
 111
 112 : histogram-by ( seq quot: ( x -- bin ) -- hashtable )
 113     [ inc-at ] sequence>hashtable ; inline
 114
 115 : histogram ( seq -- hashtable )
 116     [ ] histogram-by ;
 117
 118 : sorted-histogram ( seq -- alist )
 119     histogram sort-values ;
 120
 121 : collect-pairs ( seq quot: ( x -- v k ) -- hashtable )
 122     [ push-at ] sequence>hashtable ; inline
 123
 124 : collect-by ( seq quot: ( x -- x' ) -- hashtable )
 125     [ dup ] prepose collect-pairs ; inline
 126
 127 : mode ( seq -- x )
 128     histogram >alist
 129     [ ] [ [ [ second ] bi@ > ] most ] map-reduce first ;
 130
 131 ERROR: empty-sequence ;
 132
 133 : minmax ( seq -- min max )
 134     [
 135         empty-sequence
 136     ] [
 137         [ first dup ] keep [ [ min ] [ max ] bi-curry bi* ] each
 138     ] if-empty ;
 139
 140 : range ( seq -- x )
 141     minmax swap - ;
 142
 143 : sample-var ( seq -- x )
 144     #! normalize by N-1
 145     dup length 1 <= [
 146         drop 0
 147     ] [
 148         [ [ mean ] keep [ - sq ] with map-sum ]
 149         [ length 1 - ] bi /
 150     ] if ;
 151
 152 : var ( seq -- x )
 153     dup length 1 <= [
 154         drop 0
 155     ] [
 156         [ [ mean ] keep [ - sq ] with map-sum ]
 157         [ length ] bi /
 158     ] if ;
 159
 160 : std ( seq -- x ) var sqrt ;
 161
 162 : ste ( seq -- x ) [ std ] [ length ] bi sqrt / ;
 163
 164 : ((r)) ( mean(x) mean(y) {x} {y} -- (r) )
 165     ! finds sigma((xi-mean(x))(yi-mean(y))
 166     0 [ [ [ pick ] dip swap - ] bi@ * + ] 2reduce 2nip ;
 167
 168 : (r) ( mean(x) mean(y) {x} {y} sx sy -- r )
 169     * recip [ [ ((r)) ] keep length 1 - / ] dip * ;
 170
 171 : [r] ( {{x,y}...} -- mean(x) mean(y) {x} {y} sx sy )
 172     first2 [ [ [ mean ] bi@ ] 2keep ] 2keep [ std ] bi@ ;
 173
 174 : r ( {{x,y}...} -- r )
 175     [r] (r) ;
 176
 177 : r^2 ( {{x,y}...} -- r )
 178     r sq ;
 179
 180 : least-squares ( {{x,y}...} -- alpha beta )
 181     [r] { [ 2dup ] [ ] [ ] [ ] [ ] } spread
 182     ! stack is mean(x) mean(y) mean(x) mean(y) {x} {y} sx sy
 183     [ (r) ] 2keep ! stack is mean(x) mean(y) r sx sy
 184     swap / * ! stack is mean(x) mean(y) beta
 185     [ swapd * - ] keep ;
 186
 187 : cov ( {x} {y} -- cov )
 188     [ dup mean v-n ] bi@ v* mean ;
 189
 190 : corr ( {x} {y} -- corr )
 191      [ cov ] [ [ var ] bi@ * sqrt ] 2bi / ;
 192
 193 : cum-sum ( seq -- seq' )
 194     0 swap [ + dup ] map nip ;
 195
 196 : cum-product ( seq -- seq' )
 197     1 swap [ * dup ] map nip ;
 198
 199 : cum-min ( seq -- seq' )
 200     [ ?first ] keep [ min dup ] map nip ;
 201
 202 : cum-max ( seq -- seq' )
 203     [ ?first ] keep [ max dup ] map nip ;