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

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

import Numeric.AD.Mode
import Numeric.AD.Internal.Forward.Double

-- | Compute the directional derivative of a function given a zipped up 'Functor' of the input values and their derivatives
du
  :: Functor f
  => (f ForwardDouble -> ForwardDouble)
  -> f (Double, Double)
  -> Double
du :: forall (f :: * -> *).
Functor f =>
(f ForwardDouble -> ForwardDouble) -> f (Double, Double) -> Double
du f ForwardDouble -> ForwardDouble
f = ForwardDouble -> Double
tangent (ForwardDouble -> Double)
-> (f (Double, Double) -> ForwardDouble)
-> f (Double, Double)
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> ForwardDouble
f (f ForwardDouble -> ForwardDouble)
-> (f (Double, Double) -> f ForwardDouble)
-> f (Double, Double)
-> ForwardDouble
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Double) -> ForwardDouble)
-> f (Double, Double) -> f ForwardDouble
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Double -> Double -> ForwardDouble)
-> (Double, Double) -> ForwardDouble
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Double -> Double -> ForwardDouble
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
  => (f ForwardDouble -> ForwardDouble)
  -> f (Double, Double)
  -> (Double, Double)
du' :: forall (f :: * -> *).
Functor f =>
(f ForwardDouble -> ForwardDouble)
-> f (Double, Double) -> (Double, Double)
du' f ForwardDouble -> ForwardDouble
f = ForwardDouble -> (Double, Double)
unbundle (ForwardDouble -> (Double, Double))
-> (f (Double, Double) -> ForwardDouble)
-> f (Double, Double)
-> (Double, Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> ForwardDouble
f (f ForwardDouble -> ForwardDouble)
-> (f (Double, Double) -> f ForwardDouble)
-> f (Double, Double)
-> ForwardDouble
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Double) -> ForwardDouble)
-> f (Double, Double) -> f ForwardDouble
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Double -> Double -> ForwardDouble)
-> (Double, Double) -> ForwardDouble
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Double -> Double -> ForwardDouble
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)
  => (f ForwardDouble -> g ForwardDouble)
  -> f (Double, Double)
  -> g Double
duF :: forall (f :: * -> *) (g :: * -> *).
(Functor f, Functor g) =>
(f ForwardDouble -> g ForwardDouble)
-> f (Double, Double) -> g Double
duF f ForwardDouble -> g ForwardDouble
f = (ForwardDouble -> Double) -> g ForwardDouble -> g Double
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ForwardDouble -> Double
tangent (g ForwardDouble -> g Double)
-> (f (Double, Double) -> g ForwardDouble)
-> f (Double, Double)
-> g Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> g ForwardDouble
f (f ForwardDouble -> g ForwardDouble)
-> (f (Double, Double) -> f ForwardDouble)
-> f (Double, Double)
-> g ForwardDouble
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Double) -> ForwardDouble)
-> f (Double, Double) -> f ForwardDouble
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Double -> Double -> ForwardDouble)
-> (Double, Double) -> ForwardDouble
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Double -> Double -> ForwardDouble
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)
  => (f ForwardDouble -> g ForwardDouble)
  -> f (Double, Double)
  -> g (Double, Double)
duF' :: forall (f :: * -> *) (g :: * -> *).
(Functor f, Functor g) =>
(f ForwardDouble -> g ForwardDouble)
-> f (Double, Double) -> g (Double, Double)
duF' f ForwardDouble -> g ForwardDouble
f = (ForwardDouble -> (Double, Double))
-> g ForwardDouble -> g (Double, Double)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ForwardDouble -> (Double, Double)
unbundle (g ForwardDouble -> g (Double, Double))
-> (f (Double, Double) -> g ForwardDouble)
-> f (Double, Double)
-> g (Double, Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> g ForwardDouble
f (f ForwardDouble -> g ForwardDouble)
-> (f (Double, Double) -> f ForwardDouble)
-> f (Double, Double)
-> g ForwardDouble
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Double, Double) -> ForwardDouble)
-> f (Double, Double) -> f ForwardDouble
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Double -> Double -> ForwardDouble)
-> (Double, Double) -> ForwardDouble
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Double -> Double -> ForwardDouble
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
  :: (ForwardDouble -> ForwardDouble)
  -> Double
  -> Double
diff :: (ForwardDouble -> ForwardDouble) -> Double -> Double
diff ForwardDouble -> ForwardDouble
f Double
a = ForwardDouble -> Double
tangent (ForwardDouble -> Double) -> ForwardDouble -> Double
forall a b. (a -> b) -> a -> b
$ (ForwardDouble -> ForwardDouble) -> Double -> ForwardDouble
forall b. (ForwardDouble -> b) -> Double -> b
apply ForwardDouble -> ForwardDouble
f Double
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'
  :: (ForwardDouble -> ForwardDouble)
  -> Double
  -> (Double, Double)
diff' :: (ForwardDouble -> ForwardDouble) -> Double -> (Double, Double)
diff' ForwardDouble -> ForwardDouble
f Double
a = ForwardDouble -> (Double, Double)
unbundle (ForwardDouble -> (Double, Double))
-> ForwardDouble -> (Double, Double)
forall a b. (a -> b) -> a -> b
$ (ForwardDouble -> ForwardDouble) -> Double -> ForwardDouble
forall b. (ForwardDouble -> b) -> Double -> b
apply ForwardDouble -> ForwardDouble
f Double
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
  => (ForwardDouble -> f ForwardDouble)
  -> Double
  -> f Double
diffF :: forall (f :: * -> *).
Functor f =>
(ForwardDouble -> f ForwardDouble) -> Double -> f Double
diffF ForwardDouble -> f ForwardDouble
f Double
a = ForwardDouble -> Double
tangent (ForwardDouble -> Double) -> f ForwardDouble -> f Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ForwardDouble -> f ForwardDouble) -> Double -> f ForwardDouble
forall b. (ForwardDouble -> b) -> Double -> b
apply ForwardDouble -> f ForwardDouble
f Double
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
  => (ForwardDouble -> f ForwardDouble)
  -> Double
  -> f (Double, Double)
diffF' :: forall (f :: * -> *).
Functor f =>
(ForwardDouble -> f ForwardDouble) -> Double -> f (Double, Double)
diffF' ForwardDouble -> f ForwardDouble
f Double
a = ForwardDouble -> (Double, Double)
unbundle (ForwardDouble -> (Double, Double))
-> f ForwardDouble -> f (Double, Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ForwardDouble -> f ForwardDouble) -> Double -> f ForwardDouble
forall b. (ForwardDouble -> b) -> Double -> b
apply ForwardDouble -> f ForwardDouble
f Double
a
{-# INLINE diffF' #-}

-- | A fast, simple, transposed Jacobian computed with forward-mode AD.
jacobianT
  :: (Traversable f, Functor g)
  => (f ForwardDouble -> g ForwardDouble)
  -> f Double
  -> f (g Double)
jacobianT :: forall (f :: * -> *) (g :: * -> *).
(Traversable f, Functor g) =>
(f ForwardDouble -> g ForwardDouble) -> f Double -> f (g Double)
jacobianT f ForwardDouble -> g ForwardDouble
f = (f ForwardDouble -> g Double) -> f Double -> f (g Double)
forall (f :: * -> *) b.
Traversable f =>
(f ForwardDouble -> b) -> f Double -> f b
bind ((ForwardDouble -> Double) -> g ForwardDouble -> g Double
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ForwardDouble -> Double
tangent (g ForwardDouble -> g Double)
-> (f ForwardDouble -> g ForwardDouble)
-> f ForwardDouble
-> g Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> g ForwardDouble
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)
  => (Double -> Double -> b)
  -> (f ForwardDouble -> g ForwardDouble)
  -> f Double
  -> f (g b)
jacobianWithT :: forall (f :: * -> *) (g :: * -> *) b.
(Traversable f, Functor g) =>
(Double -> Double -> b)
-> (f ForwardDouble -> g ForwardDouble) -> f Double -> f (g b)
jacobianWithT Double -> Double -> b
g f ForwardDouble -> g ForwardDouble
f = (Double -> g ForwardDouble -> g b)
-> (f ForwardDouble -> g ForwardDouble) -> f Double -> f (g b)
forall (f :: * -> *) b c.
Traversable f =>
(Double -> b -> c) -> (f ForwardDouble -> b) -> f Double -> f c
bindWith Double -> g ForwardDouble -> g b
forall {f :: * -> *}. Functor f => Double -> f ForwardDouble -> f b
g' f ForwardDouble -> g ForwardDouble
f where
  g' :: Double -> f ForwardDouble -> f b
g' Double
a f ForwardDouble
ga = Double -> Double -> b
g Double
a (Double -> b) -> (ForwardDouble -> Double) -> ForwardDouble -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForwardDouble -> Double
tangent (ForwardDouble -> b) -> f ForwardDouble -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f ForwardDouble
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)
  => (f ForwardDouble -> g ForwardDouble)
  -> f Double
  -> g (f Double)
jacobian :: forall (f :: * -> *) (g :: * -> *).
(Traversable f, Traversable g) =>
(f ForwardDouble -> g ForwardDouble) -> f Double -> g (f Double)
jacobian f ForwardDouble -> g ForwardDouble
f f Double
as = (Double -> f Double -> f Double)
-> f (g Double) -> g Double -> g (f Double)
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 Double -> f Double) -> Double -> f Double -> f Double
forall a b. a -> b -> a
const f Double -> f Double
forall a. a -> a
id) f (g Double)
t g Double
p where
  (g Double
p, f (g Double)
t) = (f ForwardDouble -> g Double)
-> f Double -> (g Double, f (g Double))
forall (f :: * -> *) b.
Traversable f =>
(f ForwardDouble -> b) -> f Double -> (b, f b)
bind' ((ForwardDouble -> Double) -> g ForwardDouble -> g Double
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ForwardDouble -> Double
tangent (g ForwardDouble -> g Double)
-> (f ForwardDouble -> g ForwardDouble)
-> f ForwardDouble
-> g Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> g ForwardDouble
f) f Double
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)
  => (Double -> Double -> b)
  -> (f ForwardDouble -> g ForwardDouble)
  -> f Double
  -> g (f b)
jacobianWith :: forall (f :: * -> *) (g :: * -> *) b.
(Traversable f, Traversable g) =>
(Double -> Double -> b)
-> (f ForwardDouble -> g ForwardDouble) -> f Double -> g (f b)
jacobianWith Double -> Double -> b
g f ForwardDouble -> g ForwardDouble
f f Double
as = (ForwardDouble -> f b -> f b)
-> f (g b) -> g ForwardDouble -> 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) -> ForwardDouble -> 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 ForwardDouble
p where
  (g ForwardDouble
p, f (g b)
t) = (Double -> g ForwardDouble -> g b)
-> (f ForwardDouble -> g ForwardDouble)
-> f Double
-> (g ForwardDouble, f (g b))
forall (f :: * -> *) b c.
Traversable f =>
(Double -> b -> c)
-> (f ForwardDouble -> b) -> f Double -> (b, f c)
bindWith' Double -> g ForwardDouble -> g b
forall {f :: * -> *}. Functor f => Double -> f ForwardDouble -> f b
g' f ForwardDouble -> g ForwardDouble
f f Double
as
  g' :: Double -> f ForwardDouble -> f b
g' Double
a f ForwardDouble
ga = Double -> Double -> b
g Double
a (Double -> b) -> (ForwardDouble -> Double) -> ForwardDouble -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForwardDouble -> Double
tangent (ForwardDouble -> b) -> f ForwardDouble -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f ForwardDouble
ga
{-# INLINE jacobianWith #-}

-- | Compute the Jacobian using 'Forward' mode 'AD' along with the actual answer.
jacobian'
  :: (Traversable f, Traversable g)
  => (f ForwardDouble -> g ForwardDouble)
  -> f Double
  -> g (Double, f Double)
jacobian' :: forall (f :: * -> *) (g :: * -> *).
(Traversable f, Traversable g) =>
(f ForwardDouble -> g ForwardDouble)
-> f Double -> g (Double, f Double)
jacobian' f ForwardDouble -> g ForwardDouble
f f Double
as = (ForwardDouble -> f ForwardDouble -> (Double, f Double))
-> f (g ForwardDouble) -> g ForwardDouble -> g (Double, f Double)
forall (f :: * -> *) (g :: * -> *) b a c.
(Functor f, Foldable f, Traversable g) =>
(b -> f a -> c) -> f (g a) -> g b -> g c
transposeWith ForwardDouble -> f ForwardDouble -> (Double, f Double)
forall {f :: * -> *}.
Functor f =>
ForwardDouble -> f ForwardDouble -> (Double, f Double)
row f (g ForwardDouble)
t g ForwardDouble
p where
  (g ForwardDouble
p, f (g ForwardDouble)
t) = (f ForwardDouble -> g ForwardDouble)
-> f Double -> (g ForwardDouble, f (g ForwardDouble))
forall (f :: * -> *) b.
Traversable f =>
(f ForwardDouble -> b) -> f Double -> (b, f b)
bind' f ForwardDouble -> g ForwardDouble
f f Double
as
  row :: ForwardDouble -> f ForwardDouble -> (Double, f Double)
row ForwardDouble
x f ForwardDouble
as' = (ForwardDouble -> Double
primal ForwardDouble
x, ForwardDouble -> Double
tangent (ForwardDouble -> Double) -> f ForwardDouble -> f Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f ForwardDouble
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)
  => (Double -> Double -> b)
  -> (f ForwardDouble -> g ForwardDouble)
  -> f Double
  -> g (Double, f b)
jacobianWith' :: forall (f :: * -> *) (g :: * -> *) b.
(Traversable f, Traversable g) =>
(Double -> Double -> b)
-> (f ForwardDouble -> g ForwardDouble)
-> f Double
-> g (Double, f b)
jacobianWith' Double -> Double -> b
g f ForwardDouble -> g ForwardDouble
f f Double
as = (ForwardDouble -> f b -> (Double, f b))
-> f (g b) -> g ForwardDouble -> g (Double, 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 ForwardDouble -> f b -> (Double, f b)
forall {b}. ForwardDouble -> b -> (Double, b)
row f (g b)
t g ForwardDouble
p where
  (g ForwardDouble
p, f (g b)
t) = (Double -> g ForwardDouble -> g b)
-> (f ForwardDouble -> g ForwardDouble)
-> f Double
-> (g ForwardDouble, f (g b))
forall (f :: * -> *) b c.
Traversable f =>
(Double -> b -> c)
-> (f ForwardDouble -> b) -> f Double -> (b, f c)
bindWith' Double -> g ForwardDouble -> g b
forall {f :: * -> *}. Functor f => Double -> f ForwardDouble -> f b
g' f ForwardDouble -> g ForwardDouble
f f Double
as
  row :: ForwardDouble -> b -> (Double, b)
row ForwardDouble
x b
as' = (ForwardDouble -> Double
primal ForwardDouble
x, b
as')
  g' :: Double -> f ForwardDouble -> f b
g' Double
a f ForwardDouble
ga = Double -> Double -> b
g Double
a (Double -> b) -> (ForwardDouble -> Double) -> ForwardDouble -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForwardDouble -> Double
tangent (ForwardDouble -> b) -> f ForwardDouble -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f ForwardDouble
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
  => (f ForwardDouble -> ForwardDouble)
  -> f Double
  -> f Double
grad :: forall (f :: * -> *).
Traversable f =>
(f ForwardDouble -> ForwardDouble) -> f Double -> f Double
grad f ForwardDouble -> ForwardDouble
f = (f ForwardDouble -> Double) -> f Double -> f Double
forall (f :: * -> *) b.
Traversable f =>
(f ForwardDouble -> b) -> f Double -> f b
bind (ForwardDouble -> Double
tangent (ForwardDouble -> Double)
-> (f ForwardDouble -> ForwardDouble) -> f ForwardDouble -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> ForwardDouble
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
  => (f ForwardDouble -> ForwardDouble)
  -> f Double
  -> (Double, f Double)
grad' :: forall (f :: * -> *).
Traversable f =>
(f ForwardDouble -> ForwardDouble)
-> f Double -> (Double, f Double)
grad' f ForwardDouble -> ForwardDouble
f f Double
as = (ForwardDouble -> Double
primal ForwardDouble
b, ForwardDouble -> Double
tangent (ForwardDouble -> Double) -> f ForwardDouble -> f Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f ForwardDouble
bs)
    where
        (ForwardDouble
b, f ForwardDouble
bs) = (f ForwardDouble -> ForwardDouble)
-> f Double -> (ForwardDouble, f ForwardDouble)
forall (f :: * -> *) b.
Traversable f =>
(f ForwardDouble -> b) -> f Double -> (b, f b)
bind' f ForwardDouble -> ForwardDouble
f f Double
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
  => (Double -> Double -> b)
  -> (f ForwardDouble -> ForwardDouble)
  -> f Double
  -> f b
gradWith :: forall (f :: * -> *) b.
Traversable f =>
(Double -> Double -> b)
-> (f ForwardDouble -> ForwardDouble) -> f Double -> f b
gradWith Double -> Double -> b
g f ForwardDouble -> ForwardDouble
f = (Double -> Double -> b)
-> (f ForwardDouble -> Double) -> f Double -> f b
forall (f :: * -> *) b c.
Traversable f =>
(Double -> b -> c) -> (f ForwardDouble -> b) -> f Double -> f c
bindWith Double -> Double -> b
g (ForwardDouble -> Double
tangent (ForwardDouble -> Double)
-> (f ForwardDouble -> ForwardDouble) -> f ForwardDouble -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> ForwardDouble
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,[(0.0,1.0),(1.0,1.0),(2.0,1.0),(3.0,1.0),(4.0,1.0)])
gradWith'
  :: Traversable f
  => (Double -> Double -> b)
  -> (f ForwardDouble -> ForwardDouble)
  -> f Double
  -> (Double, f b)
gradWith' :: forall (f :: * -> *) b.
Traversable f =>
(Double -> Double -> b)
-> (f ForwardDouble -> ForwardDouble) -> f Double -> (Double, f b)
gradWith' Double -> Double -> b
g f ForwardDouble -> ForwardDouble
f f Double
as = (ForwardDouble -> Double
primal (ForwardDouble -> Double) -> ForwardDouble -> Double
forall a b. (a -> b) -> a -> b
$ f ForwardDouble -> ForwardDouble
f (Double -> ForwardDouble
Scalar ForwardDouble -> ForwardDouble
forall t. Mode t => Scalar t -> t
auto (Double -> ForwardDouble) -> f Double -> f ForwardDouble
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f Double
as), (Double -> Double -> b)
-> (f ForwardDouble -> Double) -> f Double -> f b
forall (f :: * -> *) b c.
Traversable f =>
(Double -> b -> c) -> (f ForwardDouble -> b) -> f Double -> f c
bindWith Double -> Double -> b
g (ForwardDouble -> Double
tangent (ForwardDouble -> Double)
-> (f ForwardDouble -> ForwardDouble) -> f ForwardDouble -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f ForwardDouble -> ForwardDouble
f) f Double
as)
{-# INLINE gradWith' #-}