1 ! Copyright (C) 2012 Doug Coleman.
2 ! See http://factorcode.org/license.txt for BSD license.
3 USING: accessors arrays assocs combinators.short-circuit fry
4 grouping kernel math math.functions math.order math.vectors
5 parser prettyprint.custom sequences sequences.deep ;
8 : flat? ( array -- ? ) [ sequence? ] none? ; inline
10 GENERIC: array-replace ( object -- shape )
14 M: object array-replace drop f ;
16 M: sequence array-replace
23 TUPLE: uniform-shape shape ;
24 C: <uniform-shape> uniform-shape
26 TUPLE: abnormal-shape shape ;
27 C: <abnormal-shape> abnormal-shape
29 GENERIC: wrap-shape ( object -- shape )
32 1array <uniform-shape> ;
34 M: sequence wrap-shape
36 [ length ] [ first ] bi 2array <uniform-shape>
41 GENERIC: shape ( array -- shape )
43 M: sequence shape array-replace wrap-shape ;
45 : ndim ( array -- n ) shape length ;
47 ERROR: no-negative-shape-components shape ;
49 : check-shape-domain ( seq -- seq )
50 dup [ 0 < ] any? [ no-negative-shape-components ] when ;
52 GENERIC: shape-capacity ( shape -- n )
54 M: sequence shape-capacity check-shape-domain product ;
56 M: uniform-shape shape-capacity
59 M: abnormal-shape shape-capacity
61 [ dup sequence? [ drop ] [ + ] if ] [ 1 + ] if*
64 ERROR: underlying-shape-mismatch underlying shape ;
66 ERROR: no-abnormally-shaped-arrays underlying shape ;
68 GENERIC: check-underlying-shape ( underlying shape -- underlying shape )
70 M: abnormal-shape check-underlying-shape
71 no-abnormally-shaped-arrays ;
73 M: uniform-shape check-underlying-shape
74 shape>> check-underlying-shape ;
76 M: sequence check-underlying-shape
77 2dup [ length ] [ shape-capacity ] bi*
78 = [ underlying-shape-mismatch ] unless ; inline
80 ERROR: shape-mismatch shaped0 shaped1 ;
84 : check-shape ( shaped-array shaped-array -- shaped-array shaped-array )
87 sequence= [ shape-mismatch ] unless ;
89 TUPLE: shaped-array underlying shape ;
90 TUPLE: row-array < shaped-array ;
91 TUPLE: col-array < shaped-array ;
93 M: shaped-array length underlying>> length ; inline
95 M: shaped-array shape shape>> ;
97 : make-shaped-array ( underlying shape class -- shaped-array )
98 [ check-underlying-shape ] dip new
100 swap >>underlying ; inline
102 : <shaped-array> ( underlying shape -- shaped-array )
103 shaped-array make-shaped-array ; inline
105 : <row-array> ( underlying shape -- shaped-array )
106 row-array make-shaped-array ; inline
108 : <col-array> ( underlying shape -- shaped-array )
109 col-array make-shaped-array ; inline
111 GENERIC: >shaped-array ( array -- shaped-array )
112 GENERIC: >row-array ( array -- shaped-array )
113 GENERIC: >col-array ( array -- shaped-array )
115 M: sequence >shaped-array
116 [ { } flatten-as ] [ shape ] bi <shaped-array> ;
118 M: shaped-array >shaped-array ;
120 M: shaped-array >row-array
121 [ underlying>> ] [ shape>> ] bi <row-array> ;
123 M: shaped-array >col-array
124 [ underlying>> ] [ shape>> ] bi <col-array> ;
126 M: sequence >col-array
127 [ flatten ] [ shape ] bi <col-array> ;
129 : shaped-unary-op ( shaped quot -- )
130 [ >shaped-array ] dip
131 [ underlying>> ] prepose
132 [ shape>> clone ] bi shaped-array boa ; inline
134 : shaped-shaped-binary-op ( shaped0 shaped1 quot -- c )
136 [ [ underlying>> ] bi@ ] prepose
137 [ drop shape>> clone ] 2bi shaped-array boa ; inline
139 : shaped+ ( a b -- c ) [ v+ ] shaped-shaped-binary-op ;
140 : shaped- ( a b -- c ) [ v- ] shaped-shaped-binary-op ;
141 : shaped*. ( a b -- c ) [ v* ] shaped-shaped-binary-op ;
143 : shaped*n ( a b -- c ) [ v*n ] curry shaped-unary-op ;
144 : n*shaped ( a b -- c ) swap shaped*n ;
146 : shaped-cos ( a -- b ) [ [ cos ] map ] shaped-unary-op ;
147 : shaped-sin ( a -- b ) [ [ sin ] map ] shaped-unary-op ;
149 : shaped-array>array ( shaped-array -- array )
150 [ underlying>> ] [ shape>> ] bi
154 [ rest-slice reverse [ group ] each ] unless-empty
157 : reshape ( shaped-array shape -- array )
158 check-underlying-shape
159 [ >shaped-array ] dip >>shape ;
161 : shaped-like ( shaped-array shape -- array )
162 [ underlying>> clone ] dip <shaped-array> ;
164 : repeated-shaped ( shape element -- shaped-array )
165 [ [ shape-capacity ] dip <array> ]
166 [ drop 1 1 pad-head ] 2bi <shaped-array> ;
168 : zeros ( shape -- shaped-array ) 0 repeated-shaped ;
170 : ones ( shape -- shaped-array ) 1 repeated-shaped ;
172 : increasing ( shape -- shaped-array )
173 [ shape-capacity <iota> >array ] [ ] bi <shaped-array> ;
175 : decreasing ( shape -- shaped-array )
176 [ shape-capacity <iota> <reversed> >array ] [ ] bi <shaped-array> ;
178 : row-length ( shape -- n ) rest-slice product ; inline
180 : column-length ( shape -- n ) first ; inline
182 : each-row ( shaped-array quot -- )
183 [ [ underlying>> ] [ shape>> row-length <groups> ] bi ] dip
186 TUPLE: transposed shaped-array ;
188 : transposed-shape ( shaped-array -- shape )
191 TUPLE: row-traverser shaped-array index ;
193 GENERIC: next-index ( object -- index )
195 SYNTAX: sa{ \ } [ >shaped-array ] parse-literal ;
197 ! M: row-array pprint* shaped-array>array pprint* ;
198 ! M: col-array pprint* shaped-array>array flip pprint* ;
199 M: shaped-array pprint-delims drop \ sa{ \ } ;
200 M: shaped-array >pprint-sequence shaped-array>array ;
201 M: shaped-array pprint* pprint-object ;
202 M: shaped-array pprint-narrow? drop f ;
204 ERROR: shaped-bounds-error seq shape ;
206 : shaped-bounds-check ( seq shaped -- seq shaped )
207 2dup shape [ < ] 2all? [ shaped-bounds-error ] unless ;
210 : calculate-row-major-index ( seq shape -- i )
211 1 [ * ] accumulate nip reverse vdot ;
213 : calculate-column-major-index ( seq shape -- i )
214 1 [ * ] accumulate nip vdot ;
216 : get-shaped-row-major ( seq shaped -- elt )
217 shaped-bounds-check [ shape calculate-row-major-index ] [ underlying>> ] bi nth ;
219 : set-shaped-row-major ( obj seq shaped -- )
220 shaped-bounds-check [ shape calculate-row-major-index ] [ underlying>> ] bi set-nth ;
222 : get-shaped-column-major ( seq shaped -- elt )
223 shaped-bounds-check [ shape calculate-column-major-index ] [ underlying>> ] bi nth ;
225 : set-shaped-column-major ( obj seq shaped -- )
226 shaped-bounds-check [ shape calculate-column-major-index ] [ underlying>> ] bi set-nth ;
229 : 2d? ( shape -- ? ) length 2 = ;
230 ERROR: 2d-expected shaped ;
231 : check-2d ( shaped -- shaped ) dup shape>> 2d? [ 2d-expected ] unless ;
233 : diagonal? ( coord -- ? ) { [ 2d? ] [ first2 = ] } 1&& ;
235 ! : definite? ( sa -- ? )
237 : shaped-each ( .. sa quot -- )
238 [ underlying>> ] dip each ; inline
240 ! : set-shaped-where ( .. elt sa quot -- )
242 ! [ underlying>> [ length <iota> ] keep zip ]
244 ! ] dip '[ _ [ _ set- ] @ ] assoc-each ; inline
246 : shaped-map! ( .. sa quot -- sa )
247 '[ _ map ] change-underlying ; inline
249 : shaped-map ( .. sa quot -- sa' )
250 [ [ underlying>> ] dip map ]
251 [ drop shape>> ] 2bi <shaped-array> ; inline
253 : pad-shapes ( sa0 sa1 -- sa0' sa1' )
255 2dup longer length '[ _ 1 pad-head ] bi@
256 [ shaped-like ] bi-curry@ bi* ;
258 : output-shape ( sa0 sa1 -- shape )
260 [ 2dup [ zero? ] either? [ max ] [ 2drop 0 ] if ] 2map ;
262 : broadcast-shape-matches? ( sa broadcast-shape -- ? )
264 { [ drop 1 = ] [ = ] } 2||
267 : broadcastable? ( sa0 sa1 -- ? )
269 [ [ shape>> ] bi@ ] [ output-shape ] 2bi
270 '[ _ broadcast-shape-matches? ] both? ;
272 TUPLE: block-array shaped shape ;
274 : <block-array> ( underlying shape -- obj )
277 : iteration-indices ( shaped -- seq )
279 cartesian-product concat
280 [ dup first array? [ first2 suffix ] when ] map
283 : map-shaped-index ( shaped quot -- shaped )
285 [ [ underlying>> ] [ shape>> iteration-indices ] bi zip ] dip map
286 ] dip swap >>underlying ; inline
288 : identity-matrix ( n -- shaped )
289 dup 2array zeros [ second first2 = 1 0 ? ] map-shaped-index ;
291 : map-strict-lower ( shaped quot -- shaped )
293 '[ first2 first2 > _ when ] map-shaped-index ; inline
295 : map-lower ( shaped quot -- shaped )
297 '[ first2 first2 >= _ when ] map-shaped-index ; inline
299 : map-strict-upper ( shaped quot -- shaped )
301 '[ first2 first2 < _ when ] map-shaped-index ; inline
303 : map-upper ( shaped quot -- shaped )
305 '[ first2 first2 <= _ when ] map-shaped-index ; inline
307 : map-diagonal ( shaped quot -- shaped )
309 '[ first2 first2 = _ when ] map-shaped-index ; inline
311 : upper ( shape obj -- shaped )
312 [ zeros check-2d ] dip '[ drop _ ] map-upper ;
314 : strict-upper ( shape obj -- shaped )
315 [ zeros check-2d ] dip '[ drop _ ] map-strict-upper ;
317 : lower ( shape obj -- shaped )
318 [ zeros check-2d ] dip '[ drop _ ] map-lower ;
320 : strict-lower ( shape obj -- shaped )
321 [ zeros check-2d ] dip '[ drop _ ] map-strict-lower ;