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