Add some utility functions for arrays

This commit is contained in:
Kiana Sheibani 2022-05-11 17:14:54 -04:00
parent e694536939
commit c44fef3c94
Signed by: toki
GPG key ID: 6CB106C25E86A9F7

View file

@ -1,41 +1,63 @@
module Data.NumIdr.Array
import Data.Vect
import Data.Reify
import Data.NumIdr.PrimArray
%default total
export
data Array : Vect rank Nat -> Type -> Type where
MkArray : PrimArray a -> Array dim a
record Array {0 rk : Nat} (s : Vect rk Nat) a where
constructor MkArray
shape : Reify s
contents : PrimArray a
mkArray : {s : _} -> PrimArray a -> Array s a
mkArray arr = MkArray (MkReify _) arr
export
size : Array dim a -> Nat
size (MkArray arr) = length arr
size : Array s a -> Nat
size arr = length arr.contents
export
fromVect : {0 dim : Vect rank Nat} -> Vect (product dim) a -> Array dim a
fromVect = MkArray . fromList . toList
shape : {0 s : Vect rk Nat} -> Array s a -> Vect rk Nat
shape arr = getReify arr.shape
export
reshape : {0 dim : Vect rank Nat} -> (0 dim' : Vect rank' Nat) ->
Array dim a -> product dim = product dim' =>
Array dim' a
reshape _ (MkArray arr) = MkArray arr
rank : Array s a -> Nat
rank arr = length $ shape arr
Vects : Vect rank Nat -> Type -> Type
export
fromVect : {s : Vect rk Nat} -> Vect (product s) a -> Array s a
fromVect = mkArray . fromList . toList
public export
Vects : Vect rk Nat -> Type -> Type
Vects [] a = a
Vects (d :: dim) a = Vect d (Vects dim a)
Vects (d :: s) a = Vect d (Vects s a)
export
fromVects : {s : _} -> Vects s a -> Array s a
fromVects = mkArray . fromList . collapse
where
collapse : {s : Vect rk Nat} -> Vects s a -> List a
collapse {s=[]} x = [x]
collapse {s=_::_} v = concat $ map collapse v
export
reshape : {0 s : Vect rk Nat} -> (s' : Vect rk' Nat) ->
Array s a -> product s = product s' =>
Array s' a
reshape _ arr = mkArray arr.contents
export
constant : a -> (dim : Vect rank Nat) -> Array dim a
constant x dim = MkArray $ create (product dim) (const x)
constant : a -> (s : Vect rk Nat) -> Array s a
constant x s = mkArray $ create (product s) (const x)
export
zeros : Num a => (dim : Vect rank Nat) -> Array dim a
zeros : Num a => (s : Vect rk Nat) -> Array s a
zeros = constant 0