1 ! Copyright (C) 2012 John Benediktsson
2 ! See http://factorcode.org/license.txt for BSD license
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 ;
18 : (stirling) ( n k -- x )
19 [ [ 1 - ] bi@ stirling ]
20 [ [ 1 - ] dip stirling ]
25 MEMO: stirling ( n k -- x )
26 2dup { [ = ] [ nip 1 = ] } 2||
27 [ 2drop 1 ] [ (stirling) ] if ;
29 :: ramanujan ( x -- y )
30 pi sqrt x e / x ^ * x 8 * 4 + x * 1 + x * 1/30 + 1/6 ^ * ;
36 : (bernoulli) ( p -- n )
38 0 [ [ nCk ] [ bernoulli * ] bi + ] with reduce
43 MEMO: bernoulli ( p -- n )
44 [ 1 ] [ (bernoulli) ] if-zero ;
46 : chi2 ( actual expected -- n )
47 0 [ dup 0 > [ [ - sq ] keep / + ] [ 2drop ] if ] 2reduce ;
52 even? [ "odd degrees of freedom" throw ] unless ;
54 : (chi2P) ( chi/2 df/2 -- p )
55 [1,b) dupd n/v cum-product swap neg e^ [ v*n sum ] keep + ;
59 : chi2P ( chi df -- p )
60 dup df-check [ 2.0 / ] [ 2 /i ] bi* (chi2P) 1.0 min ;
64 : check-jacobi ( m -- m )
65 dup { [ integer? ] [ 0 > ] [ odd? ] } 1&&
66 [ "modulus must be odd positive integer" throw ] unless ;
69 [ mod ] keep over zero? [ drop ] [
70 2dup [ sgn ] same? [ drop ] [ + ] if
76 check-jacobi [ mod' ] keep 1
80 over 8 mod' { 3 5 } member? [ neg ] when
82 2over [ 4 mod' 3 = ] both? [ neg ] when
84 ] until [ nip 1 = ] dip 0 ? ;
88 : check-legendere ( m -- m )
89 dup prime? [ "modulus must be prime positive integer" throw ] unless ;
93 : legendere ( a m -- n )
94 check-legendere jacobi ;
96 : moving-average ( seq n -- newseq )
97 <clumps> [ mean ] map ;
99 : exponential-moving-average ( seq a -- newseq )
100 [ 1 ] 2dip '[ dupd swap - _ * + dup ] map nip ;
102 : moving-median ( u n -- v )
103 <clumps> [ median ] map ;
105 : moving-supremum ( u n -- v )
106 <clumps> [ supremum ] map ;
108 : moving-infimum ( u n -- v )
109 <clumps> [ infimum ] map ;
111 : moving-sum ( u n -- v )
112 <clumps> [ sum ] map ;
114 : moving-count ( ... u n quot: ( ... elt -- ... ? ) -- ... v )
115 [ <clumps> ] [ '[ _ count ] map ] bi* ; inline
117 : nonzero ( seq -- seq' )
120 : bartlett ( n -- seq )
121 dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
122 [ iota ] [ 1 - 2 / ] bi [
123 [ recip * ] [ >= ] 2bi [ 2 swap - ] when
127 : [0,2pi] ( n -- seq )
128 [ iota ] [ 1 - 2pi swap / ] bi v*n ;
130 : hanning ( n -- seq )
131 dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
132 [0,2pi] [ cos -0.5 * 0.5 + ] map!
135 : hamming ( n -- seq )
136 dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
137 [0,2pi] [ cos -0.46 * 0.54 + ] map!
140 : blackman ( n -- seq )
141 dup 1 <= [ 1 = [ 1 1array ] [ { } ] if ] [
143 [ cos -0.5 * ] [ 2 * cos 0.08 * ] bi + 0.42 +
147 : nan-sum ( seq -- n )
148 0 [ dup fp-nan? [ drop ] [ + ] if ] binary-reduce ;
150 : nan-min ( seq -- n )
151 [ fp-nan? ] reject infimum ;
153 : nan-max ( seq -- n )
154 [ fp-nan? ] reject supremum ;
156 : fill-nans ( seq -- newseq )
158 dup fp-nan? [ drop dup ] [ nip dup ] if
162 [ 1 ] [ pi * [ sin ] [ / ] bi ] if-zero ;
164 : until-zero ( n quot -- )
165 [ dup zero? ] swap until drop ; inline
167 : cum-reduce ( seq identity quot: ( prev elt -- next ) -- result cum-result )
168 [ dup rot ] dip dup '[ _ curry dip dupd @ ] each ; inline
172 :: (gini) ( seq -- x )
173 seq natural-sort :> sorted
175 sorted 0 [ + ] cum-reduce :> ( a b )
177 1 len recip + 2 B * - ;
182 dup length 1 <= [ drop 0 ] [ (gini) ] if ;
184 : concentration-coefficient ( seq -- x )
188 [ (gini) ] [ length [ ] [ 1 - ] bi / ] bi *
191 : herfindahl ( seq -- x )
192 [ sum-of-squares ] [ sum sq ] bi / ;
194 : normalized-herfindahl ( seq -- x )
195 [ herfindahl ] [ length recip ] bi
196 [ - ] [ 1 swap - / ] bi ;
198 : exponential-index ( seq -- x )
199 dup sum '[ _ / dup ^ ] map-product ;
201 : weighted-random ( histogram -- obj )
202 unzip cum-sum [ last random ] [ bisect-left ] bi swap nth ;
204 : unique-indices ( seq -- unique indices )
205 [ members ] keep over dup length iota H{ } zip-as '[ _ at ] map ;
207 : digitize] ( seq bins -- seq' )
208 '[ _ bisect-left ] map ;
210 : digitize) ( seq bins -- seq' )
211 '[ _ bisect-right ] map ;
215 : steps ( a b length -- a b step )
216 [ 2dup swap - ] dip / ; inline
220 : linspace[a,b) ( a b length -- seq )
223 : linspace[a,b] ( a b length -- seq )
225 { [ dup 1 < ] [ 3drop { } ] }
226 { [ dup 1 = ] [ 2drop 1array ] }
227 [ 1 - steps <range> ]
230 : logspace[a,b) ( a b length base -- seq )
231 [ linspace[a,b) ] dip swap n^v ;
233 : logspace[a,b] ( a b length base -- seq )
234 [ linspace[a,b] ] dip swap n^v ;
236 : majority ( seq -- elt/f )
238 over zero? [ 2nip 1 ] [
239 pick = [ 1 + ] [ 1 - ] if
241 ] each zero? [ drop f ] when ;
243 : compression-lengths ( a b -- len(a+b) len(a) len(b) )
244 [ append ] 2keep [ >byte-array compress data>> length ] tri@ ;
246 : compression-distance ( a b -- n )
247 compression-lengths sort-pair [ - ] [ / ] bi* ;
249 : compression-dissimilarity ( a b -- n )
250 compression-lengths + / ;
252 GENERIC: round-to-even ( x -- y )
254 M: integer round-to-even ; inline
256 M: ratio round-to-even
257 >fraction [ /mod abs 2 * ] keep > [ dup 0 < -1 1 ? + ] when ;
259 M: float round-to-even
261 dup 0x1p52 <= [ 0x1p52 + 0x1p52 - ] when
263 dup -0x1p52 >= [ 0x1p52 - 0x1p52 + ] when
266 : round-to-decimal ( x n -- y )
267 10^ [ * 0.5 over 0 > [ + ] [ - ] if truncate ] [ / ] bi ;
269 : round-to-step ( x step -- y )
270 [ [ / round ] [ * ] bi ] unless-zero ;
272 GENERIC: round-away-from-zero ( x -- y )
274 M: integer round-away-from-zero ; inline
276 M: real round-away-from-zero
277 dup 0 < [ floor ] [ ceiling ] if ;
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
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
295 : kahan+ ( c sum elt -- c' sum' )
296 rot - 2dup + [ -rot [ - ] bi@ ] keep ; inline
300 : kahan-sum ( seq -- n )
301 [ 0.0 0.0 ] dip [ kahan+ ] each nip ;
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
307 ! SYNTAX: .. dup pop scan-object [a,b) suffix! ;
308 ! SYNTAX: ... dup pop scan-object [a,b] suffix! ;
310 GENERIC: sum-squares ( seq -- n )
311 M: object sum-squares [ sq ] map-sum ;
312 M: iota-tuple sum-squares
313 length 1 - [ ] [ 1 + ] [ 1/2 + ] tri * * 3 / ;
315 GENERIC: sum-cubes ( seq -- n )
316 M: object sum-cubes [ 3 ^ ] map-sum ;
317 M: iota-tuple sum-cubes sum sq ;
320 group-factors values [ 1 ] [
322 [ drop 0 ] [ length even? 1 -1 ? ] if
325 : kelly ( winning-probability odds -- fraction )
326 [ 1 + * 1 - ] [ / ] bi ;