
   0 | module Data.NumIdr.Vector
   1 | 
   2 | import Data.Vect
   3 | import Data.Permutation
   4 | import Data.NumIdr.Interfaces
   5 | import public Data.NumIdr.Array
   6 | 
   7 | %default total
   8 | 
   9 | 
  10 | ||| A vector is a rank-1 array.
  11 | public export
  12 | Vector : Nat -> Type -> Type
  13 | Vector n = Array [n]
  14 | 
  15 | %name Vector v,w,u
  16 | 
  17 | ||| The length (number of dimensions) of the vector.
  18 | public export
  19 | dim : Vector n a -> Nat
  20 | dim = head . shape
  21 | 
  22 | export
  23 | dimEq : (v : Vector n a) -> n = dim v
  24 | dimEq v = cong head $ shapeEq v
  25 | 
  26 | 
  27 | --------------------------------------------------------------------------------
  28 | -- Vector constructors
  29 | --------------------------------------------------------------------------------
  30 | 
  31 | 
  32 | ||| Construct a vector from a `Vect`.
  33 | export
  34 | vector : {default B rep : Rep} -> RepConstraint rep a => Vect n a -> Vector n a
  35 | vector v = rewrite sym (lengthCorrect v)
  36 |            in array {rep,s=[length v]} $
  37 |            rewrite lengthCorrect v in v
  38 | 
  39 | ||| Convert a vector into a `Vect`.
  40 | export
  41 | toVect : Vector n a -> Vect n a
  42 | toVect v = replace {p = flip Vect a} (believe_me $ Refl {x=()}) $
  43 |   Vect.fromList $ Prelude.toList v
  44 | 
  45 | ||| Return the `i`-th basis vector.
  46 | export
  47 | basis : {default B rep : Rep} -> RepConstraint rep a =>
  48 |           Num a => {n : _} -> (i : Fin n) -> Vector n a
  49 | basis i = fromFunction {rep} _ (\[j] => if i == j then 1 else 0)
  50 | 
  51 | ||| The first basis vector.
  52 | export
  53 | basisX : {default B rep : Rep} -> RepConstraint rep a =>
  54 |           {n : _} -> Num a => Vector (1 + n) a
  55 | basisX = basis {rep} FZ
  56 | 
  57 | ||| The second basis vector.
  58 | export
  59 | basisY : {default B rep : Rep} -> RepConstraint rep a =>
  60 |           {n : _} -> Num a => Vector (2 + n) a
  61 | basisY = basis {rep} (FS FZ)
  62 | 
  63 | ||| The third basis vector.
  64 | export
  65 | basisZ : {default B rep : Rep} -> RepConstraint rep a =>
  66 |           {n : _} -> Num a => Vector (3 + n) a
  67 | basisZ = basis {rep} (FS (FS FZ))
  68 | 
  69 | ||| Calculate the 2D unit vector with the given angle off the x-axis.
  70 | export
  71 | unit2D : {default B rep : Rep} -> RepConstraint rep Double =>
  72 |           (ang : Double) -> Vector 2 Double
  73 | unit2D ang = vector {rep} [cos ang, sin ang]
  74 | 
  75 | ||| Calculate the 3D unit vector corresponding to the given spherical coordinates,
  76 | ||| where the polar axis is the z-axis.
  77 | |||
  78 | ||| @ pol The polar angle of the vector
  79 | ||| @ az  The azimuthal angle of the vector
  80 | export
  81 | unit3D : {default B rep : Rep} -> RepConstraint rep Double =>
  82 |           (pol, az : Double) -> Vector 3 Double
  83 | unit3D pol az = vector {rep} [cos az * sin pol, sin az * sin pol, cos pol]
  84 | 
  85 | 
  86 | 
  87 | --------------------------------------------------------------------------------
  88 | -- Indexing
  89 | --------------------------------------------------------------------------------
  90 | 
  91 | infixl 10 !!
  92 | infixl 10 !?
  93 | 
  94 | ||| Index the vector at the given coordinate.
  95 | export
  96 | index : Fin n -> Vector n a -> a
  97 | index n = Array.index [n]
  98 | 
  99 | ||| Index the vector at the given coordinate.
 100 | |||
 101 | ||| This is the operator form of `index`.
 102 | export
 103 | (!!) : Vector n a -> Fin n -> a
 104 | (!!) = flip index
 105 | 
 106 | ||| Index the vector at the given coordinate, returning `Nothing` if the
 107 | ||| coordinate is out of bounds.
 108 | export
 109 | indexNB : Nat -> Vector n a -> Maybe a
 110 | indexNB n = Array.indexNB [n]
 111 | 
 112 | ||| Index the vector at the given coordinate, returning `Nothing` if the
 113 | ||| coordinate is out of bounds.
 114 | |||
 115 | ||| This is the operator form of `indexNB`.
 116 | export
 117 | (!?) : Vector n a -> Nat -> Maybe a
 118 | (!?) = flip indexNB
 119 | 
 120 | 
 121 | -- Named projections
 122 | 
 123 | ||| Return the x-coordinate (the first value) of the vector.
 124 | export
 125 | (.x) : Vector (1 + n) a -> a
 126 | (.x) = index FZ
 127 | 
 128 | ||| Return the y-coordinate (the second value) of the vector.
 129 | export
 130 | (.y) : Vector (2 + n) a -> a
 131 | (.y) = index (FS FZ)
 132 | 
 133 | ||| Return the z-coordinate (the third value) of the vector.
 134 | export
 135 | (.z) : Vector (3 + n) a -> a
 136 | (.z) = index (FS (FS FZ))
 137 | 
 138 | 
 139 | --------------------------------------------------------------------------------
 140 | -- Swizzling & permuting elements
 141 | --------------------------------------------------------------------------------
 142 | 
 143 | 
 144 | ||| Construct a vector by pulling coordinates from another vector.
 145 | export
 146 | swizzle : Vect n (Fin m) -> Vector m a -> Vector n a
 147 | swizzle p v = rewrite sym (lengthCorrect p)
 148 |               in fromFunction {rep=_} @{getRepC v} [length p] (\[i] =>
 149 |                 index (index (rewrite sym (lengthCorrect p) in i) p) v)
 150 | 
 151 | 
 152 | ||| Swap two coordinates in a vector.
 153 | export
 154 | swapCoords : (i,j : Fin n) -> Vector n a -> Vector n a
 155 | swapCoords = swapInAxis 0
 156 | 
 157 | ||| Permute the coordinates in a vector.
 158 | export
 159 | permuteCoords : Permutation n -> Vector n a -> Vector n a
 160 | permuteCoords = permuteInAxis 0
 161 | 
 162 | 
 163 | --------------------------------------------------------------------------------
 164 | -- Vector operations
 165 | --------------------------------------------------------------------------------
 166 | 
 167 | 
 168 | ||| Concatenate one vector with another.
 169 | export
 170 | (++) : Vector m a -> Vector n a -> Vector (m + n) a
 171 | (++) = concat 0
 172 | 
 173 | ||| Calculate the dot product of the two vectors.
 174 | export
 175 | dot : Num a => Vector n a -> Vector n a -> a
 176 | dot = sum .: zipWith (*)
 177 | 
 178 | 
 179 | ||| Calculate the perpendicular product of the two vectors.
 180 | export
 181 | perp : Neg a => Vector 2 a -> Vector 2 a -> a
 182 | perp a b = a.x * b.y - a.y * b.x
 183 | 
 184 | 
 185 | ||| Calculate the cross product of the two vectors.
 186 | export
 187 | cross : Neg a => Vector 3 a -> Vector 3 a -> Vector 3 a
 188 | cross v1 v2 = let [a, b, c] = toVect v1
 189 |                   [x, y, z] = toVect v2
 190 |               in  vector [b*z - c*y, c*x - a*z, a*y - b*x]
 191 | 
 192 | ||| Calculate the triple product of the three vectors.
 193 | export
 194 | triple : Neg a => Vector 3 a -> Vector 3 a -> Vector 3 a -> a
 195 | triple a b c = a `dot` (b `cross` c)
 196 |