{-# LANGUAGE RebindableSyntax #-}

-- | Inverse and implicit functions via Newton iteration on 'Diff'.
--
-- These are first-order theorems in action: the inverse-function theorem
-- says @(f⁻¹)'(f(a)) = 1/f'(a)@, and the implicit-function theorem says
-- @dy/dx = -(∂F/∂y)⁻¹ · ∂F/∂x@.  We use those derivatives (pulled back by
-- 'Diff') to drive Newton steps, and verify against exact oracles.
module Circuit.Diff.Inverse
  ( -- * Newton iteration
    newton,
    inverseN,
    implicit1N,

    -- * Primitives
    varDiff,
    constDiff,
  )
where

import Circuit.Diff (Diff (..), runDiff)
import NumHask.Algebra.Additive (Additive (..), Subtractive (..))
import NumHask.Algebra.Multiplicative (Divisive (..), Multiplicative (..))
import NumHask.Prelude

-- | The identity differentiable function @x ↦ x@.
varDiff :: Diff p a a
varDiff :: forall {k} (p :: k) a. Diff p a a
varDiff = (a -> (a, a -> a)) -> Diff p a a
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (\a
s -> (a
s, a -> a
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id))

-- | Constant differentiable function.
constDiff :: (Additive a) => b -> Diff p a b
constDiff :: forall {k} a b (p :: k). Additive a => b -> Diff p a b
constDiff b
b = (a -> (b, b -> a)) -> Diff p a b
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((b, b -> a) -> a -> (b, b -> a)
forall a b. a -> b -> a
const (b
b, a -> b -> a
forall a b. a -> b -> a
const a
forall a. Additive a => a
zero))

-- | Newton iteration for solving @f(x) = target@.
--
-- @newton f target x0 n@ takes @n@ steps starting from @x0@.
newton ::
  (Subtractive a, Divisive a) =>
  Diff p a a ->
  a ->
  a ->
  Int ->
  a
newton :: forall {k} a (p :: k).
(Subtractive a, Divisive a) =>
Diff p a a -> a -> a -> Int -> a
newton Diff p a a
f a
target a
x0 Int
n =
  let step :: a -> a
step a
x =
        let (a
y, a -> a
pb) = Diff p a a -> a -> (a, a -> a)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff p a a
f a
x
            dy :: a
dy = a -> a
pb a
forall a. Multiplicative a => a
one
         in a
x a -> a -> a
forall a. Subtractive a => a -> a -> a
- (a
y a -> a -> a
forall a. Subtractive a => a -> a -> a
- a
target) a -> a -> a
forall a. Divisive a => a -> a -> a
/ a
dy
   in ([a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
n) ((a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate a -> a
step a
x0)

-- | Newton iteration for the inverse value: find @x@ such that @f(x) = y@.
inverseN ::
  (Subtractive a, Divisive a) =>
  Diff p a a ->
  a ->
  a ->
  Int ->
  a
inverseN :: forall {k} a (p :: k).
(Subtractive a, Divisive a) =>
Diff p a a -> a -> a -> Int -> a
inverseN Diff p a a
f a
y = Diff p a a -> a -> a -> Int -> a
forall {k} a (p :: k).
(Subtractive a, Divisive a) =>
Diff p a a -> a -> a -> Int -> a
newton Diff p a a
f a
y

-- | Newton iteration for a scalar implicit equation: find @y@ such that
-- @g(y) = 0@.  The caller fixes any ambient parameters (e.g. @x@ in
-- @F(x,y)=0@) by building them into @g@ with 'constDiff.
implicit1N ::
  (Subtractive b, Divisive b) =>
  Diff p b b ->
  b ->
  Int ->
  b
implicit1N :: forall {k} b (p :: k).
(Subtractive b, Divisive b) =>
Diff p b b -> b -> Int -> b
implicit1N Diff p b b
g = Diff p b b -> b -> b -> Int -> b
forall {k} a (p :: k).
(Subtractive a, Divisive a) =>
Diff p a a -> a -> a -> Int -> a
newton Diff p b b
g b
forall a. Additive a => a
zero