0 | module Data.Tensor.Utils
  1 |
  2 | import Data.Nat -- Add import for Cast
  3 | import Data.List
  4 | import System.Random
  5 |
  6 | import Data.Tensor.Tensor
  7 | import Data.Container.Additive
  8 | import Data.Container.SubTerm
  9 | import Misc
 10 |
 11 |
 12 | {-------------------------------------------------------------------------------
 13 | {-------------------------------------------------------------------------------
 14 | This file defines common tensor utility functions
 15 | Mirrors those found in numpy/pytorch, and includes:
 16 | * zeros
 17 | * ones
 18 | * range
 19 | * size
 20 | * flatten
 21 | * oneHot
 22 | and the corresponding container variants, when they exist.
 23 |
 24 | Naming needs to be made more consistent
 25 |
 26 | -------------------------------------------------------------------------------}
 27 | -------------------------------------------------------------------------------}
 28 |
 29 | namespace CommonNames
 30 |   public export
 31 |   Scalar : (a : Type) -> Type
 32 |   Scalar a = Tensor [] a
 33 |
 34 |   public export
 35 |   Vector : (c : Axis) -> (a : Type) -> Type
 36 |   Vector c a = Tensor [c] a
 37 |   
 38 |   public export
 39 |   Matrix : (row, col : Axis) -> ConsistentWith row [col] =>
 40 |     (a : Type) -> Type
 41 |   Matrix row col a = Tensor [row, col] a
 42 |
 43 | namespace FillZerosOnes
 44 |   public export
 45 |   fill : Num a => {shape : TensorShape rank} ->
 46 |     AllC TensorMonoid shape =>
 47 |     a -> Tensor shape a
 48 |   fill x = tensorReplicate x
 49 |
 50 |   public export
 51 |   zeros : Num a => {shape : TensorShape rank} ->
 52 |     AllC TensorMonoid shape => 
 53 |     Tensor shape a
 54 |   zeros = fill (fromInteger 0)
 55 |
 56 |   public export
 57 |   ones : Num a => {shape : TensorShape rank} ->
 58 |     AllC TensorMonoid shape => 
 59 |     Tensor shape a
 60 |   ones = fill (fromInteger 1)
 61 |
 62 |   ||| An identity matrix with True on the diagonal and False elsewhere
 63 |   public export
 64 |   identityBool : {0 c : Axis} -> IsCubical c =>
 65 |     Tensor [c, c] Bool
 66 |   identityBool @{MkIsCubical _ n}
 67 |     = outerWith (==) (positions {sh=()}) (positions {sh=()})
 68 |
 69 |   ||| An identity matrix with ones on the diagonal and zeros elsewhere
 70 |   ||| Analogous to numpy.eye
 71 |   public export
 72 |   identity : {0 c : Axis} -> IsCubical c =>
 73 |     Num a => Tensor [c, c] a
 74 |   identity @{MkIsCubical _ n} = fromBool <$> identityBool
 75 |
 76 | namespace Range
 77 |   {-----
 78 |   This one is interesting, as in the cubical case it's effectively a version of 'tabulate' from Naperian functors.
 79 |   The cubical version is implemented first, and it's possible that a more general version of rangeA can be defined for container based tensors, possibly by tabulating contents of each shape respectively
 80 |   -----}
 81 |   ||| Separate implementation for the case of one vs two arguments
 82 |   ||| This allows the typechecker to more easily match the right implementation at call sites, as with only TwoArg implementation the following doesn't compile:
 83 |   ||| a = Tensor [6] Double
 84 |   ||| a = arange
 85 |   ||| Otoh, we preclude calling this without a type specified
 86 |   namespace OneArg
 87 |     ||| A range of numbers [0, stop>
 88 |     public export
 89 |     arange : {0 stop : Axis} -> IsCubical stop =>
 90 |       Cast Nat a => Tensor [stop] a
 91 |     arange @{MkIsCubical _ n} = cast . finToNat <$> positions {sh=()}
 92 |
 93 |   namespace TwoArgs
 94 |     ||| A range of numbers [start, stop>
 95 |     public export
 96 |     arangeFromTo : {default (TTInternalName ~~> 0) start : Axis} ->
 97 |       {0 stop : Axis} ->
 98 |       (cStart : IsCubical start) => (cStop : IsCubical stop) =>
 99 |       Cast Nat a => Tensor [stop.name ~~> minus (dim stop) (dim start)] a
100 |     arangeFromTo {cStart=(MkIsCubical _ n)} {cStop=(MkIsCubical _ m)}
101 |       = cast . (+n) . finToNat <$> positions {sh=()}
102 |
103 | namespace Flip
104 |   ||| Reverse a tensor along a given axis
105 |   public export
106 |   flip : {shape : TensorShape rank} ->
107 |     (axis : Fin rank) ->
108 |     IsCubical (index axis (toVect shape)) =>
109 |     Tensor shape a -> Tensor shape a
110 |
111 |
112 | namespace Concatenate
113 |   ||| Concatenate two tensors along an existing axis, the first one
114 |   ||| TODO extend to allow concatenation along an arbitrary/named axis
115 |   public export
116 |   concat : {shape : TensorShape rank} -> {l : AxisName} ->
117 |     {x, y : Axis} -> IsCubical x => IsCubical y =>
118 |     ConsistentWith (l ~~> dim x + dim y) shape =>
119 |     ConsistentWith x shape =>
120 |     ConsistentWith y shape =>
121 |     Tensor (x :: shape) a ->
122 |     Tensor (y :: shape) a ->
123 |     Tensor ((l ~~> dim x + dim y) :: shape) a
124 |   concat @{MkIsCubical _ n} @{MkIsCubical _ m} t t'
125 |     = embedTopExt $ extractTopExt t ++ extractTopExt t'
126 |
127 | namespace Size
128 |   {-----
129 |   Can we measure the size of a tensor of containers?
130 |   Likely need to impose an additional constraint that the set of positions is enumerable
131 |   -----}
132 |   ||| Number of elements in a non-cubical tensor
133 |   public export
134 |   size : {shape : TensorShape rank} ->
135 |     Tensor shape a -> Nat
136 |
137 |   namespace Cubical 
138 |     ||| Number of elements in a cubical tensor
139 |     public export
140 |     size : {shape : TensorShape rank} ->
141 |       All IsCubical (toVect shape) =>
142 |       (0 _ : Tensor shape a) -> Nat
143 |     size {shape} _ = size (toVect shape)
144 |
145 | namespace Flatten
146 |   ||| Flatten a non-cubical tensor into a list
147 |   ||| Requires that we have Foldable on all the components
148 |   ||| In general we won't know the number of elements of a non-cubical tensor at compile time
149 |   public export
150 |   flatten : {0 shape : TensorShape rank} ->
151 |     Foldable (Tensor shape) =>
152 |     Tensor shape a -> List a
153 |   flatten = toList
154 |
155 |   namespace Cubical
156 |     ||| Flatten a cubical tensor into a vector
157 |     ||| Number of elements is known at compile time
158 |     ||| Can even be zero, if any of shape elements is zero
159 |     flatten : {shape : TensorShape rank} ->
160 |       All IsCubical (conts shape) =>
161 |       Tensor shape a -> Vect (size shape) a
162 |     flatten = toVect . extMap (flattenCubical DefaultLayoutOrder) . GetT
163 |
164 | namespace Max
165 |   ||| Maximum value in a tensor
166 |   ||| Returns Nothing if the tensor is empty
167 |   public export
168 |   max : {0 shape : TensorShape rank} ->
169 |     Foldable (Tensor shape) => Ord a =>
170 |     Tensor shape a -> Maybe a
171 |   max = max . flatten
172 |
173 | namespace ArgMinMax
174 |   ||| At the moment this simply reuses the vect implementation
175 |   ||| To be revised at some point later
176 |   public export
177 |   argmax : {name : AxisName} -> {n : Nat} -> IsSucc n => Ord a =>
178 |     Tensor [name ~~> n] a -> Fin n
179 |   argmax = Vect.argmax . (#>)
180 |
181 |   ||| At the moment this simply reuses the vect implementation
182 |   ||| To be revised at some point later
183 |   public export
184 |   argmin : {name : AxisName} -> {n : Nat} -> IsSucc n => Ord a =>
185 |     Tensor [name ~~> n] a -> Fin n
186 |   argmin = Vect.argmin . (#>)
187 |
188 | namespace AllClose
189 |   ||| Scalar approximate equality, following NumPy's `isclose`:
190 |   ||| |x - y| <= atol + rtol * |y|
191 |   public export
192 |   isClose : {default 1.0e-8 atol : Double} -> {default 1.0e-5 rtol : Double} ->
193 |     Double -> Double -> Bool
194 |   isClose x y = abs (x - y) <= atol + rtol * abs y
195 |
196 |   ||| Elementwise approximate equality of tensors, following NumPy's `allclose`
197 |   public export
198 |   allClose : {0 shape : TensorShape rank} ->
199 |     Applicative (Tensor shape) => Foldable (Tensor shape) =>
200 |     {default 1.0e-8 atol : Double} -> {default 1.0e-5 rtol : Double} ->
201 |     Tensor shape Double -> Tensor shape Double -> Bool
202 |   allClose t t' = all id [| isClose {atol} {rtol} t t' |]
203 |
204 | namespace OneHot
205 |   public export
206 |   oneHot : {0 c : Axis} -> IsCubical c =>
207 |     (i : Fin (dim c)) ->
208 |     Num a =>  Tensor [c] a
209 |   oneHot @{MkIsCubical _ n} i = set zeros [i] 1 
210 |
211 | namespace Triangular
212 |   -- should we have ni,ni here, or ni,nj?
213 |   public export
214 |   cTriBool : {c : Axis} ->
215 |     (ip : InterfaceOnPositions c.cont MOrd) =>
216 |     TensorMonoid c.cont =>
217 |     (sh : c.cont.Shp) -> Tensor [c, c] Bool
218 |   cTriBool {ip = MkI p} sh
219 |     = let cPositions = positions {sh=sh}
220 |           pp : MOrd (c.cont.Pos sh) := p sh
221 |       in outerWith (flip isSubTerm) cPositions cPositions
222 |
223 |   public export
224 |   triBool : {0 c : Axis} -> IsCubical c =>
225 |     Tensor [c, c] Bool
226 |   triBool @{MkIsCubical _ n} = cTriBool ()
227 |
228 |
229 |   ||| A matrix with ones on and below the diagonal, and zeros elsewhere
230 |   ||| Analogous to numpy.tri
231 |   public export
232 |   tri : {0 c : Axis} -> IsCubical c =>
233 |     Num a => Tensor [c, c] a
234 |   tri @{MkIsCubical _ n} = fromBool <$> triBool
235 |
236 |   ||| Lower triangular part of a matrix. Elements above the diagonal are set to
237 |   ||| zero. Analogous to numpy.tril
238 |   public export
239 |   lowerTriangular : {0 c : Axis} -> IsCubical c =>
240 |     Num a => Tensor [c, c] a -> Tensor [c, c] a
241 |   lowerTriangular @{MkIsCubical _ n} = (* tri)
242 |
243 |   ||| Upper triangular part of a matrix. Elements below the diagonal are set to
244 |   ||| zero. Analogous to numpy.triu(.., k=1)
245 |   public export
246 |   upperTriangular : {0 c : Axis} -> IsCubical c =>
247 |     Num a => Tensor [c, c] a -> Tensor [c, c] a
248 |   upperTriangular @{MkIsCubical _ n} = (* ((fromBool . not) <$> triBool))
249 |
250 |   ||| Fill the elements of a tensor `t` with `fill` where `mask` is True
251 |   public export
252 |   maskedFill : {shape : TensorShape rank} ->
253 |     Num a => AllC TensorMonoid shape =>
254 |     (t : Tensor shape a) ->
255 |     (mask : Tensor shape Bool) ->
256 |     (fill : a) ->
257 |     Tensor shape a
258 |   maskedFill t mask fill = liftA2Tensor mask t <&>
259 |     (\(maskVal, tVal) => if maskVal then fill else tVal)
260 |
261 | namespace Misc
262 |   public export
263 |   sum : {shape : TensorShape rank} ->
264 |     Algebra (Tensor shape) a =>
265 |     Tensor shape a -> a
266 |   sum = reduce
267 |
268 |   public export
269 |   mean : {shape : TensorShape rank} ->
270 |     All IsCubical (toVect shape) =>
271 |     Cast Nat a =>
272 |     Fractional a => 
273 |     Algebra (Tensor shape) a =>
274 |     Tensor shape a -> a
275 |   mean t = sum t / cast (Cubical.size t)
276 |
277 |   public export
278 |   variance : {c : Axis} -> IsCubical c =>
279 |     Neg a => Fractional a => Cast Nat a =>
280 |     Tensor [c] a -> a
281 |   variance @{MkIsCubical _ n} t =
282 |     let inputMinusMean = t - pure (mean t)
283 |     in mean (inputMinusMean * inputMinusMean)
284 |
285 |   public export
286 |   cumulativeSum : {c : Axis} -> Num a =>
287 |     (isCubical : IsCubical c) =>
288 |     Tensor [c] a -> Tensor [c] a
289 |   cumulativeSum {isCubical=(MkIsCubical _ n)} t
290 |     = (#>#) (scanl1 (+)) t
291 |     
292 |     -- let tt = n -- map {f=Vect n} (scanl1 (+)) (#> t)
293 |     --       
294 |     --   in ?qwerrr -- #> ((scanl1 (+)) (#> t))  --(#>#) 
295 |
296 |
297 |
298 |
299 |
300 | namespace TensorComMonoid
301 |   ||| Pointwise commutative monoid structure on tensors, lifted from the one
302 |   ||| on the underlying type
303 |   public export
304 |   tensorComMonoid : {shape : TensorShape rank} -> AllC TensorMonoid shape =>
305 |     ComMonoid a -> ComMonoid (Tensor shape a)
306 |   tensorComMonoid (MkComMonoid p z) = MkComMonoid
307 |     (\t, t' => [| p t t' |])
308 |     (pure z)
309 |
310 | namespace Traversals
311 |   public export
312 |   inorder : Tensor [b ~> BinTreeNode] a -> Tensor [l ~> List] a
313 |   inorder = extToVector . extMap BinTreeNode.inorder . vectorToExt
314 |
315 | namespace Random
316 |   ||| Sampling a tensor is sampling each entry
317 |   ||| Probably can be done more efficiently
318 |   public export
319 |   {shape : TensorShape rank} ->
320 |   Random a =>
321 |   AllC TensorMonoid shape =>
322 |   Traversable (Tensor shape) =>
323 |   Random (Tensor shape a) where
324 |     randomIO = sequence (pure randomIO)
325 |     randomRIO (lo, hi) = traverse randomRIO [| MkPair lo hi |]
326 |
327 |
328 |   tta : Applicative (Tensor ["a" ~~> 1])
329 |   tta = %search
330 |
331 |   ttt : Traversable (Tensor ["b" ~~> 1])
332 |   ttt = %search
333 |   
334 |   ttd : Random Double
335 |   ttd = %search
336 |
337 |   randTensorShape : Random (Tensor ["a" ~~> 2, "b" ~~> 3] Double)
338 |   randTensorShape = %search
339 |
340 | -- Idris can't find the parametric randomIO interface so reimpementing here
341 | public export
342 | random : Num a => Random a => HasIO io =>
343 |   (shape : TensorShape rank) ->
344 |   All IsCubical (toVect shape) =>
345 |   Applicative (Tensor shape) => 
346 |   Traversable (Tensor shape) => 
347 |   io (Tensor shape a)
348 | random shape = sequence $ pure $ randomRIO (0, 1)