{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Scalar Taylor tower as a circuits arrow.
--
-- A value @Taylor n a b@ is a morphism from @a@ to @b@ whose values are
-- truncated Taylor series of order @n@.  The internal representation is
-- structural: scalar wires carry @n+1@ coefficients, unit wires carry @()@,
-- and product wires pair the two shapes.  This makes the cartesian instances
-- straightforward by recursion on the value shape.
--
-- The carrier is intentionally scalar-first: objects are built from 'Double'
-- and @(,)@, and the bimonoid instances are supplied for 'Double' (and the
-- unit).  The 'Traced' instance ties a lazy knot on the value shape, just as
-- the pure @(->)@ trace does; for stable feedback the coefficients resolve
-- order-by-order.
--
-- Scalar primitives ('addT', 'mulT', etc.) are implemented via
-- 'Circuit.Diff.Jet', so they inherit the correct truncated-series recurrences
-- for addition, multiplication, and elementary functions.
module Circuit.Diff.Taylor
  ( -- * Taylor arrow
    Taylor (..),
    TaylorV (..),

    -- * Scalar primitives
    constT,
    addT,
    mulT,
    sinT,
    cosT,
    expT,
    logT,

    -- * Polynomial construction and evaluation
    polyT,
    shiftT,
    evalTaylor,
    evalTaylorDerivs,

    -- * Bridge from 'Diff'
    taylorCoeffsFromDiff,
    approxTaylorFromDiff,
  )
where

import Circuit.Bimonoid (Copy (..), Discard (..), Merge (..), Zero (..))
import Circuit.Category (Category (..))
import Circuit.Channel (Channel (..), Strength (..), Traced (..))
import Circuit.Diff (Diff, runDiff)
import Circuit.Diff.Jet (Jet (..), constant, variable)
import Circuit.Tensor (Action (..), Tensor (..), Unital (..))
import Data.Proxy (Proxy (..))
import GHC.TypeNats (KnownNat, Nat, natVal, someNatVal)
import NumHask.Algebra.Additive qualified as NA
import NumHask.Algebra.Field (ExpField (..), TrigField (..))
import NumHask.Algebra.Multiplicative qualified as NM
import Prelude hiding (cos, exp, id, log, sin, (.))

-- | Values carried on a Taylor wire.
data TaylorV (n :: Nat) a where
  VT :: TaylorV n ()
  VD :: [Double] -> TaylorV n Double
  VP :: TaylorV n a -> TaylorV n b -> TaylorV n (a, b)

-- | Truncated Taylor series arrow of order @n@.
newtype Taylor (n :: Nat) a b = Taylor
  { forall (n :: Nat) a b. Taylor n a b -> TaylorV n a -> TaylorV n b
runTaylor :: TaylorV n a -> TaylorV n b
  }

-- ---------------------------------------------------------------------------
-- Scalar series helpers (delegated to Jet)
-- ---------------------------------------------------------------------------

addJets :: [Double] -> [Double] -> [Double]
addJets :: [Double] -> [Double] -> [Double]
addJets [Double]
xs [Double]
ys = Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients ([Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
xs Jet Double -> Jet Double -> Jet Double
forall a. Additive a => a -> a -> a
NA.+ [Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
ys)

mulJets :: [Double] -> [Double] -> [Double]
mulJets :: [Double] -> [Double] -> [Double]
mulJets [Double]
xs [Double]
ys = Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients ([Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
xs Jet Double -> Jet Double -> Jet Double
forall a. Multiplicative a => a -> a -> a
NM.* [Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
ys)

zeroJet :: Int -> Jet Double
zeroJet :: Int -> Jet Double
zeroJet Int
n = Int -> Double -> Jet Double
forall a. Additive a => Int -> a -> Jet a
constant Int
n Double
0

-- ---------------------------------------------------------------------------
-- Category structure
-- ---------------------------------------------------------------------------

instance Category (Taylor n) where
  id :: forall a. Taylor n a a
id = (TaylorV n a -> TaylorV n a) -> Taylor n a a
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor (\TaylorV n a
x -> TaylorV n a
x)
  {-# INLINE id #-}

  Taylor TaylorV n b -> TaylorV n c
g . :: forall b c a. Taylor n b c -> Taylor n a b -> Taylor n a c
. Taylor TaylorV n a -> TaylorV n b
f = (TaylorV n a -> TaylorV n c) -> Taylor n a c
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor (TaylorV n b -> TaylorV n c
g (TaylorV n b -> TaylorV n c)
-> (TaylorV n a -> TaylorV n b) -> TaylorV n a -> TaylorV n c
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. TaylorV n a -> TaylorV n b
f)
  {-# INLINE (.) #-}

-- ---------------------------------------------------------------------------
-- Cartesian structural maps
-- ---------------------------------------------------------------------------

instance Channel (,) (Taylor n) where
  assoc :: forall a b c. Taylor n ((a, b), c) (a, (b, c))
assoc = (TaylorV n ((a, b), c) -> TaylorV n (a, (b, c)))
-> Taylor n ((a, b), c) (a, (b, c))
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n ((a, b), c) -> TaylorV n (a, (b, c)))
 -> Taylor n ((a, b), c) (a, (b, c)))
-> (TaylorV n ((a, b), c) -> TaylorV n (a, (b, c)))
-> Taylor n ((a, b), c) (a, (b, c))
forall a b. (a -> b) -> a -> b
$ \case
    VP (VP TaylorV n a
x TaylorV n b
y) TaylorV n b
z -> TaylorV n a -> TaylorV n (b, c) -> TaylorV n (a, (b, c))
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n a
TaylorV n a
x (TaylorV n b -> TaylorV n c -> TaylorV n (b, c)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n b
TaylorV n b
y TaylorV n c
TaylorV n b
z)
  {-# INLINE assoc #-}

  assoc' :: forall a b c. Taylor n (a, (b, c)) ((a, b), c)
assoc' = (TaylorV n (a, (b, c)) -> TaylorV n ((a, b), c))
-> Taylor n (a, (b, c)) ((a, b), c)
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (a, (b, c)) -> TaylorV n ((a, b), c))
 -> Taylor n (a, (b, c)) ((a, b), c))
-> (TaylorV n (a, (b, c)) -> TaylorV n ((a, b), c))
-> Taylor n (a, (b, c)) ((a, b), c)
forall a b. (a -> b) -> a -> b
$ \case
    VP TaylorV n a
x (VP TaylorV n a
y TaylorV n b
z) -> TaylorV n (a, b) -> TaylorV n c -> TaylorV n ((a, b), c)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP (TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n a
TaylorV n a
x TaylorV n b
TaylorV n a
y) TaylorV n c
TaylorV n b
z
  {-# INLINE assoc' #-}

  slide :: forall a b c. Taylor n (a, (b, c)) (b, (a, c))
slide = (TaylorV n (a, (b, c)) -> TaylorV n (b, (a, c)))
-> Taylor n (a, (b, c)) (b, (a, c))
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (a, (b, c)) -> TaylorV n (b, (a, c)))
 -> Taylor n (a, (b, c)) (b, (a, c)))
-> (TaylorV n (a, (b, c)) -> TaylorV n (b, (a, c)))
-> Taylor n (a, (b, c)) (b, (a, c))
forall a b. (a -> b) -> a -> b
$ \case
    VP TaylorV n a
x (VP TaylorV n a
y TaylorV n b
z) -> TaylorV n b -> TaylorV n (a, c) -> TaylorV n (b, (a, c))
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n b
TaylorV n a
y (TaylorV n a -> TaylorV n c -> TaylorV n (a, c)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n a
TaylorV n a
x TaylorV n c
TaylorV n b
z)
  {-# INLINE slide #-}

instance Strength (,) (Taylor n) where
  strength :: forall b c a. Taylor n b c -> Taylor n (a, b) (a, c)
strength (Taylor TaylorV n b -> TaylorV n c
f) = (TaylorV n (a, b) -> TaylorV n (a, c)) -> Taylor n (a, b) (a, c)
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (a, b) -> TaylorV n (a, c)) -> Taylor n (a, b) (a, c))
-> (TaylorV n (a, b) -> TaylorV n (a, c)) -> Taylor n (a, b) (a, c)
forall a b. (a -> b) -> a -> b
$ \case
    VP TaylorV n a
a TaylorV n b
b -> TaylorV n a -> TaylorV n c -> TaylorV n (a, c)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n a
TaylorV n a
a (TaylorV n b -> TaylorV n c
f TaylorV n b
TaylorV n b
b)
  {-# INLINE strength #-}

instance Unital (,) (Taylor n) where
  unitl :: forall a. Taylor n (Unit (,), a) a
unitl = (TaylorV n (Unit (,), a) -> TaylorV n a)
-> Taylor n (Unit (,), a) a
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (Unit (,), a) -> TaylorV n a)
 -> Taylor n (Unit (,), a) a)
-> (TaylorV n (Unit (,), a) -> TaylorV n a)
-> Taylor n (Unit (,), a) a
forall a b. (a -> b) -> a -> b
$ \case
    VP TaylorV n a
VT TaylorV n b
x -> TaylorV n a
TaylorV n b
x
  {-# INLINE unitl #-}

  unitl' :: forall a. Taylor n a (Unit (,), a)
unitl' = (TaylorV n a -> TaylorV n (Unit (,), a))
-> Taylor n a (Unit (,), a)
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n a -> TaylorV n (Unit (,), a))
 -> Taylor n a (Unit (,), a))
-> (TaylorV n a -> TaylorV n (Unit (,), a))
-> Taylor n a (Unit (,), a)
forall a b. (a -> b) -> a -> b
$ \TaylorV n a
x -> TaylorV n () -> TaylorV n a -> TaylorV n ((), a)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT TaylorV n a
x
  {-# INLINE unitl' #-}

  unitr :: forall a. Taylor n (a, Unit (,)) a
unitr = (TaylorV n (a, Unit (,)) -> TaylorV n a)
-> Taylor n (a, Unit (,)) a
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (a, Unit (,)) -> TaylorV n a)
 -> Taylor n (a, Unit (,)) a)
-> (TaylorV n (a, Unit (,)) -> TaylorV n a)
-> Taylor n (a, Unit (,)) a
forall a b. (a -> b) -> a -> b
$ \case
    VP TaylorV n a
x TaylorV n b
VT -> TaylorV n a
TaylorV n a
x
  {-# INLINE unitr #-}

  unitr' :: forall a. Taylor n a (a, Unit (,))
unitr' = (TaylorV n a -> TaylorV n (a, Unit (,)))
-> Taylor n a (a, Unit (,))
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n a -> TaylorV n (a, Unit (,)))
 -> Taylor n a (a, Unit (,)))
-> (TaylorV n a -> TaylorV n (a, Unit (,)))
-> Taylor n a (a, Unit (,))
forall a b. (a -> b) -> a -> b
$ \TaylorV n a
x -> TaylorV n a -> TaylorV n () -> TaylorV n (a, ())
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n a
x TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT
  {-# INLINE unitr' #-}

instance Tensor (,) (Taylor n) where
  tensor :: forall a b c d.
Taylor n a b -> Taylor n c d -> Taylor n (a, c) (b, d)
tensor (Taylor TaylorV n a -> TaylorV n b
f) (Taylor TaylorV n c -> TaylorV n d
g) = (TaylorV n (a, c) -> TaylorV n (b, d)) -> Taylor n (a, c) (b, d)
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (a, c) -> TaylorV n (b, d)) -> Taylor n (a, c) (b, d))
-> (TaylorV n (a, c) -> TaylorV n (b, d)) -> Taylor n (a, c) (b, d)
forall a b. (a -> b) -> a -> b
$ \case
    VP TaylorV n a
x TaylorV n b
y -> TaylorV n b -> TaylorV n d -> TaylorV n (b, d)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP (TaylorV n a -> TaylorV n b
f TaylorV n a
TaylorV n a
x) (TaylorV n c -> TaylorV n d
g TaylorV n c
TaylorV n b
y)
  {-# INLINE tensor #-}

instance Action (,) (Taylor n) where
  braid :: forall a b. Taylor n (a, b) (b, a)
braid = (TaylorV n (a, b) -> TaylorV n (b, a)) -> Taylor n (a, b) (b, a)
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (a, b) -> TaylorV n (b, a)) -> Taylor n (a, b) (b, a))
-> (TaylorV n (a, b) -> TaylorV n (b, a)) -> Taylor n (a, b) (b, a)
forall a b. (a -> b) -> a -> b
$ \case
    VP TaylorV n a
x TaylorV n b
y -> TaylorV n b -> TaylorV n a -> TaylorV n (b, a)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n b
TaylorV n b
y TaylorV n a
TaylorV n a
x
  {-# INLINE braid #-}

-- ---------------------------------------------------------------------------
-- Trace: lazy knot on the value shape
-- ---------------------------------------------------------------------------

instance Traced (,) (Taylor n) where
  trace :: forall a b c. Taylor n (a, b) (a, c) -> Taylor n b c
trace (Taylor TaylorV n (a, b) -> TaylorV n (a, c)
body) = (TaylorV n b -> TaylorV n c) -> Taylor n b c
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n b -> TaylorV n c) -> Taylor n b c)
-> (TaylorV n b -> TaylorV n c) -> Taylor n b c
forall a b. (a -> b) -> a -> b
$ \TaylorV n b
b ->
    let VP TaylorV n a
TaylorV n a
a TaylorV n c
TaylorV n b
c = TaylorV n (a, b) -> TaylorV n (a, c)
body (TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n a
a TaylorV n b
b)
     in TaylorV n c
c
  {-# INLINE trace #-}

-- ---------------------------------------------------------------------------
-- Bimonoid structure for the unit and for scalar 'Double'
-- ---------------------------------------------------------------------------

instance Copy (Taylor n) () where
  copy :: Taylor n () ((), ())
copy = (TaylorV n () -> TaylorV n ((), ())) -> Taylor n () ((), ())
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n () -> TaylorV n ((), ())) -> Taylor n () ((), ()))
-> (TaylorV n () -> TaylorV n ((), ())) -> Taylor n () ((), ())
forall a b. (a -> b) -> a -> b
$ \TaylorV n ()
VT -> TaylorV n () -> TaylorV n () -> TaylorV n ((), ())
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT
  {-# INLINE copy #-}

instance Discard (Taylor n) () where
  discard :: Taylor n () ()
discard = (TaylorV n () -> TaylorV n ()) -> Taylor n () ()
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n () -> TaylorV n ()) -> Taylor n () ())
-> (TaylorV n () -> TaylorV n ()) -> Taylor n () ()
forall a b. (a -> b) -> a -> b
$ \TaylorV n ()
VT -> TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT
  {-# INLINE discard #-}

instance Merge (Taylor n) () where
  plus :: Taylor n ((), ()) ()
plus = (TaylorV n ((), ()) -> TaylorV n ()) -> Taylor n ((), ()) ()
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n ((), ()) -> TaylorV n ()) -> Taylor n ((), ()) ())
-> (TaylorV n ((), ()) -> TaylorV n ()) -> Taylor n ((), ()) ()
forall a b. (a -> b) -> a -> b
$ \(VP TaylorV n a
VT TaylorV n b
VT) -> TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT
  {-# INLINE plus #-}

instance Zero (Taylor n) () where
  zero :: Taylor n () ()
zero = (TaylorV n () -> TaylorV n ()) -> Taylor n () ()
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n () -> TaylorV n ()) -> Taylor n () ())
-> (TaylorV n () -> TaylorV n ()) -> Taylor n () ()
forall a b. (a -> b) -> a -> b
$ \TaylorV n ()
VT -> TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT
  {-# INLINE zero #-}

instance Copy (Taylor n) Double where
  copy :: Taylor n Double (Double, Double)
copy = (TaylorV n Double -> TaylorV n (Double, Double))
-> Taylor n Double (Double, Double)
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n (Double, Double))
 -> Taylor n Double (Double, Double))
-> (TaylorV n Double -> TaylorV n (Double, Double))
-> Taylor n Double (Double, Double)
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) -> TaylorV n Double -> TaylorV n Double -> TaylorV n (Double, Double)
forall (n :: Nat) a b.
TaylorV n a -> TaylorV n b -> TaylorV n (a, b)
VP ([Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD [Double]
xs) ([Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD [Double]
xs)
  {-# INLINE copy #-}

instance Discard (Taylor n) Double where
  discard :: Taylor n Double ()
discard = (TaylorV n Double -> TaylorV n ()) -> Taylor n Double ()
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n ()) -> Taylor n Double ())
-> (TaylorV n Double -> TaylorV n ()) -> Taylor n Double ()
forall a b. (a -> b) -> a -> b
$ \TaylorV n Double
_ -> TaylorV n ()
forall (n :: Nat). TaylorV n ()
VT
  {-# INLINE discard #-}

instance Merge (Taylor n) Double where
  plus :: Taylor n (Double, Double) Double
plus = (TaylorV n (Double, Double) -> TaylorV n Double)
-> Taylor n (Double, Double) Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (Double, Double) -> TaylorV n Double)
 -> Taylor n (Double, Double) Double)
-> (TaylorV n (Double, Double) -> TaylorV n Double)
-> Taylor n (Double, Double) Double
forall a b. (a -> b) -> a -> b
$ \(VP (VD [Double]
xs) (VD [Double]
ys)) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD ([Double] -> [Double] -> [Double]
addJets [Double]
xs [Double]
ys)
  {-# INLINE plus #-}

instance (KnownNat n) => Zero (Taylor n) Double where
  zero :: Taylor n () Double
zero = (TaylorV n () -> TaylorV n Double) -> Taylor n () Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n () -> TaylorV n Double) -> Taylor n () Double)
-> (TaylorV n () -> TaylorV n Double) -> Taylor n () Double
forall a b. (a -> b) -> a -> b
$ \TaylorV n ()
VT -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double
0)
    where
      n :: Int
n = Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy n -> Nat
forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n))
  {-# INLINE zero #-}

-- ---------------------------------------------------------------------------
-- Scalar primitives
-- ---------------------------------------------------------------------------

-- | Constant scalar morphism.
constT :: Double -> Taylor n Double Double
constT :: forall (n :: Nat). Double -> Taylor n Double Double
constT Double
c = (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double)
-> (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD (Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients (Int -> Double -> Jet Double
forall a. Additive a => Int -> a -> Jet a
constant ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double
c))
{-# INLINE constT #-}

-- | Add two scalar series.
addT :: Taylor n (Double, Double) Double
addT :: forall (n :: Nat). Taylor n (Double, Double) Double
addT = (TaylorV n (Double, Double) -> TaylorV n Double)
-> Taylor n (Double, Double) Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (Double, Double) -> TaylorV n Double)
 -> Taylor n (Double, Double) Double)
-> (TaylorV n (Double, Double) -> TaylorV n Double)
-> Taylor n (Double, Double) Double
forall a b. (a -> b) -> a -> b
$ \(VP (VD [Double]
xs) (VD [Double]
ys)) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD ([Double] -> [Double] -> [Double]
addJets [Double]
xs [Double]
ys)
{-# INLINE addT #-}

-- | Multiply two scalar series.
mulT :: Taylor n (Double, Double) Double
mulT :: forall (n :: Nat). Taylor n (Double, Double) Double
mulT = (TaylorV n (Double, Double) -> TaylorV n Double)
-> Taylor n (Double, Double) Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n (Double, Double) -> TaylorV n Double)
 -> Taylor n (Double, Double) Double)
-> (TaylorV n (Double, Double) -> TaylorV n Double)
-> Taylor n (Double, Double) Double
forall a b. (a -> b) -> a -> b
$ \(VP (VD [Double]
xs) (VD [Double]
ys)) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD ([Double] -> [Double] -> [Double]
mulJets [Double]
xs [Double]
ys)
{-# INLINE mulT #-}

-- | Sine of a scalar series.
sinT :: Taylor n Double Double
sinT :: forall (n :: Nat). Taylor n Double Double
sinT = (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double)
-> (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD (Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients (Jet Double -> Jet Double
forall a. TrigField a => a -> a
sin ([Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
xs)))
{-# INLINE sinT #-}

-- | Cosine of a scalar series.
cosT :: Taylor n Double Double
cosT :: forall (n :: Nat). Taylor n Double Double
cosT = (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double)
-> (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD (Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients (Jet Double -> Jet Double
forall a. TrigField a => a -> a
cos ([Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
xs)))
{-# INLINE cosT #-}

-- | Exponential of a scalar series.
expT :: Taylor n Double Double
expT :: forall (n :: Nat). Taylor n Double Double
expT = (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double)
-> (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD (Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients (Jet Double -> Jet Double
forall a. ExpField a => a -> a
exp ([Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
xs)))
{-# INLINE expT #-}

-- | Logarithm of a scalar series.
logT :: Taylor n Double Double
logT :: forall (n :: Nat). Taylor n Double Double
logT = (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double)
-> (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) -> [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD (Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients (Jet Double -> Jet Double
forall a. ExpField a => a -> a
log ([Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
xs)))
{-# INLINE logT #-}

-- ---------------------------------------------------------------------------
-- Polynomial construction and evaluation
-- ---------------------------------------------------------------------------

-- | Build a polynomial morphism from coefficients @[c0, c1, ..., ck]@,
-- mapping the input series @u@ to @c0 + c1*u + c2*u^2 + ... + ck*u^k@.
polyT :: [Double] -> Taylor n Double Double
polyT :: forall (n :: Nat). [Double] -> Taylor n Double Double
polyT [Double]
cs = (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double)
-> (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) ->
  let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      x :: Jet Double
x = [Double] -> Jet Double
forall a. [a] -> Jet a
Jet [Double]
xs
      go :: Double -> Jet Double -> Jet Double
go Double
c Jet Double
acc = Int -> Double -> Jet Double
forall a. Additive a => Int -> a -> Jet a
constant Int
n Double
c Jet Double -> Jet Double -> Jet Double
forall a. Additive a => a -> a -> a
NA.+ Jet Double
x Jet Double -> Jet Double -> Jet Double
forall a. Multiplicative a => a -> a -> a
NM.* Jet Double
acc
   in [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD (Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients ((Double -> Jet Double -> Jet Double)
-> Jet Double -> [Double] -> Jet Double
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Double -> Jet Double -> Jet Double
go (Int -> Jet Double
zeroJet Int
n) [Double]
cs))
{-# INLINE polyT #-}

-- | Shift the input series by a constant @x0@.
shiftT :: Double -> Taylor n Double Double
shiftT :: forall (n :: Nat). Double -> Taylor n Double Double
shiftT Double
x0 = (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall (n :: Nat) a b. (TaylorV n a -> TaylorV n b) -> Taylor n a b
Taylor ((TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double)
-> (TaylorV n Double -> TaylorV n Double) -> Taylor n Double Double
forall a b. (a -> b) -> a -> b
$ \(VD [Double]
xs) ->
  let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
   in [Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD ([Double] -> [Double] -> [Double]
addJets [Double]
xs (Jet Double -> [Double]
forall a. Jet a -> [a]
coefficients (Int -> Double -> Jet Double
forall a. Additive a => Int -> a -> Jet a
constant Int
n Double
x0)))
{-# INLINE shiftT #-}

-- | Evaluate a scalar Taylor morphism at a point and return the Taylor
-- coefficients @[f(x), f'(x), f''(x)/2!, ...]@.
evalTaylor :: forall n. (KnownNat n) => Taylor n Double Double -> Double -> [Double]
evalTaylor :: forall (n :: Nat).
KnownNat n =>
Taylor n Double Double -> Double -> [Double]
evalTaylor Taylor n Double Double
t Double
x = case Taylor n Double Double -> TaylorV n Double -> TaylorV n Double
forall (n :: Nat) a b. Taylor n a b -> TaylorV n a -> TaylorV n b
runTaylor Taylor n Double Double
t ([Double] -> TaylorV n Double
forall (n :: Nat). [Double] -> TaylorV n Double
VD [Double]
seed) of
  VD [Double]
ys -> [Double]
ys
  where
    seed :: [Double]
seed = case Proxy n -> Nat
forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n) of
      Nat
0 -> [Double
x]
      Nat
m -> Double
x Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Double
1 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate (Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Nat
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double
0
{-# INLINE evalTaylor #-}

-- | Evaluate a scalar Taylor morphism and return the raw derivatives
-- @[f(x), f'(x), f''(x), ..., f^(n)(x)]@.
evalTaylorDerivs :: forall n. (KnownNat n) => Taylor n Double Double -> Double -> [Double]
evalTaylorDerivs :: forall (n :: Nat).
KnownNat n =>
Taylor n Double Double -> Double -> [Double]
evalTaylorDerivs Taylor n Double Double
t Double
x =
  let cs :: [Double]
cs = Taylor n Double Double -> Double -> [Double]
forall (n :: Nat).
KnownNat n =>
Taylor n Double Double -> Double -> [Double]
evalTaylor Taylor n Double Double
t Double
x
      facts :: [Double]
facts = (Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Double
1.0 [Double
1.0 .. Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n]
   in (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
cs [Double]
facts
  where
    n :: Int
    n :: Int
n = Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy n -> Nat
forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n))
{-# INLINE evalTaylorDerivs #-}

-- ---------------------------------------------------------------------------
-- Bridge from a 'Diff' scalar function
-- ---------------------------------------------------------------------------

-- | Approximate Taylor coefficients of a 'Diff' scalar function at a point
-- using forward differences.
--
-- The returned list @[c0, c1, ..., ck]@ represents
-- @f(x0 + eps) = c0 + c1*eps + c2*eps^2 + ... + ck*eps^k@.
taylorCoeffsFromDiff :: Diff p Double Double -> Double -> Int -> [Double]
taylorCoeffsFromDiff :: forall {k} (p :: k).
Diff p Double Double -> Double -> Int -> [Double]
taylorCoeffsFromDiff Diff p Double Double
f Double
x0 Int
k =
  let h :: Double
h = Double
1e-4 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1.0 (Double -> Double
forall a. Num a => a -> a
abs Double
x0)
      samples :: [Double]
samples = [(Double, Double -> Double) -> Double
forall a b. (a, b) -> a
fst (Diff p Double Double -> Double -> (Double, Double -> Double)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff p Double Double
f (Double
x0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)) | Int
j <- [Int
0 .. Int
k]]
      fwdDiffs :: [[Double]]
fwdDiffs = ([Double] -> [Double]) -> [Double] -> [[Double]]
forall a. (a -> a) -> a -> [a]
iterate (\[Double]
ds -> (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
ds) [Double]
ds) [Double]
samples
      c :: Int -> Double
c Int
i = ([[Double]]
fwdDiffs [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall {a}. Fractional a => Int -> a
factD Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
i)
   in [Int -> Double
c Int
i | Int
i <- [Int
0 .. Int
k]]
  where
    factD :: Int -> a
factD Int
0 = a
1.0
    factD Int
n' = Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int
1 .. Int
n' :: Int] :: Int)

-- | Build a 'Taylor' morphism that approximates a 'Diff' scalar function
-- near @x0@.
approxTaylorFromDiff ::
  forall n p.
  (KnownNat n) =>
  Diff p Double Double ->
  Double ->
  Taylor n Double Double
approxTaylorFromDiff :: forall {k} (n :: Nat) (p :: k).
KnownNat n =>
Diff p Double Double -> Double -> Taylor n Double Double
approxTaylorFromDiff Diff p Double Double
f Double
x0 =
  let cs :: [Double]
cs = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Diff p Double Double -> Double -> Int -> [Double]
forall {k} (p :: k).
Diff p Double Double -> Double -> Int -> [Double]
taylorCoeffsFromDiff Diff p Double Double
f Double
x0 Int
n)
   in [Double] -> Taylor n Double Double
forall (n :: Nat). [Double] -> Taylor n Double Double
polyT [Double]
cs Taylor n Double Double
-> Taylor n Double Double -> Taylor n Double Double
forall b c a. Taylor n b c -> Taylor n a b -> Taylor n a c
forall k (arr :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category arr =>
arr b c -> arr a b -> arr a c
. Double -> Taylor n Double Double
forall (n :: Nat). Double -> Taylor n Double Double
shiftT Double
x0
  where
    n :: Int
n = Nat -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy n -> Nat
forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n))