numidr/src/Data/NumIdr/Array/Array.idr

613 lines
21 KiB
Idris

module Data.NumIdr.Array.Array
import Data.List
import Data.List1
import Data.Vect
import Data.Zippable
import Data.NP
import Data.NumIdr.Multiply
import Data.NumIdr.PrimArray
import Data.NumIdr.Array.Order
import Data.NumIdr.Array.Coords
%default total
||| Arrays are the central data structure of NumIdr. They are an `n`-dimensional
||| grid of values, where `n` is a value known as the *rank* of the array. Arrays
||| of rank 0 are single values, arrays of rank 1 are vectors, and arrays of rank
||| 2 are matrices.
|||
||| Each array has a *shape*, which is a vector of values giving the dimensions
||| of each axis of the array. The shape is also sometimes used to determine the
||| array's total size.
|||
||| Arrays are indexed by row first, as in the standard mathematical notation for
||| matrices. This is independent of the actual order in which they are stored; the
||| default order is row-major, but this is configurable.
|||
||| @ rk The rank of the array
||| @ s The shape of the array
||| @ a The type of the array's elements
export
data Array : (s : Vect rk Nat) -> (a : Type) -> Type where
||| Internally, arrays are stored using Idris's primitive array type. This is
||| stored along with the array's shape, and a vector of values called the
||| *strides*, which determine how indexes into the internal array should be
||| performed. This is how the order of the array is configurable.
|||
||| @ ord The order of the elements of the array
||| @ sts The strides of the array
||| @ s The shape of the array
MkArray : (ord : Order) -> (sts : Vect rk Nat) ->
(s : Vect rk Nat) -> PrimArray a -> Array s a
--------------------------------------------------------------------------------
-- Properties of arrays
--------------------------------------------------------------------------------
||| Extract the primitive array value.
export
getPrim : Array s a -> PrimArray a
getPrim (MkArray _ _ _ arr) = arr
||| The order of the elements of the array
export
getOrder : Array s a -> Order
getOrder (MkArray ord _ _ _) = ord
||| The strides of the array, returned in the same axis order as in the shape.
export
strides : Array {rk} s a -> Vect rk Nat
strides (MkArray _ sts _ _) = sts
||| The total number of elements of the array
|||
||| This is equivalent to `product s`.
export
size : Array s a -> Nat
size = length . getPrim
||| The shape of the array
export
shape : Array {rk} s a -> Vect rk Nat
shape (MkArray _ _ s _) = s
||| The rank of the array
export
rank : Array s a -> Nat
rank = length . shape
-- Get a list of all coordinates
getAllCoords' : Vect rk Nat -> List (Vect rk Nat)
getAllCoords' = traverse (\case Z => []; S n => [0..n])
getAllCoords : (s : Vect rk Nat) -> List (Coords s)
getAllCoords [] = pure []
getAllCoords (Z :: s) = []
getAllCoords (S d :: s) = [| forget (allFins d) :: getAllCoords s |]
--------------------------------------------------------------------------------
-- Shape view
--------------------------------------------------------------------------------
export
shapeEq : (arr : Array s a) -> s = shape arr
shapeEq (MkArray _ _ _ _) = Refl
public export
data ShapeView : Array s a -> Type where
Shape : (s : Vect rk Nat) -> {0 arr : Array s a} -> ShapeView arr
export
viewShape : (arr : Array s a) -> ShapeView arr
viewShape arr = rewrite shapeEq arr in
Shape (shape arr)
{arr = rewrite sym (shapeEq arr) in arr}
--------------------------------------------------------------------------------
-- Array constructors
--------------------------------------------------------------------------------
||| Create an array by repeating a single value.
|||
||| @ s The shape of the constructed array
||| @ ord The order of the constructed array
export
repeat' : (s : Vect rk Nat) -> (ord : Order) -> a -> Array s a
repeat' s ord x = MkArray ord (calcStrides ord s) s (constant (product s) x)
||| Create an array by repeating a single value.
||| To specify the order of the array, use `repeat'`.
|||
||| @ s The shape of the constructed array
export
repeat : (s : Vect rk Nat) -> a -> Array s a
repeat s = repeat' s COrder
export
zeros : Num a => (s : Vect rk Nat) -> Array s a
zeros s = repeat s 0
export
ones : Num a => (s : Vect rk Nat) -> Array s a
ones s = repeat s 1
||| Create an array given a vector of its elements. The elements of the vector
||| are arranged into the provided shape using the provided order.
|||
||| @ s The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromVect' : (s : Vect rk Nat) -> (ord : Order) -> Vect (product s) a -> Array s a
fromVect' s ord v = MkArray ord (calcStrides ord s) s (fromList $ toList v)
||| Create an array given a vector of its elements. The elements of the vector
||| are arranged into the provided shape using row-major order (the last axis is the
||| least significant).
||| To specify the order of the array, use `fromVect'`.
|||
||| @ s The shape of the constructed array
export
fromVect : (s : Vect rk Nat) -> Vect (product s) a -> Array s a
fromVect s = fromVect' s COrder
||| Create an array by taking values from a stream.
|||
||| @ s The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromStream' : (s : Vect rk Nat) -> (ord : Order) -> Stream a -> Array s a
fromStream' s ord st = MkArray ord (calcStrides ord s) s (fromList $ take (product s) st)
||| Create an array by taking values from a stream.
||| To specify the order of the array, use `fromStream'`.
|||
||| @ s The shape of the constructed array
export
fromStream : (s : Vect rk Nat) -> Stream a -> Array s a
fromStream s = fromStream' s COrder
||| Create an array given a function to generate its elements.
|||
||| @ s The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromFunctionNB' : (s : Vect rk Nat) -> (ord : Order) -> (Vect rk Nat -> a) -> Array s a
fromFunctionNB' s ord f = let sts = calcStrides ord s
in MkArray ord sts s (unsafeFromIns (product s) $
map (\is => (getLocation' sts is, f is)) $ getAllCoords' s)
||| Create an array given a function to generate its elements.
||| To specify the order of the array, use `fromFunctionNB'`.
|||
||| @ s The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromFunctionNB : (s : Vect rk Nat) -> (Vect rk Nat -> a) -> Array s a
fromFunctionNB s = fromFunctionNB' s COrder
||| Create an array given a function to generate its elements.
|||
||| @ s The shape of the constructed array
||| @ ord The order to interpret the elements
export
fromFunction' : (s : Vect rk Nat) -> (ord : Order) -> (Coords s -> a) -> Array s a
fromFunction' s ord f = let sts = calcStrides ord s
in MkArray ord sts s (unsafeFromIns (product s) $
map (\is => (getLocation sts is, f is)) $ getAllCoords s)
||| Create an array given a function to generate its elements.
||| To specify the order of the array, use `fromFunction'`.
|||
||| @ s The shape of the constructed array
export
fromFunction : (s : Vect rk Nat) -> (Coords s -> a) -> Array s a
fromFunction s = fromFunction' s COrder
||| Construct an array using a structure of nested vectors. The elements are arranged
||| to the specified order before being written.
|||
||| @ s The shape of the constructed array
||| @ ord The order of the constructed array
export
array' : (s : Vect rk Nat) -> (ord : Order) -> Vects s a -> Array s a
array' s ord v = MkArray ord sts s (unsafeFromIns (product s) ins)
where
sts : Vect rk Nat
sts = calcStrides ord s
ins : List (Nat, a)
ins = collapse $ mapWithIndex (MkPair . getLocation' sts) v
||| Construct an array using a structure of nested vectors.
||| To explicitly specify the shape and order of the array, use `array'`.
export
array : {s : Vect rk Nat} -> Vects s a -> Array s a
array v = MkArray COrder (calcStrides COrder s) s (fromList $ collapse v)
--------------------------------------------------------------------------------
-- Indexing
--------------------------------------------------------------------------------
infixl 10 !!
infixl 10 !?
infixl 11 !!..
infixl 11 !?..
||| Index the array using the given coordinates.
export
index : Coords s -> Array s a -> a
index is arr = index (getLocation (strides arr) is) (getPrim arr)
||| Index the array using the given coordinates.
|||
||| This is the operator form of `index`.
export
(!!) : Array s a -> Coords s -> a
(!!) = flip index
-- TODO: Create set/update at index functions
||| Update the value at the given coordinates using the function.
export
indexUpdate : Coords s -> (a -> a) -> Array s a -> Array s a
indexUpdate is f (MkArray ord sts s arr) =
MkArray ord sts s (updateAt (getLocation sts is) f arr)
||| Set the value at the given coordinates to the given value.
export
indexSet : Coords s -> a -> Array s a -> Array s a
indexSet is = indexUpdate is . const
||| Index the array using the given coordinates, returning `Nothing` if the
||| coordinates are out of bounds.
export
indexNB : Vect rk Nat -> Array {rk} s a -> Maybe a
indexNB is arr = if and $ map delay $ zipWith (<) is (shape arr)
then Just $ index (getLocation' (strides arr) is) (getPrim arr)
else Nothing
||| Index the array using the given coordinates, returning `Nothing` if the
||| coordinates are out of bounds.
|||
||| This is the operator form of `indexNB`.
export
(!?) : Array {rk} s a -> Vect rk Nat -> Maybe a
(!?) = flip indexNB
||| Update the value at the given coordinates using the function. The array is
||| returned unchanged if the coordinate is out of bounds.
export
indexUpdateNB : Vect rk Nat -> (a -> a) -> Array {rk} s a -> Array s a
indexUpdateNB is f (MkArray ord sts s arr) =
MkArray ord sts s (updateAt (getLocation' sts is) f arr)
||| Set the value at the given coordinates to the given value. The array is
||| returned unchanged if the coordinate is out of bounds.
export
indexSetNB : Vect rk Nat -> a -> Array {rk} s a -> Array s a
indexSetNB is = indexUpdateNB is . const
||| Index the array using the given range of coordinates, returning a new array.
export
indexRange : (rs : CoordsRange s) -> Array s a -> Array (newShape rs) a
indexRange rs arr with (viewShape arr)
_ | Shape s =
let ord = getOrder arr
sts = calcStrides ord s'
in MkArray ord sts s'
(unsafeFromIns (product s') $
map (\(is,is') => (getLocation' sts is',
index (getLocation' (strides arr) is) (getPrim arr)))
$ getCoordsList rs)
where s' : Vect ? Nat
s' = newShape rs
||| Index the array using the given range of coordinates, returning a new array.
|||
||| This is the operator form of `indexRange`.
export
(!!..) : Array s a -> (rs : CoordsRange s) -> Array (newShape rs) a
arr !!.. rs = indexRange rs arr
export
indexUnsafe : Vect rk Nat -> Array {rk} s a -> a
indexUnsafe is arr = index (getLocation' (strides arr) is) (getPrim arr)
export
(!#) : Array {rk} s a -> Vect rk Nat -> a
(!#) = flip indexUnsafe
--------------------------------------------------------------------------------
-- Operations on arrays
--------------------------------------------------------------------------------
||| Reshape the array into the given shape and reinterpret it according to
||| the given order.
|||
||| @ s' The shape to convert the array to
||| @ ord The order to reinterpret the array by
export
reshape' : (s' : Vect rk' Nat) -> (ord : Order) -> Array {rk} s a ->
(0 ok : product s = product s') => Array s' a
reshape' s' ord' arr = MkArray ord' (calcStrides ord' s') s' (getPrim arr)
||| Reshape the array into the given shape.
|||
||| @ s' The shape to convert the array to
export
reshape : (s' : Vect rk' Nat) -> Array {rk} s a ->
(0 ok : product s = product s') => Array s' a
reshape s' arr = reshape' s' (getOrder arr) arr
||| Change the internal order of the array's elements.
export
reorder : Order -> Array s a -> Array s a
reorder ord' arr with (viewShape arr)
_ | Shape s = let sts = calcStrides ord' s
in MkArray ord' sts _ (unsafeFromIns (product s) $
map (\is => (getLocation' sts is, index (getLocation' (strides arr) is) (getPrim arr))) $
getAllCoords' s)
||| Resize the array to a new shape, preserving the coordinates of the original
||| elements. New coordinates are filled with a default value.
|||
||| @ s' The shape to resize the array to
||| @ def The default value to fill the array with
export
resize : (s' : Vect rk Nat) -> (def : a) -> Array {rk} s a -> Array s' a
resize s' def arr = fromFunction' s' (getOrder arr) (fromMaybe def . (arr !?) . toNB)
||| Resize the array to a new shape, preserving the coordinates of the original
||| elements. This function requires a proof that the new shape is strictly
||| smaller than the current shape of the array.
|||
||| @ s' The shape to resize the array to
export
-- TODO: Come up with a solution that doesn't use `believe_me` or trip over some
-- weird bug in the type-checker
resizeLTE : (s' : Vect rk Nat) -> (0 ok : NP Prelude.id (zipWith LTE s' s)) =>
Array {rk} s a -> Array s' a
resizeLTE s' arr = resize s' (believe_me ()) arr
||| List all of the values in an array in row-major order.
export
elements : Array {rk} s a -> Vect (product s) a
elements (MkArray _ sts sh p) =
let elems = map (flip index p . getLocation' sts) (getAllCoords' sh)
in assert_total $ case toVect (product sh) elems of Just v => v
||| List all of the values in an array along with their coordinates.
export
enumerateNB : Array {rk} s a -> List (Vect rk Nat, a)
enumerateNB (MkArray _ sts sh p) =
map (\is => (is, index (getLocation' sts is) p)) (getAllCoords' sh)
||| List all of the values in an array along with their coordinates.
export
enumerate : Array s a -> List (Coords s, a)
enumerate arr with (viewShape arr)
_ | Shape s = map (\is => (is, index is arr)) (getAllCoords s)
||| Join two arrays along a particular axis, e.g. combining two matrices
||| vertically or horizontally. The arrays must have the same shape on all other axes.
|||
||| @ axis The axis to join the arrays on
export
concat : (axis : Fin rk) -> Array {rk} s a -> Array (replaceAt axis d s) a ->
Array (updateAt axis (+d) s) a
concat axis a b = let sA = shape a
sB = shape b
dA = index axis sA
dB = index axis sB
s = replaceAt axis (dA + dB) sA
sts = calcStrides COrder s
ins = map (mapFst $ getLocation' sts . toNB) (enumerate a)
++ map (mapFst $ getLocation' sts . updateAt axis (+dA) . toNB) (enumerate b)
-- TODO: prove that the type-level shape and `s` are equivalent
in believe_me $ MkArray COrder sts s (unsafeFromIns (product s) ins)
||| Stack multiple arrays along a new axis, e.g. stacking vectors to form a matrix.
|||
||| @ axis The axis to stack the arrays along
export
stack : {s : _} -> (axis : Fin (S rk)) -> Vect n (Array {rk} s a) -> Array (insertAt axis n s) a
stack axis arrs = rewrite sym (lengthCorrect arrs) in
fromFunction _ (\is => case getAxisInd axis (rewrite sym (lengthCorrect arrs) in is) of
(i,is') => index is' (index i arrs))
where
getAxisInd : {0 rk : _} -> {s : _} -> (ax : Fin (S rk)) -> Coords (insertAt ax n s) -> (Fin n, Coords s)
getAxisInd FZ (i :: is) = (i, is)
getAxisInd {s=_::_} (FS ax) (i :: is) = mapSnd (i::) (getAxisInd ax is)
--------------------------------------------------------------------------------
-- Implementations
--------------------------------------------------------------------------------
-- Most of these implementations apply the operation pointwise. If there are
-- multiple arrays involved with different orders, then all of the arrays are
-- reordered to match one.
export
Zippable (Array s) where
zipWith f a b with (viewShape a)
_ | Shape s = MkArray (getOrder a) (strides a) s $
if getOrder a == getOrder b
then unsafeZipWith f (getPrim a) (getPrim b)
else unsafeZipWith f (getPrim a) (getPrim $ reorder (getOrder a) b)
zipWith3 f a b c with (viewShape a)
_ | Shape s = MkArray (getOrder a) (strides a) s $
if (getOrder a == getOrder b) && (getOrder b == getOrder c)
then unsafeZipWith3 f (getPrim a) (getPrim b) (getPrim c)
else unsafeZipWith3 f (getPrim a) (getPrim $ reorder (getOrder a) b)
(getPrim $ reorder (getOrder a) c)
unzipWith f arr with (viewShape arr)
_ | Shape s = case unzipWith f (getPrim arr) of
(a, b) => (MkArray (getOrder arr) (strides arr) s a,
MkArray (getOrder arr) (strides arr) s b)
unzipWith3 f arr with (viewShape arr)
_ | Shape s = case unzipWith3 f (getPrim arr) of
(a, b, c) => (MkArray (getOrder arr) (strides arr) s a,
MkArray (getOrder arr) (strides arr) s b,
MkArray (getOrder arr) (strides arr) s c)
export
Functor (Array s) where
map f (MkArray ord sts s arr) = MkArray ord sts s (map f arr)
export
{s : _} -> Applicative (Array s) where
pure = repeat s
(<*>) = zipWith apply
export
{s : _} -> Monad (Array s) where
join arr = fromFunction s (\is => arr !! is !! is)
-- Foldable and Traversable operate on the primitive array directly. This means
-- that their operation is dependent on the order of the array.
export
Foldable (Array s) where
foldl f z = foldl f z . getPrim
foldr f z = foldr f z . getPrim
null arr = size arr == Z
toList = toList . getPrim
export
Traversable (Array s) where
traverse f (MkArray ord sts s arr) =
map (MkArray ord sts s) (traverse f arr)
export
Cast a b => Cast (Array s a) (Array s b) where
cast = map cast
export
Eq a => Eq (Array s a) where
a == b = if getOrder a == getOrder b
then unsafeEq (getPrim a) (getPrim b)
else unsafeEq (getPrim a) (getPrim $ reorder (getOrder a) b)
export
Semigroup a => Semigroup (Array s a) where
(<+>) = zipWith (<+>)
export
{s : _} -> Monoid a => Monoid (Array s a) where
neutral = repeat s neutral
-- The shape must be known at runtime here due to `fromInteger`. If `fromInteger`
-- were moved into its own interface, this constraint could be removed.
export
{s : _} -> Num a => Num (Array s a) where
(+) = zipWith (+)
(*) = zipWith (*)
fromInteger = repeat s . fromInteger
export
{s : _} -> Neg a => Neg (Array s a) where
negate = map negate
(-) = zipWith (-)
export
{s : _} -> Fractional a => Fractional (Array s a) where
recip = map recip
(/) = zipWith (/)
export
Num a => Mult a (Array {rk} s a) (Array s a) where
(*.) x = map (*x)
export
Num a => Mult (Array {rk} s a) a (Array s a) where
(*.) = flip (*.)
export
Show a => Show (Array s a) where
showPrec d arr = let orderedElems = PrimArray.toList $ getPrim $
if getOrder arr == COrder then arr else reorder COrder arr
in showCon d "array " $ concat $ insertPunct (shape arr) $ map show orderedElems
where
splitWindow : Nat -> List String -> List (List String)
splitWindow n xs = case splitAt n xs of
(xs, []) => [xs]
(l1, l2) => l1 :: splitWindow n (assert_smaller xs l2)
insertPunct : Vect rk Nat -> List String -> List String
insertPunct [] strs = strs
insertPunct [d] strs = "[" :: intersperse ", " strs `snoc` "]"
insertPunct (Z :: s) strs = ["[","]"]
insertPunct (d :: s) strs =
let secs = if null strs
then List.replicate d ("[]" :: Prelude.Nil)
else map (insertPunct s) $ splitWindow (length strs `div` d) strs
in "[" :: (concat $ intersperse [", "] secs) `snoc` "]"
--------------------------------------------------------------------------------
-- Numeric array operations
--------------------------------------------------------------------------------
||| Linearly interpolate between two arrays.
export
lerp : Neg a => a -> Array s a -> Array s a -> Array s a
lerp t a b = zipWith (+) (a *. (1 - t)) (b *. t)
||| Calculate the square of an array's Eulidean norm.
export
normSq : Num a => Array s a -> a
normSq arr = sum $ zipWith (*) arr arr
||| Calculate an array's Eucliean norm.
export
norm : Array s Double -> Double
norm = sqrt . normSq
||| Normalize the array to a norm of 1.
|||
||| If the array contains all zeros, then it is returned unchanged.
export
normalize : Array s Double -> Array s Double
normalize arr = if all (==0) arr then arr else map (/ norm arr) arr
||| Calculate the Lp-norm of an array.
export
pnorm : (p : Double) -> Array s Double -> Double
pnorm p = (`pow` recip p) . sum . map (`pow` p)