-- 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 AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Array.Internal.Ranked(
  Array(..), Vector, ShapeL,
  size, shapeL, rank,
  toList, fromList, toVector, fromVector,
  normalize,
  scalar, unScalar, constant,
  reshape, stretch, stretchOuter, transpose,
  index, pad,
  mapA, zipWithA, zipWith3A,
  append, concatOuter,
  ravel, unravel,
  window, stride, rotate,
  slice, rerank, rerank2, rev,
  reduce, foldrA, traverseA,
  allSameA,
  sumA, productA, minimumA, maximumA,
  anyA, allA,
  broadcast,
  generate, iterateN, iota,
  ) where
import Control.DeepSeq
import Data.Data(Data)
import Data.Coerce(coerce)
import qualified Data.Vector as V
import GHC.Stack
import GHC.TypeLits(KnownNat, type (+), type (<=))
import Test.QuickCheck hiding (generate)

import GHC.Generics(Generic)
import Data.Array.Internal.Dynamic()  -- Vector instance
import qualified Data.Array.Internal.RankedG as G
import Data.Array.Internal(ShapeL, Vector(..))
import Text.PrettyPrint.HughesPJClass hiding ((<>))

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

instance NFData a => NFData (Array n a)

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

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

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

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

-- | The number of elements in the array.
{-# INLINE size #-}
size :: Array n a -> Int
size :: Array n a -> Int
size = ShapeL -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (ShapeL -> Int) -> (Array n a -> ShapeL) -> Array n a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> ShapeL
forall (n :: Nat) a. Array n 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 n a -> ShapeL
shapeL :: Array n a -> ShapeL
shapeL = Array n Vector a -> ShapeL
forall (n :: Nat) (v :: * -> *) a. Array n v a -> ShapeL
G.shapeL (Array n Vector a -> ShapeL)
-> (Array n a -> Array n Vector a) -> Array n a -> ShapeL
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

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

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

-- | Convert to a list with the elements in the linearization order.
-- O(n) time.
toList :: Array n a -> [a]
toList :: Array n a -> [a]
toList = Array n Vector a -> [a]
forall (v :: * -> *) a (n :: Nat).
(Vector v, VecElem v a) =>
Array n v a -> [a]
G.toList (Array n Vector a -> [a])
-> (Array n a -> Array n Vector a) -> Array n a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: forall n a . (KnownNat n) => ShapeL -> [a] -> Array n a
fromList :: ShapeL -> [a] -> Array n a
fromList ShapeL
ss = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> ([a] -> Array n Vector a) -> [a] -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> [a] -> Array n Vector a
forall (n :: Nat) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, KnownNat n) =>
ShapeL -> [a] -> Array n 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 :: Array n a -> V.Vector a
toVector :: Array n a -> Vector a
toVector = Array n Vector a -> Vector a
forall (v :: * -> *) a (n :: Nat).
(Vector v, VecElem v a) =>
Array n v a -> v a
G.toVector (Array n Vector a -> Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: forall n a . (KnownNat n) => ShapeL -> V.Vector a -> Array n a
fromVector :: ShapeL -> Vector a -> Array n a
fromVector ShapeL
ss = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Vector a -> Array n Vector a) -> Vector a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Vector a -> Array n Vector a
forall (n :: Nat) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, KnownNat n) =>
ShapeL -> v a -> Array n 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.
normalize :: (KnownNat n) => Array n a -> Array n a
normalize :: Array n a -> Array n a
normalize = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n Vector a -> Array n Vector a
forall (v :: * -> *) a (n :: Nat).
(Vector v, VecElem v a, KnownNat n) =>
Array n v a -> Array n v a
G.normalize (Array n Vector a -> Array n Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: forall n' n a . (KnownNat n, KnownNat n') => ShapeL -> Array n a -> Array n' a
reshape :: ShapeL -> Array n a -> Array n' a
reshape ShapeL
s = Array n' Vector a -> Array n' a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n' Vector a -> Array n' a)
-> (Array n a -> Array n' Vector a) -> Array n a -> Array n' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array n Vector a -> Array n' Vector a
forall (n :: Nat) (n' :: Nat) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, KnownNat n, KnownNat n') =>
ShapeL -> Array n v a -> Array n' v a
G.reshape ShapeL
s (Array n Vector a -> Array n' Vector a)
-> (Array n a -> Array n Vector a)
-> Array n a
-> Array n' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: ShapeL -> Array n a -> Array n a
stretch :: ShapeL -> Array n a -> Array n a
stretch ShapeL
s = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array n Vector a -> Array n Vector a
forall (n :: Nat) (v :: * -> *) a.
HasCallStack =>
ShapeL -> Array n v a -> Array n v a
G.stretch ShapeL
s (Array n Vector a -> Array n Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

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

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

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

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

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

instance Functor (Array n) where
  fmap :: (a -> b) -> Array n a -> Array n b
fmap = (a -> b) -> Array n a -> Array n b
forall a b (n :: Nat). (a -> b) -> Array n a -> Array n b
mapA

instance Foldable (Array n) where
  foldr :: (a -> b -> b) -> b -> Array n a -> b
foldr = (a -> b -> b) -> b -> Array n a -> b
forall a b (n :: Nat). (a -> b -> b) -> b -> Array n a -> b
foldrA

instance Traversable (Array n) where
  traverse :: (a -> f b) -> Array n a -> f (Array n b)
traverse = (a -> f b) -> Array n a -> f (Array n b)
forall (f :: * -> *) a b (n :: Nat).
Applicative f =>
(a -> f b) -> Array n a -> f (Array n b)
traverseA

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

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

-- | Pad each dimension on the low and high side with the given value.
-- O(n) time.
pad :: (KnownNat n) => [(Int, Int)] -> a -> Array n a -> Array n a
pad :: [(Int, Int)] -> a -> Array n a -> Array n a
pad [(Int, Int)]
ps a
v = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Int)] -> a -> Array n Vector a -> Array n Vector a
forall (n :: Nat) a (v :: * -> *).
(Vector v, VecElem v a) =>
[(Int, Int)] -> a -> Array n v a -> Array n v a
G.pad [(Int, Int)]
ps a
v (Array n Vector a -> Array n Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: (KnownNat n) => [Int] -> Array n a -> Array n a
transpose :: ShapeL -> Array n a -> Array n a
transpose ShapeL
is = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array n Vector a -> Array n Vector a
forall (n :: Nat) (v :: * -> *) a.
KnownNat n =>
ShapeL -> Array n v a -> Array n v a
G.transpose ShapeL
is (Array n Vector a -> Array n Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

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

-- | Concatenate a number of arrays into a single array.
-- Fails if any, but the outer, dimensions differ.
-- O(n) time.
concatOuter :: (KnownNat n) => [Array n a] -> Array n a
concatOuter :: [Array n a] -> Array n a
concatOuter = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> ([Array n a] -> Array n Vector a) -> [Array n a] -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Array n Vector a] -> Array n Vector a
forall (v :: * -> *) a (n :: Nat).
(Vector v, VecElem v a, KnownNat n) =>
[Array n v a] -> Array n v a
G.concatOuter ([Array n Vector a] -> Array n Vector a)
-> ([Array n a] -> [Array n Vector a])
-> [Array n a]
-> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Array n a] -> [Array n 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 :: (KnownNat (1+n)) =>
         Array 1 (Array n a) -> Array (1+n) a
ravel :: Array 1 (Array n a) -> Array (1 + n) a
ravel = Array (1 + n) Vector a -> Array (1 + n) a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array (1 + n) Vector a -> Array (1 + n) a)
-> (Array 1 (Array n a) -> Array (1 + n) Vector a)
-> Array 1 (Array n a)
-> Array (1 + n) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array 1 Vector (Array n Vector a) -> Array (1 + n) Vector a
forall (v :: * -> *) (v' :: * -> *) a (n :: Nat).
(Vector v, Vector v', VecElem v a, VecElem v' (Array n v a),
 KnownNat (1 + n)) =>
Array 1 v' (Array n v a) -> Array (1 + n) v a
G.ravel (Array 1 Vector (Array n Vector a) -> Array (1 + n) Vector a)
-> (Array 1 (Array n a) -> Array 1 Vector (Array n Vector a))
-> Array 1 (Array n a)
-> Array (1 + n) Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array n a -> Array n Vector a)
-> Array 1 Vector (Array n a) -> Array 1 Vector (Array n Vector a)
forall (v :: * -> *) a b (n :: Nat).
(Vector v, VecElem v a, VecElem v b) =>
(a -> b) -> Array n v a -> Array n v b
G.mapA Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA (Array 1 Vector (Array n a) -> Array 1 Vector (Array n Vector a))
-> (Array 1 (Array n a) -> Array 1 Vector (Array n a))
-> Array 1 (Array n a)
-> Array 1 Vector (Array n Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array 1 (Array n a) -> Array 1 Vector (Array n a)
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

-- | Turn an array into a nested array, this is the inverse of 'ravel'.
-- I.e., @ravel . unravel == id@.
unravel :: Array (1+n) a -> Array 1 (Array n a)
unravel :: Array (1 + n) a -> Array 1 (Array n a)
unravel = Array 1 Vector (Array n a) -> Array 1 (Array n a)
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array 1 Vector (Array n a) -> Array 1 (Array n a))
-> (Array (1 + n) a -> Array 1 Vector (Array n a))
-> Array (1 + n) a
-> Array 1 (Array n a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array n Vector a -> Array n a)
-> Array 1 Vector (Array n Vector a) -> Array 1 Vector (Array n a)
forall (v :: * -> *) a b (n :: Nat).
(Vector v, VecElem v a, VecElem v b) =>
(a -> b) -> Array n v a -> Array n v b
G.mapA Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array 1 Vector (Array n Vector a) -> Array 1 Vector (Array n a))
-> (Array (1 + n) a -> Array 1 Vector (Array n Vector a))
-> Array (1 + n) a
-> Array 1 Vector (Array n a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (1 + n) Vector a -> Array 1 Vector (Array n Vector a)
forall (v :: * -> *) (v' :: * -> *) a (n :: Nat).
(Vector v, Vector v', VecElem v a, VecElem v' (Array n v a)) =>
Array (1 + n) v a -> Array 1 v' (Array n v a)
G.unravel (Array (1 + n) Vector a -> Array 1 Vector (Array n Vector a))
-> (Array (1 + n) a -> Array (1 + n) Vector a)
-> Array (1 + n) a
-> Array 1 Vector (Array n Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (1 + n) a -> Array (1 + n) Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: (KnownNat n, KnownNat n') => [Int] -> Array n a -> Array n' a
window :: ShapeL -> Array n a -> Array n' a
window ShapeL
ws = Array n' Vector a -> Array n' a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n' Vector a -> Array n' a)
-> (Array n a -> Array n' Vector a) -> Array n a -> Array n' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array n Vector a -> Array n' Vector a
forall (n :: Nat) (n' :: Nat) (v :: * -> *) a.
(Vector v, KnownNat n, KnownNat n') =>
ShapeL -> Array n v a -> Array n' v a
G.window ShapeL
ws (Array n Vector a -> Array n' Vector a)
-> (Array n a -> Array n Vector a)
-> Array n a
-> Array n' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: [Int] -> Array n a -> Array n a
stride :: ShapeL -> Array n a -> Array n a
stride ShapeL
ws = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array n Vector a -> Array n Vector a
forall (v :: * -> *) (n :: Nat) a.
Vector v =>
ShapeL -> Array n v a -> Array n v a
G.stride ShapeL
ws (Array n Vector a -> Array n Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: forall d p a.
          (KnownNat p, KnownNat d,
          -- Nonsense
          (d + (p + 1)) ~ ((p + d) + 1),
          (d + p) ~ (p + d),
          1 <= p + 1,
          KnownNat ((p + d) + 1),
          KnownNat (p + 1),
          KnownNat (1 + (p + 1))
          ) =>
          Int -> Array (p + d) a -> Array (p + d + 1) a
rotate :: Int -> Array (p + d) a -> Array ((p + d) + 1) a
rotate Int
k = Array ((p + d) + 1) Vector a -> Array ((p + d) + 1) a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array ((p + d) + 1) Vector a -> Array ((p + d) + 1) a)
-> (Array (p + d) a -> Array ((p + d) + 1) Vector a)
-> Array (p + d) a
-> Array ((p + d) + 1) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Array (p + d) Vector a -> Array ((p + d) + 1) Vector a
forall (d :: Nat) (p :: Nat) (v :: * -> *) a.
(KnownNat p, KnownNat d, Vector v, VecElem v a,
 (d + (p + 1)) ~ ((p + d) + 1), (d + p) ~ (p + d), 1 <= (p + 1),
 KnownNat ((p + d) + 1), KnownNat (p + 1),
 KnownNat (1 + (p + 1))) =>
Int -> Array (p + d) v a -> Array ((p + d) + 1) v a
G.rotate @d @p Int
k (Array (p + d) Vector a -> Array ((p + d) + 1) Vector a)
-> (Array (p + d) a -> Array (p + d) Vector a)
-> Array (p + d) a
-> Array ((p + d) + 1) Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (p + d) a -> Array (p + d) Vector a
forall (n :: Nat) a. Array n a -> Array n 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 array.
-- The extracted slice must fall within the array dimensions.
-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@.
-- O(1) time.
slice :: [(Int, Int)] -> Array n a -> Array n a
slice :: [(Int, Int)] -> Array n a -> Array n a
slice [(Int, Int)]
ss = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Int)] -> Array n Vector a -> Array n Vector a
forall (n :: Nat) (v :: * -> *) a.
[(Int, Int)] -> Array n v a -> Array n v a
G.slice [(Int, Int)]
ss (Array n Vector a -> Array n Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: forall n i o a b . (KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) =>
          (Array i a -> Array o b) -> Array (n+i) a -> Array (n+o) b
rerank :: (Array i a -> Array o b) -> Array (n + i) a -> Array (n + o) b
rerank Array i a -> Array o b
f = Array (n + o) Vector b -> Array (n + o) b
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array (n + o) Vector b -> Array (n + o) b)
-> (Array (n + i) a -> Array (n + o) Vector b)
-> Array (n + i) a
-> Array (n + o) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array i Vector a -> Array o Vector b)
-> Array (n + i) Vector a -> Array (n + o) Vector b
forall (n :: Nat) (i :: Nat) (o :: Nat) (v :: * -> *)
       (v' :: * -> *) a b.
(Vector v, Vector v', VecElem v a, VecElem v' b, KnownNat n,
 KnownNat o, KnownNat (n + o), KnownNat (1 + o)) =>
(Array i v a -> Array o v' b)
-> Array (n + i) v a -> Array (n + o) v' b
G.rerank (Array o b -> Array o Vector b
forall (n :: Nat) a. Array n a -> Array n Vector a
unA (Array o b -> Array o Vector b)
-> (Array i Vector a -> Array o b)
-> Array i Vector a
-> Array o Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array i a -> Array o b
f (Array i a -> Array o b)
-> (Array i Vector a -> Array i a) -> Array i Vector a -> Array o b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array i Vector a -> Array i a
forall (n :: Nat) a. Array n Vector a -> Array n a
A) (Array (n + i) Vector a -> Array (n + o) Vector b)
-> (Array (n + i) a -> Array (n + i) Vector a)
-> Array (n + i) a
-> Array (n + o) Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (n + i) a -> Array (n + i) Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: forall n i o a b c .
           (KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) =>
           (Array i a -> Array i b -> Array o c) -> Array (n+i) a -> Array (n+i) b -> Array (n+o) c
rerank2 :: (Array i a -> Array i b -> Array o c)
-> Array (n + i) a -> Array (n + i) b -> Array (n + o) c
rerank2 Array i a -> Array i b -> Array o c
f Array (n + i) a
ta Array (n + i) b
tb = Array (n + o) Vector c -> Array (n + o) c
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array (n + o) Vector c -> Array (n + o) c)
-> Array (n + o) Vector c -> Array (n + o) c
forall a b. (a -> b) -> a -> b
$ (Array i Vector a -> Array i Vector b -> Array o Vector c)
-> Array (n + i) Vector a
-> Array (n + i) Vector b
-> Array (n + o) Vector c
forall (n :: Nat) (i :: Nat) (o :: Nat) a b c (v :: * -> *).
(Vector v, VecElem v a, VecElem v b, VecElem v c, KnownNat n,
 KnownNat o, KnownNat (n + o), KnownNat (1 + o)) =>
(Array i v a -> Array i v b -> Array o v c)
-> Array (n + i) v a -> Array (n + i) v b -> Array (n + o) v c
G.rerank2 @n (\ Array i Vector a
a Array i Vector b
b -> Array o c -> Array o Vector c
forall (n :: Nat) a. Array n a -> Array n Vector a
unA (Array o c -> Array o Vector c) -> Array o c -> Array o Vector c
forall a b. (a -> b) -> a -> b
$ Array i a -> Array i b -> Array o c
f (Array i Vector a -> Array i a
forall (n :: Nat) a. Array n Vector a -> Array n a
A Array i Vector a
a) (Array i Vector b -> Array i b
forall (n :: Nat) a. Array n Vector a -> Array n a
A Array i Vector b
b)) (Array (n + i) a -> Array (n + i) Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA Array (n + i) a
ta) (Array (n + i) b -> Array (n + i) Vector b
forall (n :: Nat) a. Array n a -> Array n Vector a
unA Array (n + i) b
tb)

-- | Reverse the given dimensions, with the outermost being dimension 0.
-- O(1) time.
rev :: [Int] -> Array n a -> Array n a
rev :: ShapeL -> Array n a -> Array n a
rev ShapeL
rs = Array n Vector a -> Array n a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array n Vector a -> Array n a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> Array n Vector a -> Array n Vector a
forall (n :: Nat) (v :: * -> *) a.
ShapeL -> Array n v a -> Array n v a
G.rev ShapeL
rs (Array n Vector a -> Array n Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array n Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n 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 :: (a -> a -> a) -> a -> Array n a -> Array 0 a
reduce :: (a -> a -> a) -> a -> Array n a -> Array 0 a
reduce a -> a -> a
f a
z = Array 0 Vector a -> Array 0 a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array 0 Vector a -> Array 0 a)
-> (Array n a -> Array 0 Vector a) -> Array n a -> Array 0 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> a -> Array n Vector a -> Array 0 Vector a
forall (v :: * -> *) a (n :: Nat).
(Vector v, VecElem v a) =>
(a -> a -> a) -> a -> Array n v a -> Array 0 v a
G.reduce a -> a -> a
f a
z (Array n Vector a -> Array 0 Vector a)
-> (Array n a -> Array n Vector a) -> Array n a -> Array 0 Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

-- | Constrained version of 'foldr' for Arrays.
--
-- Note that this 'Array' actually has 'Traversable' anyway.
foldrA :: (a -> b -> b) -> b -> Array n a -> b
foldrA :: (a -> b -> b) -> b -> Array n a -> b
foldrA a -> b -> b
f b
z = (a -> b -> b) -> b -> Array n Vector a -> b
forall (v :: * -> *) a b (n :: Nat).
(Vector v, VecElem v a) =>
(a -> b -> b) -> b -> Array n v a -> b
G.foldrA a -> b -> b
f b
z (Array n Vector a -> b)
-> (Array n a -> Array n Vector a) -> Array n a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

-- | Constrained version of 'traverse' for Arrays.
--
-- Note that this 'Array' actually has 'Traversable' anyway.
traverseA :: Applicative f => (a -> f b) -> Array n a -> f (Array n b)
traverseA :: (a -> f b) -> Array n a -> f (Array n b)
traverseA a -> f b
f = (Array n Vector b -> Array n b)
-> f (Array n Vector b) -> f (Array n b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array n Vector b -> Array n b
forall (n :: Nat) a. Array n Vector a -> Array n a
A (f (Array n Vector b) -> f (Array n b))
-> (Array n a -> f (Array n Vector b))
-> Array n a
-> f (Array n b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f b) -> Array n Vector a -> f (Array n Vector b)
forall (v :: * -> *) a b (f :: * -> *) (n :: Nat).
(Vector v, VecElem v a, VecElem v b, Applicative f) =>
(a -> f b) -> Array n v a -> f (Array n v b)
G.traverseA a -> f b
f (Array n Vector a -> f (Array n Vector b))
-> (Array n a -> Array n Vector a)
-> Array n a
-> f (Array n Vector b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array n a -> Array n Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

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

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

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

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

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

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

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

-- | Test if the predicate holds for all elements.
{-# INLINE allA #-}
allA :: (a -> Bool) -> Array r a -> Bool
allA :: (a -> Bool) -> Array r a -> Bool
allA a -> Bool
p = (a -> Bool) -> Array r Vector a -> Bool
forall (v :: * -> *) a (r :: Nat).
(Vector v, VecElem v a) =>
(a -> Bool) -> Array r v a -> Bool
G.allA a -> Bool
p (Array r Vector a -> Bool)
-> (Array r a -> Array r Vector a) -> Array r a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r a -> Array r Vector a
forall (n :: Nat) a. Array n a -> Array n 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.
broadcast :: forall r' r a .
             (HasCallStack, KnownNat r, KnownNat r') =>
             [Int] -> ShapeL -> Array r a -> Array r' a
broadcast :: ShapeL -> ShapeL -> Array r a -> Array r' a
broadcast ShapeL
ds ShapeL
sh = Array r' Vector a -> Array r' a
forall (n :: Nat) a. Array n Vector a -> Array n a
A (Array r' Vector a -> Array r' a)
-> (Array r a -> Array r' Vector a) -> Array r a -> Array r' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> ShapeL -> Array r Vector a -> Array r' Vector a
forall (r' :: Nat) (r :: Nat) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, KnownNat r, KnownNat r') =>
ShapeL -> ShapeL -> Array r v a -> Array r' v a
G.broadcast ShapeL
ds ShapeL
sh (Array r Vector a -> Array r' Vector a)
-> (Array r a -> Array r Vector a)
-> Array r a
-> Array r' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r a -> Array r Vector a
forall (n :: Nat) a. Array n a -> Array n Vector a
unA

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

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

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