extra/math/extras/extras.factor

   1 ! Copyright (C) 2012 John Benediktsson
   2 ! See http://factorcode.org/license.txt for BSD license
   3
   4 USING: accessors arrays assocs assocs.extras byte-arrays
   5 combinators combinators.short-circuit compression.zlib fry
   6 grouping kernel locals math math.combinatorics math.constants
   7 math.functions math.order math.primes math.primes.factors
   8 math.ranges math.ranges.private math.statistics math.vectors
   9 memoize parser random sequences sequences.extras
  10 sequences.private sets sorting sorting.extras ;
  11
  12 IN: math.extras
  13
  14 <PRIVATE
  15
  16 DEFER: stirling
  17
  18 : (stirling) ( n k -- x )
  19     [ [ 1 - ] bi@ stirling ]
  20     [ [ 1 - ] dip stirling ]
  21     [ nip * + ] 2tri ;
  22
  23 PRIVATE>
  24
  25 MEMO: stirling ( n k -- x )
  26     2dup { [ = ] [ nip 1 = ] } 2||
  27     [ 2drop 1 ] [ (stirling) ] if ;
  28
  29 :: ramanujan ( x -- y )
  30     pi sqrt x e / x ^ * x 8 * 4 + x * 1 + x * 1/30 + 1/6 ^ * ;
  31
  32 DEFER: bernoulli
  33
  34 <PRIVATE
  35
  36 : (bernoulli) ( p -- n )
  37     [ iota ] [ 1 + ] bi [
  38         0 [ [ nCk ] [ bernoulli * ] bi + ] with reduce
  39     ] keep recip neg * ;
  40
  41 PRIVATE>
  42
  43 MEMO: bernoulli ( p -- n )
  44     [ 1 ] [ (bernoulli) ] if-zero ;
  45
  46 : chi2 ( actual expected -- n )
  47     0 [ dup 0 > [ [ - sq ] keep / + ] [ 2drop ] if ] 2reduce ;
  48
  49 <PRIVATE
  50
  51 : df-check ( df -- )
  52     even? [ "odd degrees of freedom" throw ] unless ;
  53
  54 : (chi2P) ( chi/2 df/2 -- p )
  55     [1,b) dupd n/v cum-product swap neg e^ [ v*n sum ] keep + ;
  56
  57 PRIVATE>
  58
  59 : chi2P ( chi df -- p )
  60     dup df-check [ 2.0 / ] [ 2 /i ] bi* (chi2P) 1.0 min ;
  61
  62 <PRIVATE
  63
  64 : check-jacobi ( m -- m )
  65     dup { [ integer? ] [ 0 > ] [ odd? ] } 1&&
  66     [ "modulus must be odd positive integer" throw ] unless ;
  67
  68 : mod' ( x y -- n )
  69     [ mod ] keep over zero? [ drop ] [
  70         2dup [ sgn ] same? [ drop ] [ + ] if
  71     ] if ;
  72
  73 PRIVATE>
  74
  75 : jacobi ( a m -- n )
  76     check-jacobi [ mod' ] keep 1
  77     [ pick zero? ] [
  78         [ pick even? ] [
  79             [ 2 / ] 2dip
  80             over 8 mod' { 3 5 } member? [ neg ] when
  81         ] while swapd
  82         2over [ 4 mod' 3 = ] both? [ neg ] when
  83         [ [ mod' ] keep ] dip
  84     ] until [ nip 1 = ] dip 0 ? ;
  85
  86 <PRIVATE
  87
  88 : check-legendere ( m -- m )
  89     dup prime? [ "modulus must be prime positive integer" throw ] unless ;
  90
  91 PRIVATE>
  92
  93 : legendere ( a m -- n )
  94     check-legendere jacobi ;
  95
  96 : moving-average ( seq n -- newseq )
  97     <clumps> [ mean ] map ;
  98
  99 : exponential-moving-average ( seq a -- newseq )
 100     [ 1 ] 2dip '[ dupd swap - _ * + dup ] map nip ;
 101
 102 : moving-median ( u n -- v )
 103     <clumps> [ median ] map ;
 104
 105 : moving-supremum ( u n -- v )
 106     <clumps> [ supremum ] map ;
 107
 108 : moving-infimum ( u n -- v )
 109     <clumps> [ infimum ] map ;
 110
 111 : moving-sum ( u n -- v )
 112     <clumps> [ sum ] map ;
 113
 114 : moving-count ( ... u n quot: ( ... elt -- ... ? ) -- ... v )
 115     [ <clumps> ] [ '[ _ count ] map ] bi* ; inline
 116
 117 : nonzero ( seq -- seq' )
 118     [ zero? ] reject ;
 119
 120 : bartlett ( n -- seq )
 121     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 122         [ iota ] [ 1 - 2 / ] bi [
 123             [ recip * ] [ >= ] 2bi [ 2 swap - ] when
 124         ] curry map
 125     ] if ;
 126
 127 : [0,2pi] ( n -- seq )
 128     [ iota ] [ 1 - 2pi swap / ] bi v*n ;
 129
 130 : hanning ( n -- seq )
 131     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 132         [0,2pi] [ cos -0.5 * 0.5 + ] map!
 133     ] if ;
 134
 135 : hamming ( n -- seq )
 136     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 137         [0,2pi] [ cos -0.46 * 0.54 + ] map!
 138     ] if ;
 139
 140 : blackman ( n -- seq )
 141     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 142         [0,2pi] [
 143             [ cos -0.5 * ] [ 2 * cos 0.08 * ] bi + 0.42 +
 144         ] map
 145     ] if ;
 146
 147 : nan-sum ( seq -- n )
 148     0 [ dup fp-nan? [ drop ] [ + ] if ] binary-reduce ;
 149
 150 : nan-min ( seq -- n )
 151     [ fp-nan? ] reject infimum ;
 152
 153 : nan-max ( seq -- n )
 154     [ fp-nan? ] reject supremum ;
 155
 156 : fill-nans ( seq -- newseq )
 157     [ first ] keep [
 158         dup fp-nan? [ drop dup ] [ nip dup ] if
 159     ] map nip ;
 160
 161 : sinc ( x -- y )
 162     [ 1 ] [ pi * [ sin ] [ / ] bi ] if-zero ;
 163
 164 : until-zero ( n quot -- )
 165     [ dup zero? ] swap until drop ; inline
 166
 167 : cum-reduce ( seq identity quot: ( prev elt -- next ) -- result cum-result )
 168     [ dup rot ] dip dup '[ _ curry dip dupd @ ] each ; inline
 169
 170 <PRIVATE
 171
 172 :: (gini) ( seq -- x )
 173     seq natural-sort :> sorted
 174     seq length :> len
 175     sorted 0 [ + ] cum-reduce :> ( a b )
 176     b len a * / :> B
 177     1 len recip + 2 B * - ;
 178
 179 PRIVATE>
 180
 181 : gini ( seq -- x )
 182     dup length 1 <= [ drop 0 ] [ (gini) ] if ;
 183
 184 : concentration-coefficient ( seq -- x )
 185     dup length 1 <= [
 186         drop 0
 187     ] [
 188         [ (gini) ] [ length [ ] [ 1 - ] bi / ] bi *
 189     ] if ;
 190
 191 : herfindahl ( seq -- x )
 192     [ sum-of-squares ] [ sum sq ] bi / ;
 193
 194 : normalized-herfindahl ( seq -- x )
 195     [ herfindahl ] [ length recip ] bi
 196     [ - ] [ 1 swap - / ] bi ;
 197
 198 : exponential-index ( seq -- x )
 199     dup sum '[ _ / dup ^ ] map-product ;
 200
 201 : weighted-random ( histogram -- obj )
 202     unzip cum-sum [ last random ] [ bisect-left ] bi swap nth ;
 203
 204 : unique-indices ( seq -- unique indices )
 205     [ members ] keep over dup length iota H{ } zip-as '[ _ at ] map ;
 206
 207 : digitize] ( seq bins -- seq' )
 208     '[ _ bisect-left ] map ;
 209
 210 : digitize) ( seq bins -- seq' )
 211     '[ _ bisect-right ] map ;
 212
 213 <PRIVATE
 214
 215 : steps ( a b length -- a b step )
 216     [ 2dup swap - ] dip / ; inline
 217
 218 PRIVATE>
 219
 220 : linspace[a,b) ( a b length -- seq )
 221     steps ,b) <range> ;
 222
 223 : linspace[a,b] ( a b length -- seq )
 224     {
 225         { [ dup 1 < ] [ 3drop { } ] }
 226         { [ dup 1 = ] [ 2drop 1array ] }
 227         [ 1 - steps <range> ]
 228     } cond ;
 229
 230 : logspace[a,b) ( a b length base -- seq )
 231     [ linspace[a,b) ] dip swap n^v ;
 232
 233 : logspace[a,b] ( a b length base -- seq )
 234     [ linspace[a,b] ] dip swap n^v ;
 235
 236 : majority ( seq -- elt/f )
 237     [ f 0 ] dip [
 238         over zero? [ 2nip 1 ] [
 239             pick = [ 1 + ] [ 1 - ] if
 240         ] if
 241     ] each zero? [ drop f ] when ;
 242
 243 : compression-lengths ( a b -- len(a+b) len(a) len(b) )
 244     [ append ] 2keep [ >byte-array compress data>> length ] tri@ ;
 245
 246 : compression-distance ( a b -- n )
 247     compression-lengths sort-pair [ - ] [ / ] bi* ;
 248
 249 : compression-dissimilarity ( a b -- n )
 250     compression-lengths + / ;
 251
 252 GENERIC: round-to-even ( x -- y )
 253
 254 M: integer round-to-even ; inline
 255
 256 M: ratio round-to-even
 257     >fraction [ /mod abs 2 * ] keep > [ dup 0 < -1 1 ? + ] when ;
 258
 259 M: float round-to-even
 260     dup 0 > [
 261         dup 0x1p52 <= [ 0x1p52 + 0x1p52 - ] when
 262     ] [
 263         dup -0x1p52 >= [ 0x1p52 - 0x1p52 + ] when
 264     ] if ;
 265
 266 : round-to-decimal ( x n -- y )
 267     10^ [ * 0.5 over 0 > [ + ] [ - ] if truncate ] [ / ] bi ;
 268
 269 : round-to-step ( x step -- y )
 270     [ [ / round ] [ * ] bi ] unless-zero ;
 271
 272 GENERIC: round-away-from-zero ( x -- y )
 273
 274 M: integer round-away-from-zero ; inline
 275
 276 M: real round-away-from-zero
 277     dup 0 < [ floor ] [ ceiling ] if ;
 278
 279 : monotonic-count ( seq quot: ( elt1 elt2 -- ? ) -- newseq )
 280     over empty? [ 2drop { } ] [
 281         [ 0 swap unclip-slice swap ] dip '[
 282             [ @ [ 1 + ] [ drop 0 ] if ] keep over
 283         ] { } map-as 2nip 0 prefix
 284     ] if ; inline
 285
 286 : max-monotonic-count ( seq quot: ( elt1 elt2 -- ? ) -- n )
 287     over empty? [ 2drop 0 ] [
 288         [ 0 swap unclip-slice swap 0 ] dip '[
 289             [ swapd @ [ 1 + ] [ max 0 ] if ] keep swap
 290         ] reduce nip max
 291     ] if ; inline
 292
 293 <PRIVATE
 294
 295 : kahan+ ( c sum elt -- c' sum' )
 296     rot - 2dup + [ -rot [ - ] bi@ ] keep ; inline
 297
 298 PRIVATE>
 299
 300 : kahan-sum ( seq -- n )
 301     [ 0.0 0.0 ] dip [ kahan+ ] each nip ;
 302
 303 : map-kahan-sum ( ... seq quot: ( ... elt -- ... n ) -- ... n )
 304     [ 0.0 0.0 ] 2dip [ 2dip rot kahan+ ] curry
 305     [ -rot ] prepose each nip ; inline
 306
 307 SYNTAX: .. dup pop scan-object [a,b) suffix! ;
 308
 309 SYNTAX: ... dup pop scan-object [a,b] suffix! ;
 310
 311 GENERIC: sum-squares ( seq -- n )
 312 M: object sum-squares [ sq ] map-sum ;
 313 M: iota-tuple sum-squares
 314     length 1 - [ ] [ 1 + ] [ 1/2 + ] tri * * 3 / ;
 315
 316 GENERIC: sum-cubes ( seq -- n )
 317 M: object sum-cubes [ 3 ^ ] map-sum ;
 318 M: iota-tuple sum-cubes sum sq ;
 319
 320 : mobius ( n -- x )
 321     group-factors values [ 1 ] [
 322         dup [ 1 > ] any?
 323         [ drop 0 ] [ length even? 1 -1 ? ] if
 324     ] if-empty ;
 325
 326 : kelly ( winning-probability odds -- fraction )
 327     [ 1 + * 1 - ] [ / ] bi ;