module Data.Array.Repa.Stencil
( Stencil (..)
, Boundary (..)
, makeStencil, makeStencil2
, mapStencil2, forStencil2
, mapStencilFrom2, forStencilFrom2
, stencil2)
where
import Data.Array.Repa as R
import Data.Array.Repa.Internals.Base as R
import Data.Array.Repa.Stencil.Base
import Data.Array.Repa.Stencil.Template
import Data.Array.Repa.Specialised.Dim2
import qualified Data.Array.Repa.Shape as S
import qualified Data.Vector.Unboxed as V
import Data.List as List
import GHC.Exts
import GHC.Base
import Debug.Trace
data Cursor
= Cursor Int
forStencil2
:: Elt a
=> Boundary a
-> Array DIM2 a
-> Stencil DIM2 a
-> Array DIM2 a
forStencil2 boundary arr stencil
= mapStencil2 boundary stencil arr
forStencilFrom2
:: (Elt a, Elt b)
=> Boundary a
-> Array DIM2 b
-> (b -> a)
-> Stencil DIM2 a
-> Array DIM2 a
forStencilFrom2 boundary arr from stencil
= mapStencilFrom2 boundary stencil arr from
mapStencil2
:: Elt a
=> Boundary a
-> Stencil DIM2 a
-> Array DIM2 a
-> Array DIM2 a
mapStencil2 boundary stencil arr
= mapStencilFrom2 boundary stencil arr id
mapStencilFrom2
:: (Elt a, Elt b)
=> Boundary a
-> Stencil DIM2 a
-> Array DIM2 b
-> (b -> a)
-> Array DIM2 a
mapStencilFrom2 boundary stencil@(StencilStatic sExtent zero load) arr preConvert
= let (_ :. aHeight :. aWidth) = extent arr
(_ :. sHeight :. sWidth) = sExtent
sHeight2 = sHeight `div` 2
sWidth2 = sWidth `div` 2
!xMin = sWidth2
!yMin = sHeight2
!xMax = aWidth sWidth2 1
!yMax = aHeight sHeight2 1
rectsInternal
= [ Rect (Z :. yMin :. xMin) (Z :. yMax :. xMax ) ]
inInternal (Z :. y :. x)
= x >= xMin && x <= xMax
&& y >= yMin && y <= yMax
rectsBorder
= [ Rect (Z :. 0 :. 0) (Z :. yMin 1 :. aWidth 1)
, Rect (Z :. yMax + 1 :. 0) (Z :. aHeight 1 :. aWidth 1)
, Rect (Z :. yMin :. 0) (Z :. yMax :. xMin 1)
, Rect (Z :. yMin :. xMax + 1) (Z :. yMax :. aWidth 1) ]
inBorder = not . inInternal
makeCursor' (Z :. y :. x)
= Cursor (x + y * aWidth)
shiftCursor' ix (Cursor off)
= Cursor
$ case ix of
Z :. y :. x -> off + y * aWidth + x
getInner' cur
= unsafeAppStencilCursor2 shiftCursor' stencil
arr preConvert cur
getBorder' cur
= case boundary of
BoundConst c -> c
BoundClamp -> unsafeAppStencilCursor2_clamp addDim stencil
arr preConvert cur
in Array (extent arr)
[ Region (RangeRects inBorder rectsBorder)
(GenCursor id addDim getBorder')
, Region (RangeRects inInternal rectsInternal)
(GenCursor makeCursor' shiftCursor' getInner') ]
unsafeAppStencilCursor2
:: (Elt a, Elt b)
=> (DIM2 -> Cursor -> Cursor)
-> Stencil DIM2 a
-> Array DIM2 b
-> (b -> a)
-> Cursor
-> a
unsafeAppStencilCursor2 shift
stencil@(StencilStatic sExtent zero load)
arr@(Array aExtent [Region RangeAll (GenManifest vec)]) preConvert
cur@(Cursor off)
| _ :. sHeight :. sWidth <- sExtent
, _ :. aHeight :. aWidth <- aExtent
, sHeight <= 7, sWidth <= 7
= let
getData (Cursor cur) = preConvert $ vec `V.unsafeIndex` cur
oload oy ox
= let !cur' = shift (Z :. oy :. ox) cur
in load (Z :. oy :. ox) (getData cur')
in template7x7 oload zero
unsafeAppStencilCursor2_clamp
:: forall a b. (Elt a, Elt b)
=> (DIM2 -> DIM2 -> DIM2)
-> Stencil DIM2 a
-> Array DIM2 b
-> (b -> a)
-> DIM2
-> a
unsafeAppStencilCursor2_clamp shift
stencil@(StencilStatic sExtent zero load)
arr@(Array aExtent [Region RangeAll (GenManifest vec)]) preConvert
cur
| _ :. sHeight :. sWidth <- sExtent
, _ :. aHeight :. aWidth <- aExtent
, sHeight <= 7, sWidth <= 7
= let
getData :: DIM2 -> a
getData (Z :. y :. x)
= wrapLoadX x y
wrapLoadX :: Int -> Int -> a
wrapLoadX !x !y
| x < 0 = wrapLoadY 0 y
| x >= aWidth = wrapLoadY (aWidth 1) y
| otherwise = wrapLoadY x y
wrapLoadY :: Int -> Int -> a
wrapLoadY !x !y
| y < 0 = loadXY x 0
| y >= aHeight = loadXY x (aHeight 1)
| otherwise = loadXY x y
loadXY :: Int -> Int -> a
loadXY !x !y
= preConvert $ vec `V.unsafeIndex` (x + y * aWidth)
oload oy ox
= let !cur' = shift (Z :. oy :. ox) cur
in load (Z :. oy :. ox) (getData cur')
in template7x7 oload zero
template7x7
:: (Int -> Int -> a -> a)
-> a -> a
template7x7 f zero
= f (3) (3) $ f (3) (2) $ f (3) (1) $ f (3) 0 $ f (3) 1 $ f (3) 2 $ f (3) 3
$ f (2) (3) $ f (2) (2) $ f (2) (1) $ f (2) 0 $ f (2) 1 $ f (2) 2 $ f (2) 3
$ f (1) (3) $ f (1) (2) $ f (1) (1) $ f (1) 0 $ f (1) 1 $ f (1) 2 $ f (1) 3
$ f 0 (3) $ f 0 (2) $ f 0 (1) $ f 0 0 $ f 0 1 $ f 0 2 $ f 0 3
$ f 1 (3) $ f 1 (2) $ f 1 (1) $ f 1 0 $ f 1 1 $ f 1 2 $ f 1 3
$ f 2 (3) $ f 2 (2) $ f 2 (1) $ f 2 0 $ f 2 1 $ f 2 2 $ f 2 3
$ f 3 (3) $ f 3 (2) $ f 3 (1) $ f 3 0 $ f 3 1 $ f 3 2 $ f 3 3
$ zero