-----------------------------------------------------------------------------
-- |
-- Copyright   : (c) Edward Kmett 2010-2021
-- License     : BSD3
-- Maintainer  : ekmett@gmail.com
-- Stability   : experimental
-- Portability : GHC only
--
-- Higher order derivatives via a \"dual number tower\".
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Sparse
  ( Sparse
  , auto
  -- * Sparse Gradients
  , grad
  , grad'
  , gradWith
  , gradWith'
  -- * Variadic Gradients
  -- $vgrad
  , Grad
  , vgrad
  -- * Higher-Order Gradients
  , grads
  -- * Variadic Higher-Order Gradients
  , Grads
  , vgrads

  -- * Sparse Jacobians (synonyms)
  , jacobian
  , jacobian'
  , jacobianWith
  , jacobianWith'
  , jacobians

  -- * Sparse Hessians
  , hessian
  , hessian'

  , hessianF
  , hessianF'

  ) where

import Control.Comonad
import Control.Comonad.Cofree
import Numeric.AD.Jet
import Numeric.AD.Internal.Sparse
import Numeric.AD.Internal.Combinators
import Numeric.AD.Mode

second :: (a -> b) -> (c, a) -> (c, b)
second :: forall a b c. (a -> b) -> (c, a) -> (c, b)
second a -> b
g (c
a,a
b) = (c
a, a -> b
g a
b)
{-# INLINE second #-}

grad :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> f a
grad :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> f a
grad f (Sparse a) -> Sparse a
f f a
as = f a -> Sparse a -> f a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
f b -> Sparse a -> f a
d f a
as (Sparse a -> f a) -> Sparse a -> f a
forall a b. (a -> b) -> a -> b
$ (f (Sparse a) -> Sparse a) -> f a -> Sparse a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Sparse a) -> b) -> f a -> b
apply f (Sparse a) -> Sparse a
f f a
as
{-# INLINE grad #-}

grad' :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> (a, f a)
grad' :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> (a, f a)
grad' f (Sparse a) -> Sparse a
f f a
as = f a -> Sparse a -> (a, f a)
forall (f :: * -> *) a.
(Traversable f, Num a) =>
f a -> Sparse a -> (a, f a)
d' f a
as (Sparse a -> (a, f a)) -> Sparse a -> (a, f a)
forall a b. (a -> b) -> a -> b
$ (f (Sparse a) -> Sparse a) -> f a -> Sparse a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Sparse a) -> b) -> f a -> b
apply f (Sparse a) -> Sparse a
f f a
as
{-# INLINE grad' #-}

gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (f (Sparse a) -> Sparse a) -> f a -> f b
gradWith :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Sparse a) -> Sparse a) -> f a -> f b
gradWith a -> a -> b
g f (Sparse a) -> Sparse a
f f a
as = (a -> a -> b) -> f a -> f a -> f b
forall (f :: * -> *) (g :: * -> *) a b c.
(Foldable f, Traversable g) =>
(a -> b -> c) -> f a -> g b -> g c
zipWithT a -> a -> b
g f a
as (f a -> f b) -> f a -> f b
forall a b. (a -> b) -> a -> b
$ (f (Sparse a) -> Sparse a) -> f a -> f a
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> f a
grad f (Sparse a) -> Sparse a
f f a
as
{-# INLINE gradWith #-}

gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (f (Sparse a) -> Sparse a) -> f a -> (a, f b)
gradWith' :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Sparse a) -> Sparse a) -> f a -> (a, f b)
gradWith' a -> a -> b
g f (Sparse a) -> Sparse a
f f a
as = (f a -> f b) -> (a, f a) -> (a, f b)
forall a b c. (a -> b) -> (c, a) -> (c, b)
second ((a -> a -> b) -> f a -> f a -> f b
forall (f :: * -> *) (g :: * -> *) a b c.
(Foldable f, Traversable g) =>
(a -> b -> c) -> f a -> g b -> g c
zipWithT a -> a -> b
g f a
as) ((a, f a) -> (a, f b)) -> (a, f a) -> (a, f b)
forall a b. (a -> b) -> a -> b
$ (f (Sparse a) -> Sparse a) -> f a -> (a, f a)
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> (a, f a)
grad' f (Sparse a) -> Sparse a
f f a
as
{-# INLINE gradWith' #-}

jacobian :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (f a)
jacobian :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (f a)
jacobian f (Sparse a) -> g (Sparse a)
f f a
as = f a -> Sparse a -> f a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
f b -> Sparse a -> f a
d f a
as (Sparse a -> f a) -> g (Sparse a) -> g (f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Sparse a) -> g (Sparse a)) -> f a -> g (Sparse a)
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Sparse a) -> b) -> f a -> b
apply f (Sparse a) -> g (Sparse a)
f f a
as
{-# INLINE jacobian #-}

jacobian' :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f a)
jacobian' :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f a)
jacobian' f (Sparse a) -> g (Sparse a)
f f a
as = f a -> Sparse a -> (a, f a)
forall (f :: * -> *) a.
(Traversable f, Num a) =>
f a -> Sparse a -> (a, f a)
d' f a
as (Sparse a -> (a, f a)) -> g (Sparse a) -> g (a, f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Sparse a) -> g (Sparse a)) -> f a -> g (Sparse a)
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Sparse a) -> b) -> f a -> b
apply f (Sparse a) -> g (Sparse a)
f f a
as
{-# INLINE jacobian' #-}

jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (f (Sparse a) -> g (Sparse a)) -> f a -> g (f b)
jacobianWith :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b) -> (f (Sparse a) -> g (Sparse a)) -> f a -> g (f b)
jacobianWith a -> a -> b
g f (Sparse a) -> g (Sparse a)
f f a
as = (a -> a -> b) -> f a -> f a -> f b
forall (f :: * -> *) (g :: * -> *) a b c.
(Foldable f, Traversable g) =>
(a -> b -> c) -> f a -> g b -> g c
zipWithT a -> a -> b
g f a
as (f a -> f b) -> g (f a) -> g (f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Sparse a) -> g (Sparse a)) -> f a -> g (f a)
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (f a)
jacobian f (Sparse a) -> g (Sparse a)
f f a
as
{-# INLINE jacobianWith #-}

jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f b)
jacobianWith' :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b)
-> (f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f b)
jacobianWith' a -> a -> b
g f (Sparse a) -> g (Sparse a)
f f a
as = (f a -> f b) -> (a, f a) -> (a, f b)
forall a b c. (a -> b) -> (c, a) -> (c, b)
second ((a -> a -> b) -> f a -> f a -> f b
forall (f :: * -> *) (g :: * -> *) a b c.
(Foldable f, Traversable g) =>
(a -> b -> c) -> f a -> g b -> g c
zipWithT a -> a -> b
g f a
as) ((a, f a) -> (a, f b)) -> g (a, f a) -> g (a, f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f a)
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f a)
jacobian' f (Sparse a) -> g (Sparse a)
f f a
as
{-# INLINE jacobianWith' #-}

grads :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> Cofree f a
grads :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> Cofree f a
grads f (Sparse a) -> Sparse a
f f a
as = f a -> Sparse a -> Cofree f a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
f b -> Sparse a -> Cofree f a
ds f a
as (Sparse a -> Cofree f a) -> Sparse a -> Cofree f a
forall a b. (a -> b) -> a -> b
$ (f (Sparse a) -> Sparse a) -> f a -> Sparse a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Sparse a) -> b) -> f a -> b
apply f (Sparse a) -> Sparse a
f f a
as
{-# INLINE grads #-}

jacobians :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (Cofree f a)
jacobians :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (Cofree f a)
jacobians f (Sparse a) -> g (Sparse a)
f f a
as = f a -> Sparse a -> Cofree f a
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
f b -> Sparse a -> Cofree f a
ds f a
as (Sparse a -> Cofree f a) -> g (Sparse a) -> g (Cofree f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Sparse a) -> g (Sparse a)) -> f a -> g (Sparse a)
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(f (Sparse a) -> b) -> f a -> b
apply f (Sparse a) -> g (Sparse a)
f f a
as
{-# INLINE jacobians #-}

d2 :: Functor f => Cofree f a -> f (f a)
d2 :: forall (f :: * -> *) a. Functor f => Cofree f a -> f (f a)
d2 = Jet f (f (f a)) -> f (f a)
forall (f :: * -> *) a. Jet f a -> a
headJet (Jet f (f (f a)) -> f (f a))
-> (Cofree f a -> Jet f (f (f a))) -> Cofree f a -> f (f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Jet f (f a) -> Jet f (f (f a))
forall (f :: * -> *) a. Jet f a -> Jet f (f a)
tailJet (Jet f (f a) -> Jet f (f (f a)))
-> (Cofree f a -> Jet f (f a)) -> Cofree f a -> Jet f (f (f a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Jet f a -> Jet f (f a)
forall (f :: * -> *) a. Jet f a -> Jet f (f a)
tailJet (Jet f a -> Jet f (f a))
-> (Cofree f a -> Jet f a) -> Cofree f a -> Jet f (f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cofree f a -> Jet f a
forall (f :: * -> *) a. Functor f => Cofree f a -> Jet f a
jet
{-# INLINE d2 #-}

d2' :: Functor f => Cofree f a -> (a, f (a, f a))
d2' :: forall (f :: * -> *) a. Functor f => Cofree f a -> (a, f (a, f a))
d2' (a
a :< f (Cofree f a)
as) = (a
a, (Cofree f a -> (a, f a)) -> f (Cofree f a) -> f (a, f a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(a
da :< f (Cofree f a)
das) -> (a
da, Cofree f a -> a
forall a. Cofree f a -> a
forall (w :: * -> *) a. Comonad w => w a -> a
extract (Cofree f a -> a) -> f (Cofree f a) -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Cofree f a)
das)) f (Cofree f a)
as)
{-# INLINE d2' #-}

hessian :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> f (f a)
hessian :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> f (f a)
hessian f (Sparse a) -> Sparse a
f f a
as = Cofree f a -> f (f a)
forall (f :: * -> *) a. Functor f => Cofree f a -> f (f a)
d2 (Cofree f a -> f (f a)) -> Cofree f a -> f (f a)
forall a b. (a -> b) -> a -> b
$ (f (Sparse a) -> Sparse a) -> f a -> Cofree f a
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> Cofree f a
grads f (Sparse a) -> Sparse a
f f a
as
{-# INLINE hessian #-}

hessian' :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> (a, f (a, f a))
hessian' :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> (a, f (a, f a))
hessian' f (Sparse a) -> Sparse a
f f a
as = Cofree f a -> (a, f (a, f a))
forall (f :: * -> *) a. Functor f => Cofree f a -> (a, f (a, f a))
d2' (Cofree f a -> (a, f (a, f a))) -> Cofree f a -> (a, f (a, f a))
forall a b. (a -> b) -> a -> b
$ (f (Sparse a) -> Sparse a) -> f a -> Cofree f a
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Sparse a) -> Sparse a) -> f a -> Cofree f a
grads f (Sparse a) -> Sparse a
f f a
as
{-# INLINE hessian' #-}

hessianF :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (f (f a))
hessianF :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (f (f a))
hessianF f (Sparse a) -> g (Sparse a)
f f a
as = Cofree f a -> f (f a)
forall (f :: * -> *) a. Functor f => Cofree f a -> f (f a)
d2 (Cofree f a -> f (f a)) -> g (Cofree f a) -> g (f (f a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Sparse a) -> g (Sparse a)) -> f a -> g (Cofree f a)
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (Cofree f a)
jacobians f (Sparse a) -> g (Sparse a)
f f a
as
{-# INLINE hessianF #-}

hessianF' :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f (a, f a))
hessianF' :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f (a, f a))
hessianF' f (Sparse a) -> g (Sparse a)
f f a
as = Cofree f a -> (a, f (a, f a))
forall (f :: * -> *) a. Functor f => Cofree f a -> (a, f (a, f a))
d2' (Cofree f a -> (a, f (a, f a)))
-> g (Cofree f a) -> g (a, f (a, f a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Sparse a) -> g (Sparse a)) -> f a -> g (Cofree f a)
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Sparse a) -> g (Sparse a)) -> f a -> g (Cofree f a)
jacobians f (Sparse a) -> g (Sparse a)
f f a
as
{-# INLINE hessianF' #-}

-- $vgrad
--
-- Variadic combinators for variadic mixed-mode automatic differentiation.
--
-- Unfortunately, variadicity comes at the expense of being able to use
-- quantification to avoid sensitivity confusion, so be careful when
-- counting the number of 'auto' calls you use when taking the gradient
-- of a function that takes gradients!