{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Massiv.Array.Stencil (
Stencil,
makeStencil,
getStencilSize,
getStencilCenter,
Padding (..),
noPadding,
samePadding,
mapStencil,
applyStencil,
idStencil,
sumStencil,
productStencil,
avgStencil,
maxStencil,
minStencil,
foldlStencil,
foldrStencil,
foldStencil,
dimapStencil,
lmapStencil,
rmapStencil,
module Data.Massiv.Array.Stencil.Convolution,
) where
import Data.Coerce
import Data.Massiv.Array.Delayed.Windowed
import Data.Massiv.Array.Manifest
import Data.Massiv.Array.Stencil.Convolution
import Data.Massiv.Array.Stencil.Internal
import Data.Massiv.Array.Stencil.Unsafe
import Data.Massiv.Core.Common
import Data.Semigroup
import GHC.Exts (inline)
getStencilSize :: Stencil ix e a -> Sz ix
getStencilSize :: forall ix e a. Stencil ix e a -> Sz ix
getStencilSize = forall ix e a. Stencil ix e a -> Sz ix
stencilSize
getStencilCenter :: Stencil ix e a -> ix
getStencilCenter :: forall ix e a. Stencil ix e a -> ix
getStencilCenter = forall ix e a. Stencil ix e a -> ix
stencilCenter
mapStencil
:: (Index ix, Manifest r e)
=> Border e
-> Stencil ix e a
-> Array r ix e
-> Array DW ix a
mapStencil :: forall ix r e a.
(Index ix, Manifest r e) =>
Border e -> Stencil ix e a -> Array r ix e -> Array DW ix a
mapStencil Border e
b Stencil ix e a
stencil = forall ix r e a.
(Index ix, Manifest r e) =>
Padding ix e -> Stencil ix e a -> Array r ix e -> Array DW ix a
applyStencil (forall ix e a.
Index ix =>
Stencil ix e a -> Border e -> Padding ix e
samePadding Stencil ix e a
stencil Border e
b) Stencil ix e a
stencil
{-# INLINE mapStencil #-}
data Padding ix e = Padding
{ forall ix e. Padding ix e -> Sz ix
paddingFromOrigin :: !(Sz ix)
, forall ix e. Padding ix e -> Sz ix
paddingFromBottom :: !(Sz ix)
, forall ix e. Padding ix e -> Border e
paddingWithElement :: !(Border e)
}
deriving (Padding ix e -> Padding ix e -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall ix e. (Eq ix, Eq e) => Padding ix e -> Padding ix e -> Bool
/= :: Padding ix e -> Padding ix e -> Bool
$c/= :: forall ix e. (Eq ix, Eq e) => Padding ix e -> Padding ix e -> Bool
== :: Padding ix e -> Padding ix e -> Bool
$c== :: forall ix e. (Eq ix, Eq e) => Padding ix e -> Padding ix e -> Bool
Eq, Int -> Padding ix e -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall ix e. (Index ix, Show e) => Int -> Padding ix e -> ShowS
forall ix e. (Index ix, Show e) => [Padding ix e] -> ShowS
forall ix e. (Index ix, Show e) => Padding ix e -> String
showList :: [Padding ix e] -> ShowS
$cshowList :: forall ix e. (Index ix, Show e) => [Padding ix e] -> ShowS
show :: Padding ix e -> String
$cshow :: forall ix e. (Index ix, Show e) => Padding ix e -> String
showsPrec :: Int -> Padding ix e -> ShowS
$cshowsPrec :: forall ix e. (Index ix, Show e) => Int -> Padding ix e -> ShowS
Show)
noPadding :: Index ix => Padding ix e
noPadding :: forall ix e. Index ix => Padding ix e
noPadding = forall ix e. Sz ix -> Sz ix -> Border e -> Padding ix e
Padding forall ix. Index ix => Sz ix
zeroSz forall ix. Index ix => Sz ix
zeroSz forall e. Border e
Edge
samePadding :: Index ix => Stencil ix e a -> Border e -> Padding ix e
samePadding :: forall ix e a.
Index ix =>
Stencil ix e a -> Border e -> Padding ix e
samePadding (Stencil (Sz ix
sSz) ix
sCenter (ix -> e) -> (ix -> e) -> ix -> a
_) Border e
border =
Padding
{ paddingFromOrigin :: Sz ix
paddingFromOrigin = forall ix. Index ix => ix -> Sz ix
Sz ix
sCenter
, paddingFromBottom :: Sz ix
paddingFromBottom = forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
sSz (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (forall a. Num a => a -> a -> a
+ Int
1) ix
sCenter))
, paddingWithElement :: Border e
paddingWithElement = Border e
border
}
applyStencil
:: (Index ix, Manifest r e)
=> Padding ix e
-> Stencil ix e a
-> Array r ix e
-> Array DW ix a
applyStencil :: forall ix r e a.
(Index ix, Manifest r e) =>
Padding ix e -> Stencil ix e a -> Array r ix e -> Array DW ix a
applyStencil (Padding (Sz ix
po) (Sz ix
pb) Border e
border) (Stencil Sz ix
sSz ix
sCenter (ix -> e) -> (ix -> e) -> ix -> a
stencilF) !Array r ix e
arr =
forall ix e.
Index ix =>
Array D ix e -> Window ix e -> Array DW ix e
insertWindow Array D ix a
warr Window ix a
window
where
!offset :: ix
offset = forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
sCenter ix
po
!warr :: Array D ix a
warr =
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray
(forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr)
Sz ix
sz
(forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex ((ix -> e) -> (ix -> e) -> ix -> a
stencilF (forall ix r e.
(Index ix, Manifest r e) =>
Border e -> Array r ix e -> ix -> e
borderIndex Border e
border Array r ix e
arr) (forall ix r e.
(Index ix, Manifest r e) =>
Border e -> Array r ix e -> ix -> e
borderIndex Border e
border Array r ix e
arr) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
offset))
!shrinkSz :: Sz ix
shrinkSz = forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (forall a. Num a => a -> a -> a
subtract Int
1) (forall ix. Sz ix -> ix
unSz Sz ix
sSz))
!sz :: Sz ix
sz = forall ix.
Index ix =>
(Int -> Int -> Int) -> Sz ix -> Sz ix -> Sz ix
liftSz2 (-) (forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
po (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
pb (forall ix. Sz ix -> ix
unSz (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr))))) Sz ix
shrinkSz
!wsz :: Sz ix
wsz = forall ix.
Index ix =>
(Int -> Int -> Int) -> Sz ix -> Sz ix -> Sz ix
liftSz2 (-) (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr) Sz ix
shrinkSz
!window :: Window ix a
window =
Window
{ windowStart :: ix
windowStart = ix
po
, windowSize :: Sz ix
windowSize = Sz ix
wsz
, windowIndex :: ix -> a
windowIndex = (ix -> e) -> (ix -> e) -> ix -> a
stencilF (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr) (forall ix r e.
(HasCallStack, Index ix, Manifest r e) =>
Array r ix e -> ix -> e
index' Array r ix e
arr) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
offset
, windowUnrollIx2 :: Maybe Int
windowUnrollIx2 = forall ix. Sz ix -> ix
unSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM Sz ix
sSz Dim
2
}
{-# INLINE applyStencil #-}
makeStencil
:: Index ix
=> Sz ix
-> ix
-> ((ix -> e) -> a)
-> Stencil ix e a
makeStencil :: forall ix e a.
Index ix =>
Sz ix -> ix -> ((ix -> e) -> a) -> Stencil ix e a
makeStencil !Sz ix
sSz !ix
sCenter (ix -> e) -> a
relStencil = forall ix e a.
Sz ix
-> ix -> ((ix -> e) -> (ix -> e) -> ix -> a) -> Stencil ix e a
Stencil Sz ix
sSz ix
sCenter (ix -> e) -> (ix -> e) -> ix -> a
stencil
where
stencil :: (ix -> e) -> (ix -> e) -> ix -> a
stencil ix -> e
_ ix -> e
getVal !ix
ix = forall a. a -> a
inline ((ix -> e) -> a
relStencil (ix -> e
getVal forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
ix))
{-# INLINE stencil #-}
{-# INLINE makeStencil #-}
idStencil :: Index ix => Stencil ix e e
idStencil :: forall ix e. Index ix => Stencil ix e e
idStencil = forall ix e a.
Index ix =>
Sz ix -> ix -> (ix -> (ix -> e) -> a) -> Stencil ix e a
makeUnsafeStencil forall ix. Index ix => Sz ix
oneSz forall ix. Index ix => ix
zeroIndex forall a b. (a -> b) -> a -> b
$ \ix
_ ix -> e
get -> ix -> e
get forall ix. Index ix => ix
zeroIndex
{-# INLINE idStencil #-}
foldlStencil :: Index ix => (a -> e -> a) -> a -> Sz ix -> Stencil ix e a
foldlStencil :: forall ix a e.
Index ix =>
(a -> e -> a) -> a -> Sz ix -> Stencil ix e a
foldlStencil a -> e -> a
f a
acc0 Sz ix
sz =
forall ix e a.
Index ix =>
Sz ix -> ix -> (ix -> (ix -> e) -> a) -> Stencil ix e a
makeUnsafeStencil Sz ix
sz forall ix. Index ix => ix
zeroIndex forall a b. (a -> b) -> a -> b
$ \ix
_ ix -> e
get ->
forall ix a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
iter forall ix. Index ix => ix
zeroIndex (forall ix. Sz ix -> ix
unSz Sz ix
sz) forall ix. Index ix => ix
oneIndex forall a. Ord a => a -> a -> Bool
(<) a
acc0 forall a b. (a -> b) -> a -> b
$ \ix
ix -> (a -> e -> a
`f` ix -> e
get ix
ix)
{-# INLINE foldlStencil #-}
foldrStencil :: Index ix => (e -> a -> a) -> a -> Sz ix -> Stencil ix e a
foldrStencil :: forall ix e a.
Index ix =>
(e -> a -> a) -> a -> Sz ix -> Stencil ix e a
foldrStencil e -> a -> a
f a
acc0 Sz ix
sz =
let ixStart :: ix
ixStart = forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) (forall ix. Sz ix -> ix
unSz Sz ix
sz) forall ix. Index ix => ix
oneIndex
in forall ix e a.
Index ix =>
Sz ix -> ix -> (ix -> (ix -> e) -> a) -> Stencil ix e a
makeUnsafeStencil Sz ix
sz forall ix. Index ix => ix
zeroIndex forall a b. (a -> b) -> a -> b
$ \ix
_ ix -> e
get ->
forall ix a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
iter ix
ixStart forall ix. Index ix => ix
zeroIndex (forall ix. Index ix => Int -> ix
pureIndex (-Int
1)) forall a. Ord a => a -> a -> Bool
(>=) a
acc0 forall a b. (a -> b) -> a -> b
$ \ix
ix -> e -> a -> a
f (ix -> e
get ix
ix)
{-# INLINE foldrStencil #-}
foldStencil :: (Monoid e, Index ix) => Sz ix -> Stencil ix e e
foldStencil :: forall e ix. (Monoid e, Index ix) => Sz ix -> Stencil ix e e
foldStencil = forall ix a e.
Index ix =>
(a -> e -> a) -> a -> Sz ix -> Stencil ix e a
foldlStencil forall a. Monoid a => a -> a -> a
mappend forall a. Monoid a => a
mempty
{-# INLINE foldStencil #-}
maxStencil :: (Bounded e, Ord e, Index ix) => Sz ix -> Stencil ix e e
maxStencil :: forall e ix.
(Bounded e, Ord e, Index ix) =>
Sz ix -> Stencil ix e e
maxStencil = forall c d a b ix.
(c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b
dimapStencil coerce :: forall a b. Coercible a b => a -> b
coerce forall a. Max a -> a
getMax forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e ix. (Monoid e, Index ix) => Sz ix -> Stencil ix e e
foldStencil
{-# INLINE maxStencil #-}
minStencil :: (Bounded e, Ord e, Index ix) => Sz ix -> Stencil ix e e
minStencil :: forall e ix.
(Bounded e, Ord e, Index ix) =>
Sz ix -> Stencil ix e e
minStencil = forall c d a b ix.
(c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b
dimapStencil coerce :: forall a b. Coercible a b => a -> b
coerce forall a. Min a -> a
getMin forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e ix. (Monoid e, Index ix) => Sz ix -> Stencil ix e e
foldStencil
{-# INLINE minStencil #-}
sumStencil :: (Num e, Index ix) => Sz ix -> Stencil ix e e
sumStencil :: forall e ix. (Num e, Index ix) => Sz ix -> Stencil ix e e
sumStencil = forall c d a b ix.
(c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b
dimapStencil coerce :: forall a b. Coercible a b => a -> b
coerce forall a. Sum a -> a
getSum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e ix. (Monoid e, Index ix) => Sz ix -> Stencil ix e e
foldStencil
{-# INLINE sumStencil #-}
productStencil :: (Num e, Index ix) => Sz ix -> Stencil ix e e
productStencil :: forall e ix. (Num e, Index ix) => Sz ix -> Stencil ix e e
productStencil = forall c d a b ix.
(c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b
dimapStencil coerce :: forall a b. Coercible a b => a -> b
coerce forall a. Product a -> a
getProduct forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e ix. (Monoid e, Index ix) => Sz ix -> Stencil ix e e
foldStencil
{-# INLINE productStencil #-}
avgStencil :: (Fractional e, Index ix) => Sz ix -> Stencil ix e e
avgStencil :: forall e ix. (Fractional e, Index ix) => Sz ix -> Stencil ix e e
avgStencil Sz ix
sz = forall e ix. (Num e, Index ix) => Sz ix -> Stencil ix e e
sumStencil Sz ix
sz forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
{-# INLINE avgStencil #-}