{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   : (c) Edward Kmett 2010-2021
-- License     : BSD3
-- Maintainer  : ekmett@gmail.com
-- Stability   : experimental
-- Portability : GHC only
--
-- First order dense forward mode using 'Representable' functors
--
-----------------------------------------------------------------------------

module Numeric.AD.Mode.Dense.Representable
  ( AD, Repr, auto
  -- * Dense Gradients
  , grad
  , grad'
  , gradWith
  , gradWith'

  -- * Dense Jacobians (synonyms)
  , jacobian
  , jacobian'
  , jacobianWith
  , jacobianWith'
  ) where

import Data.Functor.Rep
import Numeric.AD.Internal.Dense.Representable (Repr)
import qualified Numeric.AD.Rank1.Dense.Representable as Rank1
import Numeric.AD.Internal.Type
import Numeric.AD.Mode

-- $setup
-- >>> :set -XDeriveGeneric -XDeriveFunctor
-- >>> import GHC.Generics (Generic1)
-- >>> import Data.Distributive (Distributive (..))
-- >>> import Data.Functor.Rep (Representable, distributeRep)
-- >>> data V3 a = V3 a a a deriving (Generic1, Functor, Show)
-- >>> instance Representable V3; instance Distributive V3 where distribute = distributeRep

-- | The 'grad' function calculates the gradient of a non-scalar-to-scalar function with dense-mode AD in a single pass.
--
-- >>> grad (\(V3 x y z) -> x*y+z) (V3 1 2 3)
-- V3 2 1 1
--
grad :: (Representable f, Eq (Rep f), Num a) => (forall s. f (AD s (Repr f a)) -> AD s (Repr f a)) -> f a -> f a
grad :: forall (f :: * -> *) a.
(Representable f, Eq (Rep f), Num a) =>
(forall s. f (AD s (Repr f a)) -> AD s (Repr f a)) -> f a -> f a
grad forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f = (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
Rank1.grad (AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(AD Any (Repr f a) -> Repr f a)
-> (f (Repr f a) -> AD Any (Repr f a)) -> f (Repr f a) -> Repr f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> AD Any (Repr f a)
forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f(f (AD Any (Repr f a)) -> AD Any (Repr f a))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> AD Any (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE grad #-}

grad' :: (Representable f, Eq (Rep f), Num a) => (forall s. f (AD s (Repr f a)) -> AD s (Repr f a)) -> f a -> (a, f a)
grad' :: forall (f :: * -> *) a.
(Representable f, Eq (Rep f), Num a) =>
(forall s. f (AD s (Repr f a)) -> AD s (Repr f a))
-> f a -> (a, f a)
grad' forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f = (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)
Rank1.grad' (AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(AD Any (Repr f a) -> Repr f a)
-> (f (Repr f a) -> AD Any (Repr f a)) -> f (Repr f a) -> Repr f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> AD Any (Repr f a)
forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f(f (AD Any (Repr f a)) -> AD Any (Repr f a))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> AD Any (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE grad' #-}

gradWith :: (Representable f, Eq (Rep f), Num a) => (a -> a -> b) -> (forall s. f (AD s (Repr f a)) -> AD s (Repr f a)) -> f a -> f b
gradWith :: forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Repr f a)) -> AD s (Repr f a)) -> f a -> f b
gradWith a -> a -> b
g forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f = (a -> a -> b) -> (f (Repr f a) -> Repr f a) -> f a -> f b
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
Rank1.gradWith a -> a -> b
g (AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(AD Any (Repr f a) -> Repr f a)
-> (f (Repr f a) -> AD Any (Repr f a)) -> f (Repr f a) -> Repr f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> AD Any (Repr f a)
forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f(f (AD Any (Repr f a)) -> AD Any (Repr f a))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> AD Any (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE gradWith #-}

gradWith' :: (Representable f, Eq (Rep f), Num a) => (a -> a -> b) -> (forall s. f (AD s (Repr f a)) -> AD s (Repr f a)) -> f a -> (a, f b)
gradWith' :: forall (f :: * -> *) a b.
(Representable f, Eq (Rep f), Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Repr f a)) -> AD s (Repr f a))
-> f a
-> (a, f b)
gradWith' a -> a -> b
g forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f = (a -> a -> b) -> (f (Repr f a) -> Repr f a) -> f a -> (a, f b)
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)
Rank1.gradWith' a -> a -> b
g (AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(AD Any (Repr f a) -> Repr f a)
-> (f (Repr f a) -> AD Any (Repr f a)) -> f (Repr f a) -> Repr f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> AD Any (Repr f a)
forall s. f (AD s (Repr f a)) -> AD s (Repr f a)
f(f (AD Any (Repr f a)) -> AD Any (Repr f a))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> AD Any (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE gradWith' #-}

jacobian :: (Representable f, Eq (Rep f), Functor g, Num a) => (forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))) -> f a -> g (f a)
jacobian :: forall (f :: * -> *) (g :: * -> *) a.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a)))
-> f a -> g (f a)
jacobian forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f = (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)
Rank1.jacobian ((AD Any (Repr f a) -> Repr f a)
-> g (AD Any (Repr f a)) -> g (Repr f a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(g (AD Any (Repr f a)) -> g (Repr f a))
-> (f (Repr f a) -> g (AD Any (Repr f a)))
-> f (Repr f a)
-> g (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> g (AD Any (Repr f a))
forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f(f (AD Any (Repr f a)) -> g (AD Any (Repr f a)))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> g (AD Any (Repr f a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobian #-}

jacobian' :: (Representable f, Eq (Rep f), Functor g, Num a) => (forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))) -> f a -> g (a, f a)
jacobian' :: forall (f :: * -> *) (g :: * -> *) a.
(Representable f, Eq (Rep f), Functor g, Num a) =>
(forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a)))
-> f a -> g (a, f a)
jacobian' forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f = (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)
Rank1.jacobian' ((AD Any (Repr f a) -> Repr f a)
-> g (AD Any (Repr f a)) -> g (Repr f a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(g (AD Any (Repr f a)) -> g (Repr f a))
-> (f (Repr f a) -> g (AD Any (Repr f a)))
-> f (Repr f a)
-> g (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> g (AD Any (Repr f a))
forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f(f (AD Any (Repr f a)) -> g (AD Any (Repr f a)))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> g (AD Any (Repr f a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobian' #-}

jacobianWith :: (Representable f, Eq (Rep f), Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Repr f a)) -> g (AD s (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)
-> (forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a)))
-> f a
-> g (f b)
jacobianWith a -> a -> b
g forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f = (a -> a -> b) -> (f (Repr f a) -> g (Repr f a)) -> f a -> g (f b)
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)
Rank1.jacobianWith a -> a -> b
g ((AD Any (Repr f a) -> Repr f a)
-> g (AD Any (Repr f a)) -> g (Repr f a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(g (AD Any (Repr f a)) -> g (Repr f a))
-> (f (Repr f a) -> g (AD Any (Repr f a)))
-> f (Repr f a)
-> g (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> g (AD Any (Repr f a))
forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f(f (AD Any (Repr f a)) -> g (AD Any (Repr f a)))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> g (AD Any (Repr f a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobianWith #-}

jacobianWith' :: (Representable f, Eq (Rep f), Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Repr f a)) -> g (AD s (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)
-> (forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a)))
-> f a
-> g (a, f b)
jacobianWith' a -> a -> b
g forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f = (a -> a -> b)
-> (f (Repr f a) -> g (Repr f a)) -> f a -> g (a, f b)
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)
Rank1.jacobianWith' a -> a -> b
g ((AD Any (Repr f a) -> Repr f a)
-> g (AD Any (Repr f a)) -> g (Repr f a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Repr f a) -> Repr f a
forall s a. AD s a -> a
runAD(g (AD Any (Repr f a)) -> g (Repr f a))
-> (f (Repr f a) -> g (AD Any (Repr f a)))
-> f (Repr f a)
-> g (Repr f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Repr f a)) -> g (AD Any (Repr f a))
forall s. f (AD s (Repr f a)) -> g (AD s (Repr f a))
f(f (AD Any (Repr f a)) -> g (AD Any (Repr f a)))
-> (f (Repr f a) -> f (AD Any (Repr f a)))
-> f (Repr f a)
-> g (AD Any (Repr f a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Repr f a -> AD Any (Repr f a))
-> f (Repr f a) -> f (AD Any (Repr f a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Repr f a -> AD Any (Repr f a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobianWith' #-}