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 byte-arrays combinators
   5 combinators.short-circuit compression.zlib grouping kernel math
   6 math.bitwise math.combinatorics math.constants math.functions
   7 math.order math.primes math.primes.factors math.statistics
   8 math.vectors random ranges ranges.private sequences
   9 sequences.extras sets sorting sorting.extras ;
  10
  11 IN: math.extras
  12
  13 DEFER: stirling
  14
  15 <PRIVATE
  16
  17 : (stirling) ( n k -- x )
  18     [ [ 1 - ] bi@ stirling ]
  19     [ [ 1 - ] dip stirling ]
  20     [ nip * + ] 2tri ;
  21
  22 PRIVATE>
  23
  24 MEMO: stirling ( n k -- x )
  25     2dup { [ = ] [ nip 1 = ] } 2||
  26     [ 2drop 1 ] [ (stirling) ] if ;
  27
  28 :: ramanujan ( x -- y )
  29     pi sqrt x e / x ^ * x 8 * 4 + x * 1 + x * 1/30 + 1/6 ^ * ;
  30
  31 DEFER: bernoulli
  32
  33 <PRIVATE
  34
  35 : (bernoulli) ( p -- n )
  36     [ <iota> ] [ 1 + ] bi [
  37         0 [ [ nCk ] [ bernoulli * ] bi + ] with reduce
  38     ] keep recip neg * ;
  39
  40 PRIVATE>
  41
  42 MEMO: bernoulli ( p -- n )
  43     [ 1 ] [ (bernoulli) ] if-zero ;
  44
  45 : chi2 ( actual expected -- n )
  46     0 [ dup 0 > [ [ - sq ] keep / + ] [ 2drop ] if ] 2reduce ;
  47
  48 <PRIVATE
  49
  50 : df-check ( df -- )
  51     even? [ "odd degrees of freedom" throw ] unless ;
  52
  53 : (chi2P) ( chi/2 df/2 -- p )
  54     [1..b) dupd n/v cum-product swap neg e^ [ v*n sum ] keep + ;
  55
  56 PRIVATE>
  57
  58 : chi2P ( chi df -- p )
  59     dup df-check [ 2.0 / ] [ 2 /i ] bi* (chi2P) 1.0 min ;
  60
  61 <PRIVATE
  62
  63 : check-jacobi ( m -- m )
  64     dup { [ integer? ] [ 0 > ] [ odd? ] } 1&&
  65     [ "modulus must be odd positive integer" throw ] unless ;
  66
  67 : mod' ( x y -- n )
  68     [ mod ] keep over zero? [ drop ] [
  69         2dup [ sgn ] same? [ drop ] [ + ] if
  70     ] if ;
  71
  72 PRIVATE>
  73
  74 : jacobi ( a m -- n )
  75     check-jacobi [ mod' ] keep 1
  76     [ pick zero? ] [
  77         [ pick even? ] [
  78             [ 2 / ] 2dip
  79             over 8 mod' { 3 5 } member? [ neg ] when
  80         ] while swapd
  81         2over [ 4 mod' 3 = ] both? [ neg ] when
  82         [ [ mod' ] keep ] dip
  83     ] until [ nip 1 = ] dip 0 ? ;
  84
  85 <PRIVATE
  86
  87 : check-legendere ( m -- m )
  88     dup prime? [ "modulus must be prime positive integer" throw ] unless ;
  89
  90 PRIVATE>
  91
  92 : legendere ( a m -- n )
  93     check-legendere jacobi ;
  94
  95 : moving-average ( seq n -- newseq )
  96     <clumps> [ mean ] map ;
  97
  98 : exponential-moving-average ( seq a -- newseq )
  99     [ 1 ] 2dip '[ dupd swap - _ * + dup ] map nip ;
 100
 101 : moving-median ( u n -- v )
 102     <clumps> [ median ] map ;
 103
 104 : moving-supremum ( u n -- v )
 105     <clumps> [ supremum ] map ;
 106
 107 : moving-infimum ( u n -- v )
 108     <clumps> [ infimum ] map ;
 109
 110 : moving-sum ( u n -- v )
 111     <clumps> [ sum ] map ;
 112
 113 : moving-count ( ... u n quot: ( ... elt -- ... ? ) -- ... v )
 114     [ <clumps> ] [ '[ _ count ] map ] bi* ; inline
 115
 116 : nonzero ( seq -- seq' )
 117     [ zero? ] reject ;
 118
 119 : bartlett ( n -- seq )
 120     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 121         [ <iota> ] [ 1 - 2 / ] bi [
 122             [ recip * ] [ >= ] 2bi [ 2 swap - ] when
 123         ] curry map
 124     ] if ;
 125
 126 : [0,2pi] ( n -- seq )
 127     [ <iota> ] [ 1 - 2pi swap / ] bi v*n ;
 128
 129 : hanning ( n -- seq )
 130     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 131         [0,2pi] [ cos -0.5 * 0.5 + ] map!
 132     ] if ;
 133
 134 : hamming ( n -- seq )
 135     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 136         [0,2pi] [ cos -0.46 * 0.54 + ] map!
 137     ] if ;
 138
 139 : blackman ( n -- seq )
 140     dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
 141         [0,2pi] [
 142             [ cos -0.5 * ] [ 2 * cos 0.08 * ] bi + 0.42 +
 143         ] map
 144     ] if ;
 145
 146 : nan-sum ( seq -- n )
 147     0 [ dup fp-nan? [ drop ] [ + ] if ] binary-reduce ;
 148
 149 : nan-min ( seq -- n )
 150     [ fp-nan? ] reject infimum ;
 151
 152 : nan-max ( seq -- n )
 153     [ fp-nan? ] reject supremum ;
 154
 155 : fill-nans ( seq -- newseq )
 156     [ first ] keep [
 157         dup fp-nan? [ drop dup ] [ nip dup ] if
 158     ] map nip ;
 159
 160 : sinc ( x -- y )
 161     [ 1 ] [ pi * [ sin ] [ / ] bi ] if-zero ;
 162
 163 : cum-reduce ( seq identity quot: ( prev elt -- next ) -- result cum-result )
 164     [ dup rot ] dip dup '[ _ curry dip dupd @ ] each ; inline
 165
 166 <PRIVATE
 167
 168 :: (gini) ( seq -- x )
 169     seq natural-sort :> sorted
 170     seq length :> len
 171     sorted 0 [ + ] cum-reduce :> ( a b )
 172     b len a * / :> c
 173     1 len recip + 2 c * - ;
 174
 175 PRIVATE>
 176
 177 : gini ( seq -- x )
 178     dup length 1 <= [ drop 0 ] [ (gini) ] if ;
 179
 180 : concentration-coefficient ( seq -- x )
 181     dup length 1 <= [
 182         drop 0
 183     ] [
 184         [ (gini) ] [ length [ ] [ 1 - ] bi / ] bi *
 185     ] if ;
 186
 187 : herfindahl ( seq -- x )
 188     [ sum-of-squares ] [ sum sq ] bi / ;
 189
 190 : normalized-herfindahl ( seq -- x )
 191     [ herfindahl ] [ length recip ] bi
 192     [ - ] [ 1 swap - / ] bi ;
 193
 194 : exponential-index ( seq -- x )
 195     dup sum '[ _ / dup ^ ] map-product ;
 196
 197 : weighted-random ( histogram -- obj )
 198     unzip cum-sum [ last random ] [ bisect-left ] bi swap nth ;
 199
 200 : unique-indices ( seq -- unique indices )
 201     [ members ] keep over dup length <iota>
 202     H{ } zip-as '[ _ at ] map ;
 203
 204 : digitize] ( seq bins -- seq' )
 205     '[ _ bisect-left ] map ;
 206
 207 : digitize) ( seq bins -- seq' )
 208     '[ _ bisect-right ] map ;
 209
 210 <PRIVATE
 211
 212 : steps ( a b length -- a b step )
 213     [ 2dup swap - ] dip / ; inline
 214
 215 PRIVATE>
 216
 217 : linspace[a..b) ( a b length -- seq )
 218     steps ..b) <range> ;
 219
 220 : linspace[a..b] ( a b length -- seq )
 221     {
 222         { [ dup 1 < ] [ 3drop { } ] }
 223         { [ dup 1 = ] [ 2drop 1array ] }
 224         [ 1 - steps <range> ]
 225     } cond ;
 226
 227 : logspace[a..b) ( a b length base -- seq )
 228     [ linspace[a..b) ] dip swap n^v ;
 229
 230 : logspace[a..b] ( a b length base -- seq )
 231     [ linspace[a..b] ] dip swap n^v ;
 232
 233 : majority ( seq -- elt/f )
 234     [ f 0 ] dip [
 235         over zero? [ 2nip 1 ] [
 236             pick = [ 1 + ] [ 1 - ] if
 237         ] if
 238     ] each zero? [ drop f ] when ;
 239
 240 : compression-lengths ( a b -- len(a+b) len(a) len(b) )
 241     [ append ] 2keep [ >byte-array compress data>> length ] tri@ ;
 242
 243 : compression-distance ( a b -- n )
 244     compression-lengths sort-pair [ - ] [ / ] bi* ;
 245
 246 : compression-dissimilarity ( a b -- n )
 247     compression-lengths + / ;
 248
 249 : round-to-decimal ( x n -- y )
 250     10^ [ * 0.5 over 0 > [ + ] [ - ] if truncate ] [ / ] bi ;
 251
 252 : round-to-step ( x step -- y )
 253     [ [ / round ] [ * ] bi ] unless-zero ;
 254
 255 GENERIC: round-away-from-zero ( x -- y )
 256
 257 M: integer round-away-from-zero ; inline
 258
 259 M: real round-away-from-zero
 260     dup 0 < [ floor ] [ ceiling ] if ;
 261
 262 : monotonic-count ( seq quot: ( elt1 elt2 -- ? ) -- newseq )
 263     over empty? [ 2drop { } ] [
 264         [ 0 swap unclip-slice swap ] dip '[
 265             [ @ [ 1 + ] [ drop 0 ] if ] keep over
 266         ] { } map-as 2nip 0 prefix
 267     ] if ; inline
 268
 269 : max-monotonic-count ( seq quot: ( elt1 elt2 -- ? ) -- n )
 270     over empty? [ 2drop 0 ] [
 271         [ 0 swap unclip-slice swap 0 ] dip '[
 272             [ swapd @ [ 1 + ] [ max 0 ] if ] keep swap
 273         ] reduce nip max
 274     ] if ; inline
 275
 276 <PRIVATE
 277
 278 : kahan+ ( c sum elt -- c' sum' )
 279     rot - 2dup + [ -rot [ - ] bi@ ] keep ; inline
 280
 281 PRIVATE>
 282
 283 : kahan-sum ( seq -- n )
 284     [ 0.0 0.0 ] dip [ kahan+ ] each nip ;
 285
 286 : map-kahan-sum ( ... seq quot: ( ... elt -- ... n ) -- ... n )
 287     [ 0.0 0.0 ] 2dip [ 2dip rot kahan+ ] curry
 288     [ -rot ] prepose each nip ; inline
 289
 290 <PRIVATE
 291
 292 ! Adaptive Precision Floating-Point Arithmetic and Fast Robust Geometric Predicates
 293 ! www-2.cs.cmu.edu/afs/cs/project/quake/public/papers/robust-arithmetic.ps
 294
 295 : sort-partial ( x y -- x' y' )
 296     2dup [ abs ] bi@ < [ swap ] when ; inline
 297
 298 :: partial+ ( x y -- hi lo )
 299     x y + dup x - :> yr y yr - ; inline
 300
 301 :: partial-sums ( seq -- seq' )
 302     V{ } clone :> partials
 303     seq [
 304         0 partials [
 305             swapd sort-partial partial+ swapd
 306             [ over partials set-nth 1 + ] unless-zero
 307         ] each :> i
 308         i partials shorten
 309         [ i partials set-nth ] unless-zero
 310     ] each partials ;
 311
 312 :: sum-exact ( partials -- n )
 313     partials [ 0.0 ] [
 314         ! sum from the top, stop when sum becomes inexact
 315         [ 0.0 0.0 ] dip [
 316             nip partial+ dup 0.0 = not
 317         ] find-last drop :> ( lo n )
 318
 319         ! make half-even rounding work across multiple partials
 320         n [ 0 > ] [ f ] if* [
 321             n 1 - partials nth
 322             [ 0.0 < lo 0.0 < and ]
 323             [ 0.0 > lo 0.0 > and ] bi or [
 324                 lo 2.0 * :> y
 325                 dup y + :> x
 326                 x over - :> yr
 327                 y yr = [ drop x ] when
 328             ] when
 329         ] when
 330     ] if-empty ;
 331
 332 PRIVATE>
 333
 334 : sum-floats ( seq -- n )
 335     partial-sums sum-exact ;
 336
 337 : mobius ( n -- x )
 338     group-factors values [ 1 ] [
 339         dup [ 1 > ] any?
 340         [ drop 0 ] [ length even? 1 -1 ? ] if
 341     ] if-empty ;
 342
 343 : kelly ( winning-probability odds -- fraction )
 344     [ 1 + * 1 - ] [ / ] bi ;
 345
 346 :: integer-sqrt ( m -- n )
 347     m [ 0 ] [
 348         dup 0 < [ non-negative-integer-expected ] when
 349         bit-length 1 - 2 /i :> c
 350         1 :> a!
 351         0 :> d!
 352         c bit-length <iota> <reversed> [| s |
 353             d :> e
 354             c s neg shift d!
 355             a d e - 1 - shift
 356             m 2 c * e - d - 1 + neg shift a /i + a!
 357         ] each
 358         a a sq m > [ 1 - ] when
 359     ] if-zero ;
 360
 361 <PRIVATE
 362
 363 : reduce-evens ( value u v -- value' u' v' )
 364     [ 2dup [ even? ] both? ] [ [ 2 * ] [ 2/ ] [ 2/ ] tri* ] while ;
 365
 366 : reduce-odds ( value u v -- value' u' v' )
 367     [
 368         [ [ dup even? ] [ 2/ ] while ] bi@
 369         2dup <=> {
 370             { +eq+ [ over '[ _ * ] 2dip f ] }
 371             { +lt+ [ swap [ - ] keep t ] }
 372             { +gt+ [ [ - ] keep t ] }
 373         } case
 374     ] loop ;
 375
 376 PRIVATE>
 377
 378 : stein ( u v -- w )
 379     2dup [ zero? ] both? [ "gcd for zeros is undefined" throw ] when
 380     [ dup 0 < [ neg ] when ] bi@
 381     [ 1 ] 2dip reduce-evens reduce-odds 2drop ;
 382
 383