155 lines
3.8 KiB
Idris
155 lines
3.8 KiB
Idris
module Data.NumIdr.Vector
|
|
|
|
import Data.Vect
|
|
import Data.NumIdr.Multiply
|
|
import public Data.NumIdr.Array
|
|
|
|
%default total
|
|
|
|
|
|
||| A vector is a rank-1 array.
|
|
public export
|
|
Vector : Nat -> Type -> Type
|
|
Vector n = Array [n]
|
|
|
|
|
|
||| The length (number of dimensions) of the vector
|
|
public export
|
|
dim : Vector n a -> Nat
|
|
dim = head . shape
|
|
|
|
export
|
|
dimEq : (v : Vector n a) -> n = dim v
|
|
dimEq v = cong head $ shapeEq v
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Vector constructors
|
|
--------------------------------------------------------------------------------
|
|
|
|
|
|
||| Construct a vector from a `Vect`.
|
|
export
|
|
vector : Vect n a -> Vector n a
|
|
vector v = rewrite sym (lengthCorrect v)
|
|
in fromVect [length v] $ -- the order doesn't matter here, as
|
|
rewrite lengthCorrect v in -- there is only 1 axis
|
|
rewrite multOneLeftNeutral n in v
|
|
|
|
||| Convert a vector into a `Vect`.
|
|
export
|
|
toVect : Vector n a -> Vect n a
|
|
toVect v = believe_me $ Vect.fromList $ toList v
|
|
|
|
|
|
||| Return the `i`-th basis vector.
|
|
export
|
|
basis : Num a => {n : _} -> (i : Fin n) -> Vector n a
|
|
basis i = fromFunction _ (\[j] => if i == j then 1 else 0)
|
|
|
|
|
|
||| Calculate the 2D unit vector with the given angle off the x-axis.
|
|
export
|
|
unit2D : (ang : Double) -> Vector 2 Double
|
|
unit2D ang = vector [cos ang, sin ang]
|
|
|
|
||| Calculate the 3D unit vector corresponding to the given spherical coordinates,
|
|
||| where the polar axis is the z-axis.
|
|
|||
|
|
||| @ pol The polar angle of the vector
|
|
||| @ az The azimuthal angle of the vector
|
|
export
|
|
unit3D : (pol, az : Double) -> Vector 3 Double
|
|
unit3D pol az = vector [cos az * sin pol, sin az * sin pol, cos pol]
|
|
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Indexing
|
|
--------------------------------------------------------------------------------
|
|
|
|
infixl 10 !!
|
|
infixl 10 !?
|
|
|
|
||| Index the vector at the given coordinate.
|
|
export
|
|
index : Fin n -> Vector n a -> a
|
|
index n = Array.index [n]
|
|
|
|
||| Index the vector at the given coordinate.
|
|
|||
|
|
||| This is the operator form of `index`.
|
|
export
|
|
(!!) : Vector n a -> Fin n -> a
|
|
(!!) = flip index
|
|
|
|
||| Index the vector at the given coordinate, returning `Nothing` if the
|
|
||| coordinate is out of bounds.
|
|
export
|
|
indexNB : Nat -> Vector n a -> Maybe a
|
|
indexNB n = Array.indexNB [n]
|
|
|
|
||| Index the vector at the given coordinate, returning `Nothing` if the
|
|
||| coordinate is out of bounds.
|
|
|||
|
|
||| This is the operator form of `indexNB`.
|
|
export
|
|
(!?) : Vector n a -> Nat -> Maybe a
|
|
(!?) = flip indexNB
|
|
|
|
|
|
-- Named projections
|
|
|
|
||| Return the x-coordinate (the first value) of the vector.
|
|
export
|
|
(.x) : Vector (1 + n) a -> a
|
|
(.x) = index FZ
|
|
|
|
||| Return the y-coordinate (the second value) of the vector.
|
|
export
|
|
(.y) : Vector (2 + n) a -> a
|
|
(.y) = index (FS FZ)
|
|
|
|
||| Return the z-coordinate (the third value) of the vector.
|
|
export
|
|
(.z) : Vector (3 + n) a -> a
|
|
(.z) = index (FS (FS FZ))
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Swizzling
|
|
--------------------------------------------------------------------------------
|
|
|
|
|
|
export
|
|
swizzle : Vect n (Fin m) -> Vector m a -> Vector n a
|
|
swizzle p v = rewrite sym (lengthCorrect p)
|
|
in fromFunction [length p] (\[i] =>
|
|
index (index (rewrite sym (lengthCorrect p) in i) p) v
|
|
)
|
|
|
|
|
|
|
|
--------------------------------------------------------------------------------
|
|
-- Vector operations
|
|
--------------------------------------------------------------------------------
|
|
|
|
|
|
||| Concatenate one vector with another.
|
|
export
|
|
(++) : Vector m a -> Vector n a -> Vector (m + n) a
|
|
(++) = concat 0
|
|
|
|
|
|
||| Calculate the dot product of the two vectors.
|
|
export
|
|
dot : Num a => Vector n a -> Vector n a -> a
|
|
dot = sum .: zipWith (*)
|
|
|
|
|
|
||| Calculate the perpendicular product of the two vectors.
|
|
export
|
|
perp : Neg a => Vector 2 a -> Vector 2 a -> a
|
|
perp a b = a.x * b.y - a.y * b.x
|
|
|
|
|