derivative.egi

--
--
-- Differentiation
--
--

def ∂/∂ (f : Tensor MathExpr) (x : Tensor MathExpr) : Tensor MathExpr :=
  tensorMap2 (\f x -> ∂/∂' f x) f (flipIndices x)
  
def ∂/∂' (f : MathExpr) (!x : MathExpr) : MathExpr :=
  match f as mathExpr with
    -- symbol
    | #x -> 1
    | ?isSymbol -> 0
    -- function expression
    | func _ $args ->
       sum (map2 (\s r -> (userRefs f [s]) * ∂/∂' r x) (between 1 (length args)) args)
    -- function application
    | (apply1 #exp $g) -> exp g * ∂/∂' g x
    | (apply1 #log $g) -> 1 / g * ∂/∂' g x
    | (apply1 #sqrt $g) -> 1 / (2 * sqrt g) * ∂/∂' g x
    --| (apply2 (^) $g $h) -> f * ∂/∂' (log g * h) x
    | (apply1 #cos $g) -> (- sin g) * ∂/∂' g x
    | (apply1 #sin $g) -> cos g * ∂/∂' g x
    --| (apply1 #arccos $g) -> 1 / sqrt (1 - g ^ 2) * ∂/∂' g x
    -- | apply1 $g $a1 ->
    --   `((userRefs g [1]) a1) * ∂/∂' a1 x
    -- | apply2 $g $a1 $a2 ->
    --   `((userRefs g [1]) a1 a2) * ∂/∂' a1 x + `((userRefs g [2]) a1 a2) * ∂/∂' a2 x
    -- | apply3 $g $a1 $a2 $a3 ->
    --   `((userRefs g [1]) a1 a2 a3) * ∂/∂' a1 x + `((userRefs g [2]) a1 a2 a3) * ∂/∂' a2 x + `((userRefs g [3]) a1 a2 a3) * ∂/∂' a3 x
    -- | apply4 $g $a1 $a2 $a3 $a4 ->
    --   `((userRefs g [1]) a1 a2 a3 a4) * ∂/∂' a1 x + `((userRefs g [2]) a1 a2 a3 a4) * ∂/∂' a2 x + `((userRefs g [3]) a1 a2 a3 a4) * ∂/∂' a3 x + `((userRefs g [4]) a1 a2 a3 a4) * ∂/∂' a4 x
    -- quote
    | quote $g ->
      let g' := ∂/∂' g x
       in if isMonomial g'
            then g'
            else let d := foldl1 (\a b -> (gcd a b)) (fromPoly g')
                  in d *' (mapPoly (/' d) g')
    -- term (constant)
    | #0 -> 0
    | _ * #1 -> 0
    -- term (multiplication)
    | #1 * $fx ^ $n -> n * fx ^ (n - 1) * ∂/∂' fx x
    | $a * $fx ^ $n * $r -> a * ∂/∂' (fx ^' n) x * r + a * fx ^' n * ∂/∂' r x
    -- polynomial
    | poly $ts -> sum (map 1#(∂/∂' $1 x) ts)
    -- quotient
    | $p1 / $p2 ->
      let p1' := ∂/∂' p1 x
          p2' := ∂/∂' p2 x
       in (p1' * p2 - p2' * p1) / p2 ^ 2

def d/d : MathExpr -> MathExpr -> MathExpr := ∂/∂

def pd/pd : MathExpr -> MathExpr -> MathExpr := ∂/∂

def ∇ : Tensor MathExpr -> Vector MathExpr -> Tensor MathExpr := ∂/∂

def nabla : Tensor MathExpr -> Vector MathExpr -> Tensor MathExpr := ∇

def grad : Tensor MathExpr -> Vector MathExpr -> Tensor MathExpr := ∇

def taylorExpansion (f: MathExpr) (x: MathExpr) (a: MathExpr) : [MathExpr] := 
  multivariateTaylorExpansion f [|x|] [|a|]

def maclaurinExpansion (f: MathExpr) (x: MathExpr) : [MathExpr] := taylorExpansion f x 0

def multivariateTaylorExpansion (f: MathExpr) (xs: Vector MathExpr) (ys: Vector MathExpr) 
  : [MathExpr] :=
  withSymbols [h]
    let hs := generateTensor (\[x] -> h_x) (tensorShape xs)
     in map2
          (*)
          (map 1#(1 / fact $1) nats0)
          (map
             (compose
                1#(V.substitute xs ys $1)
                1#(V.substitute hs (withSymbols [i] xs_i - ys_i) $1))
             (iterate (compose 1#(∇ $1 xs) 1#(V.* hs $1)) f))

def multivariateMaclaurinExpansion (f: MathExpr) (xs: Vector MathExpr) : [MathExpr] :=
  multivariateTaylorExpansion f xs (tensorMap 1#0 xs)