{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Data.Massiv.Array.Ops.Fold
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Ops.Fold
  (
  -- ** Unstructured folds

  -- $unstruct_folds
    fold
  , ifoldMono
  , foldMono
  , ifoldSemi
  , foldSemi
  , foldOuterSlice
  , ifoldOuterSlice
  , foldInnerSlice
  , ifoldInnerSlice
  , minimumM
  , minimum'
  , maximumM
  , maximum'
  , sum
  , product
  , and
  , or
  , all
  , any
  , elem
  , eqArrays
  , compareArrays

  -- ** Single dimension folds
  -- *** Safe inner most
  --
  -- Folding along the inner most dimension will always be faster when compared to doing the same
  -- operation along any other dimension, this is due to the fact that inner most folds follow the
  -- memory layout of data.
  , ifoldlInner
  , foldlInner
  , ifoldrInner
  , foldrInner
  , foldInner
  -- *** Type safe within
  , ifoldlWithin
  , foldlWithin
  , ifoldrWithin
  , foldrWithin
  , foldWithin
  -- *** Partial within
  , ifoldlWithin'
  , foldlWithin'
  , ifoldrWithin'
  , foldrWithin'
  , foldWithin'

  -- ** Sequential folds

  -- $seq_folds

  , foldlS
  , foldrS
  , ifoldlS
  , ifoldrS

  -- *** Monadic
  , foldlM
  , foldrM
  , foldlM_
  , foldrM_
  , ifoldlM
  , ifoldrM
  , ifoldlM_
  , ifoldrM_

  -- *** Special folds
  , foldrFB
  , lazyFoldlS
  , lazyFoldrS

  -- ** Parallel folds

  -- $par_folds

  , foldlP
  , foldrP
  , ifoldlP
  , ifoldrP
  , ifoldlIO
  , ifoldrIO
  -- , splitReduce
  ) where

import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Ops.Construct
import Data.Massiv.Array.Ops.Fold.Internal
import Data.Massiv.Core
import Data.Massiv.Core.Common
import Prelude hiding (all, and, any, foldl, foldr, map, maximum, minimum, or, product, sum, elem)


-- | /O(n)/ - Monoidal fold over an array with an index aware function. Also known as reduce.
--
-- @since 0.2.4
ifoldMono ::
     (Index ix, Source r e, Monoid m)
  => (ix -> e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`.
  -> Array r ix e -- ^ Source array
  -> m
ifoldMono :: (ix -> e -> m) -> Array r ix e -> m
ifoldMono ix -> e -> m
f = (m -> ix -> e -> m) -> m -> (m -> m -> m) -> m -> Array r ix e -> m
forall ix r e a b.
(Index ix, Source r e) =>
(a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
ifoldlInternal (\m
a ix
ix e
e -> m
a m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` ix -> e -> m
f ix
ix e
e) m
forall a. Monoid a => a
mempty m -> m -> m
forall a. Monoid a => a -> a -> a
mappend m
forall a. Monoid a => a
mempty
{-# INLINE ifoldMono #-}


-- | /O(n)/ - Semigroup fold over an array with an index aware function.
--
-- @since 0.2.4
ifoldSemi ::
     (Index ix, Source r e, Semigroup m)
  => (ix -> e -> m) -- ^ Convert each element of an array to an appropriate `Semigroup`.
  -> m -- ^ Initial element that must be neutral to the (`<>`) function.
  -> Array r ix e -- ^ Source array
  -> m
ifoldSemi :: (ix -> e -> m) -> m -> Array r ix e -> m
ifoldSemi ix -> e -> m
f m
m = (m -> ix -> e -> m) -> m -> (m -> m -> m) -> m -> Array r ix e -> m
forall ix r e a b.
(Index ix, Source r e) =>
(a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
ifoldlInternal (\m
a ix
ix e
e -> m
a m -> m -> m
forall a. Semigroup a => a -> a -> a
<> ix -> e -> m
f ix
ix e
e) m
m m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>) m
m
{-# INLINE ifoldSemi #-}


-- | /O(n)/ - Semigroup fold over an array.
--
-- @since 0.1.6
foldSemi ::
     (Index ix, Source r e, Semigroup m)
  => (e -> m) -- ^ Convert each element of an array to an appropriate `Semigroup`.
  -> m -- ^ Initial element that must be neutral to the (`<>`) function.
  -> Array r ix e -- ^ Source array
  -> m
foldSemi :: (e -> m) -> m -> Array r ix e -> m
foldSemi e -> m
f m
m = (m -> e -> m) -> m -> (m -> m -> m) -> m -> Array r ix e -> m
forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal (\m
a e
e -> m
a m -> m -> m
forall a. Semigroup a => a -> a -> a
<> e -> m
f e
e) m
m m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>) m
m
{-# INLINE foldSemi #-}


-- | Left fold along a specified dimension with an index aware function.
--
-- @since 0.2.4
ifoldlWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) =>
  Dimension n -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin :: Dimension n
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin Dimension n
dim = Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' (Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE ifoldlWithin #-}


-- | Left fold along a specified dimension.
--
-- ====__Example__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XTypeApplications
-- >>> arr = makeArrayLinear @U Seq (Sz (2 :. 5)) id
-- >>> arr
-- Array U Seq (Sz (2 :. 5))
--   [ [ 0, 1, 2, 3, 4 ]
--   , [ 5, 6, 7, 8, 9 ]
--   ]
-- >>> foldlWithin Dim1 (flip (:)) [] arr
-- Array D Seq (Sz1 2)
--   [ [4,3,2,1,0], [9,8,7,6,5] ]
-- >>> foldlWithin Dim2 (flip (:)) [] arr
-- Array D Seq (Sz1 5)
--   [ [5,0], [6,1], [7,2], [8,3], [9,4] ]
--
-- @since 0.2.4
foldlWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) =>
  Dimension n -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin :: Dimension n
-> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin Dimension n
dim a -> e -> a
f = Dimension n
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix (n :: Nat) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin Dimension n
dim ((a -> e -> a) -> ix -> a -> e -> a
forall a b. a -> b -> a
const a -> e -> a
f)
{-# INLINE foldlWithin #-}


-- | Right fold along a specified dimension with an index aware function.
--
-- @since 0.2.4
ifoldrWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) =>
  Dimension n -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin :: Dimension n
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin Dimension n
dim = Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' (Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE ifoldrWithin #-}


-- | Right fold along a specified dimension.
--
-- @since 0.2.4
foldrWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) =>
  Dimension n -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin :: Dimension n
-> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin Dimension n
dim e -> a -> a
f = Dimension n
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix (n :: Nat) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin Dimension n
dim ((e -> a -> a) -> ix -> e -> a -> a
forall a b. a -> b -> a
const e -> a -> a
f)
{-# INLINE foldrWithin #-}


-- | Similar to `ifoldlWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
-- @since 0.2.4
ifoldlWithin' :: (HasCallStack, Index (Lower ix), Index ix, Source r e) =>
  Dim -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' :: Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' Dim
dim ix -> a -> e -> a
f a
acc0 Array r ix e
arr =
  Comp -> Sz (Lower ix) -> (Lower ix -> a) -> Array D (Lower ix) a
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
getComp Array r ix e
arr) (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
szl) ((Lower ix -> a) -> Array D (Lower ix) a)
-> (Lower ix -> a) -> Array D (Lower ix) a
forall a b. (a -> b) -> a -> b
$ \Lower ix
ixl ->
    ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
forall ix a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
iter
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim Int
0)
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
      (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1)
      Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<=)
      a
acc0
      (\ix
ix a
acc' -> ix -> a -> e -> a
f ix
ix a
acc' (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
ix))
  where
    SafeSz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
    (Int
k, Lower ix
szl) = ix -> Dim -> (Int, Lower ix)
forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
sz Dim
dim
{-# INLINE ifoldlWithin' #-}


-- | Similar to `foldlWithin`, except that dimension is specified at a value level, which means it will
-- throw an exception on an invalid dimension.
--
-- @since 0.2.4
foldlWithin' :: (HasCallStack, Index (Lower ix), Index ix, Source r e) =>
  Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' :: Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' Dim
dim a -> e -> a
f = Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' Dim
dim ((a -> e -> a) -> ix -> a -> e -> a
forall a b. a -> b -> a
const a -> e -> a
f)
{-# INLINE foldlWithin' #-}


-- | Similar to `ifoldrWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
--
-- @since 0.2.4
ifoldrWithin' :: (HasCallStack, Index (Lower ix), Index ix, Source r e) =>
  Dim -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' :: Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' Dim
dim ix -> e -> a -> a
f a
acc0 Array r ix e
arr =
  Comp -> Sz (Lower ix) -> (Lower ix -> a) -> Array D (Lower ix) a
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
getComp Array r ix e
arr) (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
szl) ((Lower ix -> a) -> Array D (Lower ix) a)
-> (Lower ix -> a) -> Array D (Lower ix) a
forall a b. (a -> b) -> a -> b
$ \Lower ix
ixl ->
    ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
forall ix a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
iter
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim Int
0)
      (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex (-Int
1))
      Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(>=)
      a
acc0
      (\ix
ix a
acc' -> ix -> e -> a -> a
f ix
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
ix) a
acc')
  where
    SafeSz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
    (Int
k, Lower ix
szl) = ix -> Dim -> (Int, Lower ix)
forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
sz Dim
dim
{-# INLINE ifoldrWithin' #-}

-- | Similar to `foldrWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
-- @since 0.2.4
foldrWithin' :: (HasCallStack, Index (Lower ix), Index ix, Source r e) =>
  Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin' :: Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin' Dim
dim e -> a -> a
f = Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' Dim
dim ((e -> a -> a) -> ix -> e -> a -> a
forall a b. a -> b -> a
const e -> a -> a
f)
{-# INLINE foldrWithin' #-}


-- | Left fold over the inner most dimension with index aware function.
--
-- @since 0.2.4
ifoldlInner :: (Index (Lower ix), Index ix, Source r e) =>
  (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlInner :: (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlInner = Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' Dim
1
{-# INLINE ifoldlInner #-}

-- | Left fold over the inner most dimension.
--
-- @since 0.2.4
foldlInner :: (Index (Lower ix), Index ix, Source r e) =>
  (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlInner :: (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlInner = Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' Dim
1
{-# INLINE foldlInner #-}

-- | Right fold over the inner most dimension with index aware function.
--
-- @since 0.2.4
ifoldrInner :: (Index (Lower ix), Index ix, Source r e) =>
  (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrInner :: (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrInner = Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' Dim
1
{-# INLINE ifoldrInner #-}

-- | Right fold over the inner most dimension.
--
-- @since 0.2.4
foldrInner :: (Index (Lower ix), Index ix, Source r e) =>
  (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrInner :: (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrInner = Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin' Dim
1
{-# INLINE foldrInner #-}

-- | Monoidal fold over the inner most dimension.
--
-- @since 0.4.3
foldInner :: (Monoid e, Index (Lower ix), Index ix, Source r e) => Array r ix e -> Array D (Lower ix) e
foldInner :: Array r ix e -> Array D (Lower ix) e
foldInner = (e -> e -> e) -> e -> Array r ix e -> Array D (Lower ix) e
forall ix r e a.
(Index (Lower ix), Index ix, Source r e) =>
(a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlInner e -> e -> e
forall a. Monoid a => a -> a -> a
mappend e
forall a. Monoid a => a
mempty
{-# INLINE foldInner #-}

-- | Monoidal fold over some internal dimension.
--
-- @since 0.4.3
foldWithin ::
     (Source r a, Monoid a, Index (Lower ix), IsIndexDimension ix n)
  => Dimension n
  -> Array r ix a
  -> Array D (Lower ix) a
foldWithin :: Dimension n -> Array r ix a -> Array D (Lower ix) a
foldWithin Dimension n
dim = Dimension n
-> (a -> a -> a) -> a -> Array r ix a -> Array D (Lower ix) a
forall ix (n :: Nat) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin Dimension n
dim a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
forall a. Monoid a => a
mempty
{-# INLINE foldWithin #-}

-- | Monoidal fold over some internal dimension. This is a pratial function and will
-- result in `IndexDimensionException` if supplied dimension is invalid.
--
-- @since 0.4.3
foldWithin' ::
     (HasCallStack, Index ix, Source r a, Monoid a, Index (Lower ix))
  => Dim
  -> Array r ix a
  -> Array D (Lower ix) a
foldWithin' :: Dim -> Array r ix a -> Array D (Lower ix) a
foldWithin' Dim
dim = Dim -> (a -> a -> a) -> a -> Array r ix a -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' Dim
dim a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
forall a. Monoid a => a
mempty
{-# INLINE foldWithin' #-}


-- | Reduce each outer slice into a monoid and mappend results together
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> import Data.Monoid (Product(..))
-- >>> arr = computeAs P $ iterateN (Sz2 2 3) (+1) (10 :: Int)
-- >>> arr
-- Array P Seq (Sz (2 :. 3))
--   [ [ 11, 12, 13 ]
--   , [ 14, 15, 16 ]
--   ]
-- >>> getProduct $ foldOuterSlice (\row -> Product (A.sum row)) arr
-- 1620
-- >>> (11 + 12 + 13) * (14 + 15 + 16) :: Int
-- 1620
--
-- @since 0.4.3
foldOuterSlice ::
     (Index ix, Index (Lower ix), Source r e, Monoid m)
  => (Array r (Lower ix) e -> m)
  -> Array r ix e
  -> m
foldOuterSlice :: (Array r (Lower ix) e -> m) -> Array r ix e -> m
foldOuterSlice Array r (Lower ix) e -> m
f = (Int -> Array r (Lower ix) e -> m) -> Array r ix e -> m
forall ix r e m.
(Index ix, Index (Lower ix), Source r e, Monoid m) =>
(Int -> Array r (Lower ix) e -> m) -> Array r ix e -> m
ifoldOuterSlice ((Array r (Lower ix) e -> m) -> Int -> Array r (Lower ix) e -> m
forall a b. a -> b -> a
const Array r (Lower ix) e -> m
f)
{-# INLINE foldOuterSlice #-}


-- | Reduce each outer slice into a monoid with an index aware function and mappend results
-- together
--
-- @since 0.4.3
ifoldOuterSlice ::
     (Index ix, Index (Lower ix), Source r e, Monoid m)
  => (Ix1 -> Array r (Lower ix) e -> m)
  -> Array r ix e
  -> m
ifoldOuterSlice :: (Int -> Array r (Lower ix) e -> m) -> Array r ix e -> m
ifoldOuterSlice Int -> Array r (Lower ix) e -> m
f Array r ix e
arr = (Int -> m) -> Array D Int Int -> m
forall ix r e m.
(Index ix, Source r e, Monoid m) =>
(e -> m) -> Array r ix e -> m
foldMono Int -> m
g (Array D Int Int -> m) -> Array D Int Int -> m
forall a b. (a -> b) -> a -> b
$ Comp -> Int -> Int -> Array D Int Int
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Int
0 Int
k
  where
    (Sz1 Int
k, Sz (Lower ix)
szL) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz (Sz ix -> (Sz Int, Sz (Lower ix)))
-> Sz ix -> (Sz Int, Sz (Lower ix))
forall a b. (a -> b) -> a -> b
$ Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
    g :: Int -> m
g Int
i = Int -> Array r (Lower ix) e -> m
f Int
i (Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
forall r e ix.
(Source r e, Index ix, Index (Lower ix)) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
unsafeOuterSlice Array r ix e
arr Sz (Lower ix)
szL Int
i)
    {-# INLINE g #-}
{-# INLINE ifoldOuterSlice #-}


-- | Reduce each inner slice into a monoid and mappend results together
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> import Data.Monoid (Product(..))
-- >>> arr = computeAs P $ iterateN (Sz2 2 3) (+1) (10 :: Int)
-- >>> arr
-- Array P Seq (Sz (2 :. 3))
--   [ [ 11, 12, 13 ]
--   , [ 14, 15, 16 ]
--   ]
-- >>> getProduct $ foldInnerSlice (\column -> Product (A.sum column)) arr
-- 19575
-- >>> (11 + 14) * (12 + 15) * (13 + 16) :: Int
-- 19575
--
-- @since 0.4.3
foldInnerSlice ::
     (Source r e, Index ix, Monoid m) => (Array D (Lower ix) e -> m) -> Array r ix e -> m
foldInnerSlice :: (Array D (Lower ix) e -> m) -> Array r ix e -> m
foldInnerSlice Array D (Lower ix) e -> m
f = (Int -> Array D (Lower ix) e -> m) -> Array r ix e -> m
forall r e ix m.
(Source r e, Index ix, Monoid m) =>
(Int -> Array D (Lower ix) e -> m) -> Array r ix e -> m
ifoldInnerSlice ((Array D (Lower ix) e -> m) -> Int -> Array D (Lower ix) e -> m
forall a b. a -> b -> a
const Array D (Lower ix) e -> m
f)
{-# INLINE foldInnerSlice #-}


-- | Reduce each inner slice into a monoid with an index aware function and mappend
-- results together
--
-- @since 0.4.3
ifoldInnerSlice ::
     (Source r e, Index ix, Monoid m) => (Ix1 -> Array D (Lower ix) e -> m) -> Array r ix e -> m
ifoldInnerSlice :: (Int -> Array D (Lower ix) e -> m) -> Array r ix e -> m
ifoldInnerSlice Int -> Array D (Lower ix) e -> m
f Array r ix e
arr = (Int -> m) -> Array D Int Int -> m
forall ix r e m.
(Index ix, Source r e, Monoid m) =>
(e -> m) -> Array r ix e -> m
foldMono Int -> m
g (Array D Int Int -> m) -> Array D Int Int -> m
forall a b. (a -> b) -> a -> b
$ Comp -> Int -> Int -> Array D Int Int
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Int
0 (Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
k)
  where
    (Sz (Lower ix)
szL, !Sz Int
k) = Sz ix -> (Sz (Lower ix), Sz Int)
forall ix. Index ix => Sz ix -> (Sz (Lower ix), Sz Int)
unsnocSz (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)
    g :: Int -> m
g Int
i = Int -> Array D (Lower ix) e -> m
f Int
i (Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
unsafeInnerSlice Array r ix e
arr Sz (Lower ix)
szL Int
i)
    {-# INLINE g #-}
{-# INLINE ifoldInnerSlice #-}

-- | /O(n)/ - Compute maximum of all elements.
--
-- @since 0.3.0
maximumM :: (MonadThrow m, Shape r ix, Source r e, Ord e) => Array r ix e -> m e
maximumM :: Array r ix e -> m e
maximumM Array r ix e
arr =
  if Array r ix e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNull Array r ix e
arr
    then SizeException -> m e
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz ix -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeEmptyException (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr))
    else let !e0 :: e
e0 = 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
forall ix. Index ix => ix
zeroIndex
          in e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> e -> m e
forall a b. (a -> b) -> a -> b
$ (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal e -> e -> e
forall a. Ord a => a -> a -> a
max e
e0 e -> e -> e
forall a. Ord a => a -> a -> a
max e
e0 Array r ix e
arr
{-# INLINE maximumM #-}


-- | /O(n)/ - Compute maximum of all elements.
--
-- @since 0.3.0
maximum' ::
     forall r ix e. (HasCallStack, Shape r ix, Source r e, Ord e)
  => Array r ix e
  -> e
maximum' :: Array r ix e -> e
maximum' = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException e -> e)
-> (Array r ix e -> Either SomeException e) -> Array r ix e -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Either SomeException e
forall (m :: * -> *) r ix e.
(MonadThrow m, Shape r ix, Source r e, Ord e) =>
Array r ix e -> m e
maximumM
{-# INLINE maximum' #-}


-- | /O(n)/ - Compute minimum of all elements.
--
-- @since 0.3.0
minimumM :: (MonadThrow m, Shape r ix, Source r e, Ord e) => Array r ix e -> m e
minimumM :: Array r ix e -> m e
minimumM Array r ix e
arr =
  if Array r ix e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNull Array r ix e
arr
    then SizeException -> m e
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz ix -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeEmptyException (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr))
    else let !e0 :: e
e0 = 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
forall ix. Index ix => ix
zeroIndex
          in e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> e -> m e
forall a b. (a -> b) -> a -> b
$ (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal e -> e -> e
forall a. Ord a => a -> a -> a
min e
e0 e -> e -> e
forall a. Ord a => a -> a -> a
min e
e0 Array r ix e
arr
{-# INLINE minimumM #-}

-- | /O(n)/ - Compute minimum of all elements.
--
-- @since 0.3.0
minimum' :: forall r ix e. (HasCallStack, Shape r ix, Source r e, Ord e) => Array r ix e -> e
minimum' :: Array r ix e -> e
minimum' = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException e -> e)
-> (Array r ix e -> Either SomeException e) -> Array r ix e -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Either SomeException e
forall (m :: * -> *) r ix e.
(MonadThrow m, Shape r ix, Source r e, Ord e) =>
Array r ix e -> m e
minimumM
{-# INLINE minimum' #-}


-- -- | /O(n)/ - Compute sum of all elements.
-- --
-- -- @since 0.1.0
-- sum' ::
--      forall r ix e. (Index ix, Source r e, Numeric r e)
--   => Array r ix e
--   -> IO e
-- sum' = splitReduce (\_ -> pure . sumArray) (\x y -> pure (x + y)) 0
-- {-# INLINE sum' #-}

-- | /O(n)/ - Compute sum of all elements.
--
-- @since 0.1.0
sum :: (Index ix, Source r e, Num e) => Array r ix e -> e
sum :: Array r ix e -> e
sum = (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal e -> e -> e
forall a. Num a => a -> a -> a
(+) e
0 e -> e -> e
forall a. Num a => a -> a -> a
(+) e
0
{-# INLINE sum #-}


-- | /O(n)/ - Compute product of all elements.
--
-- @since 0.1.0
product :: (Index ix, Source r e, Num e) => Array r ix e -> e
product :: Array r ix e -> e
product = (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal e -> e -> e
forall a. Num a => a -> a -> a
(*) e
1 e -> e -> e
forall a. Num a => a -> a -> a
(*) e
1
{-# INLINE product #-}


-- | /O(n)/ - Compute conjunction of all elements.
--
-- @since 0.1.0
and :: (Index ix, Source r Bool) => Array r ix Bool -> Bool
and :: Array r ix Bool -> Bool
and = (Bool -> Bool) -> Array r ix Bool -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
all Bool -> Bool
forall a. a -> a
id
{-# INLINE and #-}


-- | /O(n)/ - Compute disjunction of all elements.
--
-- @since 0.1.0
or :: (Index ix, Source r Bool) => Array r ix Bool -> Bool
or :: Array r ix Bool -> Bool
or = (Bool -> Bool) -> Array r ix Bool -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
any Bool -> Bool
forall a. a -> a
id
{-# INLINE or #-}


-- | /O(n)/ - Determines whether all elements of the array satisfy a predicate.
--
-- @since 0.1.0
all :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool
all :: (e -> Bool) -> Array r ix e -> Bool
all e -> Bool
f = Bool -> Bool
not (Bool -> Bool) -> (Array r ix e -> Bool) -> Array r ix e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (e -> Bool) -> Array r ix e -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (e -> Bool) -> e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Bool
f)
{-# INLINE all #-}

-- | /O(n)/ - Determines whether an element is present in the array.
--
-- @since 0.5.5
elem :: (Eq e, Index ix, Source r e) => e -> Array r ix e -> Bool
elem :: e -> Array r ix e -> Bool
elem e
e = (e -> Bool) -> Array r ix e -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
any (e
e e -> e -> Bool
forall a. Eq a => a -> a -> Bool
==)
{-# INLINE elem #-}


{- $unstruct_folds

Functions in this section will fold any `Source` array with respect to the inner
`Comp`utation strategy setting.

-}


{- $seq_folds

Functions in this section will fold any `Source` array sequentially, regardless of the inner
`Comp`utation strategy setting.

-}


{- $par_folds

__Note__ It is important to compile with @-threaded -with-rtsopts=-N@ flags, otherwise there will be
no parallelization.

Functions in this section will fold any `Source` array in parallel, regardless of the inner
`Comp`utation strategy setting. All of the parallel structured folds are performed inside `IO`
monad, because referential transparency can't generally be preserved and results will depend on the
number of cores/capabilities that computation is being performed on.

In contrast to sequential folds, each parallel folding function accepts two functions and two
initial elements as arguments. This is necessary because an array is first split into chunks, which
folded individually on separate cores with the first function, and the results of those folds are
further folded with the second function.

-}