-- Copyright 2020 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# OPTIONS_GHC -Wno-orphans #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Array.Internal.DynamicS(
  Array(..), Vector, ShapeL, V.Storable, Unbox,
  size, shapeL, rank,
  toList, fromList, toVector, fromVector,
  normalize,
  scalar, unScalar, constant,
  reshape, stretch, stretchOuter, transpose,
  index, pad,
  mapA, zipWithA, zipWith3A, zipWith4A, zipWith5A,
  append, concatOuter,
  ravel, unravel,
  window, stride, rotate,
  slice, rerank, rerank2, rev,
  reduce, foldrA, traverseA,
  allSameA,
  sumA, productA, maximumA, minimumA,
  anyA, allA,
  broadcast,
  update,
  generate, iterateN, iota,
  bitcast,
  ) where
import Control.DeepSeq
import Data.Coerce(coerce)
import Data.Data(Data)
import qualified Data.Vector.Storable as V
import Foreign.Storable(sizeOf)
import GHC.Generics(Generic)
import GHC.Stack(HasCallStack)
import Test.QuickCheck hiding (generate)
import Text.PrettyPrint.HughesPJClass hiding ((<>))

import qualified Data.Array.Internal.Dynamic as D
import qualified Data.Array.Internal.DynamicG as G
import Data.Array.Internal(T(..), ShapeL, Vector(..))

type Unbox = V.Storable

instance Vector V.Vector where
  type VecElem V.Vector = Unbox
  {-# INLINE vIndex #-}
  vIndex :: Vector a -> Int -> a
vIndex = Vector a -> Int -> a
forall a. Storable a => Vector a -> Int -> a
(V.!)
  {-# INLINE vLength #-}
  vLength :: Vector a -> Int
vLength = Vector a -> Int
forall a. Storable a => Vector a -> Int
V.length
  {-# INLINE vToList #-}
  vToList :: Vector a -> [a]
vToList = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
V.toList
  {-# INLINE vFromList #-}
  vFromList :: [a] -> Vector a
vFromList = [a] -> Vector a
forall a. Storable a => [a] -> Vector a
V.fromList
  {-# INLINE vSingleton #-}
  vSingleton :: a -> Vector a
vSingleton = a -> Vector a
forall a. Storable a => a -> Vector a
V.singleton
  {-# INLINE vReplicate #-}
  vReplicate :: Int -> a -> Vector a
vReplicate = Int -> a -> Vector a
forall a. Storable a => Int -> a -> Vector a
V.replicate
  {-# INLINE vMap #-}
  vMap :: (a -> b) -> Vector a -> Vector b
vMap = (a -> b) -> Vector a -> Vector b
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
V.map
  {-# INLINE vZipWith #-}
  vZipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
vZipWith = (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith
  {-# INLINE vZipWith3 #-}
  vZipWith3 :: (a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
vZipWith3 = (a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
forall a b c d.
(Storable a, Storable b, Storable c, Storable d) =>
(a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
V.zipWith3
  {-# INLINE vZipWith4 #-}
  vZipWith4 :: (a -> b -> c -> d -> e)
-> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
vZipWith4 = (a -> b -> c -> d -> e)
-> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
forall a b c d e.
(Storable a, Storable b, Storable c, Storable d, Storable e) =>
(a -> b -> c -> d -> e)
-> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
V.zipWith4
  {-# INLINE vZipWith5 #-}
  vZipWith5 :: (a -> b -> c -> d -> e -> f)
-> Vector a
-> Vector b
-> Vector c
-> Vector d
-> Vector e
-> Vector f
vZipWith5 = (a -> b -> c -> d -> e -> f)
-> Vector a
-> Vector b
-> Vector c
-> Vector d
-> Vector e
-> Vector f
forall a b c d e f.
(Storable a, Storable b, Storable c, Storable d, Storable e,
 Storable f) =>
(a -> b -> c -> d -> e -> f)
-> Vector a
-> Vector b
-> Vector c
-> Vector d
-> Vector e
-> Vector f
V.zipWith5
  {-# INLINE vAppend #-}
  vAppend :: Vector a -> Vector a -> Vector a
vAppend = Vector a -> Vector a -> Vector a
forall a. Storable a => Vector a -> Vector a -> Vector a
(V.++)
  {-# INLINE vConcat #-}
  vConcat :: [Vector a] -> Vector a
vConcat = [Vector a] -> Vector a
forall a. Storable a => [Vector a] -> Vector a
V.concat
  {-# INLINE vFold #-}
  vFold :: (a -> a -> a) -> a -> Vector a -> a
vFold = (a -> a -> a) -> a -> Vector a -> a
forall b a. Storable b => (a -> b -> a) -> a -> Vector b -> a
V.foldl'
  {-# INLINE vSlice #-}
  vSlice :: Int -> Int -> Vector a -> Vector a
vSlice = Int -> Int -> Vector a -> Vector a
forall a. Storable a => Int -> Int -> Vector a -> Vector a
V.slice
  {-# INLINE vSum #-}
  vSum :: Vector a -> a
vSum = Vector a -> a
forall a. (Storable a, Num a) => Vector a -> a
V.sum
  {-# INLINE vProduct #-}
  vProduct :: Vector a -> a
vProduct = Vector a -> a
forall a. (Storable a, Num a) => Vector a -> a
V.product
  {-# INLINE vMaximum #-}
  vMaximum :: Vector a -> a
vMaximum = Vector a -> a
forall a. (Storable a, Ord a) => Vector a -> a
V.maximum
  {-# INLINE vMinimum #-}
  vMinimum :: Vector a -> a
vMinimum = Vector a -> a
forall a. (Storable a, Ord a) => Vector a -> a
V.minimum
  {-# INLINE vUpdate #-}
  vUpdate :: Vector a -> [(Int, a)] -> Vector a
vUpdate = Vector a -> [(Int, a)] -> Vector a
forall a. Storable a => Vector a -> [(Int, a)] -> Vector a
(V.//)
  {-# INLINE vGenerate #-}
  vGenerate :: Int -> (Int -> a) -> Vector a
vGenerate = Int -> (Int -> a) -> Vector a
forall a. Storable a => Int -> (Int -> a) -> Vector a
V.generate
  {-# INLINE vAll #-}
  vAll :: (a -> Bool) -> Vector a -> Bool
vAll = (a -> Bool) -> Vector a -> Bool
forall a. Storable a => (a -> Bool) -> Vector a -> Bool
V.all
  {-# INLINE vAny #-}
  vAny :: (a -> Bool) -> Vector a -> Bool
vAny = (a -> Bool) -> Vector a -> Bool
forall a. Storable a => (a -> Bool) -> Vector a -> Bool
V.any

newtype Array a = A { Array a -> Array Vector a
unA :: G.Array V.Vector a }
  deriving (PrettyLevel -> [Array a] -> Doc
PrettyLevel -> Rational -> Array a -> Doc
Array a -> Doc
(PrettyLevel -> Rational -> Array a -> Doc)
-> (Array a -> Doc)
-> (PrettyLevel -> [Array a] -> Doc)
-> Pretty (Array a)
forall a. (Pretty a, Storable a) => PrettyLevel -> [Array a] -> Doc
forall a.
(Pretty a, Storable a) =>
PrettyLevel -> Rational -> Array a -> Doc
forall a. (Pretty a, Storable a) => Array a -> Doc
forall a.
(PrettyLevel -> Rational -> a -> Doc)
-> (a -> Doc) -> (PrettyLevel -> [a] -> Doc) -> Pretty a
pPrintList :: PrettyLevel -> [Array a] -> Doc
$cpPrintList :: forall a. (Pretty a, Storable a) => PrettyLevel -> [Array a] -> Doc
pPrint :: Array a -> Doc
$cpPrint :: forall a. (Pretty a, Storable a) => Array a -> Doc
pPrintPrec :: PrettyLevel -> Rational -> Array a -> Doc
$cpPrintPrec :: forall a.
(Pretty a, Storable a) =>
PrettyLevel -> Rational -> Array a -> Doc
Pretty, (forall x. Array a -> Rep (Array a) x)
-> (forall x. Rep (Array a) x -> Array a) -> Generic (Array a)
forall x. Rep (Array a) x -> Array a
forall x. Array a -> Rep (Array a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Array a) x -> Array a
forall a x. Array a -> Rep (Array a) x
$cto :: forall a x. Rep (Array a) x -> Array a
$cfrom :: forall a x. Array a -> Rep (Array a) x
Generic, Typeable (Array a)
DataType
Constr
Typeable (Array a)
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> Array a -> c (Array a))
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c (Array a))
-> (Array a -> Constr)
-> (Array a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c (Array a)))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Array a)))
-> ((forall b. Data b => b -> b) -> Array a -> Array a)
-> (forall r r'.
    (r -> r' -> r)
    -> r -> (forall d. Data d => d -> r') -> Array a -> r)
-> (forall r r'.
    (r' -> r -> r)
    -> r -> (forall d. Data d => d -> r') -> Array a -> r)
-> (forall u. (forall d. Data d => d -> u) -> Array a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> Array a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> Array a -> m (Array a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Array a -> m (Array a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Array a -> m (Array a))
-> Data (Array a)
Array a -> DataType
Array a -> Constr
(forall d. Data d => c (t d)) -> Maybe (c (Array a))
(forall b. Data b => b -> b) -> Array a -> Array a
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array a -> c (Array a)
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array a)
forall a. (Data a, Storable a) => Typeable (Array a)
forall a. (Data a, Storable a) => Array a -> DataType
forall a. (Data a, Storable a) => Array a -> Constr
forall a.
(Data a, Storable a) =>
(forall b. Data b => b -> b) -> Array a -> Array a
forall a u.
(Data a, Storable a) =>
Int -> (forall d. Data d => d -> u) -> Array a -> u
forall a u.
(Data a, Storable a) =>
(forall d. Data d => d -> u) -> Array a -> [u]
forall a r r'.
(Data a, Storable a) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
forall a r r'.
(Data a, Storable a) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
forall a (m :: * -> *).
(Data a, Storable a, Monad m) =>
(forall d. Data d => d -> m d) -> Array a -> m (Array a)
forall a (m :: * -> *).
(Data a, Storable a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array a -> m (Array a)
forall a (c :: * -> *).
(Data a, Storable a) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array a)
forall a (c :: * -> *).
(Data a, Storable a) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array a -> c (Array a)
forall a (t :: * -> *) (c :: * -> *).
(Data a, Storable a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Array a))
forall a (t :: * -> * -> *) (c :: * -> *).
(Data a, Storable a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Array a))
forall a.
Typeable a
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> a -> c a)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c a)
-> (a -> Constr)
-> (a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c a))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c a))
-> ((forall b. Data b => b -> b) -> a -> a)
-> (forall r r'.
    (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall r r'.
    (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall u. (forall d. Data d => d -> u) -> a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> Data a
forall u. Int -> (forall d. Data d => d -> u) -> Array a -> u
forall u. (forall d. Data d => d -> u) -> Array a -> [u]
forall r r'.
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
forall r r'.
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> Array a -> m (Array a)
forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> Array a -> m (Array a)
forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array a)
forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array a -> c (Array a)
forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c (Array a))
forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Array a))
$cA :: Constr
$tArray :: DataType
gmapMo :: (forall d. Data d => d -> m d) -> Array a -> m (Array a)
$cgmapMo :: forall a (m :: * -> *).
(Data a, Storable a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array a -> m (Array a)
gmapMp :: (forall d. Data d => d -> m d) -> Array a -> m (Array a)
$cgmapMp :: forall a (m :: * -> *).
(Data a, Storable a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array a -> m (Array a)
gmapM :: (forall d. Data d => d -> m d) -> Array a -> m (Array a)
$cgmapM :: forall a (m :: * -> *).
(Data a, Storable a, Monad m) =>
(forall d. Data d => d -> m d) -> Array a -> m (Array a)
gmapQi :: Int -> (forall d. Data d => d -> u) -> Array a -> u
$cgmapQi :: forall a u.
(Data a, Storable a) =>
Int -> (forall d. Data d => d -> u) -> Array a -> u
gmapQ :: (forall d. Data d => d -> u) -> Array a -> [u]
$cgmapQ :: forall a u.
(Data a, Storable a) =>
(forall d. Data d => d -> u) -> Array a -> [u]
gmapQr :: (r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
$cgmapQr :: forall a r r'.
(Data a, Storable a) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
gmapQl :: (r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
$cgmapQl :: forall a r r'.
(Data a, Storable a) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array a -> r
gmapT :: (forall b. Data b => b -> b) -> Array a -> Array a
$cgmapT :: forall a.
(Data a, Storable a) =>
(forall b. Data b => b -> b) -> Array a -> Array a
dataCast2 :: (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Array a))
$cdataCast2 :: forall a (t :: * -> * -> *) (c :: * -> *).
(Data a, Storable a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c (Array a))
dataCast1 :: (forall d. Data d => c (t d)) -> Maybe (c (Array a))
$cdataCast1 :: forall a (t :: * -> *) (c :: * -> *).
(Data a, Storable a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Array a))
dataTypeOf :: Array a -> DataType
$cdataTypeOf :: forall a. (Data a, Storable a) => Array a -> DataType
toConstr :: Array a -> Constr
$ctoConstr :: forall a. (Data a, Storable a) => Array a -> Constr
gunfold :: (forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array a)
$cgunfold :: forall a (c :: * -> *).
(Data a, Storable a) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array a)
gfoldl :: (forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array a -> c (Array a)
$cgfoldl :: forall a (c :: * -> *).
(Data a, Storable a) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array a -> c (Array a)
$cp1Data :: forall a. (Data a, Storable a) => Typeable (Array a)
Data)

instance NFData a => NFData (Array a)

instance (Show a, Unbox a) => Show (Array a) where
  showsPrec :: Int -> Array a -> ShowS
showsPrec Int
p = Int -> Array Vector a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
p (Array Vector a -> ShowS)
-> (Array a -> Array Vector a) -> Array a -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

instance (Read a, Unbox a) => Read (Array a) where
  readsPrec :: Int -> ReadS (Array a)
readsPrec Int
p String
s = [(Array Vector a -> Array a
forall a. Array Vector a -> Array a
A Array Vector a
a, String
r) | (Array Vector a
a, String
r) <- Int -> ReadS (Array Vector a)
forall a. Read a => Int -> ReadS a
readsPrec Int
p String
s]

instance Eq (G.Array V.Vector a) => Eq (Array a) where
  Array a
x == :: Array a -> Array a -> Bool
== Array a
y = Array a -> ShapeL
forall a. Array a -> ShapeL
shapeL Array a
x ShapeL -> ShapeL -> Bool
forall a. Eq a => a -> a -> Bool
== Array a -> ShapeL
forall a. Array a -> ShapeL
shapeL Array a
y Bool -> Bool -> Bool
&& Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
x Array Vector a -> Array Vector a -> Bool
forall a. Eq a => a -> a -> Bool
== Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
y
  {-# INLINE (==) #-}

instance Ord (G.Array V.Vector a) => Ord (Array a) where
  compare :: Array a -> Array a -> Ordering
compare Array a
x Array a
y = ShapeL -> ShapeL -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Array a -> ShapeL
forall a. Array a -> ShapeL
shapeL Array a
x) (Array a -> ShapeL
forall a. Array a -> ShapeL
shapeL Array a
y) Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Array Vector a -> Array Vector a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
x) (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
y)
  {-# INLINE compare #-}

-- | The number of elements in the array.
{-# INLINE size #-}
size :: Array a -> Int
size :: Array a -> Int
size = ShapeL -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (ShapeL -> Int) -> (Array a -> ShapeL) -> Array a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> ShapeL
forall a. Array a -> ShapeL
shapeL

-- | The shape of an array, i.e., a list of the sizes of its dimensions.
-- In the linearization of the array the outermost (i.e. first list element)
-- varies most slowly.
-- O(1) time.
shapeL :: Array a -> ShapeL
shapeL :: Array a -> ShapeL
shapeL = Array Vector a -> ShapeL
forall (v :: * -> *) a. Array v a -> ShapeL
G.shapeL (Array Vector a -> ShapeL)
-> (Array a -> Array Vector a) -> Array a -> ShapeL
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | The rank of an array, i.e., the number if dimensions it has.
-- O(1) time.
rank :: Array a -> Int
rank :: Array a -> Int
rank = Array Vector a -> Int
forall (v :: * -> *) a. Array v a -> Int
G.rank (Array Vector a -> Int)
-> (Array a -> Array Vector a) -> Array a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Index into an array.  Fails if the array has rank 0 or if the index is out of bounds.
-- O(1) time.
index :: (HasCallStack, Unbox a) => Array a -> Int -> Array a
index :: Array a -> Int -> Array a
index Array a
a = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Int -> Array Vector a) -> Int -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array Vector a -> Int -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v) =>
Array v a -> Int -> Array v a
G.index (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
a)

-- | Convert to a list with the elements in the linearization order.
-- O(n) time.
toList :: (HasCallStack, Unbox a) => Array a -> [a]
toList :: Array a -> [a]
toList = Array Vector a -> [a]
forall (v :: * -> *) a. (Vector v, VecElem v a) => Array v a -> [a]
G.toList (Array Vector a -> [a])
-> (Array a -> Array Vector a) -> Array a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Convert from a list with the elements given in the linearization order.
-- Fails if the given shape does not have the same number of elements as the list.
-- O(n) time.
fromList :: (HasCallStack, Unbox a) => ShapeL -> [a] -> Array a
fromList :: ShapeL -> [a] -> Array a
fromList ShapeL
ss = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> ([a] -> Array Vector a) -> [a] -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> [a] -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
ShapeL -> [a] -> Array v a
G.fromList ShapeL
ss

-- | Convert to a vector with the elements in the linearization order.
-- O(n) or O(1) time (the latter if the vector is already in the linearization order).
toVector :: (HasCallStack, Unbox a) => Array a -> V.Vector a
toVector :: Array a -> Vector a
toVector = Array Vector a -> Vector a
forall (v :: * -> *) a. (Vector v, VecElem v a) => Array v a -> v a
G.toVector (Array Vector a -> Vector a)
-> (Array a -> Array Vector a) -> Array a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Convert from a vector with the elements given in the linearization order.
-- Fails if the given shape does not have the same number of elements as the list.
-- O(1) time.
fromVector :: (HasCallStack, Unbox a) => ShapeL -> V.Vector a -> Array a
fromVector :: ShapeL -> Vector a -> Array a
fromVector ShapeL
ss = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Vector a -> Array Vector a) -> Vector a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Vector a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
ShapeL -> v a -> Array v a
G.fromVector ShapeL
ss

-- | Make sure the underlying vector is in the linearization order.
-- This is semantically an identity function, but can have big performance
-- implications.
-- O(n) or O(1) time.
{-# INLINE normalize #-}
normalize :: (Unbox a) => Array a -> Array a
normalize :: Array a -> Array a
normalize = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
Array v a -> Array v a
G.normalize (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Change the shape of an array.  Fails if the arrays have different number of elements.
-- O(n) or O(1) time.
reshape :: (HasCallStack, Unbox a) => ShapeL -> Array a -> Array a
reshape :: ShapeL -> Array a -> Array a
reshape ShapeL
s = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
ShapeL -> Array v a -> Array v a
G.reshape ShapeL
s (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Change the size of dimensions with size 1.  These dimension can be changed to any size.
-- All other dimensions must remain the same.
-- O(1) time.
stretch :: (HasCallStack) => ShapeL -> Array a -> Array a
stretch :: ShapeL -> Array a -> Array a
stretch ShapeL
s = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
HasCallStack =>
ShapeL -> Array v a -> Array v a
G.stretch ShapeL
s (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Change the size of the outermost dimension by replication.
stretchOuter :: (HasCallStack) => Int -> Array a -> Array a
stretchOuter :: Int -> Array a -> Array a
stretchOuter Int
s = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
HasCallStack =>
Int -> Array v a -> Array v a
G.stretchOuter Int
s (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Convert a value to a scalar (rank 0) array.
-- O(1) time.
scalar :: (Unbox a) => a -> Array a
scalar :: a -> Array a
scalar = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (a -> Array Vector a) -> a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Array Vector a
forall (v :: * -> *) a. (Vector v, VecElem v a) => a -> Array v a
G.scalar

-- | Convert a scalar (rank 0) array to a value.
-- O(1) time.
unScalar :: (HasCallStack, Unbox a) => Array a -> a
unScalar :: Array a -> a
unScalar = Array Vector a -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
Array v a -> a
G.unScalar (Array Vector a -> a)
-> (Array a -> Array Vector a) -> Array a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Make an array with all elements having the same value.
-- O(1) time
{-# INLINE constant #-}
constant :: (Unbox a) => ShapeL -> a -> Array a
constant :: ShapeL -> a -> Array a
constant ShapeL
sh = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (a -> Array Vector a) -> a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
ShapeL -> a -> Array v a
G.constant ShapeL
sh

-- | Map over the array elements.
-- O(n) time.
{-# INLINE mapA #-}
mapA :: (Unbox a, Unbox b) => (a -> b) -> Array a -> Array b
mapA :: (a -> b) -> Array a -> Array b
mapA a -> b
f = Array Vector b -> Array b
forall a. Array Vector a -> Array a
A (Array Vector b -> Array b)
-> (Array a -> Array Vector b) -> Array a -> Array b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Array Vector a -> Array Vector b
forall (v :: * -> *) a b.
(Vector v, VecElem v a, VecElem v b) =>
(a -> b) -> Array v a -> Array v b
G.mapA a -> b
f (Array Vector a -> Array Vector b)
-> (Array a -> Array Vector a) -> Array a -> Array Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Map over the array elements.
-- O(n) time.
{-# INLINE zipWithA #-}
zipWithA :: (HasCallStack, Unbox a, Unbox b, Unbox c) =>
            (a -> b -> c) -> Array a -> Array b -> Array c
zipWithA :: (a -> b -> c) -> Array a -> Array b -> Array c
zipWithA a -> b -> c
f Array a
a Array b
b = Array Vector c -> Array c
forall a. Array Vector a -> Array a
A (Array Vector c -> Array c) -> Array Vector c -> Array c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> Array Vector a -> Array Vector b -> Array Vector c
forall (v :: * -> *) a b c.
(HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c) =>
(a -> b -> c) -> Array v a -> Array v b -> Array v c
G.zipWithA a -> b -> c
f (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
a) (Array b -> Array Vector b
forall a. Array a -> Array Vector a
unA Array b
b)

-- | Map over the array elements.
-- O(n) time.
{-# INLINE zipWith3A #-}
zipWith3A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d) =>
             (a -> b -> c -> d) -> Array a -> Array b -> Array c -> Array d
zipWith3A :: (a -> b -> c -> d) -> Array a -> Array b -> Array c -> Array d
zipWith3A a -> b -> c -> d
f Array a
a Array b
b Array c
c = Array Vector d -> Array d
forall a. Array Vector a -> Array a
A (Array Vector d -> Array d) -> Array Vector d -> Array d
forall a b. (a -> b) -> a -> b
$ (a -> b -> c -> d)
-> Array Vector a
-> Array Vector b
-> Array Vector c
-> Array Vector d
forall (v :: * -> *) a b c d.
(HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c,
 VecElem v d) =>
(a -> b -> c -> d)
-> Array v a -> Array v b -> Array v c -> Array v d
G.zipWith3A a -> b -> c -> d
f (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
a) (Array b -> Array Vector b
forall a. Array a -> Array Vector a
unA Array b
b) (Array c -> Array Vector c
forall a. Array a -> Array Vector a
unA Array c
c)

-- | Map over the array elements.
-- O(n) time.
{-# INLINE zipWith4A #-}
zipWith4A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d, Unbox e) =>
             (a -> b -> c -> d -> e) -> Array a -> Array b -> Array c -> Array d -> Array e
zipWith4A :: (a -> b -> c -> d -> e)
-> Array a -> Array b -> Array c -> Array d -> Array e
zipWith4A a -> b -> c -> d -> e
f Array a
a Array b
b Array c
c Array d
d = Array Vector e -> Array e
forall a. Array Vector a -> Array a
A (Array Vector e -> Array e) -> Array Vector e -> Array e
forall a b. (a -> b) -> a -> b
$ (a -> b -> c -> d -> e)
-> Array Vector a
-> Array Vector b
-> Array Vector c
-> Array Vector d
-> Array Vector e
forall (v :: * -> *) a b c d e.
(HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c,
 VecElem v d, VecElem v e) =>
(a -> b -> c -> d -> e)
-> Array v a -> Array v b -> Array v c -> Array v d -> Array v e
G.zipWith4A a -> b -> c -> d -> e
f (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
a) (Array b -> Array Vector b
forall a. Array a -> Array Vector a
unA Array b
b) (Array c -> Array Vector c
forall a. Array a -> Array Vector a
unA Array c
c) (Array d -> Array Vector d
forall a. Array a -> Array Vector a
unA Array d
d)

-- | Map over the array elements.
-- O(n) time.
{-# INLINE zipWith5A #-}
zipWith5A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f) =>
             (a -> b -> c -> d -> e -> f) -> Array a -> Array b -> Array c -> Array d -> Array e -> Array f
zipWith5A :: (a -> b -> c -> d -> e -> f)
-> Array a -> Array b -> Array c -> Array d -> Array e -> Array f
zipWith5A a -> b -> c -> d -> e -> f
f Array a
a Array b
b Array c
c Array d
d Array e
e = Array Vector f -> Array f
forall a. Array Vector a -> Array a
A (Array Vector f -> Array f) -> Array Vector f -> Array f
forall a b. (a -> b) -> a -> b
$ (a -> b -> c -> d -> e -> f)
-> Array Vector a
-> Array Vector b
-> Array Vector c
-> Array Vector d
-> Array Vector e
-> Array Vector f
forall (v :: * -> *) a b c d e f.
(HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c,
 VecElem v d, VecElem v e, VecElem v f) =>
(a -> b -> c -> d -> e -> f)
-> Array v a
-> Array v b
-> Array v c
-> Array v d
-> Array v e
-> Array v f
G.zipWith5A a -> b -> c -> d -> e -> f
f (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
a) (Array b -> Array Vector b
forall a. Array a -> Array Vector a
unA Array b
b) (Array c -> Array Vector c
forall a. Array a -> Array Vector a
unA Array c
c) (Array d -> Array Vector d
forall a. Array a -> Array Vector a
unA Array d
d) (Array e -> Array Vector e
forall a. Array a -> Array Vector a
unA Array e
e)

-- | Pad each dimension on the low and high side with the given value.
-- O(n) time.
pad :: (HasCallStack, Unbox a) => [(Int, Int)] -> a -> Array a -> Array a
pad :: [(Int, Int)] -> a -> Array a -> Array a
pad [(Int, Int)]
ps a
v = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Int)] -> a -> Array Vector a -> Array Vector a
forall a (v :: * -> *).
(Vector v, VecElem v a) =>
[(Int, Int)] -> a -> Array v a -> Array v a
G.pad [(Int, Int)]
ps a
v (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Do an arbitrary array transposition.
-- Fails if the transposition argument is not a permutation of the numbers
-- [0..r-1], where r is the rank of the array.
-- O(1) time.
transpose :: (HasCallStack) => [Int] -> Array a -> Array a
transpose :: ShapeL -> Array a -> Array a
transpose ShapeL
is = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
HasCallStack =>
ShapeL -> Array v a -> Array v a
G.transpose ShapeL
is (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Append two arrays along the outermost dimension.
-- All dimensions, except the outermost, must be the same.
-- O(n) time.
append :: (HasCallStack, Unbox a) => Array a -> Array a -> Array a
append :: Array a -> Array a -> Array a
append Array a
x Array a
y = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a) -> Array Vector a -> Array a
forall a b. (a -> b) -> a -> b
$ Array Vector a -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
Array v a -> Array v a -> Array v a
G.append (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
x) (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
y)

-- | Concatenate a number of arrays into a single array.
-- Fails if any, but the outer, dimensions differ.
-- O(n) time.
concatOuter :: (HasCallStack, Unbox a) => [Array a] -> Array a
concatOuter :: [Array a] -> Array a
concatOuter = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> ([Array a] -> Array Vector a) -> [Array a] -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Array Vector a] -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
[Array v a] -> Array v a
G.concatOuter ([Array Vector a] -> Array Vector a)
-> ([Array a] -> [Array Vector a]) -> [Array a] -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Array a] -> [Array Vector a]
coerce

-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost
-- dimension of the result array.  All the arrays must have the same shape.
-- O(n) time.
ravel :: (HasCallStack, Unbox a) => D.Array (Array a) -> Array a
ravel :: Array (Array a) -> Array a
ravel = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array (Array a) -> Array Vector a)
-> Array (Array a)
-> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array Vector (Array Vector a) -> Array Vector a
forall (v :: * -> *) (v' :: * -> *) a.
(HasCallStack, Vector v, Vector v', VecElem v a,
 VecElem v' (Array v a)) =>
Array v' (Array v a) -> Array v a
G.ravel (Array Vector (Array Vector a) -> Array Vector a)
-> (Array (Array a) -> Array Vector (Array Vector a))
-> Array (Array a)
-> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array a -> Array Vector a)
-> Array Vector (Array a) -> Array Vector (Array Vector a)
forall (v :: * -> *) a b.
(Vector v, VecElem v a, VecElem v b) =>
(a -> b) -> Array v a -> Array v b
G.mapA Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA (Array Vector (Array a) -> Array Vector (Array Vector a))
-> (Array (Array a) -> Array Vector (Array a))
-> Array (Array a)
-> Array Vector (Array Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (Array a) -> Array Vector (Array a)
forall a. Array a -> Array Vector a
D.unA

-- | Turn an array into a nested array, this is the inverse of 'ravel'.
-- I.e., @ravel . unravel == id@.
unravel :: (HasCallStack, Unbox a) => Array a -> D.Array (Array a)
unravel :: Array a -> Array (Array a)
unravel = Array Vector (Array a) -> Array (Array a)
forall a. Array Vector a -> Array a
D.A (Array Vector (Array a) -> Array (Array a))
-> (Array a -> Array Vector (Array a))
-> Array a
-> Array (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array Vector a -> Array a)
-> Array Vector (Array Vector a) -> Array Vector (Array a)
forall (v :: * -> *) a b.
(Vector v, VecElem v a, VecElem v b) =>
(a -> b) -> Array v a -> Array v b
G.mapA Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector (Array Vector a) -> Array Vector (Array a))
-> (Array a -> Array Vector (Array Vector a))
-> Array a
-> Array Vector (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array Vector a -> Array Vector (Array Vector a)
forall (v :: * -> *) (v' :: * -> *) a.
(Vector v, Vector v', VecElem v a, VecElem v' (Array v a)) =>
Array v a -> Array v' (Array v a)
G.unravel (Array Vector a -> Array Vector (Array Vector a))
-> (Array a -> Array Vector a)
-> Array a
-> Array Vector (Array Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Make a window of the outermost dimensions.
-- The rank increases with the length of the window list.
-- E.g., if the shape of the array is @[10,12,8]@ and
-- the window size is @[3,3]@ then the resulting array will have shape
-- @[8,10,3,3,8]@.
-- O(1) time.
window :: (HasCallStack) => [Int] -> Array a -> Array a
window :: ShapeL -> Array a -> Array a
window ShapeL
ws = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v) =>
ShapeL -> Array v a -> Array v a
G.window ShapeL
ws (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Stride the outermost dimensions.
-- E.g., if the array shape is @[10,12,8]@ and the strides are
-- @[2,2]@ then the resulting shape will be @[5,6,8]@.
-- O(1) time.
stride :: (HasCallStack) => [Int] -> Array a -> Array a
stride :: ShapeL -> Array a -> Array a
stride ShapeL
ws = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v) =>
ShapeL -> Array v a -> Array v a
G.stride ShapeL
ws (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Rotate the array k times along the d'th dimension.
-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4,
-- the resulting shape will be @[2, 4, 3, 2]@.
rotate :: (HasCallStack, Unbox a) => Int -> Int -> Array a -> Array a
rotate :: Int -> Int -> Array a -> Array a
rotate Int
d Int
k = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
Int -> Int -> Array v a -> Array v a
G.rotate Int
d Int
k (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Extract a slice of an array.
-- The first argument is a list of (offset, length) pairs.
-- The length of the slicing argument must not exceed the rank of the arrar.
-- The extracted slice mul fall within the array dimensions.
-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@.
-- O(1) time.
slice :: (HasCallStack) => [(Int, Int)] -> Array a -> Array a
slice :: [(Int, Int)] -> Array a -> Array a
slice [(Int, Int)]
ss = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Int)] -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
HasCallStack =>
[(Int, Int)] -> Array v a -> Array v a
G.slice [(Int, Int)]
ss (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Apply a function to the subarrays /n/ levels down and make
-- the results into an array with the same /n/ outermost dimensions.
-- The /n/ must not exceed the rank of the array.
-- O(n) time.
rerank :: (HasCallStack, Unbox a, Unbox b) => Int -> (Array a -> Array b) -> Array a -> Array b
rerank :: Int -> (Array a -> Array b) -> Array a -> Array b
rerank Int
n Array a -> Array b
f = Array Vector b -> Array b
forall a. Array Vector a -> Array a
A (Array Vector b -> Array b)
-> (Array a -> Array Vector b) -> Array a -> Array b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> (Array Vector a -> Array Vector b)
-> Array Vector a
-> Array Vector b
forall (v :: * -> *) (v' :: * -> *) a b.
(HasCallStack, Vector v, Vector v', VecElem v a, VecElem v' b) =>
Int -> (Array v a -> Array v' b) -> Array v a -> Array v' b
G.rerank Int
n (Array b -> Array Vector b
forall a. Array a -> Array Vector a
unA (Array b -> Array Vector b)
-> (Array Vector a -> Array b) -> Array Vector a -> Array Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array b
f (Array a -> Array b)
-> (Array Vector a -> Array a) -> Array Vector a -> Array b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array Vector a -> Array a
forall a. Array Vector a -> Array a
A) (Array Vector a -> Array Vector b)
-> (Array a -> Array Vector a) -> Array a -> Array Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Apply a two-argument function to the subarrays /n/ levels down and make
-- the results into an array with the same /n/ outermost dimensions.
-- The /n/ must not exceed the rank of the array.
-- O(n) time.
rerank2 :: (HasCallStack, Unbox a, Unbox b, Unbox c) =>
           Int -> (Array a -> Array b -> Array c) -> Array a -> Array b -> Array c
rerank2 :: Int
-> (Array a -> Array b -> Array c) -> Array a -> Array b -> Array c
rerank2 Int
n Array a -> Array b -> Array c
f Array a
ta Array b
tb = Array Vector c -> Array c
forall a. Array Vector a -> Array a
A (Array Vector c -> Array c) -> Array Vector c -> Array c
forall a b. (a -> b) -> a -> b
$ Int
-> (Array Vector a -> Array Vector b -> Array Vector c)
-> Array Vector a
-> Array Vector b
-> Array Vector c
forall (v :: * -> *) a b c.
(HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c) =>
Int
-> (Array v a -> Array v b -> Array v c)
-> Array v a
-> Array v b
-> Array v c
G.rerank2 Int
n (\ Array Vector a
a Array Vector b
b -> Array c -> Array Vector c
forall a. Array a -> Array Vector a
unA (Array c -> Array Vector c) -> Array c -> Array Vector c
forall a b. (a -> b) -> a -> b
$ Array a -> Array b -> Array c
f (Array Vector a -> Array a
forall a. Array Vector a -> Array a
A Array Vector a
a) (Array Vector b -> Array b
forall a. Array Vector a -> Array a
A Array Vector b
b)) (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
ta) (Array b -> Array Vector b
forall a. Array a -> Array Vector a
unA Array b
tb)

-- | Reverse the given dimensions, with the outermost being dimension 0.
-- O(1) time.
rev :: [Int] -> Array a -> Array a
rev :: ShapeL -> Array a -> Array a
rev ShapeL
rs = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
HasCallStack =>
ShapeL -> Array v a -> Array v a
G.rev ShapeL
rs (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Reduce all elements of an array into a rank 0 array.
-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'.
-- O(n) time.
reduce :: (Unbox a) => (a -> a -> a) -> a -> Array a -> Array a
reduce :: (a -> a -> a) -> a -> Array a -> Array a
reduce a -> a -> a
f a
z = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> a -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
(a -> a -> a) -> a -> Array v a -> Array v a
G.reduce a -> a -> a
f a
z (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Constrained version of 'foldr' for Arrays.
foldrA :: (Unbox a) => (a -> b -> b) -> b -> Array a -> b
foldrA :: (a -> b -> b) -> b -> Array a -> b
foldrA a -> b -> b
f b
z = (a -> b -> b) -> b -> Array Vector a -> b
forall (v :: * -> *) a b.
(Vector v, VecElem v a) =>
(a -> b -> b) -> b -> Array v a -> b
G.foldrA a -> b -> b
f b
z (Array Vector a -> b)
-> (Array a -> Array Vector a) -> Array a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Constrained version of 'traverse' for Arrays.
traverseA
  :: (Unbox a, Unbox b, Applicative f) => (a -> f b) -> Array a -> f (Array b)
traverseA :: (a -> f b) -> Array a -> f (Array b)
traverseA a -> f b
f = (Array Vector b -> Array b) -> f (Array Vector b) -> f (Array b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array Vector b -> Array b
forall a. Array Vector a -> Array a
A (f (Array Vector b) -> f (Array b))
-> (Array a -> f (Array Vector b)) -> Array a -> f (Array b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f b) -> Array Vector a -> f (Array Vector b)
forall (v :: * -> *) a b (f :: * -> *).
(Vector v, VecElem v a, VecElem v b, Applicative f) =>
(a -> f b) -> Array v a -> f (Array v b)
G.traverseA a -> f b
f (Array Vector a -> f (Array Vector b))
-> (Array a -> Array Vector a) -> Array a -> f (Array Vector b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Check if all elements of the array are equal.
{-# INLINE allSameA #-}
allSameA :: (Unbox a, Eq a) => Array a -> Bool
allSameA :: Array a -> Bool
allSameA = Array Vector a -> Bool
forall (v :: * -> *) a.
(Vector v, VecElem v a, Eq a) =>
Array v a -> Bool
G.allSameA (Array Vector a -> Bool)
-> (Array a -> Array Vector a) -> Array a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

instance (Arbitrary a, Unbox a) => Arbitrary (Array a) where arbitrary :: Gen (Array a)
arbitrary = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> Gen (Array Vector a) -> Gen (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Array Vector a)
forall a. Arbitrary a => Gen a
arbitrary

-- | Sum of all elements.
{-# INLINE sumA #-}
sumA :: (Unbox a, Num a) => Array a -> a
sumA :: Array a -> a
sumA = Array Vector a -> a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Num a) =>
Array v a -> a
G.sumA (Array Vector a -> a)
-> (Array a -> Array Vector a) -> Array a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Product of all elements.
{-# INLINE productA #-}
productA :: (Unbox a, Num a) => Array a -> a
productA :: Array a -> a
productA = Array Vector a -> a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Num a) =>
Array v a -> a
G.productA (Array Vector a -> a)
-> (Array a -> Array Vector a) -> Array a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Maximum of all elements.
{-# INLINE maximumA #-}
maximumA :: (HasCallStack, Unbox a, Ord a) => Array a -> a
maximumA :: Array a -> a
maximumA = Array Vector a -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Ord a) =>
Array v a -> a
G.maximumA (Array Vector a -> a)
-> (Array a -> Array Vector a) -> Array a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Minimum of all elements.
{-# INLINE minimumA #-}
minimumA :: (HasCallStack, Unbox a, Ord a) => Array a -> a
minimumA :: Array a -> a
minimumA = Array Vector a -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Ord a) =>
Array v a -> a
G.minimumA (Array Vector a -> a)
-> (Array a -> Array Vector a) -> Array a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Test if the predicate holds for any element.
{-# INLINE anyA #-}
anyA :: Unbox a => (a -> Bool) -> Array a -> Bool
anyA :: (a -> Bool) -> Array a -> Bool
anyA a -> Bool
p = (a -> Bool) -> Array Vector a -> Bool
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
(a -> Bool) -> Array v a -> Bool
G.anyA a -> Bool
p (Array Vector a -> Bool)
-> (Array a -> Array Vector a) -> Array a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Test if the predicate holds for all elements.
{-# INLINE allA #-}
allA :: Unbox a => (a -> Bool) -> Array a -> Bool
allA :: (a -> Bool) -> Array a -> Bool
allA a -> Bool
p = (a -> Bool) -> Array Vector a -> Bool
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
(a -> Bool) -> Array v a -> Bool
G.allA a -> Bool
p (Array Vector a -> Bool)
-> (Array a -> Array Vector a) -> Array a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Put the dimensions of the argument into the specified dimensions,
-- and just replicate the data along all other dimensions.
-- The list of dimensions indicies must have the same rank as the argument array
-- and it must be strictly ascending.
{-# INLINE broadcast #-}
broadcast :: (HasCallStack, Unbox a) =>
             [Int] -> ShapeL -> Array a -> Array a
broadcast :: ShapeL -> ShapeL -> Array a -> Array a
broadcast ShapeL
ds ShapeL
sh = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A(Array Vector a -> Array a)
-> (Array a -> Array Vector a) -> Array a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> ShapeL -> Array Vector a -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
ShapeL -> ShapeL -> Array v a -> Array v a
G.broadcast ShapeL
ds ShapeL
sh (Array Vector a -> Array Vector a)
-> (Array a -> Array Vector a) -> Array a -> Array Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA

-- | Update the array at the specified indicies to the associated value.
{-# INLINE update #-}
update :: (HasCallStack, Unbox a) =>
          Array a -> [([Int], a)] -> Array a
update :: Array a -> [(ShapeL, a)] -> Array a
update Array a
a = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> ([(ShapeL, a)] -> Array Vector a) -> [(ShapeL, a)] -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array Vector a -> [(ShapeL, a)] -> Array Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a) =>
Array v a -> [(ShapeL, a)] -> Array v a
G.update (Array a -> Array Vector a
forall a. Array a -> Array Vector a
unA Array a
a)

-- | Generate an array with a function that computes the value for each index.
{-# INLINE generate #-}
generate :: (Unbox a) => ShapeL -> ([Int] -> a) -> Array a
generate :: ShapeL -> (ShapeL -> a) -> Array a
generate ShapeL
sh = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> ((ShapeL -> a) -> Array Vector a) -> (ShapeL -> a) -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> (ShapeL -> a) -> Array Vector a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> (ShapeL -> a) -> Array v a
G.generate ShapeL
sh

-- | Iterate a function n times.
{-# INLINE iterateN #-}
iterateN :: (Unbox a) =>
            Int -> (a -> a) -> a -> Array a
iterateN :: Int -> (a -> a) -> a -> Array a
iterateN Int
n a -> a
f = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (a -> Array Vector a) -> a -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (a -> a) -> a -> Array Vector a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
Int -> (a -> a) -> a -> Array v a
G.iterateN Int
n a -> a
f

-- | Generate a vector from 0 to n-1.
{-# INLINE iota #-}
iota :: (Unbox a, Enum a, Num a) => Int -> Array a
iota :: Int -> Array a
iota = Array Vector a -> Array a
forall a. Array Vector a -> Array a
A (Array Vector a -> Array a)
-> (Int -> Array Vector a) -> Int -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Array Vector a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Enum a, Num a) =>
Int -> Array v a
G.iota

-- | Convert between types by just reinterpreting the bits as another type.
-- For instance the floating point number @(1.5 :: Float)@ will convert to
-- @(0x3fc00000 :: Word32)@ since they have the same bit representation.
{-# INLINE bitcast #-}
bitcast :: forall a b . (HasCallStack, Unbox a, Unbox b) => Array a -> Array b
bitcast :: Array a -> Array b
bitcast (A (G.A ShapeL
sh (T ShapeL
ss Int
o Vector a
v)))
  | Int
sza Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
szb
  = String -> Array b
forall a. HasCallStack => String -> a
error (String -> Array b) -> String -> Array b
forall a b. (a -> b) -> a -> b
$ String
"bitcast: the types must have the same size. " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
sza, Int
szb)
  | Bool
otherwise
  = Array Vector b -> Array b
forall a. Array Vector a -> Array a
A (ShapeL -> T Vector b -> Array Vector b
forall (v :: * -> *) a. ShapeL -> T v a -> Array v a
G.A ShapeL
sh (ShapeL -> Int -> Vector b -> T Vector b
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T ShapeL
ss Int
o (Vector a -> Vector b
forall a b. (Storable a, Storable b) => Vector a -> Vector b
V.unsafeCast Vector a
v)))
  where sza :: Int
sza = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
        szb :: Int
szb = b -> Int
forall a. Storable a => a -> Int
sizeOf (b
forall a. HasCallStack => a
undefined :: b)