module Data.Yarr.Flow (
DefaultFusion(..), DefaultIFusion(..),
dzipWith, dzipWith3, D, delay,
SE, dmapElems, dmapElemsM,
dzipElems2, dzipElems2M, dzipElems3, dzipElems3M,
dzipElems, dzipElemsM,
traverse, zipElems, mapElems, mapElemsM,
cartProduct2, icartProduct2, icartProduct2M,
cartProduct3, icartProduct3, icartProduct3M,
) where
import Data.Yarr.Base
import Data.Yarr.Shape
import Data.Yarr.Fusion
import Data.Yarr.Repr.Delayed
import Data.Yarr.Repr.Separate
import Data.Yarr.Utils.FixedVector as V
traverse
:: (USource r l sh a, Shape sh')
=> (sh -> sh')
-> ((sh -> IO a) -> sh' -> IO b)
-> UArray r l sh a
-> UArray D SH sh' b
traverse transformShape newElem arr =
ShapeDelayed
(transformShape (extent arr))
(touchArray arr) (force arr)
(newElem (index arr))
zipElems
:: (Vector v a,
USource r l sh (v a), USource fr l sh b, DefaultFusion r fr l sh)
=> Fn (Dim v) a b
-> UArray r l sh (v a)
-> UArray fr l sh b
zipElems fn arr = dmap (\v -> inspect v (Fun fn)) arr
mapElems
:: (VecRegular r slr l sh v a, USource slr l sh a,
USource fslr l sh b, DefaultFusion slr fslr l sh, Vector v b)
=> (a -> b)
-> UArray r l sh (v a)
-> UArray (SE fslr) l sh (v b)
mapElems f = dmapElems (V.replicate f)
mapElemsM
:: (VecRegular r slr l sh v a, USource slr l sh a,
USource fslr l sh b, DefaultFusion slr fslr l sh, Vector v b)
=> (a -> IO b)
-> UArray r l sh (v a)
-> UArray (SE fslr) l sh (v b)
mapElemsM f = dmapElemsM (V.replicate f)
dzipWith
:: (USource r1 l sh a, DefaultFusion r1 D l sh, USource D l sh a,
USource r2 l sh b, DefaultFusion r2 D l sh, USource D l sh b,
USource D l sh c, DefaultFusion D D l sh)
=> (a -> b -> c)
-> UArray r1 l sh a
-> UArray r2 l sh b
-> UArray D l sh c
dzipWith f arr1 arr2 = dzip2 f (delay arr1) (delay arr2)
dzipWith3
:: (USource r1 l sh a, DefaultFusion r1 D l sh, USource D l sh a,
USource r2 l sh b, DefaultFusion r2 D l sh, USource D l sh b,
USource r3 l sh c, DefaultFusion r3 D l sh, USource D l sh c,
USource D l sh d, DefaultFusion D D l sh)
=> (a -> b -> c -> d)
-> UArray r1 l sh a
-> UArray r2 l sh b
-> UArray r3 l sh c
-> UArray D l sh d
dzipWith3 f arr1 arr2 arr3 = dzip3 f (delay arr1) (delay arr2) (delay arr3)
cartProduct2
:: (USource r1 l1 Dim1 a, USource r2 l2 Dim1 b)
=> (a -> b -> c)
-> UArray r1 l1 Dim1 a
-> UArray r2 l2 Dim1 b
-> UArray D SH Dim2 c
cartProduct2 f = icartProduct2M (\_ x y -> return (f x y))
icartProduct2
:: (USource r1 l1 Dim1 a, USource r2 l2 Dim1 b)
=> (Dim2 -> a -> b -> c)
-> UArray r1 l1 Dim1 a
-> UArray r2 l2 Dim1 b
-> UArray D SH Dim2 c
icartProduct2 f = icartProduct2M (\ix x y -> return (f ix x y))
icartProduct2M
:: (USource r1 l1 Dim1 a, USource r2 l2 Dim1 b)
=> (Dim2 -> a -> b -> IO c)
-> UArray r1 l1 Dim1 a
-> UArray r2 l2 Dim1 b
-> UArray D SH Dim2 c
icartProduct2M f arr1 arr2 =
ShapeDelayed
(extent arr1, extent arr2)
(touchArray arr1 >> touchArray arr2)
(force arr1 >> force arr2)
indexF
where
indexF ix@(i, j) = do
a <- arr1 `index` i
b <- arr2 `index` j
f ix a b
cartProduct3
:: (USource r1 l1 Dim1 a, USource r2 l2 Dim1 b, USource r3 l3 Dim1 c)
=> (a -> b -> c -> d)
-> UArray r1 l1 Dim1 a
-> UArray r2 l2 Dim1 b
-> UArray r3 l3 Dim1 c
-> UArray D SH Dim3 d
cartProduct3 f = icartProduct3M (\_ x y z -> return (f x y z))
icartProduct3
:: (USource r1 l1 Dim1 a, USource r2 l2 Dim1 b, USource r3 l3 Dim1 c)
=> (Dim3 -> a -> b -> c -> d)
-> UArray r1 l1 Dim1 a
-> UArray r2 l2 Dim1 b
-> UArray r3 l3 Dim1 c
-> UArray D SH Dim3 d
icartProduct3 f = icartProduct3M (\ix x y z -> return (f ix x y z))
icartProduct3M
:: (USource r1 l1 Dim1 a, USource r2 l2 Dim1 b, USource r3 l3 Dim1 c)
=> (Dim3 -> a -> b -> c -> IO d)
-> UArray r1 l1 Dim1 a
-> UArray r2 l2 Dim1 b
-> UArray r3 l3 Dim1 c
-> UArray D SH Dim3 d
icartProduct3M f arr1 arr2 arr3 =
ShapeDelayed
(extent arr1, extent arr2, extent arr3)
(touchArray arr1 >> touchArray arr2 >> touchArray arr3)
(force arr1 >> force arr2 >> force arr3)
indexF
where
indexF ix@(i, j, k) = do
a <- arr1 `index` i
b <- arr2 `index` j
c <- arr3 `index` k
f ix a b c