{-# 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-2022
-- 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.Monad as M (foldM_, forM_, unless)
import Control.Monad.ST
import Control.Scheduler (traverse_)
import Data.Bifunctor (bimap)
import Data.Foldable as F (foldl', foldrM, length, toList)
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
import Data.Massiv.Core.Common -- (size, unsafeIndex, unsafeResize, evaluate', evaluateM)
import Prelude as P hiding (
  concat,
  drop,
  mapM_,
  reverse,
  splitAt,
  take,
  traverse,
 )

-- | 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.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> m <- resizeM (Sz (3 :. 3)) $ Ix1 1 ... 9
-- >>> m
-- Array D Seq (Sz (3 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   , [ 7, 8, 9 ]
--   ]
-- >>> extractM (0 :. 1) (Sz (2 :. 2)) m
-- Array D Seq (Sz (2 :. 2))
--   [ [ 2, 3 ]
--   , [ 5, 6 ]
--   ]
-- >>> a <- resizeM (Sz (3 :> 2 :. 4)) $ Ix1 11 ... 34
-- >>> a
-- Array D Seq (Sz (3 :> 2 :. 4))
--   [ [ [ 11, 12, 13, 14 ]
--     , [ 15, 16, 17, 18 ]
--     ]
--   , [ [ 19, 20, 21, 22 ]
--     , [ 23, 24, 25, 26 ]
--     ]
--   , [ [ 27, 28, 29, 30 ]
--     , [ 31, 32, 33, 34 ]
--     ]
--   ]
-- >>> extractM (0 :> 1 :. 1) (Sz (3 :> 1 :. 2)) a
-- Array D Seq (Sz (3 :> 1 :. 2))
--   [ [ [ 16, 17 ]
--     ]
--   , [ [ 24, 25 ]
--     ]
--   , [ [ 32, 33 ]
--     ]
--   ]
--
-- @since 0.3.0
extractM
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r e)
  => ix
  -- ^ Starting index
  -> Sz ix
  -- ^ Size of the resulting array
  -> Array r ix e
  -- ^ Source array
  -> m (Array D ix e)
extractM :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> Sz ix -> Array r ix e -> m (Array D 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 D ix e -> m (Array D ix e)
forall a. a -> m a
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
$ ix -> Sz ix -> Array r ix e -> Array D ix e
forall r e ix.
(Source r e, Index ix) =>
ix -> Sz ix -> Array r ix e -> Array D ix e
unsafeExtract ix
sIx Sz ix
newSz Array r ix e
arr
  | Bool
otherwise = SizeException -> m (Array D ix e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Array D ix e))
-> SizeException -> m (Array D 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. Size r => Array r ix e -> Sz ix
forall 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. Size r => Array r ix e -> Sz ix
forall 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'
  :: forall r ix e
   . (HasCallStack, Index ix, Source r e)
  => ix
  -- ^ Starting index
  -> Sz ix
  -- ^ Size of the resulting array
  -> Array r ix e
  -- ^ Source array
  -> Array D ix e
extract' :: forall r ix e.
(HasCallStack, Index ix, Source r e) =>
ix -> Sz ix -> Array r ix e -> Array D ix e
extract' ix
sIx Sz ix
newSz = Either SomeException (Array D ix e) -> Array D ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (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
. ix -> Sz ix -> Array r ix e -> Either SomeException (Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> Sz ix -> Array r ix e -> m (Array D 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.
--
-- ====__Examples__
--
-- >>> a <- resizeM (Sz (3 :> 2 :. 4)) $ Ix1 11 ... 34
-- >>> a
-- Array D Seq (Sz (3 :> 2 :. 4))
--   [ [ [ 11, 12, 13, 14 ]
--     , [ 15, 16, 17, 18 ]
--     ]
--   , [ [ 19, 20, 21, 22 ]
--     , [ 23, 24, 25, 26 ]
--     ]
--   , [ [ 27, 28, 29, 30 ]
--     , [ 31, 32, 33, 34 ]
--     ]
--   ]
-- >>> extractFromToM (1 :> 0 :. 1) (3 :> 2 :. 4) a
-- Array D Seq (Sz (2 :> 2 :. 3))
--   [ [ [ 20, 21, 22 ]
--     , [ 24, 25, 26 ]
--     ]
--   , [ [ 28, 29, 30 ]
--     , [ 32, 33, 34 ]
--     ]
--   ]
--
-- @since 0.3.0
extractFromToM
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r e)
  => ix
  -- ^ Starting index
  -> ix
  -- ^ Index up to which elements should be extracted.
  -> Array r ix e
  -- ^ Source array.
  -> m (Array D ix e)
extractFromToM :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> ix -> Array r ix e -> m (Array D ix e)
extractFromToM ix
sIx ix
eIx = ix -> Sz ix -> Array r ix e -> m (Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> Sz ix -> Array r ix e -> m (Array D 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 `extractFromToM`, but throws an error on invalid indices.
--
-- @since 0.2.4
extractFromTo'
  :: forall r ix e
   . (HasCallStack, Index ix, Source r e)
  => ix
  -- ^ Starting index
  -> ix
  -- ^ Index up to which elmenets should be extracted.
  -> Array r ix e
  -- ^ Source array.
  -> Array D ix e
extractFromTo' :: forall r ix e.
(HasCallStack, Index ix, Source r e) =>
ix -> ix -> Array r ix e -> Array D ix e
extractFromTo' ix
sIx ix
eIx = ix -> Sz ix -> Array r ix e -> Array D ix e
forall r ix e.
(HasCallStack, Index ix, Source r e) =>
ix -> Sz ix -> Array r ix e -> Array D ix e
extract' ix
sIx (Sz ix -> Array r ix e -> Array D ix e)
-> Sz ix -> Array r ix e -> Array D 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)/ - Change the size of an array. Throws
-- `SizeElementsMismatchException` if total number of elements does not match
-- the supplied array.
--
-- @since 0.3.0
resizeM
  :: forall r ix ix' e m
   . (MonadThrow m, Index ix', Index ix, Size r)
  => Sz ix'
  -> Array r ix e
  -> m (Array r ix' e)
resizeM :: forall r ix ix' e (m :: * -> *).
(MonadThrow m, Index ix', Index ix, Size r) =>
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. Size r => Array r ix e -> Sz ix
forall 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 a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Array r ix' e -> m (Array r ix' e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Sz ix' -> Array r ix e -> Array r ix' e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index 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'
  :: forall r ix ix' e
   . (HasCallStack, Index ix', Index ix, Size r)
  => Sz ix'
  -> Array r ix e
  -> Array r ix' e
resize' :: forall r ix ix' e.
(HasCallStack, Index ix', Index ix, Size r) =>
Sz ix' -> Array r ix e -> Array r ix' e
resize' Sz ix'
sz = Either SomeException (Array r ix' e) -> Array r ix' e
forall a. HasCallStack => Either SomeException a -> a
throwEither (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 r ix ix' e (m :: * -> *).
(MonadThrow m, Index ix', Index ix, Size r) =>
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 :: forall r ix e. (Index ix, Size r) => Array r ix e -> Vector r e
flatten :: forall r ix e. (Index ix, Size r) => Array r ix e -> Vector r e
flatten Array r ix e
arr = Sz Int -> Array r ix e -> Array r Int e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index 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. Size r => Array r ix e -> Sz ix
forall 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 :: forall r e. Source r e => Matrix r e -> Matrix D e
transpose :: forall r e. Source r e => Matrix r e -> Matrix D e
transpose = Array r Ix2 e -> Array D Ix2 e
forall r ix e.
(Index (Lower ix), Index ix, Source r 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
  :: forall r ix e
   . (Index (Lower ix), Index ix, Source r e)
  => Array r ix e
  -> Array D ix e
transposeInner :: forall r ix e.
(Index (Lower ix), Index ix, Source r e) =>
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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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. (HasCallStack, 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
forall (m :: * -> *). 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
forall (m :: * -> *). 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
forall (m :: * -> *). 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
forall (m :: * -> *). 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 ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => 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. Size r => Array r ix e -> Sz ix
forall 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
forall (proxy :: * -> *). 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
  :: forall r ix e
   . (Index (Lower ix), Index ix, Source r e)
  => Array r ix e
  -> Array D ix e
transposeOuter :: forall r ix e.
(Index (Lower ix), Index ix, Source r e) =>
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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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. (HasCallStack, 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
forall (m :: * -> *). MonadThrow m => c -> 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
forall (m :: * -> *). MonadThrow m => c -> 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
forall (m :: * -> *). MonadThrow m => c -> Dim -> Int -> m c
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
forall (m :: * -> *). MonadThrow m => c -> Dim -> Int -> m c
setDimM c
ix' Dim
2 Int
n
    {-# INLINE transOuter #-}
    newVal :: ix -> e
newVal = Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => 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. Size r => Array r ix e -> Sz ix
forall 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
  :: forall n r ix e
   . (IsIndexDimension ix n, Index ix, Source r e)
  => Dimension n
  -> Array r ix e
  -> Array D ix e
reverse :: forall (n :: Natural) r ix e.
(IsIndexDimension ix n, Index ix, Source r e) =>
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.
(HasCallStack, Index ix, Source r e) =>
Dim -> Array r ix e -> Array D ix e
reverse' (Dimension n -> Dim
forall (n :: Natural). 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
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r e)
  => Dim
  -> Array r ix e
  -> m (Array D ix e)
reverseM :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
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. Size r => Array r ix e -> Sz ix
forall 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
forall (m :: * -> *). 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 a. a -> m a
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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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 ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => 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.
(HasCallStack, 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'
  :: forall r ix e
   . (HasCallStack, Index ix, Source r e)
  => Dim
  -> Array r ix e
  -> Array D ix e
reverse' :: forall r ix e.
(HasCallStack, Index ix, Source r e) =>
Dim -> Array r ix e -> Array D ix e
reverse' Dim
dim = Either SomeException (Array D ix e) -> Array D ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (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 r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r 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
   . (Manifest r e, Index ix, Source r' e, Index ix', 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 :: forall r ix e r' ix' (m :: * -> *).
(Manifest r e, Index ix, Source r' e, Index ix', MonadUnliftIO m,
 PrimMonad m, MonadThrow m) =>
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, Manifest r e, Index ix) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray (Array r' ix' e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r' ix e -> Comp
getComp Array r' ix' e
arr) Sz ix
sz (Array r' ix' e -> ix' -> m e
forall ix r e (m :: * -> *).
(Index ix, Source r 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'
  :: forall r ix ix' e
   . (HasCallStack, Source r e, Index ix, 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' :: forall r ix ix' e.
(HasCallStack, Source r e, Index ix, Index ix') =>
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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix'
sz (Array r ix e -> ix -> e
forall ix r e.
(HasCallStack, Index ix, Source r 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, Index ix, Source r1 e, Source r2 e)
  => Dim
  -> Array r1 ix e
  -> Array r2 ix e
  -> m (Array DL ix e)
appendM :: forall r1 r2 ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r1 e, Source r2 e) =>
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. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix e
arr1
      !sz2 :: Sz ix
sz2 = Array r2 ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r2 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 e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, 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 :: Loader e
      load :: Loader e
load Scheduler s ()
scheduler !Int
startAt Int -> e -> ST s ()
dlWrite Int -> Sz Int -> e -> ST s ()
_dlSet = do
        Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
          ix
-> ix -> ix -> (Int -> Int -> Bool) -> (ix -> ST s ()) -> ST s ()
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ 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 -> ST s ()) -> ST s ()) -> (ix -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
            Int -> e -> ST s ()
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 ix. Index ix => Array r1 ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e
arr1 ix
ix)
        Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
          ix
-> ix -> ix -> (Int -> Int -> Bool) -> (ix -> ST s ()) -> ST s ()
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ 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 -> ST s ()) -> ST s ()) -> (ix -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
            let i :: Int
i = ix -> Dim -> Int
forall ix. (HasCallStack, Index ix) => ix -> Dim -> Int
getDim' ix
ix Dim
n
                ix' :: ix
ix' = ix -> Dim -> Int -> ix
forall ix. (HasCallStack, 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 -> ST s ()
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 ix. Index ix => Array r2 ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => 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 a. a -> m a
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
      { dlComp :: Comp
dlComp = Array r1 ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r1 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. Strategy r => Array r ix e -> Comp
forall ix e. Array r2 ix e -> Comp
getComp Array r2 ix e
arr2
      , dlSize :: Sz ix
dlSize = Sz ix
newSz
      , dlLoad :: Loader e
dlLoad = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz Int -> e -> ST s ())
-> ST s ()
Loader e
load
      }
{-# INLINE appendM #-}

-- | Same as `appendM`, but will throw an exception in pure code on mismatched sizes.
--
-- @since 0.3.0
append'
  :: forall r1 r2 ix e
   . (HasCallStack, Index ix, Source r1 e, Source r2 e)
  => Dim
  -> Array r1 ix e
  -> Array r2 ix e
  -> Array DL ix e
append' :: forall r1 r2 ix e.
(HasCallStack, Index ix, Source r1 e, Source r2 e) =>
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 = Either SomeException (Array DL ix e) -> Array DL ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (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, Index ix, Source r1 e, Source r2 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'
  :: forall f r ix e
   . (HasCallStack, Foldable f, Index ix, Source r e)
  => Dim
  -> f (Array r ix e)
  -> Array DL ix e
concat' :: forall (f :: * -> *) r ix e.
(HasCallStack, Foldable f, Index ix, Source r e) =>
Dim -> f (Array r ix e) -> Array DL ix e
concat' Dim
n = Either SomeException (Array DL ix e) -> Array DL ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array DL ix e) -> Array DL ix e)
-> (f (Array r ix e) -> Either SomeException (Array DL ix e))
-> f (Array r ix e)
-> Array DL ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> f (Array r ix e) -> Either SomeException (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(MonadThrow m, Foldable f, Index ix, Source r e) =>
Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM Dim
n
{-# 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, Index ix, Source r e)
  => Dim
  -> f (Array r ix e)
  -> m (Array DL ix e)
concatM :: forall r ix e (f :: * -> *) (m :: * -> *).
(MonadThrow m, Foldable f, Index ix, Source r e) =>
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 a. f a -> [a]
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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array DL ix e
forall r ix e. Load 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. Size r => Array r ix e -> Sz ix
forall 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. Size r => Array r ix e -> Sz ix
forall 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)
forall (m :: * -> *).
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 a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
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)
forall (m :: * -> *).
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 e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, 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 b a. (b -> a -> b) -> b -> [a] -> b
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 :: Loader e
          load :: Loader e
load Scheduler s ()
scheduler Int
startAt Int -> e -> ST s ()
dlWrite Int -> Sz Int -> e -> ST s ()
_dlSet =
            let arrayLoader :: Int -> (Int, Array r ix e) -> ST s Int
arrayLoader !Int
kAcc (!Int
kCur, Array r ix e
arr) = do
                  Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
                    Array r ix e -> (ix -> e -> ST s ()) -> ST s ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
Array r ix a -> (ix -> a -> m b) -> m ()
iforM_ Array r ix e
arr ((ix -> e -> ST s ()) -> ST s ())
-> (ix -> e -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ix
ix e
e -> do
                      Int
i <- ix -> Dim -> ST s Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
forall (m :: * -> *). MonadThrow m => ix -> Dim -> m Int
getDimM ix
ix Dim
n
                      ix
ix' <- ix -> Dim -> Int -> ST s ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
forall (m :: * -> *). MonadThrow m => ix -> Dim -> Int -> m ix
setDimM ix
ix Dim
n (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kAcc)
                      Int -> e -> ST s ()
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 -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$! Int
kAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kCur
                {-# INLINE arrayLoader #-}
             in (Int -> (Int, Array r ix e) -> ST s Int)
-> Int -> [(Int, Array r ix e)] -> ST s ()
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) -> ST s Int
arrayLoader Int
0 ([(Int, Array r ix e)] -> ST s ())
-> [(Int, Array r ix e)] -> ST s ()
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 a. a -> m a
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{dlComp :: Comp
dlComp = Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
a Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> (Array r ix e -> Comp) -> [Array r ix e] -> Comp
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp [Array r ix e]
arrs, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: Loader e
dlLoad = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz Int -> e -> ST s ())
-> ST s ()
Loader e
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, Index (Lower ix), Source r e, Index ix)
  => Dim
  -> f (Array r (Lower ix) e)
  -> m (Array DL ix e)
stackSlicesM :: forall r ix e (f :: * -> *) (m :: * -> *).
(Foldable f, MonadThrow m, Index (Lower ix), Source r e,
 Index ix) =>
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 a. f a -> [a]
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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array DL ix e
forall r ix e. Load 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. Size r => Array r ix e -> Sz ix
forall 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 a. f a -> 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. Size r => Array r ix e -> Sz ix
forall 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 e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, 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. Size r => Array r ix e -> Sz ix
forall 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 :: Loader e
          load :: Loader e
load Scheduler s ()
scheduler Int
startAt Int -> e -> ST s ()
dlWrite Int -> Sz Int -> e -> ST s ()
_dlSet =
            let loadIndex :: Int -> Lower ix -> e -> ST s ()
loadIndex Int
k Lower ix
ix = Int -> e -> ST s ()
dlWrite (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
newSz (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, 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 -> ST s 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 -> ST s () -> ST s Int
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler ((Lower ix -> e -> ST s ()) -> Array r (Lower ix) e -> ST s ()
forall ix r a (m :: * -> *) b.
(Index ix, Source r a, Monad m) =>
(ix -> a -> m b) -> Array r ix a -> m ()
imapM_ (Int -> Lower ix -> e -> ST s ()
loadIndex Int
k) Array r (Lower ix) e
arr)
                {-# INLINE arrayLoader #-}
             in (Int -> Array r (Lower ix) e -> ST s Int)
-> Int -> f (Array r (Lower ix) e) -> ST s ()
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 -> ST s Int
arrayLoader Int
0 f (Array r (Lower ix) e)
arrsF
          {-# INLINE load #-}
      Array DL ix e -> m (Array DL ix e)
forall a. a -> m a
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{dlComp :: Comp
dlComp = (Array r (Lower ix) e -> Comp) -> [Array r (Lower ix) e] -> Comp
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Array r (Lower ix) e -> 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) e]
arrs, dlSize :: Sz ix
dlSize = Sz ix
newSz, dlLoad :: Loader e
dlLoad = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz Int -> e -> ST s ())
-> ST s ()
Loader e
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 P Seq (Sz1 3)
--   [ 1, 2, 3 ]
-- Array P Seq (Sz1 3)
--   [ 4, 5, 6 ]
-- Array P 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, Index (Lower ix), Source r e, Index ix)
  => f (Array r (Lower ix) e)
  -> m (Array DL ix e)
stackOuterSlicesM :: forall r ix e (f :: * -> *) (m :: * -> *).
(Foldable f, MonadThrow m, Index (Lower ix), Source r e,
 Index ix) =>
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, Index (Lower ix), Source r 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
forall (proxy :: * -> *). 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 D Seq (Sz1 3)
--   [ 1, 4, 7 ]
-- Array D Seq (Sz1 3)
--   [ 2, 5, 8 ]
-- Array D 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, Index (Lower ix), Source r e, Index ix)
  => f (Array r (Lower ix) e)
  -> m (Array DL ix e)
stackInnerSlicesM :: forall r ix e (f :: * -> *) (m :: * -> *).
(Foldable f, MonadThrow m, Index (Lower ix), Source r e,
 Index ix) =>
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, Index (Lower ix), Source r 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
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r e)
  => Dim
  -- ^ Dimension along which to split
  -> Int
  -- ^ Index along the dimension to split at
  -> Array r ix e
  -- ^ Source array
  -> m (Array D ix e, Array D ix e)
splitAtM :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Dim -> Int -> Array r ix e -> m (Array D ix e, Array D 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. Size r => Array r ix e -> Sz ix
forall 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
forall (m :: * -> *). 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
forall (m :: * -> *). MonadThrow m => ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim Int
i
  Array D ix e
arr1 <- ix -> ix -> Array r ix e -> m (Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> ix -> Array r ix e -> m (Array D ix e)
extractFromToM ix
forall ix. Index ix => ix
zeroIndex ix
eIx Array r ix e
arr
  Array D ix e
arr2 <- ix -> ix -> Array r ix e -> m (Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> ix -> Array r ix e -> m (Array D ix e)
extractFromToM ix
sIx ix
sz Array r ix e
arr
  (Array D ix e, Array D ix e) -> m (Array D ix e, Array D ix e)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array D ix e
arr1, Array D 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'
  :: forall r ix e
   . (HasCallStack, Index ix, Source r e)
  => Dim
  -> Int
  -> Array r ix e
  -> (Array D ix e, Array D ix e)
splitAt' :: forall r ix e.
(HasCallStack, Index ix, Source r e) =>
Dim -> Int -> Array r ix e -> (Array D ix e, Array D ix e)
splitAt' Dim
dim Int
i = Either SomeException (Array D ix e, Array D ix e)
-> (Array D ix e, Array D ix e)
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array D ix e, Array D ix e)
 -> (Array D ix e, Array D ix e))
-> (Array r ix e
    -> Either SomeException (Array D ix e, Array D ix e))
-> Array r ix e
-> (Array D ix e, Array D ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim
-> Int
-> Array r ix e
-> Either SomeException (Array D ix e, Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Dim -> Int -> Array r ix e -> m (Array D ix e, Array D ix e)
splitAtM Dim
dim Int
i
{-# INLINE splitAt' #-}

-- | Split an array in three parts across some dimension
--
-- @since 0.3.5
splitExtractM
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r 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 D ix e, Array D ix e, Array D ix e)
splitExtractM :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array D ix e, Array D ix e, Array D 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. Size r => Array r ix e -> Sz ix
forall 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
forall (m :: * -> *). 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
forall (m :: * -> *). MonadThrow m => ix -> Dim -> Int -> m ix
setDimM ix
szIx Dim
dim Int
extractSzIx1
  Array D ix e
midArr <- ix -> Sz ix -> Array r ix e -> m (Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> Sz ix -> Array r ix e -> m (Array D 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
forall (m :: * -> *). MonadThrow m => ix -> Dim -> Int -> m ix
setDimM ix
szIx Dim
dim Int
startIx1
  Array D ix e
leftArr <- ix -> Sz ix -> Array r ix e -> m (Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> Sz ix -> Array r ix e -> m (Array D 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
forall (m :: * -> *). 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 D ix e
rightArr <- ix -> ix -> Array r ix e -> m (Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
ix -> ix -> Array r ix e -> m (Array D ix e)
extractFromToM ix
rightArrStartIx ix
szIx Array r ix e
arr
  (Array D ix e, Array D ix e, Array D ix e)
-> m (Array D ix e, Array D ix e, Array D ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D ix e
leftArr, Array D ix e
midArr, Array D 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
  :: forall r r' ix e m
   . (MonadThrow m, Source r e, Source r' e, Index ix, Index (Lower ix))
  => Dim
  -> Ix1
  -> Array r' (Lower ix) e
  -> Array r ix e
  -> m (Array DL ix e)
replaceSlice :: forall r r' ix e (m :: * -> *).
(MonadThrow m, Source r e, Source r' e, Index ix,
 Index (Lower ix)) =>
Dim
-> Int
-> Array r' (Lower ix) e
-> Array r ix e
-> m (Array DL ix e)
replaceSlice Dim
dim Int
i Array r' (Lower ix) e
sl Array r ix e
arr = do
  (Array D ix e
l, Array D ix e
m, Array D ix e
r) <- Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array D ix e, Array D ix e, Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array D ix e, Array D ix e, Array D 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' ix e
m' <- Sz ix -> Array r' (Lower ix) e -> m (Array r' ix e)
forall r ix ix' e (m :: * -> *).
(MonadThrow m, Index ix', Index ix, Size r) =>
Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM (Array D ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
size Array D ix e
m) Array r' (Lower ix) e
sl
  Dim -> [Array D ix e] -> m (Array DL ix e)
forall r ix e (f :: * -> *) (m :: * -> *).
(MonadThrow m, Foldable f, Index ix, Source r e) =>
Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM Dim
dim [Array D ix e
l, Array r' ix e -> Array D ix e
forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay Array r' ix e
m', Array D 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
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r e, Load r (Lower ix) e)
  => Ix1
  -> Array r (Lower ix) e
  -> Array r ix e
  -> m (Array DL ix e)
replaceOuterSlice :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e, Load r (Lower ix) e) =>
Int -> Array r (Lower ix) e -> Array r ix e -> m (Array DL ix e)
replaceOuterSlice Int
i Array r (Lower ix) e
sl Array r ix e
arr = Dim
-> Int -> Array r (Lower ix) e -> Array r ix e -> m (Array DL ix e)
forall r r' ix e (m :: * -> *).
(MonadThrow m, Source r e, Source r' e, Index ix,
 Index (Lower ix)) =>
Dim
-> Int
-> Array 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
forall (proxy :: * -> *). proxy ix -> Dim
dimensions (Array r ix e -> Sz 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 ix e
arr)) Int
i Array 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
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r 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 :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
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 D ix e
leftArr, Array D ix e
_, Array D ix e
rightArr) <- Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array D ix e, Array D ix e, Array D ix e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Dim
-> Int
-> Sz Int
-> Array r ix e
-> m (Array D ix e, Array D ix e, Array D ix e)
splitExtractM Dim
dim Int
ix Sz Int
sz Array r ix e
arr
  Dim -> Array D ix e -> Array D ix e -> m (Array DL ix e)
forall r1 r2 ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r1 e, Source r2 e) =>
Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e)
appendM Dim
dim Array D ix e
leftArr Array D 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
  :: forall r ix e m
   . (MonadThrow m, Index ix, Index (Lower ix), Source r e)
  => Ix1
  -> Sz Ix1
  -> Array r ix e
  -> m (Array DL ix e)
deleteRowsM :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
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 r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r 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
  :: forall r ix e m
   . (MonadThrow m, Index ix, Source r e)
  => Ix1
  -> Sz Ix1
  -> Array r ix e
  -> m (Array DL ix e)
deleteColumnsM :: forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
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 r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r 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 e, Load r ix e)
  => Stride ix
  -> Array r ix e
  -> Array DL ix e
downsample :: forall r ix e.
(Source r e, Load r ix e) =>
Stride ix -> Array r ix e -> Array DL ix e
downsample Stride ix
stride Array r ix e
arr =
  DLArray{dlComp :: Comp
dlComp = Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr, dlSize :: Sz ix
dlSize = Sz ix
resultSize, dlLoad :: Loader e
dlLoad = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz Int -> e -> ST s ())
-> ST s ()
Loader e
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. Size r => Array r ix e -> Sz ix
forall 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 ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => 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 :: Loader e
    load :: Loader e
load Scheduler s ()
scheduler Int
startAt Int -> e -> ST s ()
dlWrite Int -> Sz Int -> e -> ST s ()
_ =
      Scheduler s ()
-> Int
-> Int
-> (Int -> ST s e)
-> (Int -> e -> ST s ())
-> ST s ()
forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s ()
-> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_
        Scheduler s ()
scheduler
        Int
startAt
        (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
resultSize)
        (e -> ST s e
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> ST s e) -> (Int -> e) -> Int -> ST s e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
unsafeLinearWriteWithStride)
        Int -> e -> ST s ()
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 :: forall r ix e.
Load r ix e =>
e -> Stride ix -> Array r ix e -> Array DL ix e
upsample !e
fillWith Stride ix
safeStride Array r ix e
arr =
  DLArray
    { dlComp :: Comp
dlComp = Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr
    , dlSize :: Sz ix
dlSize = Sz ix
newsz
    , dlLoad :: Loader e
dlLoad = Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz Int -> 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 ()
uWrite Int -> Sz Int -> e -> ST s ()
uSet = do
      Int -> Sz Int -> e -> ST s ()
uSet Int
startAt (Sz ix -> Sz Int
forall ix. Index ix => Sz ix -> Sz Int
toLinearSz Sz ix
newsz) e
fillWith
      Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
forall s.
Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array r ix e
arr (\Int
i -> Int -> e -> ST s ()
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 ix
sz = Array r ix e -> Sz ix
forall e. Array r ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr -- intentionally lazy in case it is used with DS
    !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
   . (Manifest r e, Index ix, Source r' e', Index ix', 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 :: forall r ix e r' ix' e' a (m :: * -> *).
(Manifest r e, Index ix, Source r' e', Index ix', 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)
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. Size r => Array r ix e -> Sz ix
forall 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, Manifest r e, Index ix) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray (Array r' ix' e' -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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 ix r e (m :: * -> *).
(Index ix, Source r 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'
  :: forall ix e r' ix' e' a
   . (HasCallStack, Source r' e', Index ix', Index ix)
  => (Sz ix' -> (Sz ix, a))
  -> (a -> (ix' -> e') -> ix -> e)
  -> Array r' ix' e'
  -> Array D ix e
transform' :: forall ix e r' ix' e' a.
(HasCallStack, Source r' e', Index ix', 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)
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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r' ix' e' -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall 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 ix r e.
(HasCallStack, Index ix, Source r 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. Size r => Array r ix e -> Sz ix
forall 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
  :: ( Manifest r e
     , Index ix
     , Source r1 e1
     , Source r2 e2
     , Index ix1
     , Index ix2
     , 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 :: forall r e ix r1 e1 r2 e2 ix1 ix2 (m :: * -> *) a.
(Manifest r e, Index ix, Source r1 e1, Source r2 e2, Index ix1,
 Index ix2, 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)
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. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> Sz ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r2 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, Manifest r e, Index ix) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray (Array r1 ix1 e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r1 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. Strategy r => Array r ix e -> Comp
forall ix e. Array r2 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 ix r e (m :: * -> *).
(Index ix, Source r e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> ix2 -> m e2
forall ix r e (m :: * -> *).
(Index ix, Source r 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'
  :: (HasCallStack, Source r1 e1, Source r2 e2, Index ix, Index ix1, Index ix2)
  => (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' :: forall r1 e1 r2 e2 ix ix1 ix2 a e.
(HasCallStack, Source r1 e1, Source r2 e2, Index ix, Index ix1,
 Index ix2) =>
(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.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r1 ix1 e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r1 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. Strategy r => Array r ix e -> Comp
forall ix e. Array r2 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 ix r e.
(HasCallStack, Index ix, Source r e) =>
Array r ix e -> ix -> e
evaluate' Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> ix2 -> e2
forall ix r e.
(HasCallStack, Index ix, Source r 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. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix1 e1
arr1) (Array r2 ix2 e2 -> Sz ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r2 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
   . (Index ix, Source r e)
  => e
  -- ^ Value to use for the grid
  -> Stride ix
  -- ^ Scaling factor
  -> Array r ix e
  -- ^ Source array
  -> Array DL ix e
zoomWithGrid :: forall r ix e.
(Index ix, Source r e) =>
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 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
Seq Sz ix
newSz (e -> Maybe e
forall a. a -> Maybe a
Just e
gridVal) Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ()
forall s. Scheduler s () -> Int -> (Int -> e -> ST s ()) -> ST s ()
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. Size r => Array r ix e -> Sz ix
forall 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 :: 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
_ Int -> e -> ST s ()
writeElement =
      Scheduler s () -> Array r ix e -> (ix -> e -> ST s ()) -> ST s ()
forall ix r e s (m :: * -> *) a.
(Index ix, Source r e, MonadPrimBase s m) =>
Scheduler s () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler s ()
scheduler Array r ix e
arr ((ix -> e -> ST s ()) -> ST s ())
-> (ix -> e -> ST s ()) -> ST s ()
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 -> ST s ()) -> Array D ix ix -> ST s ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_ (\ !ix
ix' -> Int -> e -> ST s ()
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 -> ST s ()) -> Array D ix ix -> ST s ()
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
   . (Index ix, Source r e)
  => Stride ix
  -- ^ Scaling factor
  -> Array r ix e
  -- ^ Source array
  -> Array DL ix e
zoom :: forall r ix e.
(Index ix, Source r e) =>
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 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
Seq Sz ix
newSz 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
    !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. Size r => Array r ix e -> Sz ix
forall 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 :: 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
_ Int -> e -> ST s ()
writeElement =
      Scheduler s () -> Array r ix e -> (ix -> e -> ST s ()) -> ST s ()
forall ix r e s (m :: * -> *) a.
(Index ix, Source r e, MonadPrimBase s m) =>
Scheduler s () -> Array r ix e -> (ix -> e -> m a) -> m ()
iforSchedulerM_ Scheduler s ()
scheduler Array r ix e
arr ((ix -> e -> ST s ()) -> ST s ())
-> (ix -> e -> ST s ()) -> ST s ()
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 -> ST s ()) -> Array D ix ix -> ST s ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
mapM_ (\ !ix
ix' -> Int -> e -> ST s ()
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 -> ST s ()) -> Array D ix ix -> ST s ()
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 #-}