module Data.Array.Repa.Operators.Reduction
( fold, foldAll
, sum, sumAll)
where
import Data.Array.Repa.Index
import Data.Array.Repa.Internals.Elt
import Data.Array.Repa.Internals.Base
import Data.Array.Repa.Shape as S
import qualified Data.Vector.Unboxed as V
import Prelude hiding (sum)
fold :: (Shape sh, Elt a)
=> (a -> a -> a)
-> a
-> Array (sh :. Int) a
-> Array sh a
fold f x arr
= x `seq` arr `deepSeqArray`
let sh' :. n = extent arr
elemFn i = V.foldl' f x
$ V.map (\ix -> arr ! (i :. ix))
(V.enumFromTo 0 (n 1))
in fromFunction sh' elemFn
foldAll :: (Shape sh, Elt a)
=> (a -> a -> a)
-> a
-> Array sh a
-> a
foldAll f x arr
= V.foldl' f x
$ V.map ((arr !) . (S.fromIndex (extent arr)))
$ V.enumFromTo
0
((S.size $ extent arr) 1)
sum :: (Shape sh, Elt a, Num a)
=> Array (sh :. Int) a
-> Array sh a
sum arr = fold (+) 0 arr
sumAll :: (Shape sh, Elt a, Num a)
=> Array sh a
-> a
sumAll arr
= V.foldl' (+) 0
$ V.map ((arr !) . (S.fromIndex (extent arr)))
$ V.enumFromTo
0
((S.size $ extent arr) 1)