`math/analysis/derivative.egi`

--
--
-- Differentiation
--
--

<<<<<<< HEAD
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 _ $argnames $args ->
       sum (map2 (\s r -> (userRefs f [s]) * ∂/∂' r x) (between 1 (length argnames)) 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 ∇ : MathExpr -> MathExpr -> MathExpr := ∂/∂

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

def grad : MathExpr -> MathExpr -> 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] :=
=======
-- Differentiable type class. Each CAS shape's instance owns its own
-- structural decomposition rule. Recursive sub-calls go through
-- `partialDiff` (the typeclass method) so runtime-type dispatch picks the
-- right instance again — and the dispatch is O(1) thanks to
-- IRuntimeDispatch reusing the already-evaluated argument (see
-- `evalExprShallow env (IRuntimeDispatch ...)` in Core.hs).
--
-- We don't write an explicit `instance Differentiable MathValue`: the
-- compiler emits `IRuntimeDispatch` for that case, picking the right
-- concrete instance below from the runtime CAS shape.
class Differentiable a where
  partialDiff (f: a) (x: MathValue) : MathValue

instance Differentiable Factor where
  -- Atomic shapes: a single symbol, apply1-4, quote, or func.
  -- Reachable via static dispatch from Term inst's `$fx` (typed Factor through
  -- the `^` pattern's `as (factor, integer)` signature). Symbols short-circuit
  -- here; apply / quote / func delegate to the chain dispatcher.
  partialDiff f x := match f as factor with
    | #x -> 1
    | symbol _ _ -> 0
    | _ -> chainPartialDiff f x

instance Differentiable (Term MathValue [..]) where
  -- The `^` constructor in `inductive pattern MathValue` is `(^) Factor Integer`,
  -- so `$fx` is statically Factor and `partialDiff fx x` dispatches at compile
  -- time to the Factor instance. No need for in-body atomic short-circuits.
  partialDiff f x := match f as mathValue with
    | #0 -> 0
    | _ * #1 -> 0
    | #1 * $fx ^ $n -> n * fx ^ (n - 1) * partialDiff fx x
    | $a * $fx ^ $n * $r -> a * partialDiff (fx ^' n) x * r + a * fx ^' n * partialDiff r x
    | _ -> 0

instance Differentiable (Poly MathValue [..]) where
  -- Polynomial: sum of per-term derivatives. The `poly` constructor in
  -- `inductive pattern MathValue` is `poly [Term MathValue [..]]`, so each
  -- `$t` has static type `Term MathValue [..]` and `partialDiff t x` dispatches
  -- to the Term instance at compile time — no runtime dispatch needed.
  partialDiff f x := match f as mathValue with
    | poly $ts -> sum (map (\t -> partialDiff t x) ts)
    | _ -> chainPartialDiff f x

instance Differentiable (Frac MathValue) where
  -- Quotient rule.
  partialDiff f x := match f as mathValue with
    | $p1 / $p2 ->
        let p1' := partialDiff p1 x
            p2' := partialDiff p2 x
         in (p1' * p2 - p2' * p1) / p2 ^ 2
    | _ -> chainPartialDiff f x

-- ---------------------------------------------------------------------------
-- Top-level differentiation operator (∂/∂) and the chain dispatcher.
-- ---------------------------------------------------------------------------

-- partialDiffMV is a top-level def whose static type (MathValue ->
-- MathValue -> MathValue) is concrete, so the `partialDiff` inside gets
-- expanded by the typeclass machinery (IRuntimeDispatch) at this
-- definition site. Inlining it as an anonymous lambda inside `∂/∂` would
-- leave the partialDiff call in a context whose static type is still
-- polymorphic by the time TypeClassExpand sees it.
-- Validate the complete CAS value before differentiation so an application
-- without a registered analytic derivative cannot silently become zero.
def partialDiffMV (f : MathValue) (x : MathValue) : MathValue :=
  partialDiff (requireAnalyticDerivative f x) x

def ∂/∂ (f : Tensor MathValue) (x : Tensor MathValue) : Tensor MathValue :=
  tensorMap2 partialDiffMV f (flipIndices x)

-- chainPartialDiff is the user-facing dispatcher for "non-decomposable"
-- shapes — atomic symbols, built-in apply1 derivatives, function
-- expressions and quotes. It does NOT handle term / poly / frac shapes
-- (those have their own typeclass instances above). Each
-- `declare derivative <name>` redefines this var to add new apply1
-- branches before falling back to `chainPartialDiffBuiltin` (which is
-- never redefined), so user-declared derivatives win over the built-in
-- cases without infinite recursion.
def chainPartialDiff (v : MathValue) (dx : MathValue) : MathValue :=
  chainPartialDiffBuiltin v dx

-- chainPartialDiffBuiltin handles the shapes that don't have their own
-- type-class instance: function expressions and quote nodes. apply1
-- cases are dispatched first by the redefined `chainPartialDiff` via
-- `declare derivative` rules; symbols are filtered upstream by Term
-- inst's `#dx`/`?isSymbol` short-circuits. Everything else returns 0
-- (treated as a constant, which is the correct conservative answer for
-- truly opaque values).
def chainPartialDiffBuiltin (v : MathValue) (dx : MathValue) : MathValue :=
  match v as mathValue with
    -- function expression
    | func _ $args ->
       sum (map2 (\s r -> (userRefs v [s]) * partialDiff r dx) (between 1 (length args)) args)
    -- general power u^w with a symbolic exponent: an integer exponent is
    -- a monomial and is decomposed by the Term instance, and e^w never
    -- arrives either ((^) rewrites it to exp w), so what reaches here is
    -- the opaque apply2 form.  (u^w)' = u^w * (w' * log u + w * u' / u);
    -- previously this fell to the final 0 (silently treated as constant).
    | apply2 #(^) $u $w ->
        let u' := partialDiff u dx
            w' := partialDiff w dx
         in v * (w' * log u + w * u' / u)
    -- quote
    | quote $g ->
      let g' := partialDiff g dx
       in if isMonomial g'
            then g'
            else let d := foldl1 (\a b -> (gcdForMathValue a b)) (fromPoly g')
                  in d *' (mapPoly (/' d) g')
    | _ -> 0

-- ---------------------------------------------------------------------------
-- Built-in derivatives. Each `declare derivative <name> = <expr>` rebuilds
-- `chainPartialDiff` to dispatch `apply1 #<name>` through the chain rule
-- using `<expr>` as f'(g) for `f(g(x))`. Adding new derivatives here is
-- equivalent to adding new `apply1 #<name>` cases to a hand-written matcher,
-- but cleaner and uniform with user-declared derivatives.
declare derivative sin = cos
declare derivative cos = \z -> - sin z
declare derivative exp = exp
declare derivative log = \z -> 1 / z
declare derivative sqrt = \z -> 1 / (2 * sqrt z)

-- ---------------------------------------------------------------------------
-- Aliases.
-- ---------------------------------------------------------------------------

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

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

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

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

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

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

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

def multivariateTaylorExpansion (f: MathValue) (xs: Vector MathValue) (ys: Vector MathValue)
  : [MathValue] :=
>>>>>>> cbd50c069e9b0791268a2a6bf136ea17f7f995fb
  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))

<<<<<<< HEAD
def multivariateMaclaurinExpansion (f: MathExpr) (xs: Vector MathExpr) : [MathExpr] :=
=======
def multivariateMaclaurinExpansion (f: MathValue) (xs: Vector MathValue) : [MathValue] :=
>>>>>>> cbd50c069e9b0791268a2a6bf136ea17f7f995fb
  multivariateTaylorExpansion f xs (tensorMap 1#0 xs)