{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module      : Data.Massiv.Array.Ops.Construct
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Ops.Construct
  ( -- ** With constant value
    empty
  , singleton
  , replicate
    -- ** With a function
  , makeArray
  , makeArrayLinear
  , makeArrayR
  , makeArrayLinearR
  , makeVectorR
    -- *** Iterating
  , iterateN
  , iiterateN
    -- *** Unfolding
  , unfoldlS_
  -- , unfoldlS
  , iunfoldlS_
  --, iunfoldlS
  , unfoldrS_
  --, unfoldrS
  , iunfoldrS_
  --, iunfoldrS
    -- *** Random
  , randomArray
  , randomArrayS
  , randomArrayWS
    -- *** Applicative
  , makeArrayA
  , makeArrayAR
  , makeArrayLinearA
    -- ** Enumeration
  , (...)
  , (..:)
  , range
  , rangeStepM
  , rangeStep'
  , rangeInclusive
  , rangeStepInclusiveM
  , rangeStepInclusive'
  , rangeSize
  , rangeStepSize
  , enumFromN
  , enumFromStepN
    -- ** Expansion
  , expandWithin
  , expandWithinM
  , expandWithin'
  , expandOuter
  , expandInner
  ) where

import Control.Applicative hiding (empty)
import Control.Monad (when, void)
import Control.Monad.ST
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Delayed.Push
--import Data.Massiv.Array.Delayed.Stream (unfoldr, unfoldrN)
import Data.Massiv.Array.Mutable
import Data.Massiv.Core.Common
import Prelude hiding (enumFromTo, replicate)

-- | Just like `makeArray` but with ability to specify the result representation as an
-- argument. Note the `Data.Massiv.Array.U`nboxed type constructor in the below example.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> makeArrayR U Par (Sz (2 :> 3 :. 4)) (\ (i :> j :. k) -> i * i + j * j == k * k)
-- Array U Par (Sz (2 :> 3 :. 4))
--   [ [ [ True, False, False, False ]
--     , [ False, True, False, False ]
--     , [ False, False, True, False ]
--     ]
--   , [ [ False, True, False, False ]
--     , [ False, False, False, False ]
--     , [ False, False, False, False ]
--     ]
--   ]
--
-- @since 0.1.0
makeArrayR :: Construct r ix e => r -> Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArrayR :: r -> Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArrayR r
_ = Comp -> Sz ix -> (ix -> e) -> Array r ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray
{-# INLINE makeArrayR #-}

-- | Same as `makeArrayLinear`, but with ability to supply resulting representation
--
-- @since 0.3.0
makeArrayLinearR :: Construct r ix e => r -> Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinearR :: r -> Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinearR r
_ = Comp -> Sz ix -> (Int -> e) -> Array r ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear
{-# INLINE makeArrayLinearR #-}

-- | Same as `makeArrayR`, but restricted to 1-dimensional arrays.
--
-- @since 0.1.0
makeVectorR :: Construct r Ix1 e => r -> Comp -> Sz1 -> (Ix1 -> e) -> Array r Ix1 e
makeVectorR :: r -> Comp -> Sz1 -> (Int -> e) -> Array r Int e
makeVectorR r
_ = Comp -> Sz1 -> (Int -> e) -> Array r Int e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray
{-# INLINE makeVectorR #-}


newtype STA r ix a = STA {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 :: Mutable r ix e => Sz ix -> STA r ix e -> Array r ix e
runSTA :: 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 s. ST s (Array r ix e)) -> Array r ix e
forall a. (forall s. ST s a) -> a
runST (Sz ix -> ST s (MArray (PrimState (ST s)) r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz ST s (MArray s r ix e)
-> (MArray s r ix e -> ST s (Array r ix e)) -> ST s (Array r ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray s r ix e -> ST s (Array r ix e)
forall s. MArray s r ix e -> ST s (Array r ix e)
m)
{-# INLINE runSTA  #-}

-- | Similar to `makeArray`, but construct the array sequentially using an `Applicative` interface.
--
-- /Note/ - using `Data.Massiv.Array.Mutable.generateArray` or
-- `Data.Massiv.Array.Mutable.generateArrayS` will always be faster, althought not always possible.
--
--
-- @since 0.2.6
makeArrayA ::
     forall r ix e f. (Mutable r ix e, Applicative f)
  => Sz ix
  -> (ix -> f e)
  -> f (Array r ix e)
makeArrayA :: Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA !Sz ix
sz ix -> f e
f =
  let n :: Int
n = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
      go :: Int -> f (STA r ix e)
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n =
          (e -> STA r ix e -> STA r ix e)
-> f e -> f (STA r ix e) -> f (STA r ix e)
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 s. MArray s r ix e -> ST s (Array r ix e)) -> STA r ix e
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 -> MArray (PrimState (ST s)) r ix e -> Int -> e -> ST s ()
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray s r ix e
MArray (PrimState (ST s)) r ix e
ma Int
i e
e ST s () -> ST s (Array r ix e) -> ST s (Array r ix e)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MArray s r ix e -> ST s (Array r ix e)
forall s. MArray s r ix e -> ST s (Array r ix e)
st MArray s r ix e
ma))
            (ix -> f e
f (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i))
            (Int -> f (STA r ix e)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        | Bool
otherwise = STA r ix e -> f (STA r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((forall s. MArray s r ix e -> ST s (Array r ix e)) -> STA r ix e
forall r ix a.
(forall s. MArray s r ix a -> ST s (Array r ix a)) -> STA r ix a
STA (Comp -> MArray (PrimState (ST s)) r ix e -> ST s (Array r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq))
   in Sz ix -> STA r ix e -> Array r ix e
forall r ix e.
Mutable r ix e =>
Sz ix -> STA r ix e -> Array r ix e
runSTA Sz ix
sz (STA r ix e -> Array r ix e) -> f (STA r ix e) -> f (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA r ix e)
go Int
0
{-# INLINE makeArrayA  #-}

-- | Same as `makeArrayA`, but with linear index.
--
-- @since 0.4.5
makeArrayLinearA ::
     forall r ix e f. (Mutable r ix e, Applicative f)
  => Sz ix
  -> (Int -> f e)
  -> f (Array r ix e)
makeArrayLinearA :: Sz ix -> (Int -> f e) -> f (Array r ix e)
makeArrayLinearA !Sz ix
sz Int -> f e
f =
  let n :: Int
n = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
      go :: Int -> f (STA r ix e)
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n =
          (e -> STA r ix e -> STA r ix e)
-> f e -> f (STA r ix e) -> f (STA r ix e)
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 s. MArray s r ix e -> ST s (Array r ix e)) -> STA r ix e
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 -> MArray (PrimState (ST s)) r ix e -> Int -> e -> ST s ()
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray s r ix e
MArray (PrimState (ST s)) r ix e
ma Int
i e
e ST s () -> ST s (Array r ix e) -> ST s (Array r ix e)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MArray s r ix e -> ST s (Array r ix e)
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) (Int -> f (STA r ix e)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        | Bool
otherwise = STA r ix e -> f (STA r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((forall s. MArray s r ix e -> ST s (Array r ix e)) -> STA r ix e
forall r ix a.
(forall s. MArray s r ix a -> ST s (Array r ix a)) -> STA r ix a
STA (Comp -> MArray (PrimState (ST s)) r ix e -> ST s (Array r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq))
   in Sz ix -> STA r ix e -> Array r ix e
forall r ix e.
Mutable r ix e =>
Sz ix -> STA r ix e -> Array r ix e
runSTA Sz ix
sz (STA r ix e -> Array r ix e) -> f (STA r ix e) -> f (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA r ix e)
go Int
0
{-# INLINE makeArrayLinearA  #-}


-- | Same as `makeArrayA`, but with ability to supply result array representation.
--
-- @since 0.2.6
makeArrayAR ::
     forall r ix e f. (Mutable r ix e, Applicative f)
  => r
  -> Sz ix
  -> (ix -> f e)
  -> f (Array r ix e)
makeArrayAR :: r -> Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayAR r
_ = Sz ix -> (ix -> f e) -> f (Array r ix e)
forall r ix e (f :: * -> *).
(Mutable r ix e, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
makeArrayA
{-# INLINE makeArrayAR #-}


-- | Sequentially iterate over each cell in the array in the row-major order while continuously
-- aplying the accumulator at each step.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> iterateN (Sz2 2 10) succ (10 :: Int)
-- Array DL Seq (Sz (2 :. 10))
--   [ [ 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ]
--   , [ 21, 22, 23, 24, 25, 26, 27, 28, 29, 30 ]
--   ]
--
-- @since 0.3.0
iterateN :: forall ix e . Index ix => Sz ix -> (e -> e) -> e -> Array DL ix e
iterateN :: Sz ix -> (e -> e) -> e -> Array DL ix e
iterateN Sz ix
sz e -> e
f = Sz ix -> (e -> (e, e)) -> e -> Array DL ix e
forall ix e a.
Construct DL ix e =>
Sz ix -> (a -> (e, a)) -> a -> Array DL ix e
unfoldrS_ Sz ix
sz ((e -> (e, e)) -> e -> Array DL ix e)
-> (e -> (e, e)) -> e -> Array DL ix e
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 #-}

-- | Same as `iterateN`, but with index aware function.
--
-- @since 0.3.0
iiterateN :: forall ix e . Index ix => Sz ix -> (e -> ix -> e) -> e -> Array DL ix e
iiterateN :: Sz ix -> (e -> ix -> e) -> e -> Array DL ix e
iiterateN Sz ix
sz e -> ix -> e
f = Sz ix -> (e -> ix -> (e, e)) -> e -> Array DL ix e
forall ix e a.
Construct DL ix e =>
Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e
iunfoldrS_ Sz ix
sz ((e -> ix -> (e, e)) -> e -> Array DL ix e)
-> (e -> ix -> (e, e)) -> e -> Array DL ix e
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 #-}

-- | Right unfold into a delayed load array. For the opposite direction use `unfoldlS_`.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> unfoldrS_ (Sz1 10) (\xs -> (Prelude.head xs, Prelude.tail xs)) ([10 ..] :: [Int])
-- Array DL Seq (Sz1 10)
--   [ 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 ]
--
-- @since 0.3.0
unfoldrS_ :: forall ix e a . Construct DL ix e => Sz ix -> (a -> (e, a)) -> a -> Array DL ix e
unfoldrS_ :: Sz ix -> (a -> (e, a)) -> a -> Array DL ix e
unfoldrS_ Sz ix
sz a -> (e, a)
f = Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e
forall ix e a.
Construct DL ix e =>
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_ #-}

-- | Right unfold of a delayed load array with index aware function
--
-- @since 0.3.0
iunfoldrS_ ::
     forall ix e a. Construct DL ix e
  => Sz ix
  -> (a -> ix -> (e, a))
  -> a
  -> Array DL ix e
iunfoldrS_ :: Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e
iunfoldrS_ Sz ix
sz a -> ix -> (e, a)
f a
acc0 = DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m ()
    -> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = Comp
Seq, dlSize :: Sz ix
dlSize = Sz ix
sz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ()
load}
  where
    load :: Monad m =>
      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()
    load :: Scheduler m ()
-> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ()
load Scheduler m ()
_ Int
startAt Int -> e -> m ()
dlWrite Int -> Sz1 -> e -> m ()
_ =
      m a -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m a -> m ()) -> m a -> m ()
forall a b. (a -> b) -> a -> b
$
      Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
startAt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
acc0 ((Int -> a -> m a) -> m a) -> (Int -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
acc ->
        let (e
e, a
acc') = a -> ix -> (e, a)
f a
acc (ix -> (e, a)) -> ix -> (e, a)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
startAt)
         in a
acc' a -> m () -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> m ()
dlWrite Int
i e
e
    {-# INLINE load #-}
{-# INLINE iunfoldrS_ #-}


-- | Unfold sequentially from the end. There is no way to save the accumulator after
-- unfolding is done, since resulting array is delayed, but it's possible to use
-- `Data.Massiv.Array.Mutable.unfoldlPrimM` to achive such effect.
--
-- @since 0.3.0
unfoldlS_ :: Construct DL ix e => Sz ix -> (a -> (a, e)) -> a -> Array DL ix e
unfoldlS_ :: Sz ix -> (a -> (a, e)) -> a -> Array DL ix e
unfoldlS_ Sz ix
sz a -> (a, e)
f = Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e
forall ix e a.
Construct DL ix e =>
Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e
iunfoldlS_ Sz ix
sz ((a -> (a, e)) -> ix -> a -> (a, e)
forall a b. a -> b -> a
const a -> (a, e)
f)
{-# INLINE unfoldlS_ #-}

-- | Unfold sequentially from the right with an index aware function.
--
-- @since 0.3.0
iunfoldlS_ ::
     forall ix e a. Construct DL ix e
  => Sz ix
  -> (ix -> a -> (a, e))
  -> a
  -> Array DL ix e
iunfoldlS_ :: Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e
iunfoldlS_ Sz ix
sz ix -> a -> (a, e)
f a
acc0 = DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m ()
    -> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = Comp
Seq, dlSize :: Sz ix
dlSize = Sz ix
sz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ()
load}
  where
    load :: Monad m =>
      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()
    load :: Scheduler m ()
-> Int -> (Int -> e -> m ()) -> (Int -> Sz1 -> e -> m ()) -> m ()
load Scheduler m ()
_ Int
startAt Int -> e -> m ()
dlWrite Int -> Sz1 -> e -> m ()
_ =
      m a -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m a -> m ()) -> m a -> m ()
forall a b. (a -> b) -> a -> b
$
      Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM Int
startAt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
acc0 ((Int -> a -> m a) -> m a) -> (Int -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
acc ->
        let (a
acc', e
e) = ix -> a -> (a, e)
f (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
startAt)) a
acc
         in a
acc' a -> m () -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> m ()
dlWrite Int
i e
e
    {-# INLINE load #-}
{-# INLINE iunfoldlS_ #-}


-- | Create an array with random values by using a pure splittable random number generator
-- such as one provided by either [splitmix](https://www.stackage.org/package/splitmix) or
-- [random](https://www.stackage.org/package/random) packages. If you don't have a
-- splittable generator consider using `randomArrayS` or `randomArrayWS` instead.
--
-- Because of the pure nature of the generator and its splitability we are not only able
-- to parallelize the random value generation, but also guarantee that it will be
-- deterministic, granted none of the arguments have changed.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> import System.Random.SplitMix as SplitMix
-- >>> gen = SplitMix.mkSMGen 217
-- >>> randomArray gen SplitMix.splitSMGen SplitMix.nextDouble (ParN 2) (Sz2 2 3) :: Array DL Ix2 Double
-- Array DL (ParN 2) (Sz (2 :. 3))
--   [ [ 0.7383156058619669, 0.39904053166835896, 0.5617584038393628 ]
--   , [ 0.7218718218678238, 0.7006722805067258, 0.7225894731396042 ]
--   ]
--
-- >>> import Data.Massiv.Array
-- >>> import System.Random as System
-- >>> gen = System.mkStdGen 217
-- >>> randomArray gen System.split System.random (ParN 2) (Sz2 2 3) :: Array DL Ix2 Double
-- Array DL (ParN 2) (Sz (2 :. 3))
--   [ [ 0.15191527341922206, 0.2045537167404079, 0.9635356052820256 ]
--   , [ 9.308278528094238e-2, 0.7200934018606843, 0.23173694193083583 ]
--   ]
--
-- @since 0.3.3
randomArray ::
     forall ix e g. Index ix
  => g -- ^ Initial random value generator
  -> (g -> (g, g))
     -- ^ A function that can split a generator in two independent generators
  -> (g -> (e, g))
     -- ^ A function that produces a random value and the next generator
  -> Comp -- ^ Computation strategy.
  -> Sz ix -- ^ Resulting size of the array.
  -> Array DL ix e
randomArray :: 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 = Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())
-> Array DL ix e
unsafeMakeLoadArray Comp
comp Sz ix
sz Maybe e
forall a. Maybe a
Nothing forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
  where
    !totalLength :: Int
totalLength = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
    load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
writeAt =
      Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler m () -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers Scheduler m ()
scheduler) Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
        let slackStartAt :: Int
slackStartAt = Int
slackStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt
            writeRandom :: Int -> g -> m g
writeRandom Int
k g
genII =
              let (e
e, g
genII') = g -> (e, g)
nextRandom g
genII
               in g
genII' g -> m () -> m g
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> m ()
writeAt Int
k e
e
        g
genForSlack <-
          Int
-> (Int -> Bool) -> (Int -> Int) -> g -> (Int -> g -> m g) -> m g
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
startAt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStartAt) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) g
gen ((Int -> g -> m g) -> m g) -> (Int -> g -> m g) -> m g
forall a b. (a -> b) -> a -> b
$ \Int
start g
genI -> do
            let (g
genI0, g
genI1) =
                  if Scheduler m () -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers Scheduler m ()
scheduler Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                    then (g
genI, g
genI)
                    else g -> (g, g)
splitGen g
genI
            Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
              m g -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m g -> m ()) -> m g -> m ()
forall a b. (a -> b) -> a -> b
$ Int
-> (Int -> Bool) -> (Int -> Int) -> g -> (Int -> g -> m g) -> m g
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) g
genI0 Int -> g -> m g
writeRandom
            g -> m g
forall (f :: * -> *) a. Applicative f => a -> f a
pure g
genI1
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStartAt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
          Scheduler m () -> m () -> m ()
forall (m :: * -> *). Scheduler m () -> m () -> m ()
scheduleWork_ Scheduler m ()
scheduler (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
          m g -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m g -> m ()) -> m g -> m ()
forall a b. (a -> b) -> a -> b
$ Int
-> (Int -> Bool) -> (Int -> Int) -> g -> (Int -> g -> m g) -> m g
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
slackStartAt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) g
genForSlack Int -> g -> m g
writeRandom
{-# INLINE randomArray #-}

-- | Similar to `randomArray` but performs generation sequentially, which means it doesn't
-- require splitability property. Another consequence is that it returns the new generator
-- together with /manifest/ array of random values.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> import System.Random.SplitMix as SplitMix
-- >>> gen = SplitMix.mkSMGen 217
-- >>> snd $ randomArrayS gen (Sz2 2 3) SplitMix.nextDouble :: Array P Ix2 Double
-- Array P Seq (Sz (2 :. 3))
--   [ [ 0.8878273949359751, 0.11290807610140963, 0.7383156058619669 ]
--   , [ 0.39904053166835896, 0.5617584038393628, 0.16248374266020216 ]
--   ]
--
-- >>> import Data.Massiv.Array
-- >>> import System.Random.Mersenne.Pure64 as MT
-- >>> gen = MT.pureMT 217
-- >>> snd $ randomArrayS gen (Sz2 2 3) MT.randomDouble :: Array P Ix2 Double
-- Array P Seq (Sz (2 :. 3))
--   [ [ 0.5504018416543631, 0.22504666452851707, 0.4480480867867128 ]
--   , [ 0.7139711572975297, 0.49401087853770953, 0.9397201599368645 ]
--   ]
--
-- >>> import Data.Massiv.Array
-- >>> import System.Random as System
-- >>> gen = System.mkStdGen 217
-- >>> snd $ randomArrayS gen (Sz2 2 3) System.random :: Array P Ix2 Double
-- Array P Seq (Sz (2 :. 3))
--   [ [ 0.7972230393466304, 0.4485860543300083, 0.257773196880671 ]
--   , [ 0.19115043859955794, 0.33784788936970034, 3.479381605706322e-2 ]
--   ]
--
-- @since 0.3.4
randomArrayS ::
     forall r ix e g. Mutable r ix e
  => g -- ^ Initial random value generator
  -> Sz ix -- ^ Resulting size of the array.
  -> (g -> (e, g))
     -- ^ A function that produces a random value and the next generator
  -> (g, Array r ix e)
randomArrayS :: g -> Sz ix -> (g -> (e, g)) -> (g, Array r ix e)
randomArrayS g
gen Sz ix
sz g -> (e, g)
nextRandom =
  (forall s. ST s (g, Array r ix e)) -> (g, Array r ix e)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (g, Array r ix e)) -> (g, Array r ix e))
-> (forall s. ST s (g, Array r ix e)) -> (g, Array r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> (g -> ST s (e, g)) -> g -> ST s (g, Array r ix e)
forall r ix e a (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Sz ix -> (a -> m (e, a)) -> a -> m (a, Array r ix e)
unfoldrPrimM Sz ix
sz ((e, g) -> ST s (e, g)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((e, g) -> ST s (e, g)) -> (g -> (e, g)) -> g -> ST s (e, g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> (e, g)
nextRandom) g
gen
{-# INLINE randomArrayS #-}

-- | This is a stateful approach of generating random values. If your generator is pure
-- and splittable, it is better to use `randomArray` instead, which will give you a pure,
-- deterministic and parallelizable generation of arrays. On the other hand, if your
-- generator is not thread safe, which is most likely the case, instead of using some sort
-- of global mutex, `WorkerStates` allows you to keep track of individual state per worker
-- (thread), which fits parallelization of random value generation perfectly. All that
-- needs to be done is generators need to be initialized once per worker and then they can
-- be reused as many times as necessary.
--
-- ==== __Examples__
--
-- In the example below we take a stateful random generator from
-- [wmc-random](https://www.stackage.org/package/mwc-random), which is not thread safe,
-- and safely parallelize it by giving each thread it's own generator:
--
-- > λ> import Data.Massiv.Array
-- > λ> import System.Random.MWC (createSystemRandom, uniformR)
-- > λ> import System.Random.MWC.Distributions (standard)
-- > λ> gens <- initWorkerStates Par (\_ -> createSystemRandom)
-- > λ> randomArrayWS gens (Sz2 2 3) standard :: IO (Array P Ix2 Double)
-- > Array P Par (Sz (2 :. 3))
-- >   [ [ -0.9066144845415213, 0.5264323240310042, -1.320943607597422 ]
-- >   , [ -0.6837929005619592, -0.3041255565826211, 6.53353089112833e-2 ]
-- >   ]
-- > λ> randomArrayWS gens (Sz1 10) (uniformR (0, 9)) :: IO (Array P Ix1 Int)
-- > Array P Par (Sz1 10)
-- >   [ 3, 6, 1, 2, 1, 7, 6, 0, 8, 8 ]
--
-- @since 0.3.4
randomArrayWS ::
     forall r ix e g m. (Mutable r ix e, MonadUnliftIO m, PrimMonad m)
  => WorkerStates g -- ^ Use `initWorkerStates` to initialize you per thread generators
  -> Sz ix -- ^ Resulting size of the array
  -> (g -> m e) -- ^ Generate the value using the per thread generator.
  -> m (Array r ix e)
randomArrayWS :: WorkerStates g -> Sz ix -> (g -> m e) -> m (Array r ix e)
randomArrayWS WorkerStates g
states Sz ix
sz g -> m e
genRandom = WorkerStates g -> Sz ix -> (Int -> g -> m e) -> m (Array r ix e)
forall r ix e s (m :: * -> *).
(Mutable r ix e, MonadUnliftIO m, PrimMonad m) =>
WorkerStates s -> Sz ix -> (Int -> s -> m e) -> m (Array r ix e)
generateArrayLinearWS WorkerStates g
states Sz ix
sz ((g -> m e) -> Int -> g -> m e
forall a b. a -> b -> a
const g -> m e
genRandom)
{-# INLINE randomArrayWS #-}

infix 4 ..., ..:

-- | Handy synonym for @`rangeInclusive` `Seq`@. Similar to @..@ for list.
--
-- >>> Ix1 4 ... 10
-- Array D Seq (Sz1 7)
--   [ 4, 5, 6, 7, 8, 9, 10 ]
--
-- @since 0.3.0
(...) :: Index ix => ix -> ix -> Array D ix ix
... :: ix -> ix -> Array D ix ix
(...) = Comp -> ix -> ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
rangeInclusive Comp
Seq
{-# INLINE (...) #-}

-- | Handy synonym for @`range` `Seq`@
--
-- >>> Ix1 4 ..: 10
-- Array D Seq (Sz1 6)
--   [ 4, 5, 6, 7, 8, 9 ]
--
-- @since 0.3.0
(..:) :: Index ix => ix -> ix -> Array D ix ix
..: :: ix -> ix -> Array D ix ix
(..:) = Comp -> ix -> ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range Comp
Seq
{-# INLINE (..:) #-}


-- prop> range comp from to == rangeStep comp from 1 to
--
-- | Create an array of indices with a range from start to finish (not-including), where indices are
-- incremeted by one.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> range Seq (Ix1 1) 6
-- Array D Seq (Sz1 5)
--   [ 1, 2, 3, 4, 5 ]
-- >>> fromIx2 <$> range Seq (-1) (2 :. 2)
-- Array D Seq (Sz (3 :. 3))
--   [ [ (-1,-1), (-1,0), (-1,1) ]
--   , [ (0,-1), (0,0), (0,1) ]
--   , [ (1,-1), (1,0), (1,1) ]
--   ]
--
-- @since 0.1.0
range :: Index ix => Comp -> ix -> ix -> Array D ix ix
range :: Comp -> ix -> ix -> Array D ix ix
range Comp
comp !ix
from !ix
to = Comp -> ix -> Sz ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> Sz ix -> Array D ix ix
rangeSize Comp
comp ix
from (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
to ix
from))
{-# INLINE range #-}

-- | Same as `range`, but with a custom step.
--
-- /__Throws Exceptions__/: `IndexZeroException`
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> rangeStepM Seq (Ix1 1) 2 8
-- Array D Seq (Sz1 4)
--   [ 1, 3, 5, 7 ]
-- >>> rangeStepM Seq (Ix1 1) 0 8
-- *** Exception: IndexZeroException: 0
--
-- @since 0.3.0
rangeStepM :: (Index ix, MonadThrow m) =>
              Comp -- ^ Computation strategy
           -> ix -- ^ Start
           -> ix -- ^ Step (Can't have zeros)
           -> ix -- ^ End
           -> m (Array D ix ix)
rangeStepM :: Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepM Comp
comp !ix
from !ix
step !ix
to
  | (Bool -> Int -> Bool) -> Bool -> ix -> Bool
forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
foldlIndex (\Bool
acc Int
i -> Bool
acc Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) Bool
False ix
step = IndexException -> m (Array D ix ix)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m (Array D ix ix))
-> IndexException -> m (Array D ix ix)
forall a b. (a -> b) -> a -> b
$ ix -> IndexException
forall ix. Index ix => ix -> IndexException
IndexZeroException ix
step
  | Bool
otherwise =
    let dist :: ix
dist = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
to ix
from
        sz :: ix
sz = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Integral a => a -> a -> a
div ix
dist ix
step
        r :: ix
r = (Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex Int -> Int
forall a. Num a => a -> a
signum (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Integral a => a -> a -> a
mod ix
dist ix
step
     in Array D ix ix -> m (Array D ix ix)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D ix ix -> m (Array D ix ix))
-> Array D ix ix -> m (Array D ix ix)
forall a b. (a -> b) -> a -> b
$ Comp -> ix -> ix -> Sz ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> ix -> Sz ix -> Array D ix ix
rangeStepSize Comp
comp ix
from ix
step (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
sz ix
r))
{-# INLINE rangeStepM #-}

-- | Same as `rangeStepM`, but will throw an error whenever @step@ contains zeros.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> rangeStep' Seq (Ix1 1) 2 6
-- Array D Seq (Sz1 3)
--   [ 1, 3, 5 ]
--
-- @since 0.3.0
rangeStep' :: Index ix => Comp -> ix -> ix -> ix -> Array D ix ix
rangeStep' :: Comp -> ix -> ix -> ix -> Array D ix ix
rangeStep' Comp
comp ix
from ix
step = (SomeException -> Array D ix ix)
-> (Array D ix ix -> Array D ix ix)
-> Either SomeException (Array D ix ix)
-> Array D ix ix
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D ix ix
forall a e. Exception e => e -> a
throw Array D ix ix -> Array D ix ix
forall a. a -> a
id  (Either SomeException (Array D ix ix) -> Array D ix ix)
-> (ix -> Either SomeException (Array D ix ix))
-> ix
-> Array D ix ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> ix -> ix -> ix -> Either SomeException (Array D ix ix)
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' #-}

-- | Just like `range`, except the finish index is included.
--
-- @since 0.3.0
rangeInclusive :: Index ix => Comp -> ix -> ix -> Array D ix ix
rangeInclusive :: Comp -> ix -> ix -> Array D ix ix
rangeInclusive Comp
comp ix
ixFrom ix
ixTo =
  Comp -> ix -> Sz ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> Sz ix -> Array D ix ix
rangeSize Comp
comp ix
ixFrom (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ix
ixTo) ix
ixFrom))
{-# INLINE rangeInclusive #-}


-- | Just like `rangeStep`, except the finish index is included.
--
-- @since 0.3.0
rangeStepInclusiveM :: (MonadThrow m, Index ix) => Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepInclusiveM :: Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepInclusiveM Comp
comp ix
ixFrom ix
step ix
ixTo = Comp -> ix -> ix -> ix -> m (Array D ix ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
Comp -> ix -> ix -> ix -> m (Array D ix ix)
rangeStepM Comp
comp ix
ixFrom ix
step ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+) ix
ixTo)
{-# INLINE rangeStepInclusiveM #-}

-- | Just like `range`, except the finish index is included.
--
-- @since 0.3.1
rangeStepInclusive' :: Index ix => Comp -> ix -> ix -> ix -> Array D ix ix
rangeStepInclusive' :: Comp -> ix -> ix -> ix -> Array D ix ix
rangeStepInclusive' Comp
comp ix
ixFrom ix
step = (SomeException -> Array D ix ix)
-> (Array D ix ix -> Array D ix ix)
-> Either SomeException (Array D ix ix)
-> Array D ix ix
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D ix ix
forall a e. Exception e => e -> a
throw Array D ix ix -> Array D ix ix
forall a. a -> a
id  (Either SomeException (Array D ix ix) -> Array D ix ix)
-> (ix -> Either SomeException (Array D ix ix))
-> ix
-> Array D ix ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> ix -> ix -> ix -> Either SomeException (Array D ix ix)
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' #-}


-- | Create an array of specified size with indices starting with some index at position @0@ and
-- incremented by @1@ until the end of the array is reached
--
-- @since 0.3.0
rangeSize :: Index ix =>
             Comp -- ^ Computation strategy
          -> ix -- ^ @x@ - start value
          -> Sz ix -- ^ @sz@ - Size of resulting array
          -> Array D ix ix
rangeSize :: Comp -> ix -> Sz ix -> Array D ix ix
rangeSize Comp
comp !ix
from !Sz ix
sz = Comp -> Sz ix -> (ix -> ix) -> Array D ix ix
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
from)
{-# INLINE rangeSize #-}

-- | Same as `rangeSize`, but with ability to specify the step.
--
-- @since 0.3.0
rangeStepSize :: Index ix =>
                 Comp -- ^ Computation strategy
              -> ix -- ^ @x@ - start value
              -> ix -- ^ @delta@ - step value
              -> Sz ix -- ^ @sz@ - Size of resulting array
              -> Array D ix ix
rangeStepSize :: Comp -> ix -> ix -> Sz ix -> Array D ix ix
rangeStepSize Comp
comp !ix
from !ix
step !Sz ix
sz =
  Comp -> Sz ix -> (ix -> ix) -> Array D ix ix
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
from (ix -> ix) -> (ix -> ix) -> ix -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ix
step)
{-# INLINE rangeStepSize #-}


-- | Same as `enumFromStepN` with step @dx = 1@.
--
-- /Related/: `Data.Massiv.Vector.senumFromN`, `Data.Massiv.Vector.senumFromStepN`,
-- `enumFromStepN`, `rangeSize`, `rangeStepSize`, `range`
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> enumFromN Seq (5 :: Double) 3
-- Array D Seq (Sz1 3)
--   [ 5.0, 6.0, 7.0 ]
--
-- __/Similar/__:
--
-- [@Prelude.`Prelude.enumFromTo`@] Very similar to @[i .. i + n - 1]@, except that
-- `senumFromN` is faster, but it only works for `Num` and not for `Enum` elements
--
-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromN`@]
--
-- @since 0.1.0
enumFromN :: Num e =>
             Comp
          -> e -- ^ @x@ - start value
          -> Sz1 -- ^ @n@ - length of resulting vector.
          -> Vector D e
enumFromN :: Comp -> e -> Sz1 -> Vector D e
enumFromN Comp
comp !e
from !Sz1
sz = Comp -> Sz1 -> (Int -> e) -> Vector D e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz1
sz ((Int -> e) -> Vector D e) -> (Int -> e) -> Vector D e
forall a b. (a -> b) -> a -> b
$ \ Int
i -> e
from e -> e -> e
forall a. Num a => a -> a -> a
+ Int -> e
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
{-# INLINE enumFromN #-}


-- | Enumerate from a starting number @x@ exactly @n@ times with a custom step value
-- @dx@. Unlike `Data.Massiv.Vector.senumFromStepN`, there is no dependency on neigboring
-- elements therefore `enumFromStepN` is parallelizable.
--
-- /Related/: `Data.Massiv.Vector.senumFromN`, `Data.Massiv.Vector.senumFromStepN`,
-- `enumFromN`, `rangeSize`, `rangeStepSize`, `range`, `rangeStepM`
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> enumFromStepN Seq 1 (0.1 :: Double) 5
-- Array D Seq (Sz1 5)
--   [ 1.0, 1.1, 1.2, 1.3, 1.4 ]
-- >>> enumFromStepN Seq (-pi :: Float) (pi/4) 9
-- Array D Seq (Sz1 9)
--   [ -3.1415927, -2.3561945, -1.5707964, -0.78539824, 0.0, 0.78539824, 1.5707963, 2.3561947, 3.1415927 ]
--
-- __/Similar/__:
--
-- [@Prelude.`Prelude.enumFrom`@] Similar to @take n [x, x + dx ..]@, except that
-- `enumFromStepN` is parallelizable and it only works for `Num` and not for `Enum`
-- elements. Floating point value will be slightly different as well.
--
-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromStepN`@] Similar in the
-- outcome, but very different in the way it works.
--
--
-- @since 0.1.0
enumFromStepN :: Num e =>
                 Comp
              -> e -- ^ @x@ - start number
              -> e -- ^ @dx@ - step number
              -> Sz1 -- ^ @n@ - length of resulting vector
              -> Vector D e
enumFromStepN :: Comp -> e -> e -> Sz1 -> Vector D e
enumFromStepN Comp
comp !e
from !e
step !Sz1
sz = Comp -> Sz1 -> (Int -> e) -> Vector D e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz1
sz ((Int -> e) -> Vector D e) -> (Int -> e) -> Vector D e
forall a b. (a -> b) -> a -> b
$ \ Int
i -> e
from e -> e -> e
forall a. Num a => a -> a -> a
+ Int -> e
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i e -> e -> e
forall a. Num a => a -> a -> a
* e
step
{-# INLINE enumFromStepN #-}


-- | Function that expands an array to one with a higher dimension.
--
-- This is useful for constructing arrays where there is shared computation
-- between multiple cells.  The makeArray method of constructing arrays:
--
-- > makeArray :: Construct r ix e => Comp -> ix -> (ix -> e) -> Array r ix e
--
-- ...runs a function @ix -> e@ at every array index. This is inefficient if
-- there is a substantial amount of repeated computation that could be shared
-- while constructing elements on the same dimension. The expand functions make
-- this possible. First you construct an @Array r (Lower ix) a@ of one fewer
-- dimensions where @a@ is something like @`Array` r `Ix1` a@ or @`Array` r `Ix2` a@. Then
-- you use 'expandWithin' and a creation function @a -> Int -> b@ to create an
-- @`Array` `D` `Ix2` b@ or @`Array` `D` `Ix3` b@ respectfully.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> a = makeArrayR U Seq (Sz1 6) (+10) -- Imagine (+10) is some expensive function
-- >>> a
-- Array U Seq (Sz1 6)
--   [ 10, 11, 12, 13, 14, 15 ]
-- >>> expandWithin Dim1 5 (\ e j -> (j + 1) * 100 + e) a :: Array D Ix2 Int
-- Array D Seq (Sz (6 :. 5))
--   [ [ 110, 210, 310, 410, 510 ]
--   , [ 111, 211, 311, 411, 511 ]
--   , [ 112, 212, 312, 412, 512 ]
--   , [ 113, 213, 313, 413, 513 ]
--   , [ 114, 214, 314, 414, 514 ]
--   , [ 115, 215, 315, 415, 515 ]
--   ]
-- >>> expandWithin Dim2 5 (\ e j -> (j + 1) * 100 + e) a :: Array D Ix2 Int
-- Array D Seq (Sz (5 :. 6))
--   [ [ 110, 111, 112, 113, 114, 115 ]
--   , [ 210, 211, 212, 213, 214, 215 ]
--   , [ 310, 311, 312, 313, 314, 315 ]
--   , [ 410, 411, 412, 413, 414, 415 ]
--   , [ 510, 511, 512, 513, 514, 515 ]
--   ]
--
-- @since 0.2.6
expandWithin ::
     forall ix e r n a. (IsIndexDimension ix n, Manifest r (Lower ix) a)
  => Dimension n
  -> Sz1
  -> (a -> Ix1 -> e)
  -> Array r (Lower ix) a
  -> Array D ix e
expandWithin :: 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 =
  Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz ((ix -> e) -> Array D ix e) -> (ix -> e) -> Array D ix e
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
    let (Int
i, Lower ix
ixl) = ix -> Dimension n -> (Int, Lower ix)
forall ix (n :: Nat).
IsIndexDimension ix n =>
ix -> Dimension n -> (Int, Lower ix)
pullOutDimension ix
ix Dimension n
dim
     in a -> Int -> e
f (Array r (Lower ix) a -> Lower ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
  where
    szl :: Lower ix
szl = Sz (Lower ix) -> Lower ix
forall ix. Sz ix -> ix
unSz (Array r (Lower ix) a -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr)
    sz :: Sz ix
sz = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (Lower ix -> Dimension n -> Int -> ix
forall ix (n :: Nat).
IsIndexDimension ix n =>
Lower ix -> Dimension n -> Int -> ix
insertDimension Lower ix
szl Dimension n
dim Int
k)
{-# INLINE expandWithin #-}

-- | Similar to `expandWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
-- @since 0.2.6
expandWithin'
  :: (Index ix, Manifest r (Lower ix) a)
  => Dim
  -> Sz1
  -> (a -> Ix1 -> b)
  -> Array r (Lower ix) a
  -> Array D ix b
expandWithin' :: Dim
-> Sz1 -> (a -> Int -> b) -> Array r (Lower ix) a -> Array D ix b
expandWithin' Dim
dim Sz1
k a -> Int -> b
f Array r (Lower ix) a
arr = (SomeException -> Array D ix b)
-> (Array D ix b -> Array D ix b)
-> Either SomeException (Array D ix b)
-> Array D ix b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D ix b
forall a e. Exception e => e -> a
throw Array D ix b -> Array D ix b
forall a. a -> a
id (Either SomeException (Array D ix b) -> Array D ix b)
-> Either SomeException (Array D ix b) -> Array D ix b
forall a b. (a -> b) -> a -> b
$ Dim
-> Sz1
-> (a -> Int -> b)
-> Array r (Lower ix) a
-> Either SomeException (Array D ix b)
forall ix r a (m :: * -> *) b.
(Index ix, Manifest r (Lower ix) 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
{-# INLINE expandWithin' #-}

-- | Similar to `expandWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
-- @since 0.4.0
expandWithinM
  :: (Index ix, Manifest r (Lower ix) a, MonadThrow m)
  => Dim
  -> Sz1
  -> (a -> Ix1 -> b)
  -> Array r (Lower ix) a
  -> m (Array D ix b)
expandWithinM :: 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 <- Sz (Lower ix) -> Dim -> Sz1 -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz1 -> m (Sz ix)
insertSzM (Array r (Lower ix) a -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr) Dim
dim Sz1
k
  Array D ix b -> m (Array D ix b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D ix b -> m (Array D ix b))
-> Array D ix b -> m (Array D ix b)
forall a b. (a -> b) -> a -> b
$
    Comp -> Sz ix -> (ix -> b) -> Array D ix b
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz ((ix -> b) -> Array D ix b) -> (ix -> b) -> Array D ix b
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
      let (Int
i, Lower ix
ixl) = ix -> Dim -> (Int, Lower ix)
forall ix. Index ix => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
ix Dim
dim -- dim has been checked above
       in a -> Int -> b
f (Array r (Lower ix) a -> Lower ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
{-# INLINE expandWithinM #-}

-- | Similar to `expandWithin`, except it uses the outermost dimension.
--
-- @since 0.2.6
expandOuter
  :: (Index ix, Manifest r (Lower ix) a)
  => Sz1
  -> (a -> Ix1 -> b)
  -> Array r (Lower ix) a
  -> Array D ix b
expandOuter :: 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 =
  Comp -> Sz ix -> (ix -> b) -> Array D ix b
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz ((ix -> b) -> Array D ix b) -> (ix -> b) -> Array D ix b
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
    let (Int
i, Lower ix
ixl) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
     in a -> Int -> b
f (Array r (Lower ix) a -> Lower ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
  where
    szl :: Sz (Lower ix)
szl = Array r (Lower ix) a -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr
    sz :: Sz ix
sz = Sz1 -> Sz (Lower ix) -> Sz ix
forall ix. Index ix => Sz1 -> Sz (Lower ix) -> Sz ix
consSz Sz1
k Sz (Lower ix)
szl
{-# INLINE expandOuter #-}

-- | Similar to `expandWithin`, except it uses the innermost dimension.
--
-- @since 0.2.6
expandInner
  :: (Index ix, Manifest r (Lower ix) a)
  => Sz1
  -> (a -> Ix1 -> b)
  -> Array r (Lower ix) a
  -> Array D ix b
expandInner :: 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 =
  Comp -> Sz ix -> (ix -> b) -> Array D ix b
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r (Lower ix) a
arr) Sz ix
sz ((ix -> b) -> Array D ix b) -> (ix -> b) -> Array D ix b
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
    let (Lower ix
ixl, Int
i) = ix -> (Lower ix, Int)
forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
ix
     in a -> Int -> b
f (Array r (Lower ix) a -> Lower ix -> a
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r (Lower ix) a
arr Lower ix
ixl) Int
i
  where
    szl :: Sz (Lower ix)
szl = Array r (Lower ix) a -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) a
arr
    sz :: Sz ix
sz = Sz (Lower ix) -> Sz1 -> Sz ix
forall ix. Index ix => Sz (Lower ix) -> Sz1 -> Sz ix
snocSz Sz (Lower ix)
szl Sz1
k
{-# INLINE expandInner #-}