module Data.Array.Repa.Operators.Mapping
( map
, zipWith
, (+^)
, (-^)
, (*^)
, (/^))
where
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 qualified Prelude as P
import Prelude (($), (.), (+), (*), (+), (/), ())
map :: (Shape sh, Elt a, Elt b)
=> (a -> b)
-> Array sh a
-> Array sh b
map f (Array sh regions)
= Array sh (mapRegions regions)
where
mapRegions rs
= case rs of
[] -> []
[r] -> [mapRegion r]
[r1, r2] -> [mapRegion r1, mapRegion r2]
[r1, r2, r3] -> [mapRegion r1, mapRegion r2, mapRegion r3]
[r1, r2, r3, r4] -> [mapRegion r1, mapRegion r2, mapRegion r3, mapRegion r4]
_ -> mapRegions' rs
mapRegions' rs
= case rs of
[] -> []
(r : rs') -> mapRegion r : mapRegions' rs'
mapRegion (Region range gen)
= Region range (mapGen gen)
mapGen gen
= case gen of
GenManifest vec
-> GenCursor
P.id
addDim
(\ix -> f $ V.unsafeIndex vec $ S.toIndex sh ix)
GenCursor makeCursor shiftCursor loadElem
-> GenCursor makeCursor shiftCursor (f . loadElem)
zipWith :: (Shape sh, Elt a, Elt b, Elt c)
=> (a -> b -> c)
-> Array sh a
-> Array sh b
-> Array sh c
zipWith f arr1 arr2
| Array sh2 [_] <- arr1
, Array sh1 [ Region g21 (GenCursor make21 _ load21)
, Region g22 (GenCursor make22 _ load22)] <- arr2
= let
load21' ix = f (arr1 `unsafeIndex` ix) (load21 $ make21 ix)
load22' ix = f (arr1 `unsafeIndex` ix) (load22 $ make22 ix)
in Array (S.intersectDim sh1 sh2)
[ Region g21 (GenCursor P.id addDim load21')
, Region g22 (GenCursor P.id addDim load22') ]
| P.otherwise
= let
getElem' ix = f (arr1 `unsafeIndex` ix) (arr2 `unsafeIndex` ix)
in fromFunction
(S.intersectDim (extent arr1) (extent arr2))
getElem'
(+^) = zipWith (+)
(-^) = zipWith ()
(*^) = zipWith (*)
(/^) = zipWith (/)