{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Massiv.Array.Ops.Construct (
empty,
singleton,
replicate,
makeArray,
makeArrayLinear,
makeArrayR,
makeArrayLinearR,
makeVectorR,
iterateN,
iiterateN,
unfoldlS_,
iunfoldlS_,
unfoldrS_,
iunfoldrS_,
makeSplitSeedArray,
uniformArray,
uniformRangeArray,
randomArray,
randomArrayS,
randomArrayWS,
makeArrayA,
makeArrayAR,
makeArrayLinearA,
(...),
(..:),
range,
rangeStepM,
rangeStep',
rangeInclusive,
rangeStepInclusiveM,
rangeStepInclusive',
rangeSize,
rangeStepSize,
enumFromN,
enumFromStepN,
expandWithin,
expandWithinM,
expandWithin',
expandOuter,
expandInner,
) where
import Control.Applicative hiding (empty)
import Control.Monad (void, when)
import Control.Monad.ST
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Delayed.Push
import Data.Massiv.Array.Mutable
import Data.Massiv.Core.Common
import System.Random.Stateful
import Prelude hiding (enumFromTo, replicate)
makeArrayR :: Load r ix e => r -> Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArrayR :: forall r ix e.
Load r ix e =>
r -> Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArrayR r
_ = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray
{-# INLINE makeArrayR #-}
makeArrayLinearR :: Load r ix e => r -> Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinearR :: forall r ix e.
Load r ix e =>
r -> Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinearR r
_ = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear
{-# INLINE makeArrayLinearR #-}
makeVectorR :: Load r Ix1 e => r -> Comp -> Sz1 -> (Ix1 -> e) -> Vector r e
makeVectorR :: forall r e.
Load r Int e =>
r -> Comp -> Sz1 -> (Int -> e) -> Vector r e
makeVectorR r
_ = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray
{-# INLINE makeVectorR #-}
newtype STA r ix a = STA {forall r ix a.
STA r ix a -> forall s. MArray s r ix a -> ST s (Array r ix a)
_runSTA :: forall s. MArray s r ix a -> ST s (Array r ix a)}
runSTA :: (Manifest r e, Index ix) => Sz ix -> STA r ix e -> Array r ix e
runSTA :: forall r e ix.
(Manifest r e, Index ix) =>
Sz ix -> STA r ix e -> Array r ix e
runSTA !Sz ix
sz (STA forall s. MArray s r ix e -> ST s (Array r ix e)
m) = forall a. (forall s. ST s a) -> a
runST (forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s. MArray s r ix e -> ST s (Array r ix e)
m)
{-# INLINE runSTA #-}
makeArrayA
:: forall r ix e f
. (Manifest r e, Index ix, Applicative f)
=> Sz ix
-> (ix -> f e)
-> f (Array r ix e)
makeArrayA :: forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA sz :: Sz ix
sz@(Sz ix
n) ix -> f e
f =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall r e ix.
(Manifest r e, Index ix) =>
Sz ix -> STA r ix e -> Array r ix e
runSTA Sz ix
sz) forall a b. (a -> b) -> a -> b
$
forall ix (f :: * -> *) a.
Index ix =>
ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f a
-> (ix -> f a -> f a)
-> f a
iterF forall ix. Index ix => ix
zeroIndex ix
n forall ix. Index ix => ix
oneIndex forall a. Ord a => a -> a -> Bool
(<) (forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall r ix a.
(forall s. MArray s r ix a -> ST s (Array r ix a)) -> STA r ix a
STA (forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq))) forall a b. (a -> b) -> a -> b
$ \ix
ix f (STA r ix e)
g ->
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\e
e (STA forall s. MArray s r ix e -> ST s (Array r ix e)
st) -> forall r ix a.
(forall s. MArray s r ix a -> ST s (Array r ix a)) -> STA r ix a
STA (\MArray s r ix e
ma -> forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> e -> m ()
unsafeWrite MArray s r ix e
ma ix
ix e
e forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. MArray s r ix e -> ST s (Array r ix e)
st MArray s r ix e
ma)) (ix -> f e
f ix
ix) f (STA r ix e)
g
{-# INLINE makeArrayA #-}
makeArrayLinearA
:: forall r ix e f
. (Manifest r e, Index ix, Applicative f)
=> Sz ix
-> (Int -> f e)
-> f (Array r ix e)
makeArrayLinearA :: forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (Int -> f e) -> f (Array r ix e)
makeArrayLinearA !Sz ix
sz Int -> f e
f =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall r e ix.
(Manifest r e, Index ix) =>
Sz ix -> STA r ix e -> Array r ix e
runSTA Sz ix
sz) forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
0 (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (forall a. Num a => a -> a -> a
+ Int
1) (forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall r ix a.
(forall s. MArray s r ix a -> ST s (Array r ix a)) -> STA r ix a
STA (forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq))) forall a b. (a -> b) -> a -> b
$ \Int
i ->
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\e
e (STA forall s. MArray s r ix e -> ST s (Array r ix e)
st) -> forall r ix a.
(forall s. MArray s r ix a -> ST s (Array r ix a)) -> STA r ix a
STA (\MArray s r ix e
ma -> forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray s r ix e
ma Int
i e
e forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. MArray s r ix e -> ST s (Array r ix e)
st MArray s r ix e
ma)) (Int -> f e
f Int
i)
{-# INLINE makeArrayLinearA #-}
makeArrayAR
:: forall r ix e f
. (Manifest r e, Index ix, Applicative f)
=> r
-> Sz ix
-> (ix -> f e)
-> f (Array r ix e)
makeArrayAR :: forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
r -> Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayAR r
_ = forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA
{-# INLINE makeArrayAR #-}
iterateN :: forall ix e. Index ix => Sz ix -> (e -> e) -> e -> Array DL ix e
iterateN :: forall ix e. Index ix => Sz ix -> (e -> e) -> e -> Array DL ix e
iterateN Sz ix
sz e -> e
f = forall ix e a.
Index ix =>
Sz ix -> (a -> (e, a)) -> a -> Array DL ix e
unfoldrS_ Sz ix
sz forall a b. (a -> b) -> a -> b
$ \e
a -> let !a' :: e
a' = e -> e
f e
a in (e
a', e
a')
{-# INLINE iterateN #-}
iiterateN :: forall ix e. Index ix => Sz ix -> (e -> ix -> e) -> e -> Array DL ix e
iiterateN :: forall ix e.
Index ix =>
Sz ix -> (e -> ix -> e) -> e -> Array DL ix e
iiterateN Sz ix
sz e -> ix -> e
f = forall ix e a.
Index ix =>
Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e
iunfoldrS_ Sz ix
sz forall a b. (a -> b) -> a -> b
$ \e
a ix
ix -> let !a' :: e
a' = e -> ix -> e
f e
a ix
ix in (e
a', e
a')
{-# INLINE iiterateN #-}
unfoldrS_
:: forall ix e a
. Index ix
=> Sz ix
-> (a -> (e, a))
-> a
-> Array DL ix e
unfoldrS_ :: forall ix e a.
Index ix =>
Sz ix -> (a -> (e, a)) -> a -> Array DL ix e
unfoldrS_ Sz ix
sz a -> (e, a)
f = forall ix e a.
Index ix =>
Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e
iunfoldrS_ Sz ix
sz (\a
a ix
_ -> a -> (e, a)
f a
a)
{-# INLINE unfoldrS_ #-}
iunfoldrS_
:: forall ix e a
. Index ix
=> Sz ix
-> (a -> ix -> (e, a))
-> a
-> Array DL ix e
iunfoldrS_ :: forall ix e a.
Index ix =>
Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e
iunfoldrS_ Sz ix
sz a -> ix -> (e, a)
f a
acc0 = DLArray{dlComp :: Comp
dlComp = Comp
Seq, dlSize :: Sz ix
dlSize = Sz ix
sz, dlLoad :: Loader e
dlLoad = Loader e
load}
where
load :: Loader e
load :: Loader e
load Scheduler s ()
_ Int
startAt Int -> e -> ST s ()
dlWrite Int -> Sz1 -> e -> ST s ()
_ =
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM RowMajor
defRowMajor Int
startAt Sz ix
sz forall ix. Index ix => ix
zeroIndex forall ix. Index ix => Stride ix
oneStride a
acc0 forall a b. (a -> b) -> a -> b
$ \ !Int
i !ix
ix !a
acc ->
case a -> ix -> (e, a)
f a
acc ix
ix of
(e
e, !a
acc') -> a
acc' forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
dlWrite Int
i e
e
{-# INLINE load #-}
{-# INLINE iunfoldrS_ #-}
unfoldlS_ :: Index ix => Sz ix -> (a -> (a, e)) -> a -> Array DL ix e
unfoldlS_ :: forall ix a e.
Index ix =>
Sz ix -> (a -> (a, e)) -> a -> Array DL ix e
unfoldlS_ Sz ix
sz a -> (a, e)
f = forall ix e a.
Index ix =>
Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e
iunfoldlS_ Sz ix
sz (forall a b. a -> b -> a
const a -> (a, e)
f)
{-# INLINE unfoldlS_ #-}
iunfoldlS_
:: forall ix e a
. Index ix
=> Sz ix
-> (ix -> a -> (a, e))
-> a
-> Array DL ix e
iunfoldlS_ :: forall ix e a.
Index ix =>
Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e
iunfoldlS_ Sz ix
sz ix -> a -> (a, e)
f a
acc0 = DLArray{dlComp :: Comp
dlComp = Comp
Seq, dlSize :: Sz ix
dlSize = Sz ix
sz, dlLoad :: Loader e
dlLoad = Loader e
load}
where
load :: Loader e
load :: Loader e
load Scheduler s ()
_ Int
startAt Int -> e -> ST s ()
dlWrite Int -> Sz1 -> e -> ST s ()
_ =
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM Int
startAt (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz forall a. Num a => a -> a -> a
+ Int
startAt) (forall a. Num a => a -> a -> a
+ Int
1) a
acc0 forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
acc ->
let (a
acc', e
e) = ix -> a -> (a, e)
f (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz (Int
i forall a. Num a => a -> a -> a
- Int
startAt)) a
acc
in a
acc' forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
dlWrite Int
i e
e
{-# INLINE load #-}
{-# INLINE iunfoldlS_ #-}
randomArray
:: forall ix e g
. Index ix
=> g
-> (g -> (g, g))
-> (g -> (e, g))
-> Comp
-> Sz ix
-> Array DL ix e
randomArray :: forall ix e g.
Index ix =>
g
-> (g -> (g, g)) -> (g -> (e, g)) -> Comp -> Sz ix -> Array DL ix e
randomArray g
gen g -> (g, g)
splitGen g -> (e, g)
nextRandom Comp
comp Sz ix
sz = forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall s.
Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ())
-> Array DL ix e
unsafeMakeLoadArray Comp
comp Sz ix
sz forall a. Maybe a
Nothing forall s. Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ()
load
where
!totalLength :: Int
totalLength = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
load :: forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()
load :: forall s. Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ()
load Scheduler s ()
scheduler Int
startAt Int -> e -> ST s ()
writeAt =
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
let slackStartAt :: Int
slackStartAt = Int
slackStart forall a. Num a => a -> a -> a
+ Int
startAt
writeRandom :: Int -> g -> ST s g
writeRandom Int
k g
genII =
let (e
e, g
genII') = g -> (e, g)
nextRandom g
genII
in g
genII' forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
writeAt Int
k e
e
g
genForSlack <-
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
startAt (forall a. Ord a => a -> a -> Bool
< Int
slackStartAt) (forall a. Num a => a -> a -> a
+ Int
chunkLength) g
gen forall a b. (a -> b) -> a -> b
$ \Int
start g
genI -> do
let (g
genI0, g
genI1) =
if forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler forall a. Eq a => a -> a -> Bool
== Int
1
then (g
genI, g
genI)
else g -> (g, g)
splitGen g
genI
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
start (forall a. Ord a => a -> a -> Bool
< Int
start forall a. Num a => a -> a -> a
+ Int
chunkLength) (forall a. Num a => a -> a -> a
+ Int
1) g
genI0 Int -> g -> ST s g
writeRandom
forall (f :: * -> *) a. Applicative f => a -> f a
pure g
genI1
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
slackStartAt (forall a. Ord a => a -> a -> Bool
< Int
totalLength forall a. Num a => a -> a -> a
+ Int
startAt) (forall a. Num a => a -> a -> a
+ Int
1) g
genForSlack Int -> g -> ST s g
writeRandom
{-# INLINE randomArray #-}
makeSplitSeedArray
:: forall ix e g it
. (Iterator it, Index ix)
=> it
-> g
-> (g -> (g, g))
-> Comp
-> Sz ix
-> (Ix1 -> ix -> g -> (e, g))
-> Array DL ix e
makeSplitSeedArray :: forall ix e g it.
(Iterator it, Index ix) =>
it
-> g
-> (g -> (g, g))
-> Comp
-> Sz ix
-> (Int -> ix -> g -> (e, g))
-> Array DL ix e
makeSplitSeedArray it
it g
seed g -> (g, g)
splitSeed Comp
comp Sz ix
sz Int -> ix -> g -> (e, g)
genFunc =
DLArray{dlComp :: Comp
dlComp = Comp
comp, dlSize :: Sz ix
dlSize = Sz ix
sz, dlLoad :: Loader e
dlLoad = Loader e
load}
where
load :: Loader e
load :: Loader e
load Scheduler s ()
scheduler Int
startAt Int -> e -> ST s ()
writeAt Int -> Sz1 -> e -> ST s ()
_ =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullAccST_ it
it Scheduler s ()
scheduler Int
startAt Sz ix
sz g
seed (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> (g, g)
splitSeed) forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix g
g ->
case Int -> ix -> g -> (e, g)
genFunc (Int
i forall a. Num a => a -> a -> a
- Int
startAt) ix
ix g
g of
(e
x, g
g') -> g
g' forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
writeAt Int
i e
x
{-# INLINE load #-}
{-# INLINE makeSplitSeedArray #-}
uniformArray
:: forall ix e g
. (Index ix, RandomGen g, Uniform e)
=> g
-> Comp
-> Sz ix
-> Array DL ix e
uniformArray :: forall ix e g.
(Index ix, RandomGen g, Uniform e) =>
g -> Comp -> Sz ix -> Array DL ix e
uniformArray g
gen = forall ix e g.
Index ix =>
g
-> (g -> (g, g)) -> (g -> (e, g)) -> Comp -> Sz ix -> Array DL ix e
randomArray g
gen forall g. RandomGen g => g -> (g, g)
split forall g a. (RandomGen g, Uniform a) => g -> (a, g)
uniform
{-# INLINE uniformArray #-}
uniformRangeArray
:: forall ix e g
. (Index ix, RandomGen g, UniformRange e)
=> g
-> (e, e)
-> Comp
-> Sz ix
-> Array DL ix e
uniformRangeArray :: forall ix e g.
(Index ix, RandomGen g, UniformRange e) =>
g -> (e, e) -> Comp -> Sz ix -> Array DL ix e
uniformRangeArray g
gen (e, e)
r = forall ix e g.
Index ix =>
g
-> (g -> (g, g)) -> (g -> (e, g)) -> Comp -> Sz ix -> Array DL ix e
randomArray g
gen forall g. RandomGen g => g -> (g, g)
split (forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (e, e)
r)
{-# INLINE uniformRangeArray #-}
randomArrayS
:: forall r ix e g
. (Manifest r e, Index ix)
=> g
-> Sz ix
-> (g -> (e, g))
-> (g, Array r ix e)
randomArrayS :: forall r ix e g.
(Manifest r e, Index ix) =>
g -> Sz ix -> (g -> (e, g)) -> (g, Array r ix e)
randomArrayS g
gen Sz ix
sz g -> (e, g)
nextRandom =
forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (a -> m (e, a)) -> a -> m (a, Array r ix e)
unfoldrPrimM Sz ix
sz (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> (e, g)
nextRandom) g
gen
{-# INLINE randomArrayS #-}
randomArrayWS
:: forall r ix e g m
. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m)
=> WorkerStates g
-> Sz ix
-> (g -> m e)
-> m (Array r ix e)
randomArrayWS :: forall r ix e g (m :: * -> *).
(Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) =>
WorkerStates g -> Sz ix -> (g -> m e) -> m (Array r ix e)
randomArrayWS WorkerStates g
states Sz ix
sz g -> m e
genRandom = forall r ix e s (m :: * -> *).
(Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) =>
WorkerStates s -> Sz ix -> (Int -> s -> m e) -> m (Array r ix e)
generateArrayLinearWS WorkerStates g
states Sz ix
sz (forall a b. a -> b -> a
const g -> m e
genRandom)
{-# INLINE randomArrayWS #-}
infix 4 ..., ..:
(...) :: Index ix => ix -> ix -> Array D ix ix
... :: forall ix. Index ix => ix -> ix -> Array D ix ix
(...) = forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
rangeInclusive Comp
Seq
{-# INLINE (...) #-}
(..:) :: Index ix => ix -> ix -> Array D ix ix
..: :: forall ix. Index ix => ix -> ix -> Array D ix ix
(..:) = forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range Comp
Seq
{-# INLINE (..:) #-}
range :: Index ix => Comp -> ix -> ix -> Array D ix ix
range :: forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range Comp
comp !ix
from !ix
to = forall ix. Index ix => Comp -> ix -> Sz ix -> Array D ix ix
rangeSize Comp
comp ix
from (forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
to ix
from))
{-# INLINE range #-}
rangeStepM
:: forall ix m
. (Index ix, MonadThrow m)
=> Comp
-> ix
-> ix
-> ix
-> m (Array D ix ix)
rangeStepM :: forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepM Comp
comp !ix
from !ix
step !ix
to
| forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
foldlIndex (\Bool
acc Int
i -> Bool
acc Bool -> Bool -> Bool
|| Int
i forall a. Eq a => a -> a -> Bool
== Int
0) Bool
False ix
step = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => ix -> IndexException
IndexZeroException ix
step
| Bool
otherwise =
let dist :: ix
dist = forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
to ix
from
sz :: ix
sz = forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Integral a => a -> a -> a
div ix
dist ix
step
r :: ix
r = forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex forall a. Num a => a -> a
signum forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Integral a => a -> a -> a
mod ix
dist ix
step
in forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Comp -> ix -> ix -> Sz ix -> Array D ix ix
rangeStepSize Comp
comp ix
from ix
step (forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
sz ix
r))
{-# INLINE rangeStepM #-}
rangeStep' :: (HasCallStack, Index ix) => Comp -> ix -> ix -> ix -> Array D ix ix
rangeStep' :: forall ix.
(HasCallStack, Index ix) =>
Comp -> ix -> ix -> ix -> Array D ix ix
rangeStep' Comp
comp ix
from ix
step = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepM Comp
comp ix
from ix
step
{-# INLINE rangeStep' #-}
rangeInclusive :: Index ix => Comp -> ix -> ix -> Array D ix ix
rangeInclusive :: forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
rangeInclusive Comp
comp ix
ixFrom ix
ixTo =
forall ix. Index ix => Comp -> ix -> Sz ix -> Array D ix ix
rangeSize Comp
comp ix
ixFrom (forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (forall a. Num a => a -> a -> a
+ Int
1) ix
ixTo) ix
ixFrom))
{-# INLINE rangeInclusive #-}
rangeStepInclusiveM :: (MonadThrow m, Index ix) => Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepInclusiveM :: forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepInclusiveM Comp
comp ix
ixFrom ix
step ix
ixTo = forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepM Comp
comp ix
ixFrom ix
step (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int
1 forall a. Num a => a -> a -> a
+) ix
ixTo)
{-# INLINE rangeStepInclusiveM #-}
rangeStepInclusive' :: (HasCallStack, Index ix) => Comp -> ix -> ix -> ix -> Array D ix ix
rangeStepInclusive' :: forall ix.
(HasCallStack, Index ix) =>
Comp -> ix -> ix -> ix -> Array D ix ix
rangeStepInclusive' Comp
comp ix
ixFrom ix
step = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepInclusiveM Comp
comp ix
ixFrom ix
step
{-# INLINE rangeStepInclusive' #-}
rangeSize
:: Index ix
=> Comp
-> ix
-> Sz ix
-> Array D ix ix
rangeSize :: forall ix. Index ix => Comp -> ix -> Sz ix -> Array D ix ix
rangeSize Comp
comp !ix
from !Sz ix
sz = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
from)
{-# INLINE rangeSize #-}
rangeStepSize
:: Index ix
=> Comp
-> ix
-> ix
-> Sz ix
-> Array D ix ix
rangeStepSize :: forall ix. Index ix => Comp -> ix -> ix -> Sz ix -> Array D ix ix
rangeStepSize Comp
comp !ix
from !ix
step !Sz ix
sz =
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
from 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
step)
{-# INLINE rangeStepSize #-}
enumFromN
:: Num e
=> Comp
-> e
-> Sz1
-> Vector D e
enumFromN :: forall e. Num e => Comp -> e -> Sz1 -> Vector D e
enumFromN Comp
comp !e
from !Sz1
sz = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz1
sz forall a b. (a -> b) -> a -> b
$ \Int
i -> e
from forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
{-# INLINE enumFromN #-}
enumFromStepN
:: Num e
=> Comp
-> e
-> e
-> Sz1
-> Vector D e
enumFromStepN :: forall e. Num e => Comp -> e -> e -> Sz1 -> Vector D e
enumFromStepN Comp
comp !e
from !e
step !Sz1
sz = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz1
sz forall a b. (a -> b) -> a -> b
$ \Int
i -> e
from forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i forall a. Num a => a -> a -> a
* e
step
{-# INLINE enumFromStepN #-}
expandWithin
:: forall n ix e r a
. (IsIndexDimension ix n, Index (Lower ix), Manifest r a)
=> Dimension n
-> Sz1
-> (a -> Ix1 -> e)
-> Array r (Lower ix) a
-> Array D ix e
expandWithin :: forall (n :: Natural) ix e r a.
(IsIndexDimension ix n, Index (Lower ix), Manifest r a) =>
Dimension n
-> Sz1 -> (a -> Int -> e) -> Array r (Lower ix) a -> Array D ix e
expandWithin Dimension n
dim (Sz Int
k) a -> Int -> e
f Array r (Lower ix) a
arr =
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz forall a b. (a -> b) -> a -> b
$ \ix
ix ->
let (Int
i, Lower ix
ixl) = forall ix (n :: Natural).
IsIndexDimension ix n =>
ix -> Dimension n -> (Int, Lower ix)
pullOutDimension ix
ix Dimension n
dim
in a -> Int -> e
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
where
szl :: Lower ix
szl = forall ix. Sz ix -> ix
unSz (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr)
sz :: Sz ix
sz = forall ix. ix -> Sz ix
SafeSz (forall ix (n :: Natural).
IsIndexDimension ix n =>
Lower ix -> Dimension n -> Int -> ix
insertDimension Lower ix
szl Dimension n
dim Int
k)
{-# INLINE expandWithin #-}
expandWithin'
:: forall r ix a b
. (HasCallStack, Index ix, Index (Lower ix), Manifest r a)
=> Dim
-> Sz1
-> (a -> Ix1 -> b)
-> Array r (Lower ix) a
-> Array D ix b
expandWithin' :: forall r ix a b.
(HasCallStack, Index ix, Index (Lower ix), Manifest r a) =>
Dim
-> Sz1 -> (a -> Int -> b) -> Array r (Lower ix) a -> Array D ix b
expandWithin' Dim
dim Sz1
k a -> Int -> b
f = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix a b (m :: * -> *).
(Index ix, Index (Lower ix), Manifest r a, MonadThrow m) =>
Dim
-> Sz1
-> (a -> Int -> b)
-> Array r (Lower ix) a
-> m (Array D ix b)
expandWithinM Dim
dim Sz1
k a -> Int -> b
f
{-# INLINE expandWithin' #-}
expandWithinM
:: forall r ix a b m
. (Index ix, Index (Lower ix), Manifest r a, MonadThrow m)
=> Dim
-> Sz1
-> (a -> Ix1 -> b)
-> Array r (Lower ix) a
-> m (Array D ix b)
expandWithinM :: forall r ix a b (m :: * -> *).
(Index ix, Index (Lower ix), Manifest r a, MonadThrow m) =>
Dim
-> Sz1
-> (a -> Int -> b)
-> Array r (Lower ix) a
-> m (Array D ix b)
expandWithinM Dim
dim Sz1
k a -> Int -> b
f Array r (Lower ix) a
arr = do
Sz ix
sz <- forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz1 -> m (Sz ix)
insertSzM (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr) Dim
dim Sz1
k
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz forall a b. (a -> b) -> a -> b
$ \ix
ix ->
let (Int
i, Lower ix
ixl) = forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
ix Dim
dim
in a -> Int -> b
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
{-# INLINE expandWithinM #-}
expandOuter
:: forall r ix a b
. (Index ix, Index (Lower ix), Manifest r a)
=> Sz1
-> (a -> Ix1 -> b)
-> Array r (Lower ix) a
-> Array D ix b
expandOuter :: forall r ix a b.
(Index ix, Index (Lower ix), Manifest r a) =>
Sz1 -> (a -> Int -> b) -> Array r (Lower ix) a -> Array D ix b
expandOuter Sz1
k a -> Int -> b
f Array r (Lower ix) a
arr =
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz forall a b. (a -> b) -> a -> b
$ \ix
ix ->
let (Int
i, Lower ix
ixl) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
in a -> Int -> b
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
where
szl :: Sz (Lower ix)
szl = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr
sz :: Sz ix
sz = forall ix. Index ix => Sz1 -> Sz (Lower ix) -> Sz ix
consSz Sz1
k Sz (Lower ix)
szl
{-# INLINE expandOuter #-}
expandInner
:: forall r ix a b
. (Index ix, Index (Lower ix), Manifest r a)
=> Sz1
-> (a -> Ix1 -> b)
-> Array r (Lower ix) a
-> Array D ix b
expandInner :: forall r ix a b.
(Index ix, Index (Lower ix), Manifest r a) =>
Sz1 -> (a -> Int -> b) -> Array r (Lower ix) a -> Array D ix b
expandInner Sz1
k a -> Int -> b
f Array r (Lower ix) a
arr =
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz forall a b. (a -> b) -> a -> b
$ \ix
ix ->
let (Lower ix
ixl, Int
i) = forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
ix
in a -> Int -> b
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
where
szl :: Sz (Lower ix)
szl = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr
sz :: Sz ix
sz = forall ix. Index ix => Sz (Lower ix) -> Sz1 -> Sz ix
snocSz Sz (Lower ix)
szl Sz1
k
{-# INLINE expandInner #-}