-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2021
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- Forward mode automatic differentiation
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Forward
  ( Forward
  , auto
  -- * Gradient
  , grad
  , grad'
  , gradWith
  , gradWith'
  -- * Jacobian
  , jacobian
  , jacobian'
  , jacobianWith
  , jacobianWith'
  -- * Transposed Jacobian
  , jacobianT
  , jacobianWithT
  -- * Hessian Product
  , hessianProduct
  , hessianProduct'
  -- * Derivatives
  , diff
  , diff'
  , diffF
  , diffF'
  -- * Directional Derivatives
  , du
  , du'
  , duF
  , duF'
  ) where

import Numeric.AD.Internal.Forward
import Numeric.AD.Internal.On
import Numeric.AD.Mode

-- | Compute the directional derivative of a function given a zipped up 'Functor' of the input values and their derivatives
du
  :: (Functor f, Num a)
  => (f (Forward a) -> Forward a)
  -> f (a, a)
  -> a
du :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(f (Forward a) -> Forward a) -> f (a, a) -> a
du f (Forward a) -> Forward a
f = Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> a) -> (f (a, a) -> Forward a) -> f (a, a) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> Forward a
f (f (Forward a) -> Forward a)
-> (f (a, a) -> f (Forward a)) -> f (a, a) -> Forward a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, a) -> Forward a) -> f (a, a) -> f (Forward a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> a -> Forward a) -> (a, a) -> Forward a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Forward a
forall a. a -> a -> Forward a
bundle)
{-# INLINE du #-}

-- | Compute the answer and directional derivative of a function given a zipped up 'Functor' of the input values and their derivatives
du'
  :: (Functor f, Num a)
  => (f (Forward a) -> Forward a)
  -> f (a, a)
  -> (a, a)
du' :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(f (Forward a) -> Forward a) -> f (a, a) -> (a, a)
du' f (Forward a) -> Forward a
f = Forward a -> (a, a)
forall a. Num a => Forward a -> (a, a)
unbundle (Forward a -> (a, a))
-> (f (a, a) -> Forward a) -> f (a, a) -> (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> Forward a
f (f (Forward a) -> Forward a)
-> (f (a, a) -> f (Forward a)) -> f (a, a) -> Forward a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, a) -> Forward a) -> f (a, a) -> f (Forward a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> a -> Forward a) -> (a, a) -> Forward a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Forward a
forall a. a -> a -> Forward a
bundle)
{-# INLINE du' #-}

-- | Compute a vector of directional derivatives for a function given a zipped up 'Functor' of the input values and their derivatives.
duF
  :: (Functor f, Functor g, Num a)
  => (f (Forward a) -> g (Forward a))
  -> f (a, a)
  -> g a
duF :: forall (f :: * -> *) (g :: * -> *) a.
(Functor f, Functor g, Num a) =>
(f (Forward a) -> g (Forward a)) -> f (a, a) -> g a
duF f (Forward a) -> g (Forward a)
f = (Forward a -> a) -> g (Forward a) -> g a
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Forward a -> a
forall a. Num a => Forward a -> a
tangent (g (Forward a) -> g a)
-> (f (a, a) -> g (Forward a)) -> f (a, a) -> g a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> g (Forward a)
f (f (Forward a) -> g (Forward a))
-> (f (a, a) -> f (Forward a)) -> f (a, a) -> g (Forward a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, a) -> Forward a) -> f (a, a) -> f (Forward a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> a -> Forward a) -> (a, a) -> Forward a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Forward a
forall a. a -> a -> Forward a
bundle)
{-# INLINE duF #-}

-- | Compute a vector of answers and directional derivatives for a function given a zipped up 'Functor' of the input values and their derivatives.
duF'
  :: (Functor f, Functor g, Num a)
  => (f (Forward a) -> g (Forward a))
  -> f (a, a)
  -> g (a, a)
duF' :: forall (f :: * -> *) (g :: * -> *) a.
(Functor f, Functor g, Num a) =>
(f (Forward a) -> g (Forward a)) -> f (a, a) -> g (a, a)
duF' f (Forward a) -> g (Forward a)
f = (Forward a -> (a, a)) -> g (Forward a) -> g (a, a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Forward a -> (a, a)
forall a. Num a => Forward a -> (a, a)
unbundle (g (Forward a) -> g (a, a))
-> (f (a, a) -> g (Forward a)) -> f (a, a) -> g (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> g (Forward a)
f (f (Forward a) -> g (Forward a))
-> (f (a, a) -> f (Forward a)) -> f (a, a) -> g (Forward a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, a) -> Forward a) -> f (a, a) -> f (Forward a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> a -> Forward a) -> (a, a) -> Forward a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Forward a
forall a. a -> a -> Forward a
bundle)
{-# INLINE duF' #-}

-- | The 'diff' function calculates the first derivative of a scalar-to-scalar function by forward-mode 'AD'
--
-- >>> diff sin 0
-- 1.0
diff
  :: Num a
  => (Forward a -> Forward a)
  -> a
  -> a
diff :: forall a. Num a => (Forward a -> Forward a) -> a -> a
diff Forward a -> Forward a
f a
a = Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> a) -> Forward a -> a
forall a b. (a -> b) -> a -> b
$ (Forward a -> Forward a) -> a -> Forward a
forall a b. Num a => (Forward a -> b) -> a -> b
apply Forward a -> Forward a
f a
a
{-# INLINE diff #-}

-- | The 'diff'' function calculates the result and first derivative of scalar-to-scalar function by 'Forward' mode 'AD'
--
-- @
-- 'diff'' 'sin' == 'sin' 'Control.Arrow.&&&' 'cos'
-- 'diff'' f = f 'Control.Arrow.&&&' d f
-- @
--
-- >>> diff' sin 0
-- (0.0,1.0)
--
-- >>> diff' exp 0
-- (1.0,1.0)

diff'
  :: Num a
  => (Forward a -> Forward a)
  -> a
  -> (a, a)
diff' :: forall a. Num a => (Forward a -> Forward a) -> a -> (a, a)
diff' Forward a -> Forward a
f a
a = Forward a -> (a, a)
forall a. Num a => Forward a -> (a, a)
unbundle (Forward a -> (a, a)) -> Forward a -> (a, a)
forall a b. (a -> b) -> a -> b
$ (Forward a -> Forward a) -> a -> Forward a
forall a b. Num a => (Forward a -> b) -> a -> b
apply Forward a -> Forward a
f a
a
{-# INLINE diff' #-}

-- | The 'diffF' function calculates the first derivatives of scalar-to-nonscalar function by 'Forward' mode 'AD'
--
-- >>> diffF (\a -> [sin a, cos a]) 0
-- [1.0,-0.0]
diffF
  :: (Functor f, Num a)
  => (Forward a -> f (Forward a))
  -> a
  -> f a
diffF :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(Forward a -> f (Forward a)) -> a -> f a
diffF Forward a -> f (Forward a)
f a
a = Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> a) -> f (Forward a) -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Forward a -> f (Forward a)) -> a -> f (Forward a)
forall a b. Num a => (Forward a -> b) -> a -> b
apply Forward a -> f (Forward a)
f a
a
{-# INLINE diffF #-}

-- | The 'diffF'' function calculates the result and first derivatives of a scalar-to-non-scalar function by 'Forward' mode 'AD'
--
-- >>> diffF' (\a -> [sin a, cos a]) 0
-- [(0.0,1.0),(1.0,-0.0)]
diffF'
  :: (Functor f, Num a)
  => (Forward a -> f (Forward a))
  -> a
  -> f (a, a)
diffF' :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(Forward a -> f (Forward a)) -> a -> f (a, a)
diffF' Forward a -> f (Forward a)
f a
a = Forward a -> (a, a)
forall a. Num a => Forward a -> (a, a)
unbundle (Forward a -> (a, a)) -> f (Forward a) -> f (a, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Forward a -> f (Forward a)) -> a -> f (Forward a)
forall a b. Num a => (Forward a -> b) -> a -> b
apply Forward a -> f (Forward a)
f a
a
{-# INLINE diffF' #-}

-- | A fast, simple, transposed Jacobian computed with forward-mode AD.
jacobianT
  :: (Traversable f, Functor g, Num a)
  => (f (Forward a) -> g (Forward a))
  -> f a
  -> f (g a)
jacobianT :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Forward a) -> g (Forward a)) -> f a -> f (g a)
jacobianT f (Forward a) -> g (Forward a)
f = (f (Forward a) -> g a) -> f a -> f (g a)
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Forward a) -> b) -> f a -> f b
bind ((Forward a -> a) -> g (Forward a) -> g a
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Forward a -> a
forall a. Num a => Forward a -> a
tangent (g (Forward a) -> g a)
-> (f (Forward a) -> g (Forward a)) -> f (Forward a) -> g a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> g (Forward a)
f)
{-# INLINE jacobianT #-}

-- | A fast, simple, transposed Jacobian computed with 'Forward' mode 'AD' that combines the output with the input.
jacobianWithT
  :: (Traversable f, Functor g, Num a)
  => (a -> a -> b)
  -> (f (Forward a) -> g (Forward a))
  -> f a
  -> f (g b)
jacobianWithT :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b) -> (f (Forward a) -> g (Forward a)) -> f a -> f (g b)
jacobianWithT a -> a -> b
g f (Forward a) -> g (Forward a)
f = (a -> g (Forward a) -> g b)
-> (f (Forward a) -> g (Forward a)) -> f a -> f (g b)
forall (f :: * -> *) a b c.
(Traversable f, Num a) =>
(a -> b -> c) -> (f (Forward a) -> b) -> f a -> f c
bindWith a -> g (Forward a) -> g b
forall {f :: * -> *}. Functor f => a -> f (Forward a) -> f b
g' f (Forward a) -> g (Forward a)
f where
  g' :: a -> f (Forward a) -> f b
g' a
a f (Forward a)
ga = a -> a -> b
g a
a (a -> b) -> (Forward a -> a) -> Forward a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> b) -> f (Forward a) -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Forward a)
ga
{-# INLINE jacobianWithT #-}

-- | Compute the Jacobian using 'Forward' mode 'AD'. This must transpose the result, so 'jacobianT' is faster and allows more result types.
--
--
-- >>> jacobian (\[x,y] -> [y,x,x+y,x*y,exp x * sin y]) [pi,1]
-- [[0.0,1.0],[1.0,0.0],[1.0,1.0],[1.0,3.141592653589793],[19.472221418841606,12.502969588876512]]
jacobian
  :: (Traversable f, Traversable g, Num a)
  => (f (Forward a) -> g (Forward a))
  -> f a
  -> g (f a)
jacobian :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Traversable g, Num a) =>
(f (Forward a) -> g (Forward a)) -> f a -> g (f a)
jacobian f (Forward a) -> g (Forward a)
f f a
as = (a -> f a -> f a) -> f (g a) -> g a -> g (f a)
forall (f :: * -> *) (g :: * -> *) b a c.
(Functor f, Foldable f, Traversable g) =>
(b -> f a -> c) -> f (g a) -> g b -> g c
transposeWith ((f a -> f a) -> a -> f a -> f a
forall a b. a -> b -> a
const f a -> f a
forall a. a -> a
id) f (g a)
t g a
p where
  (g a
p, f (g a)
t) = (f (Forward a) -> g a) -> f a -> (g a, f (g a))
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Forward a) -> b) -> f a -> (b, f b)
bind' ((Forward a -> a) -> g (Forward a) -> g a
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Forward a -> a
forall a. Num a => Forward a -> a
tangent (g (Forward a) -> g a)
-> (f (Forward a) -> g (Forward a)) -> f (Forward a) -> g a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> g (Forward a)
f) f a
as
{-# INLINE jacobian #-}

-- | Compute the Jacobian using 'Forward' mode 'AD' and combine the output with the input. This must transpose the result, so 'jacobianWithT' is faster, and allows more result types.
jacobianWith
  :: (Traversable f, Traversable g, Num a)
  => (a -> a -> b)
  -> (f (Forward a) -> g (Forward a))
  -> f a
  -> g (f b)
jacobianWith :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Traversable g, Num a) =>
(a -> a -> b) -> (f (Forward a) -> g (Forward a)) -> f a -> g (f b)
jacobianWith a -> a -> b
g f (Forward a) -> g (Forward a)
f f a
as = (Forward a -> f b -> f b) -> f (g b) -> g (Forward a) -> g (f b)
forall (f :: * -> *) (g :: * -> *) b a c.
(Functor f, Foldable f, Traversable g) =>
(b -> f a -> c) -> f (g a) -> g b -> g c
transposeWith ((f b -> f b) -> Forward a -> f b -> f b
forall a b. a -> b -> a
const f b -> f b
forall a. a -> a
id) f (g b)
t g (Forward a)
p where
  (g (Forward a)
p, f (g b)
t) = (a -> g (Forward a) -> g b)
-> (f (Forward a) -> g (Forward a))
-> f a
-> (g (Forward a), f (g b))
forall (f :: * -> *) a b c.
(Traversable f, Num a) =>
(a -> b -> c) -> (f (Forward a) -> b) -> f a -> (b, f c)
bindWith' a -> g (Forward a) -> g b
forall {f :: * -> *}. Functor f => a -> f (Forward a) -> f b
g' f (Forward a) -> g (Forward a)
f f a
as
  g' :: a -> f (Forward a) -> f b
g' a
a f (Forward a)
ga = a -> a -> b
g a
a (a -> b) -> (Forward a -> a) -> Forward a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> b) -> f (Forward a) -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Forward a)
ga
{-# INLINE jacobianWith #-}

-- | Compute the Jacobian using 'Forward' mode 'AD' along with the actual answer.
jacobian'
  :: (Traversable f, Traversable g, Num a)
  => (f (Forward a) -> g (Forward a))
  -> f a
  -> g (a, f a)
jacobian' :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Traversable g, Num a) =>
(f (Forward a) -> g (Forward a)) -> f a -> g (a, f a)
jacobian' f (Forward a) -> g (Forward a)
f f a
as = (Forward a -> f (Forward a) -> (a, f a))
-> f (g (Forward a)) -> g (Forward a) -> g (a, f a)
forall (f :: * -> *) (g :: * -> *) b a c.
(Functor f, Foldable f, Traversable g) =>
(b -> f a -> c) -> f (g a) -> g b -> g c
transposeWith Forward a -> f (Forward a) -> (a, f a)
forall {f :: * -> *} {a} {b}.
(Functor f, Num a, Num b) =>
Forward a -> f (Forward b) -> (a, f b)
row f (g (Forward a))
t g (Forward a)
p where
  (g (Forward a)
p, f (g (Forward a))
t) = (f (Forward a) -> g (Forward a))
-> f a -> (g (Forward a), f (g (Forward a)))
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Forward a) -> b) -> f a -> (b, f b)
bind' f (Forward a) -> g (Forward a)
f f a
as
  row :: Forward a -> f (Forward b) -> (a, f b)
row Forward a
x f (Forward b)
as' = (Forward a -> a
forall a. Num a => Forward a -> a
primal Forward a
x, Forward b -> b
forall a. Num a => Forward a -> a
tangent (Forward b -> b) -> f (Forward b) -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Forward b)
as')
{-# INLINE jacobian' #-}

-- | Compute the Jacobian using 'Forward' mode 'AD' combined with the input using a user specified function, along with the actual answer.
jacobianWith'
  :: (Traversable f, Traversable g, Num a)
  => (a -> a -> b)
  -> (f (Forward a) -> g (Forward a))
  -> f a
  -> g (a, f b)
jacobianWith' :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Traversable g, Num a) =>
(a -> a -> b)
-> (f (Forward a) -> g (Forward a)) -> f a -> g (a, f b)
jacobianWith' a -> a -> b
g f (Forward a) -> g (Forward a)
f f a
as = (Forward a -> f b -> (a, f b))
-> f (g b) -> g (Forward a) -> g (a, f b)
forall (f :: * -> *) (g :: * -> *) b a c.
(Functor f, Foldable f, Traversable g) =>
(b -> f a -> c) -> f (g a) -> g b -> g c
transposeWith Forward a -> f b -> (a, f b)
forall {a} {b}. Num a => Forward a -> b -> (a, b)
row f (g b)
t g (Forward a)
p where
  (g (Forward a)
p, f (g b)
t) = (a -> g (Forward a) -> g b)
-> (f (Forward a) -> g (Forward a))
-> f a
-> (g (Forward a), f (g b))
forall (f :: * -> *) a b c.
(Traversable f, Num a) =>
(a -> b -> c) -> (f (Forward a) -> b) -> f a -> (b, f c)
bindWith' a -> g (Forward a) -> g b
forall {f :: * -> *}. Functor f => a -> f (Forward a) -> f b
g' f (Forward a) -> g (Forward a)
f f a
as
  row :: Forward a -> b -> (a, b)
row Forward a
x b
as' = (Forward a -> a
forall a. Num a => Forward a -> a
primal Forward a
x, b
as')
  g' :: a -> f (Forward a) -> f b
g' a
a f (Forward a)
ga = a -> a -> b
g a
a (a -> b) -> (Forward a -> a) -> Forward a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> b) -> f (Forward a) -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Forward a)
ga
{-# INLINE jacobianWith' #-}

-- | Compute the gradient of a function using forward mode AD.
--
-- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Reverse.grad' for @n@ inputs, in exchange for better space utilization.
grad
  :: (Traversable f, Num a)
  => (f (Forward a) -> Forward a)
  -> f a
  -> f a
grad :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Forward a) -> Forward a) -> f a -> f a
grad f (Forward a) -> Forward a
f = (f (Forward a) -> a) -> f a -> f a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Forward a) -> b) -> f a -> f b
bind (Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> a)
-> (f (Forward a) -> Forward a) -> f (Forward a) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> Forward a
f)
{-# INLINE grad #-}

-- | Compute the gradient and answer to a function using forward mode AD.
--
-- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Reverse.grad'' for @n@ inputs, in exchange for better space utilization.
grad'
  :: (Traversable f, Num a)
  => (f (Forward a) -> Forward a)
  -> f a
  -> (a, f a)
grad' :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Forward a) -> Forward a) -> f a -> (a, f a)
grad' f (Forward a) -> Forward a
f f a
as = (Forward a -> a
forall a. Num a => Forward a -> a
primal Forward a
b, Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> a) -> f (Forward a) -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Forward a)
bs) where
  (Forward a
b, f (Forward a)
bs) = (f (Forward a) -> Forward a) -> f a -> (Forward a, f (Forward a))
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Forward a) -> b) -> f a -> (b, f b)
bind' f (Forward a) -> Forward a
f f a
as
{-# INLINE grad' #-}

-- | Compute the gradient of a function using forward mode AD and combine the result with the input using a user-specified function.
--
-- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Reverse.gradWith' for @n@ inputs, in exchange for better space utilization.
gradWith
  :: (Traversable f, Num a)
  => (a -> a -> b)
  -> (f (Forward a) -> Forward a)
  -> f a
  -> f b
gradWith :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Forward a) -> Forward a) -> f a -> f b
gradWith a -> a -> b
g f (Forward a) -> Forward a
f = (a -> a -> b) -> (f (Forward a) -> a) -> f a -> f b
forall (f :: * -> *) a b c.
(Traversable f, Num a) =>
(a -> b -> c) -> (f (Forward a) -> b) -> f a -> f c
bindWith a -> a -> b
g (Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> a)
-> (f (Forward a) -> Forward a) -> f (Forward a) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> Forward a
f)
{-# INLINE gradWith #-}

-- | Compute the gradient of a function using forward mode AD and the answer, and combine the result with the input using a
-- user-specified function.
--
-- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Reverse.gradWith'' for @n@ inputs, in exchange for better space utilization.
--
-- >>> gradWith' (,) sum [0..4]
-- (10,[(0,1),(1,1),(2,1),(3,1),(4,1)])
gradWith'
  :: (Traversable f, Num a)
  => (a -> a -> b)
  -> (f (Forward a) -> Forward a)
  -> f a
  -> (a, f b)
gradWith' :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Forward a) -> Forward a) -> f a -> (a, f b)
gradWith' a -> a -> b
g f (Forward a) -> Forward a
f f a
as = (Forward a -> a
forall a. Num a => Forward a -> a
primal (Forward a -> a) -> Forward a -> a
forall a b. (a -> b) -> a -> b
$ f (Forward a) -> Forward a
f (a -> Forward a
forall a. a -> Forward a
Lift (a -> Forward a) -> f a -> f (Forward a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
as), (a -> a -> b) -> (f (Forward a) -> a) -> f a -> f b
forall (f :: * -> *) a b c.
(Traversable f, Num a) =>
(a -> b -> c) -> (f (Forward a) -> b) -> f a -> f c
bindWith a -> a -> b
g (Forward a -> a
forall a. Num a => Forward a -> a
tangent (Forward a -> a)
-> (f (Forward a) -> Forward a) -> f (Forward a) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Forward a) -> Forward a
f) f a
as)
{-# INLINE gradWith' #-}

-- | Compute the product of a vector with the Hessian using forward-on-forward-mode AD.
--
hessianProduct
  :: (Traversable f, Num a)
  => (f (On (Forward (Forward a))) -> On (Forward (Forward a)))
  -> f (a, a)
  -> f a
hessianProduct :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (On (Forward (Forward a))) -> On (Forward (Forward a)))
-> f (a, a) -> f a
hessianProduct f (On (Forward (Forward a))) -> On (Forward (Forward a))
f = (f (Forward a) -> f (Forward a)) -> f (a, a) -> f a
forall (f :: * -> *) (g :: * -> *) a.
(Functor f, Functor g, Num a) =>
(f (Forward a) -> g (Forward a)) -> f (a, a) -> g a
duF ((f (Forward a) -> f (Forward a)) -> f (a, a) -> f a)
-> (f (Forward a) -> f (Forward a)) -> f (a, a) -> f a
forall a b. (a -> b) -> a -> b
$ (f (Forward (Forward a)) -> Forward (Forward a))
-> f (Forward a) -> f (Forward a)
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Forward a) -> Forward a) -> f a -> f a
grad ((f (Forward (Forward a)) -> Forward (Forward a))
 -> f (Forward a) -> f (Forward a))
-> (f (Forward (Forward a)) -> Forward (Forward a))
-> f (Forward a)
-> f (Forward a)
forall a b. (a -> b) -> a -> b
$ On (Forward (Forward a)) -> Forward (Forward a)
forall t. On t -> t
off (On (Forward (Forward a)) -> Forward (Forward a))
-> (f (Forward (Forward a)) -> On (Forward (Forward a)))
-> f (Forward (Forward a))
-> Forward (Forward a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (On (Forward (Forward a))) -> On (Forward (Forward a))
f (f (On (Forward (Forward a))) -> On (Forward (Forward a)))
-> (f (Forward (Forward a)) -> f (On (Forward (Forward a))))
-> f (Forward (Forward a))
-> On (Forward (Forward a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Forward (Forward a) -> On (Forward (Forward a)))
-> f (Forward (Forward a)) -> f (On (Forward (Forward a)))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Forward (Forward a) -> On (Forward (Forward a))
forall t. t -> On t
On
{-# INLINE hessianProduct #-}

-- | Compute the gradient and hessian product using forward-on-forward-mode AD.
hessianProduct'
  :: (Traversable f, Num a)
  => (f (On (Forward (Forward a))) -> On (Forward (Forward a)))
  -> f (a, a) -> f (a, a)
hessianProduct' :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (On (Forward (Forward a))) -> On (Forward (Forward a)))
-> f (a, a) -> f (a, a)
hessianProduct' f (On (Forward (Forward a))) -> On (Forward (Forward a))
f = (f (Forward a) -> f (Forward a)) -> f (a, a) -> f (a, a)
forall (f :: * -> *) (g :: * -> *) a.
(Functor f, Functor g, Num a) =>
(f (Forward a) -> g (Forward a)) -> f (a, a) -> g (a, a)
duF' ((f (Forward a) -> f (Forward a)) -> f (a, a) -> f (a, a))
-> (f (Forward a) -> f (Forward a)) -> f (a, a) -> f (a, a)
forall a b. (a -> b) -> a -> b
$ (f (Forward (Forward a)) -> Forward (Forward a))
-> f (Forward a) -> f (Forward a)
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Forward a) -> Forward a) -> f a -> f a
grad ((f (Forward (Forward a)) -> Forward (Forward a))
 -> f (Forward a) -> f (Forward a))
-> (f (Forward (Forward a)) -> Forward (Forward a))
-> f (Forward a)
-> f (Forward a)
forall a b. (a -> b) -> a -> b
$ On (Forward (Forward a)) -> Forward (Forward a)
forall t. On t -> t
off (On (Forward (Forward a)) -> Forward (Forward a))
-> (f (Forward (Forward a)) -> On (Forward (Forward a)))
-> f (Forward (Forward a))
-> Forward (Forward a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (On (Forward (Forward a))) -> On (Forward (Forward a))
f (f (On (Forward (Forward a))) -> On (Forward (Forward a)))
-> (f (Forward (Forward a)) -> f (On (Forward (Forward a))))
-> f (Forward (Forward a))
-> On (Forward (Forward a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Forward (Forward a) -> On (Forward (Forward a)))
-> f (Forward (Forward a)) -> f (On (Forward (Forward a)))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Forward (Forward a) -> On (Forward (Forward a))
forall t. t -> On t
On
{-# INLINE hessianProduct' #-}