714 lines
26 KiB
Idris
714 lines
26 KiB
Idris
module Data.NumIdr.Array.Array
|
|
|
|
import Data.List
|
|
import Data.Vect
|
|
import Data.Zippable
|
|
import Data.Permutation
|
|
import Data.NumIdr.Interfaces
|
|
import Data.NumIdr.PrimArray
|
|
import Data.NumIdr.Array.Rep
|
|
import Data.NumIdr.Array.Coords
|
|
|
|
%default total
|
|
|
|
|
|
||| The type of an array.
|
|
|||
|
|
||| 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.
|
|
|||
|
|
||| @ 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 via one of a handful of representations.
|
|
|||
|
|
||| @ s The shape of the array
|
|
||| @ rep The internal representation of the array
|
|
||| @ rc A witness that the element type satisfies the representation constraint
|
|
MkArray : (rep : Rep) -> (rc : RepConstraint rep a) => (s : Vect rk Nat) ->
|
|
PrimArray rep s a @{rc} -> Array s a
|
|
|
|
%name Array arr
|
|
|
|
|
|
export
|
|
unsafeMkArray : (rep : Rep) -> (rc : RepConstraint rep a) => (s : Vect rk Nat) ->
|
|
PrimArray rep s a @{rc} -> Array s a
|
|
unsafeMkArray = MkArray
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Properties of arrays
|
|
--------------------------------------------------------------------------------
|
|
|
|
|
|
||| 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
|
|
|
|
||| The internal representation of the array.
|
|
export
|
|
getRep : Array s a -> Rep
|
|
getRep (MkArray rep _ _) = rep
|
|
|
|
||| The representation constraint of the array.
|
|
export
|
|
getRepC : (arr : Array s a) -> RepConstraint (getRep arr) a
|
|
getRepC (MkArray _ @{rc} _ _) = rc
|
|
|
|
||| Extract the primitive backend array.
|
|
export
|
|
getPrim : (arr : Array s a) -> PrimArray (getRep arr) s a @{getRepC arr}
|
|
getPrim (MkArray _ _ pr) = pr
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Shape view
|
|
--------------------------------------------------------------------------------
|
|
|
|
export
|
|
shapeEq : (arr : Array s a) -> s = shape arr
|
|
shapeEq (MkArray _ _ _) = Refl
|
|
|
|
|
|
||| A view for extracting the shape of an array.
|
|
public export
|
|
data ShapeView : Array s a -> Type where
|
|
Shape : (s : Vect rk Nat) -> {0 arr : Array s a} -> ShapeView arr
|
|
|
|
||| The covering function for the view `ShapeView`. This function takes an array
|
|
||| of type `Array s a` and returns `Shape s`.
|
|
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
|
|
||| @ rep The internal representation of the constructed array
|
|
export
|
|
repeat : {default B rep : Rep} -> RepConstraint rep a =>
|
|
(s : Vect rk Nat) -> a -> Array s a
|
|
repeat s x = MkArray rep s (PrimArray.constant s x)
|
|
|
|
||| Create an array filled with zeros.
|
|
|||
|
|
||| @ s The shape of the constructed array
|
|
||| @ rep The internal representation of the constructed array
|
|
export
|
|
zeros : {default B rep : Rep} -> RepConstraint rep a =>
|
|
Num a => (s : Vect rk Nat) -> Array s a
|
|
zeros {rep} s = repeat {rep} s 0
|
|
|
|
||| Create an array filled with ones.
|
|
|||
|
|
||| @ s The shape of the constructed array
|
|
||| @ rep The internal representation of the constructed array
|
|
export
|
|
ones : {default B rep : Rep} -> RepConstraint rep a =>
|
|
Num a => (s : Vect rk Nat) -> Array s a
|
|
ones {rep} s = repeat {rep} 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
|
|
||| @ rep The internal representation of the constructed array
|
|
export
|
|
fromVect : {default B rep : Rep} -> LinearRep rep => RepConstraint rep a =>
|
|
(s : Vect rk Nat) -> Vect (product s) a -> Array s a
|
|
fromVect {rep} s v = MkArray rep s (PrimArray.fromList {rep} s $ toList v)
|
|
|
|
|
|
||| Create an array by taking values from a stream.
|
|
|||
|
|
||| @ s The shape of the constructed array
|
|
||| @ rep The internal representation of the constructed array
|
|
export
|
|
fromStream : {default B rep : Rep} -> LinearRep rep => RepConstraint rep a =>
|
|
(s : Vect rk Nat) -> Stream a -> Array s a
|
|
fromStream {rep} s str = fromVect {rep} s (take _ str)
|
|
|
|
||| Create an array given a function to generate its elements.
|
|
|||
|
|
||| @ s The shape of the constructed array
|
|
||| @ rep The internal representation of the constructed array
|
|
export
|
|
fromFunctionNB : {default B rep : Rep} -> RepConstraint rep a =>
|
|
(s : Vect rk Nat) -> (Vect rk Nat -> a) -> Array s a
|
|
fromFunctionNB s f = MkArray rep s (PrimArray.fromFunctionNB s f)
|
|
|
|
||| Create an array given a function to generate its elements.
|
|
|||
|
|
||| @ s The shape of the constructed array
|
|
||| @ rep The internal representation of the constructed array
|
|
export
|
|
fromFunction : {default B rep : Rep} -> RepConstraint rep a =>
|
|
(s : Vect rk Nat) -> (Coords s -> a) -> Array s a
|
|
fromFunction s f = MkArray rep s (PrimArray.fromFunction s f)
|
|
|
|
||| 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' : {default B rep : Rep} -> RepConstraint rep a =>
|
|
(s : Vect rk Nat) -> Vects s a -> Array s a
|
|
array' s v = MkArray rep s (fromVects s v)
|
|
|
|
||| Construct an array using a structure of nested vectors.
|
|
||| To explicitly specify the shape and order of the array, use `array'`.
|
|
export
|
|
array : {default B rep : Rep} -> RepConstraint rep a =>
|
|
{s : Vect rk Nat} -> Vects s a -> Array s a
|
|
array {rep} = array' {rep} _
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Indexing
|
|
--------------------------------------------------------------------------------
|
|
|
|
infixl 10 !!
|
|
infixl 10 !?
|
|
infixl 10 !#
|
|
infixl 11 !!..
|
|
infixl 11 !?..
|
|
infixl 11 !#..
|
|
|
|
|
|
||| Index the array using the given coordinates.
|
|
export
|
|
index : Coords s -> Array s a -> a
|
|
index is (MkArray _ _ arr) = PrimArray.index is arr
|
|
|
|
||| Index the array using the given coordinates.
|
|
|||
|
|
||| This is the operator form of `index`.
|
|
export %inline
|
|
(!!) : Array s a -> Coords s -> a
|
|
arr !! is = index is arr
|
|
|
|
||| Update the entry at the given coordinates using the function.
|
|
export
|
|
indexUpdate : Coords s -> (a -> a) -> Array s a -> Array s a
|
|
indexUpdate is f (MkArray rep @{rc} s arr) =
|
|
MkArray rep @{rc} s (indexUpdate @{rc} is f arr)
|
|
|
|
||| Set the entry 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 range of coordinates, returning a new array.
|
|
export
|
|
indexRange : (rs : CoordsRange s) -> Array s a -> Array (newShape rs) a
|
|
indexRange rs (MkArray rep @{rc} s arr) =
|
|
MkArray rep @{rc} _ (PrimArray.indexRange @{rc} rs arr)
|
|
|
|
||| Index the array using the given range of coordinates, returning a new array.
|
|
|||
|
|
||| This is the operator form of `indexRange`.
|
|
export %inline
|
|
(!!..) : Array s a -> (rs : CoordsRange s) -> Array (newShape rs) a
|
|
arr !!.. rs = indexRange rs arr
|
|
|
|
||| Set the sub-array at the given range of coordinates to the given array.
|
|
export
|
|
indexSetRange : (rs : CoordsRange s) -> Array (newShape rs) a ->
|
|
Array s a -> Array s a
|
|
indexSetRange rs (MkArray _ _ rpl) (MkArray rep s arr) =
|
|
MkArray rep s (PrimArray.indexSetRange {rep} rs (convertRepPrim rpl) arr)
|
|
|
|
|
|
||| Update the sub-array at the given range of coordinates by applying
|
|
||| a function to it.
|
|
export
|
|
indexUpdateRange : (rs : CoordsRange s) ->
|
|
(Array (newShape rs) a -> Array (newShape rs) a) ->
|
|
Array s a -> Array s a
|
|
indexUpdateRange rs f arr = indexSetRange rs (f $ arr !!.. rs) arr
|
|
|
|
|
|
||| 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 (MkArray rep @{rc} s arr) =
|
|
map (\is => index is (MkArray rep @{rc} s arr)) (validateCoords s is)
|
|
|
|
||| Index the array using the given coordinates, returning `Nothing` if the
|
|
||| coordinates are out of bounds.
|
|
|||
|
|
||| This is the operator form of `indexNB`.
|
|
export %inline
|
|
(!?) : Array {rk} s a -> Vect rk Nat -> Maybe a
|
|
arr !? is = indexNB is arr
|
|
|
|
||| Update the entry at the given coordinates using the function. `Nothing` is
|
|
||| returned if the coordinates are out of bounds.
|
|
export
|
|
indexUpdateNB : Vect rk Nat -> (a -> a) -> Array {rk} s a -> Maybe (Array s a)
|
|
indexUpdateNB is f (MkArray rep @{rc} s arr) =
|
|
map (\is => indexUpdate is f (MkArray rep @{rc} s arr)) (validateCoords s is)
|
|
|
|
||| Set the entry at the given coordinates using the function. `Nothing` is
|
|
||| returned if the coordinates are out of bounds.
|
|
export
|
|
indexSetNB : Vect rk Nat -> a -> Array {rk} s a -> Maybe (Array s a)
|
|
indexSetNB is = indexUpdateNB is . const
|
|
|
|
|
|
||| Index the array using the given range of coordinates, returning a new array.
|
|
||| If any of the given indices are out of bounds, then `Nothing` is returned.
|
|
export
|
|
indexRangeNB : (rs : Vect rk CRangeNB) -> Array s a -> Maybe (Array (newShape s rs) a)
|
|
indexRangeNB rs (MkArray rep @{rc} s arr) =
|
|
map (\rs => believe_me $ Array.indexRange rs (MkArray rep @{rc} s arr)) (validateCRange s rs)
|
|
|
|
||| Index the array using the given range of coordinates, returning a new array.
|
|
||| If any of the given indices are out of bounds, then `Nothing` is returned.
|
|
|||
|
|
||| This is the operator form of `indexRangeNB`.
|
|
export %inline
|
|
(!?..) : Array s a -> (rs : Vect rk CRangeNB) -> Maybe (Array (newShape s rs) a)
|
|
arr !?.. rs = indexRangeNB rs arr
|
|
|
|
|
|
||| Index the array using the given coordinates.
|
|
||| WARNING: This function does not perform any bounds check on its inputs.
|
|
||| Misuse of this function can easily break memory safety.
|
|
export %unsafe
|
|
indexUnsafe : Vect rk Nat -> Array {rk} s a -> a
|
|
indexUnsafe is (MkArray _ _ arr) = PrimArray.indexUnsafe is arr
|
|
|
|
||| Index the array using the given coordinates.
|
|
||| WARNING: This function does not perform any bounds check on its inputs.
|
|
||| Misuse of this function can easily break memory safety.
|
|
|||
|
|
||| This is the operator form of `indexUnsafe`.
|
|
export %inline %unsafe
|
|
(!#) : Array {rk} s a -> Vect rk Nat -> a
|
|
arr !# is = indexUnsafe is arr
|
|
|
|
|
|
||| Index the array using the given range of coordinates, returning a new array.
|
|
||| WARNING: This function does not perform any bounds check on its inputs.
|
|
||| Misuse of this function can easily break memory safety.
|
|
export %unsafe
|
|
indexRangeUnsafe : (rs : Vect rk CRangeNB) -> Array s a -> Array (newShape s rs) a
|
|
indexRangeUnsafe rs (MkArray rep @{rc} s arr) =
|
|
believe_me $ Array.indexRange (assertCRange s rs) (MkArray rep @{rc} s arr)
|
|
|
|
||| Index the array using the given range of coordinates, returning a new array.
|
|
||| WARNING: This function does not perform any bounds check on its inputs.
|
|
||| Misuse of this function can easily break memory safety.
|
|
|||
|
|
||| This is the operator form of `indexRangeUnsafe`.
|
|
export %inline %unsafe
|
|
(!#..) : Array s a -> (rs : Vect rk CRangeNB) -> Array (newShape s rs) a
|
|
arr !#.. is = indexRangeUnsafe is arr
|
|
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Operations on arrays
|
|
--------------------------------------------------------------------------------
|
|
|
|
||| Map a function over an array.
|
|
|||
|
|
||| You should almost always use `map` instead; only use this function if you
|
|
||| know what you are doing!
|
|
export
|
|
mapArray' : (a -> a) -> Array s a -> Array s a
|
|
mapArray' f (MkArray rep _ arr) = MkArray rep _ (mapPrim f arr)
|
|
|
|
||| Map a function over an array.
|
|
|||
|
|
||| You should almost always use `map` instead; only use this function if you
|
|
||| know what you are doing!
|
|
export
|
|
mapArray : (a -> b) -> (arr : Array s a) -> RepConstraint (getRep arr) b => Array s b
|
|
mapArray f (MkArray rep _ arr) @{rc} = MkArray rep @{rc} _ (mapPrim f arr)
|
|
|
|
||| Combine two arrays of the same element type using a binary function.
|
|
|||
|
|
||| You should almost always use `zipWith` instead; only use this function if
|
|
||| you know what you are doing!
|
|
export
|
|
zipWithArray' : (a -> a -> a) -> Array s a -> Array s a -> Array s a
|
|
zipWithArray' {s} f a b with (viewShape a)
|
|
_ | Shape s = MkArray (mergeRep (getRep a) (getRep b))
|
|
@{mergeRepConstraint (getRepC a) (getRepC b)} s
|
|
$ PrimArray.fromFunctionNB @{mergeRepConstraint (getRepC a) (getRepC b)} _
|
|
(\is => f (a !# is) (b !# is))
|
|
|
|
||| Combine two arrays using a binary function.
|
|
|||
|
|
||| You should almost always use `zipWith` instead; only use this function if
|
|
||| you know what you are doing!
|
|
export
|
|
zipWithArray : (a -> b -> c) -> (arr : Array s a) -> (arr' : Array s b) ->
|
|
RepConstraint (mergeRep (getRep arr) (getRep arr')) c => Array s c
|
|
zipWithArray {s} f a b @{rc} with (viewShape a)
|
|
_ | Shape s = MkArray (mergeRep (getRep a) (getRep b)) @{rc} s
|
|
$ PrimArray.fromFunctionNB _ (\is => f (a !# is) (b !# is))
|
|
|
|
|
|
||| Reshape the array into the given shape.
|
|
|||
|
|
||| @ s' The shape to convert the array to
|
|
export
|
|
reshape : (s' : Vect rk' Nat) -> (arr : Array {rk} s a) -> LinearRep (getRep arr) =>
|
|
(0 ok : product s = product s') => Array s' a
|
|
reshape s' (MkArray rep _ arr) = MkArray rep s' (PrimArray.reshape s' arr)
|
|
|
|
||| Change the internal representation of the array's elements.
|
|
export
|
|
convertRep : (rep : Rep) -> RepConstraint rep a => Array s a -> Array s a
|
|
convertRep rep (MkArray _ s arr) = MkArray rep s (convertRepPrim arr)
|
|
|
|
||| 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 {rep=getRep arr} @{getRepC arr} s' (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
|
|
-- HACK: 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 : All 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 along with their coordinates.
|
|
export
|
|
enumerateNB : Array {rk} s a -> List (Vect rk Nat, a)
|
|
enumerateNB (MkArray _ s arr) =
|
|
map (\is => (is, PrimArray.indexUnsafe is arr)) (getAllCoords' s)
|
|
|
|
||| List all of the values in an array along with their coordinates.
|
|
export
|
|
enumerate : Array s a -> List (Coords s, a)
|
|
enumerate {s} arr with (viewShape arr)
|
|
_ | Shape s = map (\is => (is, index is arr)) (getAllCoords s)
|
|
|
|
||| List all of the values in an array in row-major order.
|
|
export
|
|
elements : Array {rk} s a -> Vect (product s) a
|
|
elements (MkArray _ s arr) =
|
|
believe_me $ Vect.fromList $
|
|
map (flip PrimArray.indexUnsafe arr) (getAllCoords' s)
|
|
|
|
||| Join two arrays along a particular axis, e.g. combining two matrices
|
|
||| vertically or horizontally. All other axes of the arrays must have the
|
|
||| same dimensions.
|
|
|||
|
|
||| @ 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 {s,d} axis a b with (viewShape a, viewShape b)
|
|
_ | (Shape s, Shape (replaceAt axis d s)) =
|
|
believe_me $ Array.fromFunctionNB {rep=mergeRep (getRep a) (getRep b)}
|
|
@{mergeRepConstraint (getRepC a) (getRepC b)} (updateAt axis (+ index axis (shape b)) s)
|
|
(\is => let limit = index axis s
|
|
in if index axis is < limit then a !# is else b !# updateAt axis (`minus` limit) is)
|
|
|
|
||| 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)
|
|
|
|
||| Join the axes of a nested array structure to form a single array.
|
|
export
|
|
joinAxes : {s' : _} -> Array s (Array s' a) -> Array (s ++ s') a
|
|
joinAxes {s} arr with (viewShape arr)
|
|
_ | Shape s = fromFunctionNB (s ++ s') (\is => arr !# takeUpTo s is !# dropUpTo s is)
|
|
where
|
|
takeUpTo : Vect rk Nat -> Vect (rk + rk') Nat -> Vect rk Nat
|
|
takeUpTo [] ys = []
|
|
takeUpTo (x::xs) (y::ys) = y :: takeUpTo xs ys
|
|
|
|
dropUpTo : Vect rk Nat -> Vect (rk + rk') Nat -> Vect rk' Nat
|
|
dropUpTo [] ys = ys
|
|
dropUpTo (x::xs) (y::ys) = dropUpTo xs ys
|
|
|
|
|
|
||| Split an array into a nested array structure along the specified axes.
|
|
export
|
|
splitAxes : (rk : Nat) -> {0 rk' : Nat} -> {s : _} ->
|
|
Array {rk=rk+rk'} s a -> Array (take {m=rk'} rk s) (Array (drop {m=rk'} rk s) a)
|
|
splitAxes _ {s} arr = fromFunctionNB _ (\is => fromFunctionNB _ (\is' => arr !# (is ++ is')))
|
|
|
|
||| Construct the transpose of an array by reversing the order of its axes.
|
|
export
|
|
transpose : Array s a -> Array (reverse s) a
|
|
transpose {s} arr with (viewShape arr)
|
|
_ | Shape s = fromFunctionNB _ (\is => arr !# reverse is)
|
|
|
|
||| Construct the transpose of an array by reversing the order of its axes.
|
|
|||
|
|
||| This is the postfix form of `transpose`.
|
|
export
|
|
(.T) : Array s a -> Array (reverse s) a
|
|
(.T) = transpose
|
|
|
|
|
|
||| Swap two axes in an array.
|
|
export
|
|
swapAxes : (i,j : Fin rk) -> Array s a -> Array (swapElems i j s) a
|
|
swapAxes {s} i j arr with (viewShape arr)
|
|
_ | Shape s = fromFunctionNB _ (\is => arr !# swapElems i j is)
|
|
|
|
||| Apply a permutation to the axes of an array.
|
|
export
|
|
permuteAxes : (p : Permutation rk) -> Array s a -> Array (permuteVect p s) a
|
|
permuteAxes {s} p arr with (viewShape arr)
|
|
_ | Shape s = fromFunctionNB _ (\is => arr !# permuteVect p s)
|
|
|
|
||| Swap two coordinates along a specific axis (e.g. swapping two rows in a matrix).
|
|
|||
|
|
||| @ axis The axis to swap the coordinates along. Slices of the array
|
|
||| perpendicular to this axis are taken when swapping.
|
|
export
|
|
swapInAxis : (axis : Fin rk) -> (i,j : Fin (index axis s)) -> Array s a -> Array s a
|
|
swapInAxis {s} axis i j arr with (viewShape arr)
|
|
_ | Shape s = fromFunctionNB _ (\is => arr !# updateAt axis (swapValues i j) is)
|
|
|
|
||| Permute the coordinates along a specific axis (e.g. permuting the rows in
|
|
||| a matrix).
|
|
|||
|
|
||| @ axis The axis to permute the coordinates along. Slices of the array
|
|
||| perpendicular to this axis are taken when permuting.
|
|
export
|
|
permuteInAxis : (axis : Fin rk) -> Permutation (index axis s) -> Array s a -> Array s a
|
|
permuteInAxis {s} axis p arr with (viewShape arr)
|
|
_ | Shape s = fromFunctionNB _ (\is => arr !# updateAt axis (permuteValues p) is)
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Implementations
|
|
--------------------------------------------------------------------------------
|
|
|
|
export
|
|
Zippable (Array s) where
|
|
zipWith {s} f a b with (viewShape a)
|
|
_ | Shape s = MkArray (mergeRepNC (getRep a) (getRep b))
|
|
@{mergeNCRepConstraint} s
|
|
$ PrimArray.fromFunctionNB @{mergeNCRepConstraint} _
|
|
(\is => f (a !# is) (b !# is))
|
|
|
|
zipWith3 {s} f a b c with (viewShape a)
|
|
_ | Shape s = MkArray (mergeRepNC (mergeRep (getRep a) (getRep b)) (getRep c))
|
|
@{mergeNCRepConstraint} s
|
|
$ PrimArray.fromFunctionNB @{mergeNCRepConstraint} _
|
|
(\is => f (a !# is) (b !# is) (c !# is))
|
|
|
|
unzipWith {s} f arr with (viewShape arr)
|
|
_ | Shape s =
|
|
let rep : Rep
|
|
rep = forceRepNC $ getRep arr
|
|
in (MkArray rep
|
|
@{forceRepConstraint} s
|
|
$ PrimArray.fromFunctionNB @{forceRepConstraint} _
|
|
(\is => fst $ f (arr !# is)),
|
|
MkArray rep
|
|
@{forceRepConstraint} s
|
|
$ PrimArray.fromFunctionNB @{forceRepConstraint} _
|
|
(\is => snd $ f (arr !# is)))
|
|
|
|
unzipWith3 {s} f arr with (viewShape arr)
|
|
_ | Shape s =
|
|
let rep : Rep
|
|
rep = forceRepNC $ getRep arr
|
|
in (MkArray rep
|
|
@{forceRepConstraint} s
|
|
$ PrimArray.fromFunctionNB @{forceRepConstraint} _
|
|
(\is => fst $ f (arr !# is)),
|
|
MkArray rep
|
|
@{forceRepConstraint} s
|
|
$ PrimArray.fromFunctionNB @{forceRepConstraint} _
|
|
(\is => fst $ snd $ f (arr !# is)),
|
|
MkArray rep
|
|
@{forceRepConstraint} s
|
|
$ PrimArray.fromFunctionNB @{forceRepConstraint} _
|
|
(\is => snd $ snd $ f (arr !# is)))
|
|
|
|
|
|
export
|
|
Functor (Array s) where
|
|
map f (MkArray rep @{rc} s arr) = MkArray (forceRepNC rep) @{forceRepConstraint} s
|
|
(mapPrim @{forceRepConstraint} @{forceRepConstraint} f
|
|
$ convertRepPrim @{rc} @{forceRepConstraint} 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 internal representation of the
|
|
-- array.
|
|
|
|
export
|
|
Foldable (Array s) where
|
|
foldl f z (MkArray _ _ arr) = PrimArray.foldl f z arr
|
|
foldr f z (MkArray _ _ arr) = PrimArray.foldr f z arr
|
|
null (MkArray _ s _) = isZero (product s)
|
|
|
|
|
|
export
|
|
Traversable (Array s) where
|
|
traverse f (MkArray rep @{rc} s arr) =
|
|
map (MkArray (forceRepNC rep) @{forceRepConstraint} s)
|
|
(PrimArray.traverse {rep=forceRepNC rep}
|
|
@{%search} @{forceRepConstraint} @{forceRepConstraint} f
|
|
(convertRepPrim @{rc} @{forceRepConstraint} 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 = and $ zipWith (delay .: (==)) (convertRep D a) (convertRep D b)
|
|
|
|
export
|
|
Semigroup a => Semigroup (Array s a) where
|
|
(<+>) = zipWithArray' (<+>)
|
|
|
|
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
|
|
(+) = zipWithArray' (+)
|
|
(*) = zipWithArray' (*)
|
|
|
|
fromInteger = repeat s . fromInteger
|
|
|
|
export
|
|
{s : _} -> Neg a => Neg (Array s a) where
|
|
negate = mapArray' negate
|
|
(-) = zipWithArray' (-)
|
|
|
|
export
|
|
{s : _} -> Fractional a => Fractional (Array s a) where
|
|
recip = mapArray' recip
|
|
(/) = zipWithArray' (/)
|
|
|
|
|
|
export
|
|
Num a => Mult a (Array {rk} s a) (Array s a) where
|
|
(*.) x = mapArray' (*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 = toList $ elements 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 = zipWithArray' (+) (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)
|