{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
-- |
-- Module      : Data.Massiv.Array.Ops.Transform
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Ops.Transform
  ( -- ** Transpose
    transpose
  , transposeInner
  , transposeOuter
  -- ** Reverse
  , reverse
  , reverse'
  , reverseM
  -- ** Backpermute
  , backpermuteM
  , backpermute'
  -- ** Resize
  , resizeM
  , resize'
  , flatten
  -- ** Extract
  , extractM
  , extract'
  , extractFromToM
  , extractFromTo'
  , deleteRowsM
  , deleteColumnsM
  , deleteRegionM
  -- ** Append/Split
  , appendOuterM
  , appendM
  , append'
  , concatOuterM
  , concatM
  , concat'
  , stackSlicesM
  , stackOuterSlicesM
  , stackInnerSlicesM
  , splitAtM
  , splitAt'
  , splitExtractM
  , replaceSlice
  , replaceOuterSlice
  -- ** Upsample/Downsample
  , upsample
  , downsample
  -- ** Zoom
  , zoom
  , zoomWithGrid
  -- ** Transform
  , transformM
  , transform'
  , transform2M
  , transform2'
  ) where

import Control.Scheduler (traverse_)
import Control.Monad as M (foldM_, unless, forM_)
import Data.Bifunctor (bimap)
import Data.Foldable as F (foldl', foldrM, toList, length)
import qualified Data.List as L (uncons)
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Delayed.Push
import Data.Massiv.Array.Mutable
import Data.Massiv.Array.Ops.Construct
import Data.Massiv.Array.Ops.Map
import Data.Massiv.Core.Common
import Prelude as P hiding (concat, splitAt, traverse, mapM_, reverse, take, drop)


-- | Extract a sub-array from within a larger source array. Array that is being extracted must be
-- fully encapsulated in a source array, otherwise `SizeSubregionException` will be thrown.
extractM :: (MonadThrow m, Extract r ix e)
         => ix -- ^ Starting index
         -> Sz ix -- ^ Size of the resulting array
         -> Array r ix e -- ^ Source array
         -> m (Array (R r) ix e)
extractM :: ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM !ix
sIx !Sz ix
newSz !Array r ix e
arr
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz1 ix
sIx Bool -> Bool -> Bool
&& Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
eIx1 ix
sIx Bool -> Bool -> Bool
&& Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz1 ix
eIx =
    Array (R r) ix e -> m (Array (R r) ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array (R r) ix e -> m (Array (R r) ix e))
-> Array (R r) ix e -> m (Array (R r) ix e)
forall a b. (a -> b) -> a -> b
$ ix -> Sz ix -> Array r ix e -> Array (R r) ix e
forall r ix e.
Extract r ix e =>
ix -> Sz ix -> Array r ix e -> Array (R r) ix e
unsafeExtract ix
sIx Sz ix
newSz Array r ix e
arr
  | Bool
otherwise = SizeException -> m (Array (R r) ix e)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> m (Array (R r) ix e))
-> SizeException -> m (Array (R r) ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> ix -> Sz ix -> SizeException
SizeSubregionException (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr) ix
sIx Sz ix
newSz
  where
    sz1 :: Sz ix
sz1 = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((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) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)))
    eIx1 :: Sz ix
eIx1 = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((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
eIx)
    eIx :: ix
eIx = (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
sIx (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
newSz
{-# INLINE extractM #-}

-- | Same as `extractM`, but will throw a runtime exception from pure code if supplied dimensions
-- are incorrect.
--
-- @since 0.1.0
extract' :: Extract r ix e
        => ix -- ^ Starting index
        -> Sz ix -- ^ Size of the resulting array
        -> Array r ix e -- ^ Source array
        -> Array (R r) ix e
extract' :: ix -> Sz ix -> Array r ix e -> Array (R r) ix e
extract' ix
sIx Sz ix
newSz = (SomeException -> Array (R r) ix e)
-> (Array (R r) ix e -> Array (R r) ix e)
-> Either SomeException (Array (R r) ix e)
-> Array (R r) ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array (R r) ix e
forall a e. Exception e => e -> a
throw Array (R r) ix e -> Array (R r) ix e
forall a. a -> a
id (Either SomeException (Array (R r) ix e) -> Array (R r) ix e)
-> (Array r ix e -> Either SomeException (Array (R r) ix e))
-> Array r ix e
-> Array (R r) ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix
-> Sz ix -> Array r ix e -> Either SomeException (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
sIx Sz ix
newSz
{-# INLINE extract' #-}


-- | Similar to `extractM`, except it takes starting and ending index. Result array will not include
-- the ending index.
--
-- @since 0.3.0
extractFromToM :: (MonadThrow m, Extract r ix e) =>
                  ix -- ^ Starting index
               -> ix -- ^ Index up to which elements should be extracted.
               -> Array r ix e -- ^ Source array.
               -> m (Array (R r) ix e)
extractFromToM :: ix -> ix -> Array r ix e -> m (Array (R r) ix e)
extractFromToM ix
sIx ix
eIx = ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
sIx (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
eIx ix
sIx))
{-# INLINE extractFromToM #-}

-- | Same as `extractFromTo`, but throws an error on invalid indices.
--
-- @since 0.2.4
extractFromTo' :: Extract r ix e =>
                 ix -- ^ Starting index
              -> ix -- ^ Index up to which elmenets should be extracted.
              -> Array r ix e -- ^ Source array.
              -> Array (R r) ix e
extractFromTo' :: ix -> ix -> Array r ix e -> Array (R r) ix e
extractFromTo' ix
sIx ix
eIx = ix -> Sz ix -> Array r ix e -> Array (R r) ix e
forall r ix e.
Extract r ix e =>
ix -> Sz ix -> Array r ix e -> Array (R r) ix e
extract' ix
sIx (Sz ix -> Array r ix e -> Array (R r) ix e)
-> Sz ix -> Array r ix e -> Array (R r) ix e
forall a b. (a -> b) -> a -> b
$ 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
eIx ix
sIx)
{-# INLINE extractFromTo' #-}


-- | /O(1)/ - Changes the shape of an array. Returns `Nothing` if total
-- number of elements does not match the source array.
--
-- @since 0.3.0
resizeM ::
     (MonadThrow m, Index ix', Load r ix e, Resize r ix)
  => Sz ix'
  -> Array r ix e
  -> m (Array r ix' e)
resizeM :: Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM Sz ix'
sz Array r ix e
arr = Sz ix -> Sz ix' -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr) Sz ix'
sz m () -> m (Array r ix' e) -> m (Array r ix' e)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Array r ix' e -> m (Array r ix' e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Sz ix' -> Array r ix e -> Array r ix' e
forall r ix ix' e.
(Resize r ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix'
sz Array r ix e
arr)
{-# INLINE resizeM #-}

-- | Same as `resizeM`, but will throw an error if supplied dimensions are incorrect.
--
-- @since 0.1.0
resize' :: (Index ix', Load r ix e, Resize r ix) => Sz ix' -> Array r ix e -> Array r ix' e
resize' :: Sz ix' -> Array r ix e -> Array r ix' e
resize' Sz ix'
sz = (SomeException -> Array r ix' e)
-> (Array r ix' e -> Array r ix' e)
-> Either SomeException (Array r ix' e)
-> Array r ix' e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array r ix' e
forall a e. Exception e => e -> a
throw Array r ix' e -> Array r ix' e
forall a. a -> a
id (Either SomeException (Array r ix' e) -> Array r ix' e)
-> (Array r ix e -> Either SomeException (Array r ix' e))
-> Array r ix e
-> Array r ix' e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix' -> Array r ix e -> Either SomeException (Array r ix' e)
forall (m :: * -> *) ix' r ix e.
(MonadThrow m, Index ix', Load r ix e, Resize r ix) =>
Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM Sz ix'
sz
{-# INLINE resize' #-}

-- | /O(1)/ - Reduce a multi-dimensional array into a flat vector
--
-- @since 0.3.1
flatten :: (Load r ix e, Resize r ix) => Array r ix e -> Array r Ix1 e
flatten :: Array r ix e -> Array r Int e
flatten Array r ix e
arr = Sz Int -> Array r ix e -> Array r Int e
forall r ix ix' e.
(Resize r ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr))) Array r ix e
arr
{-# INLINE flatten #-}


-- | Transpose a 2-dimensional array
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> arr = makeArrayLinearR D Seq (Sz (2 :. 3)) id
-- >>> arr
-- Array D Seq (Sz (2 :. 3))
--   [ [ 0, 1, 2 ]
--   , [ 3, 4, 5 ]
--   ]
-- >>> transpose arr
-- Array D Seq (Sz (3 :. 2))
--   [ [ 0, 3 ]
--   , [ 1, 4 ]
--   , [ 2, 5 ]
--   ]
--
-- @since 0.1.0
transpose :: Source r Ix2 e => Array r Ix2 e -> Array D Ix2 e
transpose :: Array r Ix2 e -> Array D Ix2 e
transpose = Array r Ix2 e -> Array D Ix2 e
forall ix r' e.
(Index (Lower ix), Source r' ix e) =>
Array r' ix e -> Array D ix e
transposeInner
{-# INLINE [1] transpose #-}

{-# RULES
"transpose . transpose" [~1] forall arr . transpose (transpose arr) = delay arr
"transposeInner . transposeInner" [~1] forall arr . transposeInner (transposeInner arr) = delay arr
"transposeOuter . transposeOuter" [~1] forall arr . transposeOuter (transposeOuter arr) = delay arr
 #-}


-- | Transpose inner two dimensions of at least rank-2 array.
--
-- ===__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> arr = makeArrayLinearR U Seq (Sz (2 :> 3 :. 4)) id
-- >>> arr
-- Array U Seq (Sz (2 :> 3 :. 4))
--   [ [ [ 0, 1, 2, 3 ]
--     , [ 4, 5, 6, 7 ]
--     , [ 8, 9, 10, 11 ]
--     ]
--   , [ [ 12, 13, 14, 15 ]
--     , [ 16, 17, 18, 19 ]
--     , [ 20, 21, 22, 23 ]
--     ]
--   ]
-- >>> transposeInner arr
-- Array D Seq (Sz (3 :> 2 :. 4))
--   [ [ [ 0, 1, 2, 3 ]
--     , [ 12, 13, 14, 15 ]
--     ]
--   , [ [ 4, 5, 6, 7 ]
--     , [ 16, 17, 18, 19 ]
--     ]
--   , [ [ 8, 9, 10, 11 ]
--     , [ 20, 21, 22, 23 ]
--     ]
--   ]
--
-- @since 0.1.0
transposeInner :: (Index (Lower ix), Source r' ix e)
               => Array r' ix e -> Array D ix e
transposeInner :: Array r' ix e -> Array D ix e
transposeInner !Array r' ix e
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' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr) Sz ix
newsz ix -> e
newVal
  where
    transInner :: ix -> ix
transInner !ix
ix =
      (SomeException -> ix)
-> (ix -> ix) -> Either SomeException ix -> ix
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> ix
forall e a. Exception e => e -> a
throwImpossible ix -> ix
forall a. a -> a
id (Either SomeException ix -> ix) -> Either SomeException ix -> ix
forall a b. (a -> b) -> a -> b
$ do
        Int
n <- ix -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM ix
ix Dim
dix
        Int
m <- ix -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM ix
ix (Dim
dix Dim -> Dim -> Dim
forall a. Num a => a -> a -> a
- Dim
1)
        ix
ix' <- ix -> Dim -> Int -> Either SomeException ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
ix Dim
dix Int
m
        ix -> Dim -> Int -> Either SomeException ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
ix' (Dim
dix Dim -> Dim -> Dim
forall a. Num a => a -> a -> a
- Dim
1) Int
n
    {-# INLINE transInner #-}
    newVal :: ix -> e
newVal = Array r' ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r' ix e
arr (ix -> e) -> (ix -> ix) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix
transInner
    {-# INLINE newVal #-}
    !newsz :: Sz ix
newsz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (ix -> ix
transInner (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r' ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix e
arr)))
    !dix :: Dim
dix = Sz ix -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions Sz ix
newsz
{-# INLINE [1] transposeInner #-}

-- | Transpose outer two dimensions of at least rank-2 array.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XTypeApplications
-- >>> arr = makeArrayLinear @U Seq (Sz (2 :> 3 :. 4)) id
-- >>> arr
-- Array U Seq (Sz (2 :> 3 :. 4))
--   [ [ [ 0, 1, 2, 3 ]
--     , [ 4, 5, 6, 7 ]
--     , [ 8, 9, 10, 11 ]
--     ]
--   , [ [ 12, 13, 14, 15 ]
--     , [ 16, 17, 18, 19 ]
--     , [ 20, 21, 22, 23 ]
--     ]
--   ]
-- >>> transposeOuter arr
-- Array D Seq (Sz (2 :> 4 :. 3))
--   [ [ [ 0, 4, 8 ]
--     , [ 1, 5, 9 ]
--     , [ 2, 6, 10 ]
--     , [ 3, 7, 11 ]
--     ]
--   , [ [ 12, 16, 20 ]
--     , [ 13, 17, 21 ]
--     , [ 14, 18, 22 ]
--     , [ 15, 19, 23 ]
--     ]
--   ]
--
--
-- @since 0.1.0
transposeOuter :: (Index (Lower ix), Source r' ix e)
               => Array r' ix e -> Array D ix e
transposeOuter :: Array r' ix e -> Array D ix e
transposeOuter !Array r' ix e
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' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr) Sz ix
newsz ix -> e
newVal
  where
    transOuter :: c -> c
transOuter !c
ix =
      (SomeException -> c) -> (c -> c) -> Either SomeException c -> c
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> c
forall e a. Exception e => e -> a
throwImpossible c -> c
forall a. a -> a
id (Either SomeException c -> c) -> Either SomeException c -> c
forall a b. (a -> b) -> a -> b
$ do
        Int
n <- c -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM c
ix Dim
1
        Int
m <- c -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM c
ix Dim
2
        c
ix' <- c -> Dim -> Int -> Either SomeException c
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM c
ix Dim
1 Int
m
        c -> Dim -> Int -> Either SomeException c
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM c
ix' Dim
2 Int
n
    {-# INLINE transOuter #-}
    newVal :: ix -> e
newVal = Array r' ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r' ix e
arr (ix -> e) -> (ix -> ix) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix
forall c. Index c => c -> c
transOuter
    {-# INLINE newVal #-}
    !newsz :: Sz ix
newsz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (ix -> ix
forall c. Index c => c -> c
transOuter (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r' ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix e
arr)))
{-# INLINE [1] transposeOuter #-}

-- | Reverse an array along some dimension. Dimension supplied is checked at compile time.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> arr = makeArrayLinear Seq (Sz2 4 5) (+10) :: Array D Ix2 Int
-- >>> arr
-- Array D Seq (Sz (4 :. 5))
--   [ [ 10, 11, 12, 13, 14 ]
--   , [ 15, 16, 17, 18, 19 ]
--   , [ 20, 21, 22, 23, 24 ]
--   , [ 25, 26, 27, 28, 29 ]
--   ]
-- >>> A.reverse Dim1 arr
-- Array D Seq (Sz (4 :. 5))
--   [ [ 14, 13, 12, 11, 10 ]
--   , [ 19, 18, 17, 16, 15 ]
--   , [ 24, 23, 22, 21, 20 ]
--   , [ 29, 28, 27, 26, 25 ]
--   ]
-- >>> A.reverse Dim2 arr
-- Array D Seq (Sz (4 :. 5))
--   [ [ 25, 26, 27, 28, 29 ]
--   , [ 20, 21, 22, 23, 24 ]
--   , [ 15, 16, 17, 18, 19 ]
--   , [ 10, 11, 12, 13, 14 ]
--   ]
--
-- @since 0.4.1
reverse :: (IsIndexDimension ix n, Source r ix e) => Dimension n -> Array r ix e -> Array D ix e
reverse :: Dimension n -> Array r ix e -> Array D ix e
reverse Dimension n
dim = Dim -> Array r ix e -> Array D ix e
forall r ix e. Source r ix e => Dim -> Array r ix e -> Array D ix e
reverse' (Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE reverse #-}

-- | Similarly to `reverse`, flip an array along a particular dimension, but throws
-- `IndexDimensionException` for an incorrect dimension.
--
-- @since 0.4.1
reverseM :: (MonadThrow m, Source r ix e) => Dim -> Array r ix e -> m (Array D ix e)
reverseM :: Dim -> Array r ix e -> m (Array D ix e)
reverseM Dim
dim Array r ix e
arr = do
  let sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
  Int
k <- ix -> Dim -> m Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) Dim
dim
  Array D ix e -> m (Array D ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D ix e -> m (Array D ix e))
-> Array D ix e -> m (Array D ix e)
forall a b. (a -> b) -> a -> b
$ 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 ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
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 ->
    Array r ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ((Int, ix) -> ix
forall a b. (a, b) -> b
snd ((Int, ix) -> ix) -> (Int, ix) -> ix
forall a b. (a -> b) -> a -> b
$ ix -> Dim -> (Int -> Int) -> (Int, ix)
forall ix. Index ix => ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' ix
ix Dim
dim (\Int
i -> Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
{-# INLINE reverseM #-}

-- | Reverse an array along some dimension. Same as `reverseM`, but throws the
-- `IndexDimensionException` from pure code.
--
-- @since 0.4.1
reverse' :: Source r ix e => Dim -> Array r ix e -> Array D ix e
reverse' :: Dim -> Array r ix e -> Array D ix e
reverse' Dim
dim = (SomeException -> Array D ix e)
-> (Array D ix e -> Array D ix e)
-> Either SomeException (Array D ix e)
-> Array D ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D ix e
forall a e. Exception e => e -> a
throw Array D ix e -> Array D ix e
forall a. a -> a
id (Either SomeException (Array D ix e) -> Array D ix e)
-> (Array r ix e -> Either SomeException (Array D ix e))
-> Array r ix e
-> Array D ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> Array r ix e -> Either SomeException (Array D ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Source r ix e) =>
Dim -> Array r ix e -> m (Array D ix e)
reverseM Dim
dim
{-# INLINE reverse' #-}

-- | Rearrange elements of an array into a new one by using a function that maps indices of the
-- newly created one into the old one. This function can throw `IndexOutOfBoundsException`.
--
-- ===__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XTypeApplications
-- >>> arr = makeArrayLinear @D Seq (Sz (2 :> 3 :. 4)) id
-- >>> arr
-- Array D Seq (Sz (2 :> 3 :. 4))
--   [ [ [ 0, 1, 2, 3 ]
--     , [ 4, 5, 6, 7 ]
--     , [ 8, 9, 10, 11 ]
--     ]
--   , [ [ 12, 13, 14, 15 ]
--     , [ 16, 17, 18, 19 ]
--     , [ 20, 21, 22, 23 ]
--     ]
--   ]
-- >>> backpermuteM @U (Sz (4 :. 2)) (\(i :. j) -> j :> j :. i) arr
-- Array U Seq (Sz (4 :. 2))
--   [ [ 0, 16 ]
--   , [ 1, 17 ]
--   , [ 2, 18 ]
--   , [ 3, 19 ]
--   ]
--
-- @since 0.3.0
backpermuteM ::
     forall r ix e r' ix' m.
     (Mutable r ix e, Source r' ix' e, MonadUnliftIO m, PrimMonad m, MonadThrow m)
  => Sz ix -- ^ Size of the result array
  -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one.
  -> Array r' ix' e -- ^ Source array.
  -> m (Array r ix e)
backpermuteM :: Sz ix -> (ix -> ix') -> Array r' ix' e -> m (Array r ix e)
backpermuteM Sz ix
sz ix -> ix'
ixF !Array r' ix' e
arr = Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, PrimMonad m, Mutable r ix e) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray (Array r' ix' e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e
arr) Sz ix
sz (Array r' ix' e -> ix' -> m e
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r' ix' e
arr (ix' -> m e) -> (ix -> ix') -> ix -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix'
ixF)
{-# INLINE backpermuteM #-}

-- | Similar to `backpermuteM`, with a few notable differences:
--
-- * Creates a delayed array, instead of manifest, therefore it can be fused
-- * Respects computation strategy, so it can be parallelized
-- * Throws a runtime `IndexOutOfBoundsException` from pure code.
--
-- @since 0.3.0
backpermute' :: (Source r' ix' e, Index ix) =>
                Sz ix -- ^ Size of the result array
             -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one.
             -> Array r' ix' e -- ^ Source array.
             -> Array D ix e
backpermute' :: Sz ix -> (ix -> ix') -> Array r' ix' e -> Array D ix e
backpermute' Sz ix
sz ix -> ix'
ixF !Array r' ix' e
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' ix' e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e
arr) Sz ix
sz (Array r' ix' e -> ix' -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r' ix' e
arr (ix' -> e) -> (ix -> ix') -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix'
ixF)
{-# INLINE backpermute' #-}


-- | Append two arrays together along a particular dimension. Sizes of both arrays must match, with
-- an allowed exception of the dimension they are being appended along, otherwise `Nothing` is
-- returned.
--
-- ====__Examples__
--
-- Append two 2D arrays along both dimensions. Note that they do agree on inner dimensions.
--
-- >>> import Data.Massiv.Array
-- >>> arrA = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('A', i, j))
-- >>> arrB = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('B', i, j))
-- >>> appendM 1 arrA arrB
-- Array DL Seq (Sz (2 :. 6))
--   [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('B',0,0), ('B',0,1), ('B',0,2) ]
--   , [ ('A',1,0), ('A',1,1), ('A',1,2), ('B',1,0), ('B',1,1), ('B',1,2) ]
--   ]
-- >>> appendM 2 arrA arrB
-- Array DL Seq (Sz (4 :. 3))
--   [ [ ('A',0,0), ('A',0,1), ('A',0,2) ]
--   , [ ('A',1,0), ('A',1,1), ('A',1,2) ]
--   , [ ('B',0,0), ('B',0,1), ('B',0,2) ]
--   , [ ('B',1,0), ('B',1,1), ('B',1,2) ]
--   ]
--
-- Now appending arrays with different sizes:
--
-- >>> arrC = makeArrayR U Seq (Sz (2 :. 4)) (\(i :. j) -> ('C', i, j))
-- >>> appendM 1 arrA arrC
-- Array DL Seq (Sz (2 :. 7))
--   [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('C',0,0), ('C',0,1), ('C',0,2), ('C',0,3) ]
--   , [ ('A',1,0), ('A',1,1), ('A',1,2), ('C',1,0), ('C',1,1), ('C',1,2), ('C',1,3) ]
--   ]
-- >>> appendM 2 arrA arrC
-- *** Exception: SizeMismatchException: (Sz (2 :. 3)) vs (Sz (2 :. 4))
--
-- @since 0.3.0
appendM ::
     forall r1 r2 ix e m. (MonadThrow m, Source r1 ix e, Source r2 ix e)
  => Dim
  -> Array r1 ix e
  -> Array r2 ix e
  -> m (Array DL ix e)
appendM :: Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e)
appendM Dim
n !Array r1 ix e
arr1 !Array r2 ix e
arr2 = do
  let !sz1 :: Sz ix
sz1 = Array r1 ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix e
arr1
      !sz2 :: Sz ix
sz2 = Array r2 ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix e
arr2
  (Sz Int
k1, Sz (Lower ix)
szl1) <- Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM Sz ix
sz1 Dim
n
  (Sz Int
k2, Sz (Lower ix)
szl2) <- Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM Sz ix
sz2 Dim
n
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz (Lower ix)
szl1 Sz (Lower ix) -> Sz (Lower ix) -> Bool
forall a. Eq a => a -> a -> Bool
== Sz (Lower ix)
szl2) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ SizeException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> m ()) -> SizeException -> m ()
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz ix
sz1 Sz ix
sz2
  let !k1' :: Int
k1' = Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
k1
  Sz ix
newSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM Sz (Lower ix)
szl1 Dim
n (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int
k1' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
k2))
  let load :: Monad n =>
        Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()
      load :: Scheduler n ()
-> Int
-> (Int -> e -> n ())
-> (Int -> Sz Int -> e -> n ())
-> n ()
load Scheduler n ()
scheduler !Int
startAt Int -> e -> n ()
dlWrite Int -> Sz Int -> e -> n ()
_dlSet = do
        Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$
          ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> n ()) -> n ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ ix
forall ix. Index ix => ix
zeroIndex (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz1) (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((ix -> n ()) -> n ()) -> (ix -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
            Int -> e -> n ()
dlWrite (Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix) (Array r1 ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e
arr1 ix
ix)
        Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$
          ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> n ()) -> n ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ ix
forall ix. Index ix => ix
zeroIndex (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz2) (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((ix -> n ()) -> n ()) -> (ix -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
            let i :: Int
i = ix -> Dim -> Int
forall ix. Index ix => ix -> Dim -> Int
getDim' ix
ix Dim
n
                ix' :: ix
ix' = ix -> Dim -> Int -> ix
forall ix. Index ix => ix -> Dim -> Int -> ix
setDim' ix
ix Dim
n (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k1')
             in Int -> e -> n ()
dlWrite (Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') (Array r2 ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e
arr2 ix
ix)
      {-# INLINE load #-}
  Array DL ix e -> m (Array DL ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array DL ix e -> m (Array DL ix e))
-> Array DL ix e -> m (Array DL ix e)
forall a b. (a -> b) -> a -> b
$
    DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m ()
    -> Int
    -> (Int -> e -> m ())
    -> (Int -> Sz Int -> e -> m ())
    -> m ())
-> Array DL ix e
DLArray
      {dlComp :: Comp
dlComp = Array r1 ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r1 ix e
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r2 ix e
arr2, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
{-# INLINE appendM #-}


-- | Same as `appendM`, but will throw an exception in pure code on mismatched sizes.
--
-- @since 0.3.0
append' :: (Source r1 ix e, Source r2 ix e) =>
           Dim -> Array r1 ix e -> Array r2 ix e -> Array DL ix e
append' :: Dim -> Array r1 ix e -> Array r2 ix e -> Array DL ix e
append' Dim
dim Array r1 ix e
arr1 Array r2 ix e
arr2 = (SomeException -> Array DL ix e)
-> (Array DL ix e -> Array DL ix e)
-> Either SomeException (Array DL ix e)
-> Array DL ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array DL ix e
forall a e. Exception e => e -> a
throw Array DL ix e -> Array DL ix e
forall a. a -> a
id (Either SomeException (Array DL ix e) -> Array DL ix e)
-> Either SomeException (Array DL ix e) -> Array DL ix e
forall a b. (a -> b) -> a -> b
$ Dim
-> Array r1 ix e
-> Array r2 ix e
-> Either SomeException (Array DL ix e)
forall r1 r2 ix e (m :: * -> *).
(MonadThrow m, Source r1 ix e, Source r2 ix e) =>
Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e)
appendM Dim
dim Array r1 ix e
arr1 Array r2 ix e
arr2
{-# INLINE append' #-}

-- | Concat many arrays together along some dimension.
--
-- @since 0.3.0
concat' :: (Foldable f, Source r ix e) => Dim -> f (Array r ix e) -> Array DL ix e
concat' :: Dim -> f (Array r ix e) -> Array DL ix e
concat' Dim
n f (Array r ix e)
arrs = (SomeException -> Array DL ix e)
-> (Array DL ix e -> Array DL ix e)
-> Either SomeException (Array DL ix e)
-> Array DL ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array DL ix e
forall a e. Exception e => e -> a
throw Array DL ix e -> Array DL ix e
forall a. a -> a
id (Either SomeException (Array DL ix e) -> Array DL ix e)
-> Either SomeException (Array DL ix e) -> Array DL ix e
forall a b. (a -> b) -> a -> b
$ Dim -> f (Array r ix e) -> Either SomeException (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(MonadThrow m, Foldable f, Source r ix e) =>
Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM Dim
n f (Array r ix e)
arrs
{-# INLINE concat' #-}

-- | Concatenate many arrays together along some dimension. It is important that all sizes are
-- equal, with an exception of the dimensions along which concatenation happens.
--
-- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException`
--
-- @since 0.3.0
concatM ::
     forall r ix e f m. (MonadThrow m, Foldable f, Source r ix e)
  => Dim
  -> f (Array r ix e)
  -> m (Array DL ix e)
concatM :: Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM Dim
n !f (Array r ix e)
arrsF =
  case [Array r ix e] -> Maybe (Array r ix e, [Array r ix e])
forall a. [a] -> Maybe (a, [a])
L.uncons (f (Array r ix e) -> [Array r ix e]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList f (Array r ix e)
arrsF) of
    Maybe (Array r ix e, [Array r ix e])
Nothing -> Array DL ix e -> m (Array DL ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array DL ix e
forall r ix e. Construct r ix e => Array r ix e
empty
    Just (Array r ix e
a, [Array r ix e]
arrs) -> do
      let sz :: ix
sz = Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
a)
          szs :: [ix]
szs = Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Sz ix -> ix) -> (Array r ix e -> Sz ix) -> Array r ix e -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size (Array r ix e -> ix) -> [Array r ix e] -> [ix]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Array r ix e]
arrs
      (Int
k, Lower ix
szl) <- ix -> Dim -> m (Int, Lower ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
sz Dim
n
      -- / remove the dimension out of all sizes along which concatenation will happen
      ([Int]
ks, [Lower ix]
szls) <-
        (ix -> ([Int], [Lower ix]) -> m ([Int], [Lower ix]))
-> ([Int], [Lower ix]) -> [ix] -> m ([Int], [Lower ix])
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
F.foldrM (\ !ix
csz ([Int]
ks, [Lower ix]
szls) -> (Int -> [Int])
-> (Lower ix -> [Lower ix])
-> (Int, Lower ix)
-> ([Int], [Lower ix])
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ks) (Lower ix -> [Lower ix] -> [Lower ix]
forall a. a -> [a] -> [a]
: [Lower ix]
szls) ((Int, Lower ix) -> ([Int], [Lower ix]))
-> m (Int, Lower ix) -> m ([Int], [Lower ix])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ix -> Dim -> m (Int, Lower ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
csz Dim
n) ([], []) [ix]
szs
      -- / make sure to fail as soon as at least one of the arrays has a mismatching inner size
      ((ix, Lower ix) -> m ()) -> [(ix, Lower ix)] -> m ()
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
traverse_
        (\(ix
sz', Lower ix
_) -> SizeException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ix
sz) (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ix
sz')))
        (((ix, Lower ix) -> Bool) -> [(ix, Lower ix)] -> [(ix, Lower ix)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile ((Lower ix -> Lower ix -> Bool
forall a. Eq a => a -> a -> Bool
== Lower ix
szl) (Lower ix -> Bool)
-> ((ix, Lower ix) -> Lower ix) -> (ix, Lower ix) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix, Lower ix) -> Lower ix
forall a b. (a, b) -> b
snd) ([(ix, Lower ix)] -> [(ix, Lower ix)])
-> [(ix, Lower ix)] -> [(ix, Lower ix)]
forall a b. (a -> b) -> a -> b
$ [ix] -> [Lower ix] -> [(ix, Lower ix)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip [ix]
szs [Lower ix]
szls)
      let kTotal :: Sz Int
kTotal = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int -> Sz Int) -> Int -> Sz Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Int -> [Int] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
k [Int]
ks
      Sz ix
newSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
szl) Dim
n Sz Int
kTotal
      let load :: Monad n =>
            Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()
          load :: Scheduler n ()
-> Int
-> (Int -> e -> n ())
-> (Int -> Sz Int -> e -> n ())
-> n ()
load Scheduler n ()
scheduler Int
startAt Int -> e -> n ()
dlWrite Int -> Sz Int -> e -> n ()
_dlSet =
            let arrayLoader :: Int -> (Int, Array r ix e) -> n Int
arrayLoader !Int
kAcc (Int
kCur, Array r ix e
arr) = do
                  Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler (n () -> n ()) -> n () -> n ()
forall a b. (a -> b) -> a -> b
$
                    Array r ix e -> (ix -> e -> n ()) -> n ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
Array r ix a -> (ix -> a -> m b) -> m ()
iforM_ Array r ix e
arr ((ix -> e -> n ()) -> n ()) -> (ix -> e -> n ()) -> n ()
forall a b. (a -> b) -> a -> b
$ \ix
ix e
e ->
                      let i :: Int
i = ix -> Dim -> Int
forall ix. Index ix => ix -> Dim -> Int
getDim' ix
ix Dim
n
                          ix' :: ix
ix' = ix -> Dim -> Int -> ix
forall ix. Index ix => ix -> Dim -> Int -> ix
setDim' ix
ix Dim
n (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kAcc)
                       in Int -> e -> n ()
dlWrite (Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') e
e
                  Int -> n Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
kAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kCur)
             in (Int -> (Int, Array r ix e) -> n Int)
-> Int -> [(Int, Array r ix e)] -> n ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
M.foldM_ Int -> (Int, Array r ix e) -> n Int
arrayLoader Int
0 ([(Int, Array r ix e)] -> n ()) -> [(Int, Array r ix e)] -> n ()
forall a b. (a -> b) -> a -> b
$ (Int
k, Array r ix e
a) (Int, Array r ix e)
-> [(Int, Array r ix e)] -> [(Int, Array r ix e)]
forall a. a -> [a] -> [a]
: [Int] -> [Array r ix e] -> [(Int, Array r ix e)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip [Int]
ks [Array r ix e]
arrs
          {-# INLINE load #-}
      Array DL ix e -> m (Array DL ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array DL ix e -> m (Array DL ix e))
-> Array DL ix e -> m (Array DL ix e)
forall a b. (a -> b) -> a -> b
$
        DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m ()
    -> Int
    -> (Int -> e -> m ())
    -> (Int -> Sz Int -> e -> m ())
    -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = (Array r ix e -> Comp) -> f (Array r ix e) -> Comp
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp f (Array r ix e)
arrsF, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
{-# INLINE concatM #-}


-- | Stack slices on top of each other along the specified dimension.
--
-- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException`
--
-- ====__Examples__
--
-- Here are the three different ways to stack up two 2D Matrix pages into a 3D array.
--
-- >>> import Data.Massiv.Array as A
-- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int
-- >>> y = compute (iterateN 3 succ 9) :: Matrix P Int
-- >>> x
-- Array P Seq (Sz (3 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   , [ 7, 8, 9 ]
--   ]
-- >>> y
-- Array P Seq (Sz (3 :. 3))
--   [ [ 10, 11, 12 ]
--   , [ 13, 14, 15 ]
--   , [ 16, 17, 18 ]
--   ]
-- >>> stackSlicesM 1 [x, y] :: IO (Array DL Ix3 Int)
-- Array DL Seq (Sz (3 :> 3 :. 2))
--   [ [ [ 1, 10 ]
--     , [ 2, 11 ]
--     , [ 3, 12 ]
--     ]
--   , [ [ 4, 13 ]
--     , [ 5, 14 ]
--     , [ 6, 15 ]
--     ]
--   , [ [ 7, 16 ]
--     , [ 8, 17 ]
--     , [ 9, 18 ]
--     ]
--   ]
-- >>> stackSlicesM 2 [x, y] :: IO (Array DL Ix3 Int)
-- Array DL Seq (Sz (3 :> 2 :. 3))
--   [ [ [ 1, 2, 3 ]
--     , [ 10, 11, 12 ]
--     ]
--   , [ [ 4, 5, 6 ]
--     , [ 13, 14, 15 ]
--     ]
--   , [ [ 7, 8, 9 ]
--     , [ 16, 17, 18 ]
--     ]
--   ]
-- >>> stackSlicesM 3 [x, y] :: IO (Array DL Ix3 Int)
-- Array DL Seq (Sz (2 :> 3 :. 3))
--   [ [ [ 1, 2, 3 ]
--     , [ 4, 5, 6 ]
--     , [ 7, 8, 9 ]
--     ]
--   , [ [ 10, 11, 12 ]
--     , [ 13, 14, 15 ]
--     , [ 16, 17, 18 ]
--     ]
--   ]
--
-- @since 0.5.4
stackSlicesM ::
     forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix)
  => Dim
  -> f (Array r (Lower ix) e)
  -> m (Array DL ix e)
stackSlicesM :: Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
stackSlicesM Dim
dim !f (Array r (Lower ix) e)
arrsF = do
  case [Array r (Lower ix) e]
-> Maybe (Array r (Lower ix) e, [Array r (Lower ix) e])
forall a. [a] -> Maybe (a, [a])
L.uncons (f (Array r (Lower ix) e) -> [Array r (Lower ix) e]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList f (Array r (Lower ix) e)
arrsF) of
    Maybe (Array r (Lower ix) e, [Array r (Lower ix) e])
Nothing -> Array DL ix e -> m (Array DL ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array DL ix e
forall r ix e. Construct r ix e => Array r ix e
empty
    Just (Array r (Lower ix) e
a, [Array r (Lower ix) e]
arrs) -> do
      let sz :: Sz (Lower ix)
sz = Array r (Lower ix) e -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) e
a
          len :: Sz Int
len = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (f (Array r (Lower ix) e) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
F.length f (Array r (Lower ix) e)
arrsF)
      -- / make sure all arrays have the same size
      f (Array r (Lower ix) e) -> (Array r (Lower ix) e -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
M.forM_ f (Array r (Lower ix) e)
arrsF ((Array r (Lower ix) e -> m ()) -> m ())
-> (Array r (Lower ix) e -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Array r (Lower ix) e
arr ->
         Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz (Lower ix)
sz Sz (Lower ix) -> Sz (Lower ix) -> Bool
forall a. Eq a => a -> a -> Bool
== Array r (Lower ix) e -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) e
arr) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ SizeException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz (Lower ix) -> Sz (Lower ix) -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz (Lower ix)
sz (Array r (Lower ix) e -> Sz (Lower ix)
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r (Lower ix) e
arr))
      Sz ix
newSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM Sz (Lower ix)
sz Dim
dim Sz Int
len
      let load :: Monad n =>
            Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()
          load :: Scheduler n ()
-> Int
-> (Int -> e -> n ())
-> (Int -> Sz Int -> e -> n ())
-> n ()
load Scheduler n ()
scheduler Int
startAt Int -> e -> n ()
dlWrite Int -> Sz Int -> e -> n ()
_dlSet =
            let loadIndex :: Int -> Lower ix -> e -> n ()
loadIndex Int
k Lower ix
ix = Int -> e -> n ()
dlWrite (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz (Lower ix -> Dim -> Int -> ix
forall ix. Index ix => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ix Dim
dim Int
k) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt)
                arrayLoader :: Int -> Array r (Lower ix) e -> n Int
arrayLoader !Int
k Array r (Lower ix) e
arr = (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> n () -> n Int
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Scheduler n () -> n () -> n ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler n ()
scheduler ((Lower ix -> e -> n ()) -> Array r (Lower ix) e -> n ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(ix -> a -> m b) -> Array r ix a -> m ()
imapM_ (Int -> Lower ix -> e -> n ()
loadIndex Int
k) Array r (Lower ix) e
arr)
                {-# INLINE arrayLoader #-}
             in (Int -> Array r (Lower ix) e -> n Int)
-> Int -> f (Array r (Lower ix) e) -> n ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
M.foldM_ Int -> Array r (Lower ix) e -> n Int
arrayLoader Int
0 f (Array r (Lower ix) e)
arrsF
          {-# INLINE load #-}
      Array DL ix e -> m (Array DL ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array DL ix e -> m (Array DL ix e))
-> Array DL ix e -> m (Array DL ix e)
forall a b. (a -> b) -> a -> b
$
        DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m ()
    -> Int
    -> (Int -> e -> m ())
    -> (Int -> Sz Int -> e -> m ())
    -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = (Array r (Lower ix) e -> Comp) -> [Array r (Lower ix) e] -> Comp
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Array r (Lower ix) e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp [Array r (Lower ix) e]
arrs, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
{-# INLINE stackSlicesM #-}

-- | Specialized `stackSlicesM` to handling stacking from the outside. It is the inverse of
-- `Data.Massiv.Array.outerSlices`.
--
-- /__Exceptions__/: `SizeMismatchException`
--
-- ====__Examples__
--
-- In this example we stack vectors as row of a matrix from top to bottom:
--
-- >>> import Data.Massiv.Array as A
-- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int
-- >>> x
-- Array P Seq (Sz (3 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   , [ 7, 8, 9 ]
--   ]
-- >>> rows = outerSlices x
-- >>> A.mapM_ print rows
-- Array M Seq (Sz1 3)
--   [ 1, 2, 3 ]
-- Array M Seq (Sz1 3)
--   [ 4, 5, 6 ]
-- Array M Seq (Sz1 3)
--   [ 7, 8, 9 ]
-- >>> stackOuterSlicesM rows :: IO (Matrix DL Int)
-- Array DL Seq (Sz (3 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   , [ 7, 8, 9 ]
--   ]
--
-- @since 0.5.4
stackOuterSlicesM ::
     forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix)
  => f (Array r (Lower ix) e)
  -> m (Array DL ix e)
stackOuterSlicesM :: f (Array r (Lower ix) e) -> m (Array DL ix e)
stackOuterSlicesM = Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix) =>
Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
stackSlicesM (Proxy ix -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions (Proxy ix
forall k (t :: k). Proxy t
Proxy :: Proxy ix))
{-# INLINE stackOuterSlicesM #-}

-- | Specialized `stackSlicesM` to handling stacking from the inside. It is the inverse of
-- `Data.Massiv.Array.innerSlices`.
--
-- /__Exceptions__/: `SizeMismatchException`
--
-- ====__Examples__
--
-- In this example we stack vectors as columns of a matrix from left to right:
--
-- >>> import Data.Massiv.Array as A
-- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int
-- >>> x
-- Array P Seq (Sz (3 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   , [ 7, 8, 9 ]
--   ]
-- >>> columns = innerSlices x
-- >>> A.mapM_ print columns
-- Array M Seq (Sz1 3)
--   [ 1, 4, 7 ]
-- Array M Seq (Sz1 3)
--   [ 2, 5, 8 ]
-- Array M Seq (Sz1 3)
--   [ 3, 6, 9 ]
-- >>> stackInnerSlicesM columns :: IO (Matrix DL Int)
-- Array DL Seq (Sz (3 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   , [ 7, 8, 9 ]
--   ]
--
-- @since 0.5.4
stackInnerSlicesM ::
     forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix)
  => f (Array r (Lower ix) e)
  -> m (Array DL ix e)
stackInnerSlicesM :: f (Array r (Lower ix) e) -> m (Array DL ix e)
stackInnerSlicesM = Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix) =>
Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e)
stackSlicesM Dim
1
{-# INLINE stackInnerSlicesM #-}


-- | /O(1)/ - Split an array into two at an index along a specified dimension.
--
-- /Related/: `splitAt'`, `splitExtractM`, `Data.Massiv.Vector.sliceAt'`, `Data.Massiv.Vector.sliceAtM`
--
-- /__Exceptions__/: `IndexDimensionException`, `SizeSubregionException`
--
-- @since 0.3.0
splitAtM ::
     (MonadThrow m, Extract r ix e)
  => Dim -- ^ Dimension along which to split
  -> Int -- ^ Index along the dimension to split at
  -> Array r ix e -- ^ Source array
  -> m (Array (R r) ix e, Array (R r) ix e)
splitAtM :: Dim
-> Int -> Array r ix e -> m (Array (R r) ix e, Array (R r) ix e)
splitAtM Dim
dim Int
i Array r ix e
arr = do
  let Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
  ix
eIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
sz Dim
dim Int
i
  ix
sIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim Int
i
  Array (R r) ix e
arr1 <- ix -> ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> ix -> Array r ix e -> m (Array (R r) ix e)
extractFromToM ix
forall ix. Index ix => ix
zeroIndex ix
eIx Array r ix e
arr
  Array (R r) ix e
arr2 <- ix -> ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> ix -> Array r ix e -> m (Array (R r) ix e)
extractFromToM ix
sIx ix
sz Array r ix e
arr
  (Array (R r) ix e, Array (R r) ix e)
-> m (Array (R r) ix e, Array (R r) ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array (R r) ix e
arr1, Array (R r) ix e
arr2)
{-# INLINE splitAtM #-}

-- | /O(1)/ - Split an array into two at an index along a specified dimension. Throws an
-- error for a wrong dimension or incorrect indices.
--
-- /Related/: `splitAtM`, `splitExtractM`, `Data.Massiv.Vector.sliceAt'`, `Data.Massiv.Vector.sliceAtM`
--
-- ==== __Examples__
--
--
-- @since 0.1.0
splitAt' :: Extract r ix e =>
            Dim -> Int -> Array r ix e -> (Array (R r) ix e, Array (R r) ix e)
splitAt' :: Dim -> Int -> Array r ix e -> (Array (R r) ix e, Array (R r) ix e)
splitAt' Dim
dim Int
i Array r ix e
arr = (SomeException -> (Array (R r) ix e, Array (R r) ix e))
-> ((Array (R r) ix e, Array (R r) ix e)
    -> (Array (R r) ix e, Array (R r) ix e))
-> Either SomeException (Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> (Array (R r) ix e, Array (R r) ix e)
forall a e. Exception e => e -> a
throw (Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e)
forall a. a -> a
id (Either SomeException (Array (R r) ix e, Array (R r) ix e)
 -> (Array (R r) ix e, Array (R r) ix e))
-> Either SomeException (Array (R r) ix e, Array (R r) ix e)
-> (Array (R r) ix e, Array (R r) ix e)
forall a b. (a -> b) -> a -> b
$ Dim
-> Int
-> Array r ix e
-> Either SomeException (Array (R r) ix e, Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
Dim
-> Int -> Array r ix e -> m (Array (R r) ix e, Array (R r) ix e)
splitAtM Dim
dim Int
i Array r ix e
arr
{-# INLINE splitAt' #-}


-- | Split an array in three parts across some dimension
--
-- @since 0.3.5
splitExtractM ::
     (MonadThrow m, Extract r ix e, Source (R r) ix e)
  => Dim -- ^ Dimension along which to do the extraction
  -> Ix1 -- ^ Start index along the dimension that needs to be extracted
  -> Sz Ix1 -- ^ Size of the extracted array along the dimension that it will be extracted
  -> Array r ix e
  -> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
splitExtractM :: Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
splitExtractM Dim
dim Int
startIx1 (Sz Int
extractSzIx1) Array r ix e
arr = do
  let Sz ix
szIx = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
  ix
midStartIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim Int
startIx1
  ix
midExtractSzIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
szIx Dim
dim Int
extractSzIx1
  Array (R r) ix e
midArr <- ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
midStartIx (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
midExtractSzIx) Array r ix e
arr
  ix
leftArrSzIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
szIx Dim
dim Int
startIx1
  Array (R r) ix e
leftArr <- ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> Sz ix -> Array r ix e -> m (Array (R r) ix e)
extractM ix
forall ix. Index ix => ix
zeroIndex (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
leftArrSzIx) Array r ix e
arr
  ix
rightArrStartIx <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim (Int
startIx1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
extractSzIx1)
  Array (R r) ix e
rightArr <- ix -> ix -> Array r ix e -> m (Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e) =>
ix -> ix -> Array r ix e -> m (Array (R r) ix e)
extractFromToM ix
rightArrStartIx ix
szIx Array r ix e
arr
  (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array (R r) ix e
leftArr, Array (R r) ix e
midArr, Array (R r) ix e
rightArr)
{-# INLINE splitExtractM #-}



-- | Replace a slice of an array with another one
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> arr = makeArrayR U Seq (Sz3 3 4 5) fromIx3
-- >>> arr' = makeArrayR U Seq (Sz3 3 4 5) (fromIx3 . liftIndex (* 100))
-- >>> replaceSlice 2 1 (arr' <!> (2, 3)) arr
-- Array DL Seq (Sz (3 :> 4 :. 5))
--   [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3), (0,0,4) ]
--     , [ (0,300,0), (0,300,100), (0,300,200), (0,300,300), (0,300,400) ]
--     , [ (0,2,0), (0,2,1), (0,2,2), (0,2,3), (0,2,4) ]
--     , [ (0,3,0), (0,3,1), (0,3,2), (0,3,3), (0,3,4) ]
--     ]
--   , [ [ (1,0,0), (1,0,1), (1,0,2), (1,0,3), (1,0,4) ]
--     , [ (100,300,0), (100,300,100), (100,300,200), (100,300,300), (100,300,400) ]
--     , [ (1,2,0), (1,2,1), (1,2,2), (1,2,3), (1,2,4) ]
--     , [ (1,3,0), (1,3,1), (1,3,2), (1,3,3), (1,3,4) ]
--     ]
--   , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3), (2,0,4) ]
--     , [ (200,300,0), (200,300,100), (200,300,200), (200,300,300), (200,300,400) ]
--     , [ (2,2,0), (2,2,1), (2,2,2), (2,2,3), (2,2,4) ]
--     , [ (2,3,0), (2,3,1), (2,3,2), (2,3,3), (2,3,4) ]
--     ]
--   ]
--
-- @since 0.6.1
replaceSlice ::
     ( MonadThrow m
     , Extract r ix e
     , Source (R r) ix e
     , Load (R r) (Lower ix) e
     , Resize (R r) (Lower ix)
     )
  => Dim
  -> Ix1
  -> Array (R r) (Lower ix) e
  -> Array r ix e
  -> m (Array DL ix e)
replaceSlice :: Dim
-> Int
-> Array (R r) (Lower ix) e
-> Array r ix e
-> m (Array DL ix e)
replaceSlice Dim
dim Int
i Array (R r) (Lower ix) e
sl Array r ix e
arr = do
  (Array (R r) ix e
l, Array (R r) ix e
m, Array (R r) ix e
r) <- Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e) =>
Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
splitExtractM Dim
dim Int
i (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz Int
1) Array r ix e
arr
  Array (R r) ix e
m' <- Sz ix -> Array (R r) (Lower ix) e -> m (Array (R r) ix e)
forall (m :: * -> *) ix' r ix e.
(MonadThrow m, Index ix', Load r ix e, Resize r ix) =>
Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM (Array (R r) ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array (R r) ix e
m) Array (R r) (Lower ix) e
sl
  Dim -> [Array (R r) ix e] -> m (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(MonadThrow m, Foldable f, Source r ix e) =>
Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM Dim
dim [Array (R r) ix e
l, Array (R r) ix e
m', Array (R r) ix e
r]
{-# INLINE replaceSlice #-}


-- | Replace an outer slice of an array with another one
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> arr = makeArrayR U Seq (Sz3 3 4 5) fromIx3
-- >>> arr' = makeArrayR U Seq (Sz3 3 4 5) (fromIx3 . liftIndex (* 100))
-- >>> replaceOuterSlice 1 (arr' !> 2) arr
-- Array DL Seq (Sz (3 :> 4 :. 5))
--   [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3), (0,0,4) ]
--     , [ (0,1,0), (0,1,1), (0,1,2), (0,1,3), (0,1,4) ]
--     , [ (0,2,0), (0,2,1), (0,2,2), (0,2,3), (0,2,4) ]
--     , [ (0,3,0), (0,3,1), (0,3,2), (0,3,3), (0,3,4) ]
--     ]
--   , [ [ (200,0,0), (200,0,100), (200,0,200), (200,0,300), (200,0,400) ]
--     , [ (200,100,0), (200,100,100), (200,100,200), (200,100,300), (200,100,400) ]
--     , [ (200,200,0), (200,200,100), (200,200,200), (200,200,300), (200,200,400) ]
--     , [ (200,300,0), (200,300,100), (200,300,200), (200,300,300), (200,300,400) ]
--     ]
--   , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3), (2,0,4) ]
--     , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3), (2,1,4) ]
--     , [ (2,2,0), (2,2,1), (2,2,2), (2,2,3), (2,2,4) ]
--     , [ (2,3,0), (2,3,1), (2,3,2), (2,3,3), (2,3,4) ]
--     ]
--   ]
--
-- @since 0.6.1
replaceOuterSlice ::
     ( MonadThrow m
     , Extract r ix e
     , Source (R r) ix e
     , Load (R r) (Lower ix) e
     , Resize (R r) (Lower ix)
     )
  => Ix1
  -> Array (R r) (Lower ix) e
  -> Array r ix e
  -> m (Array DL ix e)
replaceOuterSlice :: Int
-> Array (R r) (Lower ix) e -> Array r ix e -> m (Array DL ix e)
replaceOuterSlice Int
i Array (R r) (Lower ix) e
sl Array r ix e
arr = Dim
-> Int
-> Array (R r) (Lower ix) e
-> Array r ix e
-> m (Array DL ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e,
 Load (R r) (Lower ix) e, Resize (R r) (Lower ix)) =>
Dim
-> Int
-> Array (R r) (Lower ix) e
-> Array r ix e
-> m (Array DL ix e)
replaceSlice (Sz ix -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)) Int
i Array (R r) (Lower ix) e
sl Array r ix e
arr
{-# INLINE replaceOuterSlice #-}


-- | Delete a region from an array along the specified dimension.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> arr = fromIx3 <$> (0 :> 0 :. 0 ..: 3 :> 2 :. 6)
-- >>> deleteRegionM 1 2 3 arr
-- Array DL Seq (Sz (3 :> 2 :. 3))
--   [ [ [ (0,0,0), (0,0,1), (0,0,5) ]
--     , [ (0,1,0), (0,1,1), (0,1,5) ]
--     ]
--   , [ [ (1,0,0), (1,0,1), (1,0,5) ]
--     , [ (1,1,0), (1,1,1), (1,1,5) ]
--     ]
--   , [ [ (2,0,0), (2,0,1), (2,0,5) ]
--     , [ (2,1,0), (2,1,1), (2,1,5) ]
--     ]
--   ]
-- >>> v = Ix1 0 ... 10
-- >>> deleteRegionM 1 3 5 v
-- Array DL Seq (Sz1 6)
--   [ 0, 1, 2, 8, 9, 10 ]
--
-- @since 0.3.5
deleteRegionM ::
     (MonadThrow m, Extract r ix e, Source (R r) ix e)
  => Dim -- ^ Along which axis should the removal happen
  -> Ix1 -- ^ At which index to start dropping slices
  -> Sz Ix1 -- ^ Number of slices to drop
  -> Array r ix e -- ^ Array that will have it's subarray removed
  -> m (Array DL ix e)
deleteRegionM :: Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRegionM Dim
dim Int
ix Sz Int
sz Array r ix e
arr = do
  (Array (R r) ix e
leftArr, Array (R r) ix e
_, Array (R r) ix e
rightArr) <- Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e) =>
Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e)
splitExtractM Dim
dim Int
ix Sz Int
sz Array r ix e
arr
  Dim -> Array (R r) ix e -> Array (R r) ix e -> m (Array DL ix e)
forall r1 r2 ix e (m :: * -> *).
(MonadThrow m, Source r1 ix e, Source r2 ix e) =>
Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e)
appendM Dim
dim Array (R r) ix e
leftArr Array (R r) ix e
rightArr
{-# INLINE deleteRegionM #-}

-- | Similar to `deleteRegionM`, but drop a specified number of rows from an array that
-- has at least 2 dimensions.
--
-- ====__Example__
--
-- >>> import Data.Massiv.Array
-- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6)
-- >>> arr
-- Array D Seq (Sz (3 :. 6))
--   [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ]
--   , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ]
--   , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ]
--   ]
-- >>> deleteRowsM 1 1 arr
-- Array DL Seq (Sz (2 :. 6))
--   [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ]
--   , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ]
--   ]
--
-- @since 0.3.5
deleteRowsM ::
     (MonadThrow m, Extract r ix e, Source (R r) ix e, Index (Lower ix))
  => Ix1
  -> Sz Ix1
  -> Array r ix e
  -> m (Array DL ix e)
deleteRowsM :: Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRowsM = Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e) =>
Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRegionM Dim
2
{-# INLINE deleteRowsM #-}

-- | Similar to `deleteRegionM`, but drop a specified number of columns an array.
--
-- ====__Example__
--
-- >>> import Data.Massiv.Array
-- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6)
-- >>> arr
-- Array D Seq (Sz (3 :. 6))
--   [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ]
--   , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ]
--   , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ]
--   ]
-- >>> deleteColumnsM 2 3 arr
-- Array DL Seq (Sz (3 :. 3))
--   [ [ (0,0), (0,1), (0,5) ]
--   , [ (1,0), (1,1), (1,5) ]
--   , [ (2,0), (2,1), (2,5) ]
--   ]
--
-- @since 0.3.5
deleteColumnsM ::
     (MonadThrow m, Extract r ix e, Source (R r) ix e)
  => Ix1
  -> Sz Ix1
  -> Array r ix e
  -> m (Array DL ix e)
deleteColumnsM :: Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteColumnsM = Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Extract r ix e, Source (R r) ix e) =>
Dim -> Int -> Sz Int -> Array r ix e -> m (Array DL ix e)
deleteRegionM Dim
1
{-# INLINE deleteColumnsM #-}


-- | Discard elements from the source array according to the stride.
--
-- @since 0.3.0
downsample ::
     forall r ix e. Source r ix e
  => Stride ix
  -> Array r ix e
  -> Array DL ix e
downsample :: Stride ix -> Array r ix e -> Array DL ix e
downsample Stride ix
stride Array r ix e
arr =
  DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m ()
    -> Int
    -> (Int -> e -> m ())
    -> (Int -> Sz Int -> e -> m ())
    -> m ())
-> Array DL ix e
DLArray {dlComp :: Comp
dlComp = Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr, dlSize :: Sz ix
dlSize = Sz ix
resultSize, dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load}
  where
    resultSize :: Sz ix
resultSize = Stride ix -> Sz ix -> Sz ix
forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride ix
stride (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)
    strideIx :: ix
strideIx = Stride ix -> ix
forall ix. Stride ix -> ix
unStride Stride ix
stride
    unsafeLinearWriteWithStride :: Int -> e
unsafeLinearWriteWithStride =
      Array r ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr (ix -> e) -> (Int -> ix) -> Int -> e
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
strideIx (ix -> ix) -> (Int -> ix) -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
resultSize
    {-# INLINE unsafeLinearWriteWithStride #-}
    load :: Monad m =>
      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()
    load :: Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
dlWrite Int -> Sz Int -> e -> m ()
_ =
      Scheduler m ()
-> Int -> Int -> (Int -> m e) -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *) b c.
Monad m =>
Scheduler m ()
-> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_
        Scheduler m ()
scheduler
        Int
startAt
        (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
resultSize)
        (e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> (Int -> e) -> Int -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
unsafeLinearWriteWithStride)
        Int -> e -> m ()
dlWrite
    {-# INLINE load #-}
{-# INLINE downsample #-}


-- | Insert the same element into a `Load`able array according to the supplied stride.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> arr = iterateN (Sz2 3 2) succ (0 :: Int)
-- >>> arr
-- Array DL Seq (Sz (3 :. 2))
--   [ [ 1, 2 ]
--   , [ 3, 4 ]
--   , [ 5, 6 ]
--   ]
-- >>> upsample 0 (Stride (2 :. 3)) arr
-- Array DL Seq (Sz (6 :. 6))
--   [ [ 1, 0, 0, 2, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0 ]
--   , [ 3, 0, 0, 4, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0 ]
--   , [ 5, 0, 0, 6, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0 ]
--   ]
-- >>> upsample 9 (Stride (1 :. 2)) arr
-- Array DL Seq (Sz (3 :. 4))
--   [ [ 1, 9, 2, 9 ]
--   , [ 3, 9, 4, 9 ]
--   , [ 5, 9, 6, 9 ]
--   ]
--
-- @since 0.3.0
upsample ::
     forall r ix e. Load r ix e
  => e -- ^ Element to use for filling the newly added cells
  -> Stride ix -- ^ Fill cells according to this stride
  -> Array r ix e -- ^ Array that will have cells added to
  -> Array DL ix e
upsample :: e -> Stride ix -> Array r ix e -> Array DL ix e
upsample !e
fillWith Stride ix
safeStride Array r ix e
arr =
  DLArray :: forall ix e.
Comp
-> Sz ix
-> (forall (m :: * -> *).
    Monad m =>
    Scheduler m ()
    -> Int
    -> (Int -> e -> m ())
    -> (Int -> Sz Int -> e -> m ())
    -> m ())
-> Array DL ix e
DLArray
    { dlComp :: Comp
dlComp = Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr
    , dlSize :: Sz ix
dlSize = Sz ix
newsz
    , dlLoad :: forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> e -> m ())
-> m ()
dlLoad = forall (m :: * -> *).
Monad m =>
Scheduler m ()
-> Int
-> (Int -> e -> m ())
-> (Int -> Sz Int -> 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 -> Sz Int -> e -> m ())
-> m ()
load Scheduler m ()
scheduler Int
startAt Int -> e -> m ()
uWrite Int -> Sz Int -> e -> m ()
uSet = do
      Int -> Sz Int -> e -> m ()
uSet Int
startAt (Sz ix -> Sz Int
forall ix. Index ix => Sz ix -> Sz Int
toLinearSz Sz ix
newsz) e
fillWith
      Scheduler m () -> Array r ix e -> (Int -> e -> m ()) -> m ()
forall r ix e (m :: * -> *).
(Load r ix e, Monad m) =>
Scheduler m () -> Array r ix e -> (Int -> e -> m ()) -> m ()
loadArrayM Scheduler m ()
scheduler Array r ix e
arr (\Int
i -> Int -> e -> m ()
uWrite (Int -> Int
adjustLinearStride (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
startAt)))
    {-# INLINE load #-}
    adjustLinearStride :: Int -> Int
adjustLinearStride = Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newsz (ix -> Int) -> (Int -> ix) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> ix
timesStride (ix -> ix) -> (Int -> ix) -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz
    {-# INLINE adjustLinearStride #-}
    timesStride :: ix -> ix
timesStride !ix
ix = (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
stride ix
ix
    {-# INLINE timesStride #-}
    !stride :: ix
stride = Stride ix -> ix
forall ix. Stride ix -> ix
unStride Stride ix
safeStride
    !sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
    !newsz :: Sz ix
newsz = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> ix
timesStride (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz)
{-# INLINE upsample #-}


-- | General array transformation, that forces computation and produces a manifest array.
--
-- @since 0.3.0
transformM ::
     forall r ix e r' ix' e' a m.
     (Mutable r ix e, Source r' ix' e', MonadUnliftIO m, PrimMonad m, MonadThrow m)
  => (Sz ix' -> m (Sz ix, a))
  -> (a -> (ix' -> m e') -> ix -> m e)
  -> Array r' ix' e'
  -> m (Array r ix e)
transformM :: (Sz ix' -> m (Sz ix, a))
-> (a -> (ix' -> m e') -> ix -> m e)
-> Array r' ix' e'
-> m (Array r ix e)
transformM Sz ix' -> m (Sz ix, a)
getSzM a -> (ix' -> m e') -> ix -> m e
getM Array r' ix' e'
arr = do
  (Sz ix
sz, a
a) <- Sz ix' -> m (Sz ix, a)
getSzM (Array r' ix' e' -> Sz ix'
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix' e'
arr)
  Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, PrimMonad m, Mutable r ix e) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray (Array r' ix' e' -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e'
arr) Sz ix
sz (a -> (ix' -> m e') -> ix -> m e
getM a
a (Array r' ix' e' -> ix' -> m e'
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r' ix' e'
arr))
{-# INLINE transformM #-}


-- | General array transformation
--
-- @since 0.3.0
transform' ::
     (Source r' ix' e', Index ix)
  => (Sz ix' -> (Sz ix, a))
  -> (a -> (ix' -> e') -> ix -> e)
  -> Array r' ix' e'
  -> Array D ix e
transform' :: (Sz ix' -> (Sz ix, a))
-> (a -> (ix' -> e') -> ix -> e) -> Array r' ix' e' -> Array D ix e
transform' Sz ix' -> (Sz ix, a)
getSz a -> (ix' -> e') -> ix -> e
get Array r' ix' e'
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' ix' e' -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix' e'
arr) Sz ix
sz (a -> (ix' -> e') -> ix -> e
get a
a (Array r' ix' e' -> ix' -> e'
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r' ix' e'
arr))
  where
    (Sz ix
sz, a
a) = Sz ix' -> (Sz ix, a)
getSz (Array r' ix' e' -> Sz ix'
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix' e'
arr)
{-# INLINE transform' #-}

-- | Same as `transformM`, but operates on two arrays
--
-- @since 0.3.0
transform2M ::
     (Mutable r ix e, Source r1 ix1 e1, Source r2 ix2 e2, MonadUnliftIO m, PrimMonad m, MonadThrow m)
  => (Sz ix1 -> Sz ix2 -> m (Sz ix, a))
  -> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e)
  -> Array r1 ix1 e1
  -> Array r2 ix2 e2
  -> m (Array r ix e)
transform2M :: (Sz ix1 -> Sz ix2 -> m (Sz ix, a))
-> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e)
-> Array r1 ix1 e1
-> Array r2 ix2 e2
-> m (Array r ix e)
transform2M Sz ix1 -> Sz ix2 -> m (Sz ix, a)
getSzM a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e
getM Array r1 ix1 e1
arr1 Array r2 ix2 e2
arr2 = do
  (Sz ix
sz, a
a) <- Sz ix1 -> Sz ix2 -> m (Sz ix, a)
getSzM (Array r1 ix1 e1 -> Sz ix1
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> Sz ix2
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix2 e2
arr2)
  Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, PrimMonad m, Mutable r ix e) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray (Array r1 ix1 e1 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r1 ix1 e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix2 e2 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r2 ix2 e2
arr2) Sz ix
sz (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e
getM a
a (Array r1 ix1 e1 -> ix1 -> m e1
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> ix2 -> m e2
forall r ix e (m :: * -> *).
(Source r ix e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r2 ix2 e2
arr2))
{-# INLINE transform2M #-}


-- | Same as `transform'`, but operates on two arrays
--
-- @since 0.3.0
transform2' ::
     (Source r1 ix1 e1, Source r2 ix2 e2, Index ix)
  => (Sz ix1 -> Sz ix2 -> (Sz ix, a))
  -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e)
  -> Array r1 ix1 e1
  -> Array r2 ix2 e2
  -> Array D ix e
transform2' :: (Sz ix1 -> Sz ix2 -> (Sz ix, a))
-> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e)
-> Array r1 ix1 e1
-> Array r2 ix2 e2
-> Array D ix e
transform2' Sz ix1 -> Sz ix2 -> (Sz ix, a)
getSz a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e
get Array r1 ix1 e1
arr1 Array r2 ix2 e2
arr2 =
  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 r1 ix1 e1 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r1 ix1 e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix2 e2 -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r2 ix2 e2
arr2) Sz ix
sz (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e
get a
a (Array r1 ix1 e1 -> ix1 -> e1
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> ix2 -> e2
forall r ix e. Source r ix e => Array r ix e -> ix -> e
evaluate' Array r2 ix2 e2
arr2))
  where
    (Sz ix
sz, a
a) = Sz ix1 -> Sz ix2 -> (Sz ix, a)
getSz (Array r1 ix1 e1 -> Sz ix1
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> Sz ix2
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r2 ix2 e2
arr2)
{-# INLINE transform2' #-}



-- | Replicate each element of the array by a factor in stride along each dimension and surround each
-- such group with a box of supplied grid value. It will essentially zoom up an array and create a
-- grid around each element from the original array. Very useful for zooming up images to inspect
-- individual pixels.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> arr = resize' (Sz2 3 2) (Ix1 1 ... 6)
-- >>> arr
-- Array D Seq (Sz (3 :. 2))
--   [ [ 1, 2 ]
--   , [ 3, 4 ]
--   , [ 5, 6 ]
--   ]
-- >>> zoomWithGrid 0 (Stride (2 :. 3)) arr
-- Array DL Seq (Sz (10 :. 9))
--   [ [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ]
--   , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ]
--   , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ]
--   , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
--   ]
--
-- @since 0.3.1
zoomWithGrid ::
     forall r ix e. Source r ix e
  => e -- ^ Value to use for the grid
  -> Stride ix -- ^ Scaling factor
  -> Array r ix e -- ^ Source array
  -> Array DL ix e
zoomWithGrid :: e -> Stride ix -> Array r ix e -> Array DL ix e
zoomWithGrid e
gridVal (Stride ix
zoomFactor) Array r ix e
arr = 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
Seq Sz ix
newSz (e -> Maybe e
forall a. a -> Maybe a
Just e
gridVal) forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
  where
    !kx :: ix
kx = (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
zoomFactor
    !lastNewIx :: ix
lastNewIx = (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
kx (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)
    !newSz :: Sz ix
newSz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((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
lastNewIx)
    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
    load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
_ Int -> e -> m ()
writeElement =
      Scheduler m () -> Array r ix e -> (ix -> e -> m ()) -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
Scheduler m () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler m ()
scheduler Array r ix e
arr ((ix -> e -> m ()) -> m ()) -> (ix -> e -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !ix
ix !e
e ->
        let !kix :: ix
kix = (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
ix ix
kx
         in (ix -> m ()) -> Array D ix ix -> m ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_ (\ !ix
ix' -> Int -> e -> m ()
writeElement (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') e
e) (Array D ix ix -> m ()) -> Array D ix ix -> m ()
forall a b. (a -> b) -> a -> b
$
            Comp -> ix -> ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range Comp
Seq ((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
kix) ((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
kix ix
kx)
    {-# INLINE load #-}
{-# INLINE zoomWithGrid #-}

-- | Increaze the size of the array accoridng to the stride multiplier while replicating
-- the same element to fill the neighbors. It is exactly the same as `zoomWithGrid`, but
-- without the grid.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> arr = resize' (Sz3 1 3 2) (Ix1 1 ... 6)
-- >>> arr
-- Array D Seq (Sz (1 :> 3 :. 2))
--   [ [ [ 1, 2 ]
--     , [ 3, 4 ]
--     , [ 5, 6 ]
--     ]
--   ]
-- >>> zoom (Stride (2 :> 2 :. 3)) arr
-- Array DL Seq (Sz (2 :> 6 :. 6))
--   [ [ [ 1, 1, 1, 2, 2, 2 ]
--     , [ 1, 1, 1, 2, 2, 2 ]
--     , [ 3, 3, 3, 4, 4, 4 ]
--     , [ 3, 3, 3, 4, 4, 4 ]
--     , [ 5, 5, 5, 6, 6, 6 ]
--     , [ 5, 5, 5, 6, 6, 6 ]
--     ]
--   , [ [ 1, 1, 1, 2, 2, 2 ]
--     , [ 1, 1, 1, 2, 2, 2 ]
--     , [ 3, 3, 3, 4, 4, 4 ]
--     , [ 3, 3, 3, 4, 4, 4 ]
--     , [ 5, 5, 5, 6, 6, 6 ]
--     , [ 5, 5, 5, 6, 6, 6 ]
--     ]
--   ]
--
-- @since 0.4.4
zoom ::
     forall r ix e. Source r ix e
  => Stride ix -- ^ Scaling factor
  -> Array r ix e -- ^ Source array
  -> Array DL ix e
zoom :: Stride ix -> Array r ix e -> Array DL ix e
zoom (Stride ix
zoomFactor) Array r ix e
arr = 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
Seq Sz ix
newSz Maybe e
forall a. Maybe a
Nothing forall (m :: * -> *).
Monad m =>
Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load
  where
    !lastNewIx :: ix
lastNewIx = (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
zoomFactor (ix -> ix) -> ix -> ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)
    !newSz :: Sz ix
newSz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
lastNewIx
    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
    load :: Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()
load Scheduler m ()
scheduler Int
_ Int -> e -> m ()
writeElement =
      Scheduler m () -> Array r ix e -> (ix -> e -> m ()) -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
Scheduler m () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler m ()
scheduler Array r ix e
arr ((ix -> e -> m ()) -> m ()) -> (ix -> e -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !ix
ix !e
e ->
        let !kix :: ix
kix = (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
ix ix
zoomFactor
         in (ix -> m ()) -> Array D ix ix -> m ()
forall r ix a (m :: * -> *) b.
(Source r ix a, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_ (\ !ix
ix' -> Int -> e -> m ()
writeElement (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz ix
ix') e
e) (Array D ix ix -> m ()) -> Array D ix ix -> m ()
forall a b. (a -> b) -> a -> b
$
            Comp -> ix -> ix -> Array D ix ix
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range Comp
Seq ix
kix ((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
kix ix
zoomFactor)
    {-# INLINE load #-}
{-# INLINE zoom #-}