{-# 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-2022
-- 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_,
  iunfoldlS_,
  unfoldrS_,
  iunfoldrS_,
  makeSplitSeedArray,

  -- *** Random
  uniformArray,
  uniformRangeArray,
  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 (void, when)
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 System.Random.Stateful
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 :: 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
_ = Comp -> Sz ix -> (ix -> e) -> Array r ix e
forall r ix e.
Load 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 :: 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
_ = Comp -> Sz ix -> (Int -> e) -> Array r ix e
forall r ix e.
Load 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 :: 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
_ = Comp -> Sz1 -> (Int -> e) -> Array r Int e
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 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 e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, 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 a b. ST s a -> (a -> ST s b) -> ST s b
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
   . (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 =
  (STA r ix e -> Array r ix e) -> f (STA r ix e) -> f (Array r ix e)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Sz ix -> STA r ix e -> Array r ix e
forall r e ix.
(Manifest r e, Index ix) =>
Sz ix -> STA r ix e -> Array r ix e
runSTA Sz ix
sz) (f (STA r ix e) -> f (Array r ix e))
-> f (STA r ix e) -> f (Array r ix e)
forall a b. (a -> b) -> a -> b
$
    ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f (STA r ix e)
-> (ix -> f (STA r ix e) -> f (STA r ix e))
-> f (STA r ix e)
forall ix (f :: * -> *) a.
Index ix =>
ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f a
-> (ix -> f a -> f a)
-> f a
forall (f :: * -> *) a.
ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f a
-> (ix -> f a -> f a)
-> f a
iterF ix
forall ix. Index ix => ix
zeroIndex ix
n ix
forall ix. Index ix => ix
oneIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) (STA r ix e -> f (STA r ix e)
forall a. a -> f a
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 e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq))) ((ix -> f (STA r ix e) -> f (STA r ix e)) -> f (STA r ix e))
-> (ix -> f (STA r ix e) -> f (STA r ix e)) -> f (STA r ix e)
forall a b. (a -> b) -> a -> b
$ \ix
ix f (STA r ix e)
g ->
      (e -> STA r ix e -> STA r ix e)
-> f e -> f (STA r ix e) -> f (STA r ix e)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
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 -> ix -> e -> ST s ()
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
MArray (PrimState (ST s)) r ix e
ma ix
ix e
e ST s () -> ST s (Array r ix e) -> ST s (Array r ix e)
forall a b. ST s a -> ST s b -> ST s b
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 ix
ix) f (STA r ix e)
g
{-# INLINE makeArrayA #-}

-- | Same as `makeArrayA`, but with linear index.
--
-- @since 0.4.5
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 =
  (STA r ix e -> Array r ix e) -> f (STA r ix e) -> f (Array r ix e)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Sz ix -> STA r ix e -> Array r ix e
forall r e ix.
(Manifest r e, Index ix) =>
Sz ix -> STA r ix e -> Array r ix e
runSTA Sz ix
sz) (f (STA r ix e) -> f (Array r ix e))
-> f (STA r ix e) -> f (Array r ix e)
forall a b. (a -> b) -> a -> b
$
    Int
-> (Int -> Bool)
-> (Int -> Int)
-> f (STA r ix e)
-> (Int -> f (STA r ix e) -> f (STA r ix e))
-> f (STA r ix e)
forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
0 (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
1) (STA r ix e -> f (STA r ix e)
forall a. a -> f a
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 e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq))) ((Int -> f (STA r ix e) -> f (STA r ix e)) -> f (STA r ix e))
-> (Int -> f (STA r ix e) -> f (STA r ix e)) -> f (STA r ix e)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      (e -> STA r ix e -> STA r ix e)
-> f e -> f (STA r ix e) -> f (STA r ix e)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
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 e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, 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 a b. ST s a -> ST s b -> ST s b
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)
{-# INLINE makeArrayLinearA #-}

-- | Same as `makeArrayA`, but with ability to supply result array representation.
--
-- @since 0.2.6
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
_ = Sz ix -> (ix -> f e) -> f (Array r ix e)
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 #-}

-- | 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 :: forall ix e. Index ix => 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.
Index ix =>
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 :: forall ix e.
Index ix =>
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.
Index ix =>
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
   . 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 = Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e
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_ #-}

-- | Right unfold of a delayed load array with index aware function
--
-- @since 0.3.0
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 = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
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 ()
_ =
      ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$
        RowMajor
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> ST s a)
-> ST s a
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
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
RowMajor
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM RowMajor
defRowMajor Int
startAt Sz ix
sz ix
forall ix. Index ix => ix
zeroIndex Stride ix
forall ix. Index ix => Stride ix
oneStride a
acc0 ((Int -> ix -> a -> ST s a) -> ST s a)
-> (Int -> ix -> a -> ST s a) -> ST s a
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' a -> ST s () -> ST s a
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
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 achieve such effect.
--
-- @since 0.3.0
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 = Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e
forall ix e a.
Index ix =>
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
   . 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 = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
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 ()
_ =
      ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$
        Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> a -> ST s a)
-> ST s 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 -> ST s a) -> ST s a) -> (Int -> a -> ST s a) -> ST s 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 -> ST s () -> ST s a
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
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.
--
-- __Note__: Starting with massiv-1.1.0 this function will be deprecated in
-- favor of a more general `genSplitArray`
--
-- ==== __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 Random
-- >>> gen = Random.mkStdGen 217
-- >>> randomArray gen Random.split Random.random (ParN 2) (Sz2 2 3) :: Array DL Ix2 Double
-- Array DL (ParN 2) (Sz (2 :. 3))
--   [ [ 0.2616843941380331, 0.600959468331641, 0.4382415961606372 ]
--   , [ 0.27812817813217605, 0.2993277194932741, 0.2774105268603957 ]
--   ]
--
-- @since 1.0.0
randomArray
  :: forall ix e g
   . Index ix
  => g
  -- ^ Initial random value generator
  -> (g -> (g, g))
  -- ^ A function that can split a generator into two independent
  -- generators. It will only be called if supplied computation strategy
  -- needs more than one worker threads.
  -> (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 :: 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 = Comp
-> Sz ix
-> Maybe e
-> (forall s.
    Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ())
-> Array DL ix e
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 Maybe e
forall a. Maybe a
Nothing Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ()
forall s. Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ()
load
  where
    !totalLength :: Int
totalLength = Sz ix -> Int
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 =
      Int -> Int -> (Int -> Int -> ST s ()) -> ST s ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) Int
totalLength ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
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 -> ST s g
writeRandom Int
k g
genII =
              let (e
e, g
genII') = g -> (e, g)
nextRandom g
genII
               in g
genII' g -> ST s () -> ST s g
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
writeAt Int
k e
e
        g
genForSlack <-
          Int
-> (Int -> Bool)
-> (Int -> Int)
-> g
-> (Int -> g -> ST s g)
-> ST s 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 -> ST s g) -> ST s g) -> (Int -> g -> ST s g) -> ST s g
forall a b. (a -> b) -> a -> b
$ \Int
start g
genI -> do
            let (g
genI0, g
genI1) =
                  if Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
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 s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
              ST s g -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s g -> ST s ()) -> ST s g -> ST s ()
forall a b. (a -> b) -> a -> b
$
                Int
-> (Int -> Bool)
-> (Int -> Int)
-> g
-> (Int -> g -> ST s g)
-> ST s 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 -> ST s g
writeRandom
            g -> ST s g
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure g
genI1
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
          Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
            ST s g -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s g -> ST s ()) -> ST s g -> ST s ()
forall a b. (a -> b) -> a -> b
$
              Int
-> (Int -> Bool)
-> (Int -> Int)
-> g
-> (Int -> g -> ST s g)
-> ST s 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 -> ST s g
writeRandom
{-# INLINE randomArray #-}

-- | Create a delayed array with an initial seed and a splitting function. It is
-- somewhat similar to `iunfoldlS_` function, but it is capable of parallelizing
-- computation and iterating over the array accoriding to the supplied
-- `Iterator`. Upon parallelization every job will get the second part of the
-- result produced by the split function, while the first part will be used for
-- subsequent splits. This function is similar to
-- `Data.Massiv.Array.Manifest.generateSplitSeedArray`
--
-- @since 1.0.2
makeSplitSeedArray
  :: forall ix e g it
   . (Iterator it, Index ix)
  => it
  -- ^ Iterator
  -> g
  -- ^ Initial seed
  -> (g -> (g, g))
  -- ^ A function that can split a seed into two independent seeds. It will
  -- be called the same number of times as the number of jobs that will get
  -- scheduled during parallelization. Eg. only once for the sequential case.
  -> Comp
  -- ^ Computation strategy.
  -> Sz ix
  -- ^ Resulting size of the array.
  -> (Ix1 -> ix -> g -> (e, g))
  -- ^ A function that produces a value and the next seed. It takes both
  -- versions of the index, in linear and in multi-dimensional forms, as well as
  -- the current seeding value.
  -> 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 = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
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 ()
_ =
      it
-> Scheduler s ()
-> Int
-> Sz ix
-> g
-> (g -> ST s (g, g))
-> (Int -> ix -> g -> ST s g)
-> 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 ()
forall ix s a.
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 ((g, g) -> ST s (g, g)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((g, g) -> ST s (g, g)) -> (g -> (g, g)) -> g -> ST s (g, g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. g -> (g, g)
splitSeed) ((Int -> ix -> g -> ST s g) -> ST s ())
-> (Int -> ix -> g -> ST s g) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix g
g ->
        case Int -> ix -> g -> (e, g)
genFunc (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
startAt) ix
ix g
g of
          (e
x, g
g') -> g
g' g -> ST s () -> ST s g
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> e -> ST s ()
writeAt Int
i e
x
    {-# INLINE load #-}
{-# INLINE makeSplitSeedArray #-}

-- | Generate a random array where all elements are sampled from a uniform distribution.
--
-- @since 1.0.0
uniformArray
  :: forall ix e g
   . (Index ix, RandomGen g, Uniform e)
  => g
  -- ^ Initial random value generator.
  -> Comp
  -- ^ Computation strategy.
  -> Sz ix
  -- ^ Resulting size of the array.
  -> 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 = g
-> (g -> (g, g)) -> (g -> (e, g)) -> Comp -> Sz ix -> Array DL ix e
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)
forall g. RandomGen g => g -> (g, g)
split g -> (e, g)
forall g a. (RandomGen g, Uniform a) => g -> (a, g)
uniform
{-# INLINE uniformArray #-}

-- | Same as `uniformArray`, but will generate values in a supplied range.
--
-- @since 1.0.0
uniformRangeArray
  :: forall ix e g
   . (Index ix, RandomGen g, UniformRange e)
  => g
  -- ^ Initial random value generator.
  -> (e, e)
  -- ^ Inclusive range in which values will be generated in.
  -> Comp
  -- ^ Computation strategy.
  -> Sz ix
  -- ^ Resulting size of the array.
  -> 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 = g
-> (g -> (g, g)) -> (g -> (e, g)) -> Comp -> Sz ix -> Array DL ix e
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)
forall g. RandomGen g => g -> (g, g)
split ((e, e) -> g -> (e, g)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (e, e)
r)
{-# INLINE uniformRangeArray #-}

-- | 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.11217260506402493, 0.8870919238985904, 0.2616843941380331 ]
--   , [ 0.600959468331641, 0.4382415961606372, 0.8375162573397977 ]
--   ]
--
-- @since 0.3.4
randomArrayS
  :: forall r ix e g
   . (Manifest r e, Index ix)
  => 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 :: 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 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 :: * -> *).
(Manifest r e, Index ix, 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 a. a -> ST s a
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 number 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. There is a caveat
-- of course, statistical independence will depend on the entropy in your initial seeds,
-- so do not use the example below verbatim, since initial seeds are sequential numbers.
--
-- >>> import Data.Massiv.Array as A
-- >>> import System.Random.MWC as MWC (initialize)
-- >>> import System.Random.Stateful (uniformRM)
-- >>> import Control.Scheduler (initWorkerStates, getWorkerId)
-- >>> :set -XTypeApplications
-- >>> gens <- initWorkerStates Par (MWC.initialize . A.toPrimitiveVector . A.singleton @P @Ix1 . fromIntegral . getWorkerId)
-- >>> randomArrayWS gens (Sz2 2 3) (uniformRM (0, 9)) :: IO (Matrix P Double)
-- Array P Par (Sz (2 :. 3))
--   [ [ 8.999240522095299, 6.832223390653755, 3.065728078741671 ]
--   , [ 7.242581103346686, 2.4565807301968623, 0.4514262066689775 ]
--   ]
-- >>> randomArrayWS gens (Sz1 6) (uniformRM (0, 9)) :: IO (Vector P Int)
-- Array P Par (Sz1 6)
--   [ 8, 8, 7, 1, 1, 2 ]
--
-- @since 0.3.4
randomArrayWS
  :: forall r ix e g m
   . (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m)
  => WorkerStates g
  -- ^ Use `Control.Scheduler.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 :: 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 = WorkerStates g -> Sz ix -> (Int -> g -> m e) -> m (Array r ix e)
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 ((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
... :: forall ix. Index 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
..: :: forall ix. Index 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 :: forall ix. Index ix => 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
  :: forall ix m
   . (Index ix, MonadThrow m)
  => Comp
  -- ^ Computation strategy
  -> ix
  -- ^ Start
  -> ix
  -- ^ Step. Negative and positive values are ok, but can't have zeros
  -> ix
  -- ^ End
  -> 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
  | (Bool -> Int -> Bool) -> Bool -> ix -> Bool
forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
forall a. (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 e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, 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 a. a -> m a
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' :: (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 = Either SomeException (Array D ix ix) -> Array D ix ix
forall a. HasCallStack => Either SomeException a -> a
throwEither (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 :: forall ix. Index ix => 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 `rangeStepM`, 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 :: 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 = 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' :: (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 = Either SomeException (Array D ix ix) -> Array D ix ix
forall a. HasCallStack => Either SomeException a -> a
throwEither (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 :: forall ix. Index ix => 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.
Load 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 :: forall ix. Index ix => 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.
Load 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
-- `enumFromN` 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 :: forall e. Num e => Comp -> e -> Sz1 -> Vector D e
enumFromN Comp
comp !e
from !Sz1
sz = Comp -> Sz1 -> (Int -> e) -> Array D Int e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz1
sz ((Int -> e) -> Array D Int e) -> (Int -> e) -> Array D Int 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 :: forall e. Num e => Comp -> e -> e -> Sz1 -> Vector D e
enumFromStepN Comp
comp !e
from !e
step !Sz1
sz = Comp -> Sz1 -> (Int -> e) -> Array D Int e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz1
sz ((Int -> e) -> Array D Int e) -> (Int -> e) -> Array D Int 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 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 =
  Comp -> Sz ix -> (ix -> e) -> Array D ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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 :: Natural).
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 ix. Index ix => Array r ix a -> ix -> a
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 = Sz (Lower ix) -> Lower ix
forall ix. Sz ix -> ix
unSz (Array r (Lower ix) a -> Sz (Lower ix)
forall r ix e. Size r => Array r ix e -> Sz ix
forall 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 :: Natural).
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'
  :: 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 = Either SomeException (Array D ix b) -> Array D ix b
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array D ix b) -> Array D ix b)
-> (Array r (Lower ix) a -> Either SomeException (Array D ix b))
-> Array r (Lower ix) a
-> Array D ix b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim
-> Sz1
-> (a -> Int -> b)
-> Array r (Lower ix) a
-> Either SomeException (Array D ix b)
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' #-}

-- | 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
  :: 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 <- 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. Size r => Array r ix e -> Sz ix
forall 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 a. a -> m a
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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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. (HasCallStack, 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 ix. Index ix => Array r ix a -> ix -> a
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 #-}

-- | Similar to `expandWithin`, except it uses the outermost dimension.
--
-- @since 0.2.6
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 =
  Comp -> Sz ix -> (ix -> b) -> Array D ix b
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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 ix. Index ix => Array r ix a -> ix -> a
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 = Array r (Lower ix) a -> Sz (Lower ix)
forall r ix e. Size r => Array r ix e -> Sz ix
forall 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
  :: 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 =
  Comp -> Sz ix -> (ix -> b) -> Array D ix b
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r (Lower ix) a -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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 ix. Index ix => Array r ix a -> ix -> a
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 = Array r (Lower ix) a -> Sz (Lower ix)
forall r ix e. Size r => Array r ix e -> Sz ix
forall 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 #-}