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