{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   : (c) Edward Kmett 2010-2021
-- License     : BSD3
-- Maintainer  : ekmett@gmail.com
-- Stability   : experimental
-- Portability : GHC only
--
-- Dense forward mode automatic differentiation with representable functors.
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Dense.Representable
  ( Repr
  , auto
  -- * Sparse Gradients
  , grad
  , grad'
  , gradWith
  , gradWith'
  -- * Sparse Jacobians (synonyms)
  , jacobian
  , jacobian'
  , jacobianWith
  , jacobianWith'
  ) where

import Data.Functor.Rep
import Numeric.AD.Internal.Dense.Representable
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
  :: (Representable f, Eq (Rep f), Num a)
  => (f (Repr f a) -> Repr f a)
  -> f a
  -> f a
grad :: forall (f :: * -> *) a.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> Repr f a) -> f a -> f a
grad f (Repr f a) -> Repr f a
f f a
as = f a -> Repr f a -> f a
forall (f :: * -> *) a. f a -> Repr f a -> f a
ds (a -> f a
forall (f :: * -> *) a. Representable f => a -> f a
pureRep a
0) (Repr f a -> f a) -> Repr f a -> f a
forall a b. (a -> b) -> a -> b
$ (f (Repr f a) -> Repr f a) -> f a -> Repr f a
forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> b) -> f a -> b
apply f (Repr f a) -> Repr f a
f f a
as
{-# INLINE grad #-}

grad'
  :: (Representable f, Eq (Rep f), Num a)
  => (f (Repr f a) -> Repr f a)
  -> f a
  -> (a, f a)
grad' :: forall (f :: * -> *) a.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> Repr f a) -> f a -> (a, f a)
grad' f (Repr f a) -> Repr f a
f f a
as = f a -> Repr f a -> (a, f a)
forall a (f :: * -> *). Num a => f a -> Repr f a -> (a, f a)
ds' (a -> f a
forall (f :: * -> *) a. Representable f => a -> f a
pureRep a
0) (Repr f a -> (a, f a)) -> Repr f a -> (a, f a)
forall a b. (a -> b) -> a -> b
$ (f (Repr f a) -> Repr f a) -> f a -> Repr f a
forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> b) -> f a -> b
apply f (Repr f a) -> Repr f a
f f a
as
{-# INLINE grad' #-}

gradWith
  :: (Representable f, Eq (Rep f), Num a)
  => (a -> a -> b)
  -> (f (Repr f a) -> Repr f a)
  -> f a
  -> f b
gradWith :: forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(a -> a -> b) -> (f (Repr f a) -> Repr f a) -> f a -> f b
gradWith a -> a -> b
g f (Repr f a) -> Repr f a
f f a
as = (a -> a -> b) -> f a -> f a -> f b
forall (f :: * -> *) a b c.
Representable f =>
(a -> b -> c) -> f a -> f b -> f c
liftR2 a -> a -> b
g f a
as (f a -> f b) -> f a -> f b
forall a b. (a -> b) -> a -> b
$ (f (Repr f a) -> Repr f a) -> f a -> f a
forall (f :: * -> *) a.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> Repr f a) -> f a -> f a
grad f (Repr f a) -> Repr f a
f f a
as
{-# INLINE gradWith #-}

gradWith'
  :: (Representable f, Eq (Rep f), Num a)
  => (a -> a -> b)
  -> (f (Repr f a) -> Repr f a)
  -> f a
  -> (a, f b)
gradWith' :: forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(a -> a -> b) -> (f (Repr f a) -> Repr f a) -> f a -> (a, f b)
gradWith' a -> a -> b
g f (Repr f a) -> Repr f 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 :: * -> *) a b c.
Representable f =>
(a -> b -> c) -> f a -> f b -> f c
liftR2 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 (Repr f a) -> Repr f a) -> f a -> (a, f a)
forall (f :: * -> *) a.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> Repr f a) -> f a -> (a, f a)
grad' f (Repr f a) -> Repr f a
f f a
as
{-# INLINE gradWith' #-}

jacobian
  :: (Representable f, Eq (Rep f), Functor g, Num a)
  => (f (Repr f a) -> g (Repr f a))
  -> f a
  -> g (f a)
jacobian :: forall (f :: * -> *) (g :: * -> *) a.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(f (Repr f a) -> g (Repr f a)) -> f a -> g (f a)
jacobian f (Repr f a) -> g (Repr f a)
f f a
as = f a -> Repr f a -> f a
forall (f :: * -> *) a. f a -> Repr f a -> f a
ds (a
0 a -> f a -> f a
forall a b. a -> f b -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f a
as) (Repr f a -> f a) -> g (Repr f a) -> g (f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Repr f a) -> g (Repr f a)) -> f a -> g (Repr f a)
forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> b) -> f a -> b
apply f (Repr f a) -> g (Repr f a)
f f a
as
{-# INLINE jacobian #-}

jacobian'
  :: (Representable f, Eq (Rep f), Functor g, Num a)
  => (f (Repr f a) -> g (Repr f a))
  -> f a
  -> g (a, f a)
jacobian' :: forall (f :: * -> *) (g :: * -> *) a.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(f (Repr f a) -> g (Repr f a)) -> f a -> g (a, f a)
jacobian' f (Repr f a) -> g (Repr f a)
f f a
as = f a -> Repr f a -> (a, f a)
forall a (f :: * -> *). Num a => f a -> Repr f a -> (a, f a)
ds' (a
0 a -> f a -> f a
forall a b. a -> f b -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f a
as) (Repr f a -> (a, f a)) -> g (Repr f a) -> g (a, f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f (Repr f a) -> g (Repr f a)) -> f a -> g (Repr f a)
forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(f (Repr f a) -> b) -> f a -> b
apply f (Repr f a) -> g (Repr f a)
f f a
as
{-# INLINE jacobian' #-}

jacobianWith
  :: (Representable f, Eq (Rep f), Functor g, Num a)
  => (a -> a -> b)
  -> (f (Repr f a) -> g (Repr f a))
  -> f a
  -> g (f b)
jacobianWith :: forall (f :: * -> *) (g :: * -> *) a b.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(a -> a -> b) -> (f (Repr f a) -> g (Repr f a)) -> f a -> g (f b)
jacobianWith a -> a -> b
g f (Repr f a) -> g (Repr f a)
f f a
as = (a -> a -> b) -> f a -> f a -> f b
forall (f :: * -> *) a b c.
Representable f =>
(a -> b -> c) -> f a -> f b -> f c
liftR2 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 (Repr f a) -> g (Repr f a)) -> f a -> g (f a)
forall (f :: * -> *) (g :: * -> *) a.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(f (Repr f a) -> g (Repr f a)) -> f a -> g (f a)
jacobian f (Repr f a) -> g (Repr f a)
f f a
as
{-# INLINE jacobianWith #-}

jacobianWith'
  :: (Representable f, Eq (Rep f), Functor g, Num a)
  => (a -> a -> b)
  -> (f (Repr f a) -> g (Repr f a))
  -> f a
  -> g (a, f b)
jacobianWith' :: forall (f :: * -> *) (g :: * -> *) a b.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(a -> a -> b)
-> (f (Repr f a) -> g (Repr f a)) -> f a -> g (a, f b)
jacobianWith' a -> a -> b
g f (Repr f a) -> g (Repr f 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 :: * -> *) a b c.
Representable f =>
(a -> b -> c) -> f a -> f b -> f c
liftR2 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 (Repr f a) -> g (Repr f a)) -> f a -> g (a, f a)
forall (f :: * -> *) (g :: * -> *) a.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(f (Repr f a) -> g (Repr f a)) -> f a -> g (a, f a)
jacobian' f (Repr f a) -> g (Repr f a)
f f a
as
{-# INLINE jacobianWith' #-}