Aptwe/src/Examples/Differentiation.idr at main · CyberCat-Institute/Aptwe · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
module Examples.Differentiation

import Data.List.Quantifiers
import IxUtils
import Kernel.Kinds
import Kernel.Types
import Builtins.Types
import Interpreter.Builtins.Types
import Interpreter.Types
import Builtins.Terms
import Kernel.Structure
import Kernel.Terms
import Interpreter.Terms

doubleConst : Double -> Term [] (BaseTy Real)
doubleConst x = BaseTerm $ Builtin $ \[] => (x, \X => [])

times : Term [BaseTy Real, BaseTy Real] (BaseTy Real)
times = BaseTerm $ Builtin $ \[x, y] => (x * y, \X => [X, X])

sin : Term [BaseTy Real] (BaseTy Real)
sin = BaseTerm $ Builtin $ \[x] => (sin x, \X => [X])

cos : Term [BaseTy Real] (BaseTy Real)
cos = BaseTerm $ Builtin $ \[x] => (cos x, \X => [X])

Mono : Ty (True, True) -> Ty (True, True)
Mono a = Tensor a (Not a)

diff : Term [BaseTy Real] (BaseTy Real) -> Term [BaseTy Real] (BaseTy Real)
    -> Term [Mono (BaseTy Real)] (Mono (BaseTy Real))
diff f df = TensorElim Var
          $ Rename (Copy Z $ Insert Id Z $ Insert Id Z $ Empty)
          $ TensorIntro f
          $ NotIntroCov
          $ Rename (Insert Id (S (S Z)) $ Insert Id (S Z) $ Insert Id Z $ Empty)
          $ NotElim Var
          $ Let df
          $ times

dsin : Term [Mono (BaseTy Real)] (Mono (BaseTy Real))
dsin = diff sin cos

dtimes : Term [Mono (BaseTy Real), Mono (BaseTy Real)] (Mono (BaseTy Real))
dtimes = TensorElim Var
       $ Rename (Insert Id (S (S Z)) $ Copy Z $ Insert Id Z $ Insert Id Z $ Empty)
       $ TensorElim Var
       $ Rename (Insert Id (S (S Z)) $ Copy Z $ Insert Id (S (S (S Z))) $ Insert Id Z $ Insert Id Z $ Insert Id Z $ Empty)
       $ TensorIntro times
       $ NotIntroCov
       $ Rename (Insert Id (S Z) $ Insert Id (S (S Z)) $ Copy Z $ Insert Id (S (S Z)) $ Insert Id Z $ Insert Id Z $ Empty)
       $ NotElim Var
       $ UnitElim (NotElim Var $ times)
       $ times

dsquare : Term [Mono (BaseTy Real)] (Mono (BaseTy Real))
dsquare = Rename (Copy Z $ Insert Id Z $ Empty)
         $ dtimes

-- x sin x^2
example : Term [Mono (BaseTy Real)] (Mono (BaseTy Real))
example = Rename (Copy Z $ Insert Id Z $ Empty)
        $ Let (Let dsquare dsin)
        $ dtimes

dtest : Term [Mono (BaseTy Real)] (Mono (BaseTy Real)) -> Double -> (Double, Double -> Double)
dtest t x = let ((y, X), k) = eval t [(x, X)]
             in (y, \dy => let [(X, dx)] = k (X, dy)
                            in dx)