diff --git a/CHANGELOG.markdown b/CHANGELOG.markdown --- a/CHANGELOG.markdown +++ b/CHANGELOG.markdown @@ -1,3 +1,8 @@ +4.2 +--- +* Removed broken `Directed` mode. +* Added `Numeric.AD.Rank1` combinators and moved most infinitesimal handling back out of the modes and into an `AD` wrapper. + 4.1 --- * Fixed a bug in the type of `conjugateGradientAscent` and `conjugateGradientDescent` that prevent users from being able to ever call it. diff --git a/README.markdown b/README.markdown --- a/README.markdown +++ b/README.markdown @@ -9,13 +9,13 @@ This library contains at its core a single implementation that describes how to compute the partial derivatives of a wide array of primitive operations. It then exposes an API that enables a user to safely combine them using standard higher-order functions, just as you would with any other Haskell numerical type. -There are several ways to compose these individual [Jacobian matrices](http://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant). We hide the choice used by the API behind an explicit "Mode" type-class and universal quantification. This prevents the end user from exploiting the properties of an individual mode, and thereby potentially violating invariants or [confusing infinitesimals](http://conway.rutgers.edu/~ccshan/wiki/blog/posts/Differentiation/). +There are several ways to compose these individual [Jacobian matrices](http://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant). We hide the choice used by the API behind an explicit "Mode" type-class and universal quantification. This prevents users from [confusing infinitesimals](http://conway.rutgers.edu/~ccshan/wiki/blog/posts/Differentiation/). If you want to risk infinitesimal confusion in order to get greater flexibility in how you curry, flip and generally combine the differential operators, then the `Rank1.*` modules are probably your cup of tea. Features -------- * Provides forward- and reverse- mode AD combinators with a common API. - * Type-level "branding" is used to both prevent the end user from confusing infinitesimals and to limit unsafe access to the implementation details of each mode. + * Optional type-level "branding" is available to prevent the end user from confusing infinitesimals * Each mode has a separate module full of combinators, with a consistent look and feel. Examples @@ -94,12 +94,12 @@ * `Numeric.AD` computes using whichever mode or combination thereof is suitable to each individual combinator. This mode is the default, re-exported by `Numeric.AD` * `Numeric.AD.Mode.Forward` provides basic forward-mode AD. It is good for computing simple derivatives. * `Numeric.AD.Mode.Sparse` computes a sparse forward-mode AD tower. It is good for higher derivatives or large numbers of outputs. - * `Numeric.AD.Mode.Reverse` computes with reverse-mode AD. It is good for computing a few outputs given many inputs. - * `Numeric.AD.Mode.Wengert` computes with reverse-mode AD. It is good for computing a few outputs given many inputs, when not using sparks. + * `Numeric.AD.Mode.Kahn` computes with reverse-mode AD. It is good for computing a few outputs given many inputs. + * `Numeric.AD.Mode.Reverse` computes with reverse-mode AD. It is good for computing a few outputs given many inputs, when not using sparks heavily. * `Numeric.AD.Mode.Tower` computes a dense forward-mode AD tower useful for higher derivatives of single input functions. - * `Numeric.AD.Newton` provides a number of combinators for root finding using Newton's method with quadratic convergence. * `Numeric.AD.Halley` provides a number of combinators for root finding using Halley's method with cubic convergence. + * `Numeric.AD.Rank1.*` provides combinators for AD that are strictly rank-1. This makes it easier to flip and contort them with higher order functions at the expense of type safety when it comes to infinitsimal confusion. ### Combinators diff --git a/ad.cabal b/ad.cabal --- a/ad.cabal +++ b/ad.cabal @@ -1,5 +1,5 @@ name: ad -version: 4.1 +version: 4.2 license: BSD3 license-File: LICENSE copyright: (c) Edward Kmett 2010-2014, @@ -110,6 +110,7 @@ erf >= 2.0 && < 2.1, free >= 4.6.1 && < 5, mtl >= 2 && < 2.2, + nats >= 0.1.2 && < 1, reflection >= 1.4 && < 2, tagged >= 0.7 && < 1, template-haskell, @@ -117,21 +118,7 @@ exposed-modules: Numeric.AD - Numeric.AD.Halley - Numeric.AD.Jacobian - Numeric.AD.Jet - Numeric.AD.Newton - - Numeric.AD.Mode - Numeric.AD.Mode.Directed - Numeric.AD.Mode.Forward - Numeric.AD.Mode.Forward.Double - Numeric.AD.Mode.Kahn - Numeric.AD.Mode.Reverse - Numeric.AD.Mode.Tower - Numeric.AD.Mode.Sparse - Numeric.AD.Internal.Dense Numeric.AD.Internal.Forward Numeric.AD.Internal.Forward.Double @@ -142,6 +129,23 @@ Numeric.AD.Internal.Reverse Numeric.AD.Internal.Sparse Numeric.AD.Internal.Tower + Numeric.AD.Internal.Type + Numeric.AD.Jacobian + Numeric.AD.Jet + Numeric.AD.Mode + Numeric.AD.Mode.Forward + Numeric.AD.Mode.Forward.Double + Numeric.AD.Mode.Kahn + Numeric.AD.Mode.Reverse + Numeric.AD.Mode.Sparse + Numeric.AD.Mode.Tower + Numeric.AD.Newton + Numeric.AD.Rank1.Forward + Numeric.AD.Rank1.Forward.Double + Numeric.AD.Rank1.Sparse + Numeric.AD.Rank1.Tower + Numeric.AD.Rank1.Kahn + Numeric.AD.Rank1.Newton other-modules: Numeric.AD.Internal.Combinators @@ -155,6 +159,7 @@ -- Verify the results of the examples test-suite doctests + default-language: Haskell2010 type: exitcode-stdio-1.0 main-is: doctests.hs build-depends: @@ -169,6 +174,7 @@ hs-source-dirs: tests benchmark blackscholes + default-language: Haskell2010 type: exitcode-stdio-1.0 main-is: BlackScholes.hs hs-source-dirs: bench diff --git a/src/Numeric/AD.hs b/src/Numeric/AD.hs --- a/src/Numeric/AD.hs +++ b/src/Numeric/AD.hs @@ -41,10 +41,10 @@ ----------------------------------------------------------------------------- module Numeric.AD - ( + ( AD -- * AD modes - Mode(auto) + , Mode(auto) , Scalar -- * Gradients (Reverse Mode) @@ -129,7 +129,6 @@ , conjugateGradientAscent ) where -import Control.Applicative import Data.Functor.Compose import Data.Traversable (Traversable) import Data.Reflection (Reifies) @@ -139,94 +138,66 @@ import Numeric.AD.Internal.Reverse (Reverse, Tape) import Numeric.AD.Internal.Sparse (Sparse, Grads, vgrads) +import Numeric.AD.Internal.Type import Numeric.AD.Mode +import qualified Numeric.AD.Rank1.Forward as Forward1 import Numeric.AD.Mode.Forward ( diff, diff', diffF, diffF' , du, du', duF, duF' - , jacobianT, jacobianWithT ) + , jacobianT, jacobianWithT + ) import Numeric.AD.Mode.Tower ( diffsF, diffs0F, diffs, diffs0 , taylor, taylor0, maclaurin, maclaurin0 - , dus, dus0, dusF, dus0F ) + , dus, dus0, dusF, dus0F + ) import qualified Numeric.AD.Mode.Reverse as Reverse import Numeric.AD.Mode.Reverse - ( grad, grad', gradWith, gradWith') + ( grad, grad', gradWith, gradWith' + , jacobian, jacobian', jacobianWith, jacobianWith' + ) -- temporary until we make a full sparse mode -import qualified Numeric.AD.Mode.Sparse as Sparse +import qualified Numeric.AD.Rank1.Sparse as Sparse1 import Numeric.AD.Mode.Sparse - ( grads, jacobians, hessian', hessianF') + ( grads, jacobians, hessian', hessianF' + ) import Numeric.AD.Newton --- | Calculate the Jacobian of a non-scalar-to-non-scalar function, automatically choosing between sparse and Reverse mode AD. --- --- If you know that you have relatively many outputs per input, consider using 'Numeric.AD.Sparse.jacobian'. --- --- >>> 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, Functor g, Num a) => (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (f a) -jacobian f bs = snd <$> jacobian' f bs -{-# INLINE jacobian #-} - --- | Calculate both the answer and Jacobian of a non-scalar-to-non-scalar function, using reverse-mode AD. --- --- If you have relatively many outputs per input, consider using 'Numeric.AD.Sparse.jacobian''. -jacobian' :: (Traversable f, Functor g, Num a) => (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (a, f a) -jacobian' = Reverse.jacobian' -{-# INLINE jacobian' #-} - --- | @'jacobianWith' g f@ calculates the Jacobian of a non-scalar-to-non-scalar function, using Reverse mode AD. --- --- The resulting Jacobian matrix is then recombined element-wise with the input using @g@. --- --- If you know that you have relatively many outputs per input, consider using 'Numeric.AD.Sparse.jacobianWith'. -jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (f b) -jacobianWith g f bs = snd <$> jacobianWith' g f bs -{-# INLINE jacobianWith #-} - --- | @'jacobianWith'' g f@ calculates the answer and Jacobian of a non-scalar-to-non-scalar function, using Reverse mode AD. --- --- The resulting Jacobian matrix is then recombined element-wise with the input using @g@. --- --- If you know that you have relatively many outputs per input, consider using 'Numeric.AD.Sparse.jacobianWith''. -jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (a, f b) -jacobianWith' = Reverse.jacobianWith' -{-# INLINE jacobianWith' #-} - -- | @'hessianProduct' f wv@ computes the product of the hessian @H@ of a non-scalar-to-scalar function @f@ at @w = 'fst' <$> wv@ with a vector @v = snd <$> wv@ using \"Pearlmutter\'s method\" from , which states: -- -- > H v = (d/dr) grad_w (w + r v) | r = 0 -- -- Or in other words, we take the directional derivative of the gradient. The gradient is calculated in reverse mode, then the directional derivative is calculated in forward mode. -- -hessianProduct :: (Traversable f, Num a) => (forall s s'. Reifies s Tape => f (On (Reverse (Forward a s') s)) -> On (Reverse (Forward a s') s)) -> f (a, a) -> f a -hessianProduct f = duF (grad (off . f . fmap On)) +hessianProduct :: (Traversable f, Num a) => (forall s. Reifies s Tape => f (On (Reverse s (Forward a))) -> On (Reverse s (Forward a))) -> f (a, a) -> f a +hessianProduct f = Forward1.duF (grad (off . f . fmap On)) -- | @'hessianProduct'' f wv@ computes both the gradient of a non-scalar-to-scalar @f@ at @w = 'fst' <$> wv@ and the product of the hessian @H@ at @w@ with a vector @v = snd <$> wv@ using \"Pearlmutter's method\". The outputs are returned wrapped in the same functor. -- -- > H v = (d/dr) grad_w (w + r v) | r = 0 -- -- Or in other words, we return the gradient and the directional derivative of the gradient. The gradient is calculated in reverse mode, then the directional derivative is calculated in forward mode. -hessianProduct' :: (Traversable f, Num a) => (forall s s'. Reifies s Tape => f (On (Reverse (Forward a s') s)) -> On (Reverse (Forward a s') s)) -> f (a, a) -> f (a, a) -hessianProduct' f = duF' (grad (off . f . fmap On)) +hessianProduct' :: (Traversable f, Num a) => (forall s. Reifies s Tape => f (On (Reverse s (Forward a))) -> On (Reverse s (Forward a))) -> f (a, a) -> f (a, a) +hessianProduct' f = Forward1.duF' (grad (off . f . fmap On)) -- | Compute the Hessian via the Jacobian of the gradient. gradient is computed in reverse mode and then the Jacobian is computed in sparse (forward) mode. -- -- >>> hessian (\[x,y] -> x*y) [1,2] -- [[0,1],[1,0]] -hessian :: (Traversable f, Num a) => (forall s s'. Reifies s Tape => f (On (Reverse (Sparse a s') s)) -> On (Reverse (Sparse a s') s)) -> f a -> f (f a) -hessian f = Sparse.jacobian (grad (off . f . fmap On)) +hessian :: (Traversable f, Num a) => (forall s. Reifies s Tape => f (On (Reverse s (Sparse a))) -> On (Reverse s (Sparse a))) -> f a -> f (f a) +hessian f = Sparse1.jacobian (grad (off . f . fmap On)) -- | Compute the order 3 Hessian tensor on a non-scalar-to-non-scalar function using 'Sparse'-on-'Reverse' -- -- >>> hessianF (\[x,y] -> [x*y,x+y,exp x*cos y]) [1,2] -- [[[0.0,1.0],[1.0,0.0]],[[0.0,0.0],[0.0,0.0]],[[-1.1312043837568135,-2.4717266720048188],[-2.4717266720048188,1.1312043837568135]]] -hessianF :: (Traversable f, Functor g, Num a) => (forall s s'. Reifies s Tape => f (On (Reverse (Sparse a s') s)) -> g (On (Reverse (Sparse a s') s))) -> f a -> g (f (f a)) -hessianF f as = getCompose $ Sparse.jacobian (Compose . Reverse.jacobian (fmap off . f . fmap On)) as +hessianF :: (Traversable f, Functor g, Num a) => (forall s. Reifies s Tape => f (On (Reverse s (Sparse a))) -> g (On (Reverse s (Sparse a)))) -> f a -> g (f (f a)) +hessianF f as = getCompose $ Sparse1.jacobian (Compose . Reverse.jacobian (fmap off . f . fmap On)) as -- $vgrad -- diff --git a/src/Numeric/AD/Halley.hs b/src/Numeric/AD/Halley.hs --- a/src/Numeric/AD/Halley.hs +++ b/src/Numeric/AD/Halley.hs @@ -24,13 +24,12 @@ , extremum ) where -import Prelude hiding (all) +import Prelude import Numeric.AD.Internal.Forward (Forward) import Numeric.AD.Internal.On import Numeric.AD.Internal.Tower (Tower) -import Numeric.AD.Mode -import Numeric.AD.Mode.Tower (diffs0) -import Numeric.AD.Mode.Forward (diff) -- , diff') +import Numeric.AD.Internal.Type (AD(..)) +import qualified Numeric.AD.Rank1.Halley as Rank1 -- $setup -- >>> import Data.Complex @@ -47,11 +46,8 @@ -- -- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1) -- 0.0 :+ 1.0 -findZero :: (Fractional a, Eq a) => (forall s. Tower a s -> Tower a s) -> a -> [a] -findZero f = go where - go x = x : if x == xn then [] else go xn where - (y:y':y'':_) = diffs0 f x - xn = x - 2*y*y'/(2*y'*y'-y*y'') +findZero :: (Fractional a, Eq a) => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> [a] +findZero f = Rank1.findZero (runAD.f.AD) {-# INLINE findZero #-} -- | The 'inverse' function inverts a scalar function using @@ -61,8 +57,8 @@ -- -- Note: the @take 10 $ inverse sqrt 1 (sqrt 10)@ example that works for Newton's method -- fails with Halley's method because the preconditions do not hold! -inverse :: (Fractional a, Eq a) => (forall s. Tower a s -> Tower a s) -> a -> a -> [a] -inverse f x0 y = findZero (\x -> f x - auto y) x0 +inverse :: (Fractional a, Eq a) => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a -> [a] +inverse f = Rank1.inverse (runAD.f.AD) {-# INLINE inverse #-} -- | The 'fixedPoint' function find a fixedpoint of a scalar @@ -74,8 +70,8 @@ -- -- >>> last $ take 10 $ fixedPoint cos 1 -- 0.7390851332151607 -fixedPoint :: (Fractional a, Eq a) => (forall s. Tower a s -> Tower a s) -> a -> [a] -fixedPoint f = findZero (\x -> f x - x) +fixedPoint :: (Fractional a, Eq a) => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> [a] +fixedPoint f = Rank1.fixedPoint (runAD.f.AD) {-# INLINE fixedPoint #-} @@ -86,6 +82,6 @@ -- -- >>> take 10 $ extremum cos 1 -- [1.0,0.29616942658570555,4.59979519460002e-3,1.6220740159042513e-8,0.0] -extremum :: (Fractional a, Eq a) => (forall s s'. On (Forward (Tower a s') s) -> On (Forward (Tower a s') s)) -> a -> [a] -extremum f = findZero (diff (off . f . On)) +extremum :: (Fractional a, Eq a) => (forall s. AD s (On (Forward (Tower a))) -> AD s (On (Forward (Tower a)))) -> a -> [a] +extremum f = Rank1.extremum (runAD.f.AD) {-# INLINE extremum #-} diff --git a/src/Numeric/AD/Internal/Dense.hs b/src/Numeric/AD/Internal/Dense.hs --- a/src/Numeric/AD/Internal/Dense.hs +++ b/src/Numeric/AD/Internal/Dense.hs @@ -51,50 +51,47 @@ import Numeric.AD.Jacobian import Numeric.AD.Mode -data Dense f a s +data Dense f a = Lift !a | Dense !a (f a) | Zero -type instance Scalar (Dense f a s) = a - -instance Show a => Show (Dense f a s) where +instance Show a => Show (Dense f a) where showsPrec d (Lift a) = showsPrec d a showsPrec d (Dense a _) = showsPrec d a showsPrec _ Zero = showString "0" -ds :: f a -> Dense f a s -> f a +ds :: f a -> Dense f a -> f a ds _ (Dense _ da) = da ds z _ = z {-# INLINE ds #-} -ds' :: Num a => f a -> Dense f a s -> (a, f a) +ds' :: Num a => f a -> Dense f a -> (a, f a) ds' _ (Dense a da) = (a, da) ds' z (Lift a) = (a, z) ds' z Zero = (0, z) {-# INLINE ds' #-} -- Bind variables and count inputs -vars :: (Traversable f, Num a) => f a -> f (Dense f a s) +vars :: (Traversable f, Num a) => f a -> f (Dense f a) vars as = snd $ mapAccumL outer (0 :: Int) as where outer !i a = (i + 1, Dense a $ snd $ mapAccumL (inner i) 0 as) inner !i !j _ = (j + 1, if i == j then 1 else 0) {-# INLINE vars #-} -apply :: (Traversable f, Num a) => (f (Dense f a s) -> b) -> f a -> b +apply :: (Traversable f, Num a) => (f (Dense f a) -> b) -> f a -> b apply f as = f (vars as) {-# INLINE apply #-} -primal :: Num a => Dense f a s -> a +primal :: Num a => Dense f a -> a primal Zero = 0 primal (Lift a) = a primal (Dense a _) = a -instance (Num a, Traversable f) => Mode (Dense f a s) where +instance (Num a, Traversable f) => Mode (Dense f a) where + type Scalar (Dense f a) = a auto = Lift zero = Zero - - _ *^ Zero = Zero a *^ Lift b = Lift (a * b) a *^ Dense b db = Dense (a * b) $ fmap (a*) db @@ -105,7 +102,7 @@ Lift a ^/ b = Lift (a / b) Dense a da ^/ b = Dense (a / b) $ fmap (/b) da -(<+>) :: (Traversable f, Num a) => Dense f a s -> Dense f a s -> Dense f a s +(<+>) :: (Traversable f, Num a) => Dense f a -> Dense f a -> Dense f a Zero <+> a = a a <+> Zero = a Lift a <+> Lift b = Lift (a + b) @@ -113,14 +110,14 @@ Dense a da <+> Lift b = Dense (a + b) da Dense a da <+> Dense b db = Dense (a + b) $ zipWithT (+) da db -(<**>) :: (Traversable f, Floating a) => Dense f a s -> Dense f a s -> Dense f a s +(<**>) :: (Traversable f, Floating a) => Dense f a -> Dense f a -> Dense f a Zero <**> y = auto (0 ** primal y) _ <**> Zero = auto 1 x <**> Lift y = lift1 (**y) (\z -> y *^ z ** Id (y - 1)) x x <**> y = lift2_ (**) (\z xi yi -> (yi * z / xi, z * log xi)) x y -instance (Traversable f, Num a) => Jacobian (Dense f a s) where - type D (Dense f a s) = Id a s +instance (Traversable f, Num a) => Jacobian (Dense f a) where + type D (Dense f a) = Id a unary f _ Zero = Lift (f 0) unary f _ (Lift b) = Lift (f b) unary f (Id dadb) (Dense b db) = Dense (f b) (fmap (dadb *) db) @@ -181,7 +178,7 @@ (Id dadb, Id dadc) = df (Id a) (Id b) (Id c) productRule dbi dci = dadb * dbi + dci * dadc -#define BODY1(x) (Traversable f, x) +#define BODY1(x) (Traversable f, x) #define BODY2(x,y) (Traversable f, x, y) -#define HEAD Dense f a s +#define HEAD Dense f a #include "instances.h" diff --git a/src/Numeric/AD/Internal/Forward.hs b/src/Numeric/AD/Internal/Forward.hs --- a/src/Numeric/AD/Internal/Forward.hs +++ b/src/Numeric/AD/Internal/Forward.hs @@ -52,40 +52,40 @@ #endif -- | 'Forward' mode AD -data Forward a s +data Forward a = Forward !a a | Lift !a | Zero deriving (Show, Data, Typeable) -type instance Scalar (Forward a s) = a - -- | Calculate the 'tangent' using forward mode AD. -tangent :: Num a => Forward a s -> a +tangent :: Num a => Forward a -> a tangent (Forward _ da) = da tangent _ = 0 {-# INLINE tangent #-} -unbundle :: Num a => Forward a s -> (a, a) +unbundle :: Num a => Forward a -> (a, a) unbundle (Forward a da) = (a, da) unbundle Zero = (0,0) unbundle (Lift a) = (a, 0) {-# INLINE unbundle #-} -bundle :: a -> a -> Forward a s +bundle :: a -> a -> Forward a bundle = Forward {-# INLINE bundle #-} -apply :: Num a => (Forward a s -> b) -> a -> b +apply :: Num a => (Forward a -> b) -> a -> b apply f a = f (bundle a 1) {-# INLINE apply #-} -primal :: Num a => Forward a s -> a +primal :: Num a => Forward a -> a primal (Forward a _) = a primal (Lift a) = a primal Zero = 0 -instance Num a => Mode (Forward a s) where +instance Num a => Mode (Forward a) where + type Scalar (Forward a) = a + auto = Lift zero = Zero @@ -107,7 +107,7 @@ Lift a ^/ b = Lift (a / b) Zero ^/ _ = Zero -(<+>) :: Num a => Forward a s -> Forward a s -> Forward a s +(<+>) :: Num a => Forward a -> Forward a -> Forward a Zero <+> a = a a <+> Zero = a Forward a da <+> Forward b db = Forward (a + b) (da + db) @@ -115,14 +115,14 @@ Lift a <+> Forward b db = Forward (a + b) db Lift a <+> Lift b = Lift (a + b) -(<**>) :: Floating a => Forward a s -> Forward a s -> Forward a s +(<**>) :: Floating a => Forward a -> Forward a -> Forward a Zero <**> y = auto (0 ** primal y) _ <**> Zero = auto 1 x <**> Lift y = lift1 (**y) (\z -> y *^ z ** Id (y - 1)) x x <**> y = lift2_ (**) (\z xi yi -> (yi * z / xi, z * log xi)) x y -instance Num a => Jacobian (Forward a s) where - type D (Forward a s) = Id a s +instance Num a => Jacobian (Forward a) where + type D (Forward a) = Id a unary f (Id dadb) (Forward b db) = Forward (f b) (dadb * db) unary f _ (Lift b) = Lift (f b) @@ -175,27 +175,27 @@ (Id dadb, Id dadc) = df (Id a) (Id b) (Id c) da = dadb * db + dc * dadc -#define HEAD Forward a s +#define HEAD Forward a #include "instances.h" -bind :: (Traversable f, Num a) => (f (Forward a s) -> b) -> f a -> f b +bind :: (Traversable f, Num a) => (f (Forward a) -> b) -> f a -> f b bind f as = snd $ mapAccumL outer (0 :: Int) as where outer !i _ = (i + 1, f $ snd $ mapAccumL (inner i) 0 as) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) -bind' :: (Traversable f, Num a) => (f (Forward a s) -> b) -> f a -> (b, f b) +bind' :: (Traversable f, Num a) => (f (Forward a) -> b) -> f a -> (b, f b) bind' f as = dropIx $ mapAccumL outer (0 :: Int, b0) as where outer (!i, _) _ = let b = f $ snd $ mapAccumL (inner i) (0 :: Int) as in ((i + 1, b), b) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) b0 = f (auto <$> as) dropIx ((_,b),bs) = (b,bs) -bindWith :: (Traversable f, Num a) => (a -> b -> c) -> (f (Forward a s) -> b) -> f a -> f c +bindWith :: (Traversable f, Num a) => (a -> b -> c) -> (f (Forward a) -> b) -> f a -> f c bindWith g f as = snd $ mapAccumL outer (0 :: Int) as where outer !i a = (i + 1, g a $ f $ snd $ mapAccumL (inner i) 0 as) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) -bindWith' :: (Traversable f, Num a) => (a -> b -> c) -> (f (Forward a s) -> b) -> f a -> (b, f c) +bindWith' :: (Traversable f, Num a) => (a -> b -> c) -> (f (Forward a) -> b) -> f a -> (b, f c) bindWith' g f as = dropIx $ mapAccumL outer (0 :: Int, b0) as where outer (!i, _) a = let b = f $ snd $ mapAccumL (inner i) (0 :: Int) as in ((i + 1, b), g a b) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) diff --git a/src/Numeric/AD/Internal/Forward/Double.hs b/src/Numeric/AD/Internal/Forward/Double.hs --- a/src/Numeric/AD/Internal/Forward/Double.hs +++ b/src/Numeric/AD/Internal/Forward/Double.hs @@ -44,25 +44,26 @@ import Numeric.AD.Jacobian import Numeric.AD.Mode -data ForwardDouble a = ForwardDouble { primal, tangent :: {-# UNPACK #-} !Double } +data ForwardDouble = ForwardDouble { primal, tangent :: {-# UNPACK #-} !Double } deriving (Read, Show) -type instance Scalar (ForwardDouble s) = Double - -unbundle :: ForwardDouble s -> (Double, Double) +unbundle :: ForwardDouble -> (Double, Double) unbundle (ForwardDouble a da) = (a, da) {-# INLINE unbundle #-} -bundle :: Double -> Double -> ForwardDouble s +bundle :: Double -> Double -> ForwardDouble bundle = ForwardDouble {-# INLINE bundle #-} -apply :: (ForwardDouble s -> b) -> Double -> b +apply :: (ForwardDouble -> b) -> Double -> b apply f a = f (bundle a 1) {-# INLINE apply #-} -instance Mode (ForwardDouble s) where +instance Mode ForwardDouble where + type Scalar ForwardDouble = Double + auto = flip ForwardDouble 0 + zero = ForwardDouble 0 0 isKnownZero (ForwardDouble 0 0) = True @@ -72,16 +73,14 @@ isKnownConstant _ = False a *^ ForwardDouble b db = ForwardDouble (a * b) (a * db) - ForwardDouble a da ^* b = ForwardDouble (a * b) (da * b) - ForwardDouble a da ^/ b = ForwardDouble (a / b) (da / b) -(<+>) :: ForwardDouble s -> ForwardDouble s -> ForwardDouble s +(<+>) :: ForwardDouble -> ForwardDouble -> ForwardDouble ForwardDouble a da <+> ForwardDouble b db = ForwardDouble (a + b) (da + db) -instance Jacobian (ForwardDouble s) where - type D (ForwardDouble s) = Id Double s +instance Jacobian ForwardDouble where + type D ForwardDouble = Id Double unary f (Id dadb) (ForwardDouble b db) = ForwardDouble (f b) (dadb * db) @@ -104,13 +103,13 @@ (Id dadb, Id dadc) = df (Id a) (Id b) (Id c) da = dadb * db + dc * dadc -instance Eq (ForwardDouble s) where +instance Eq ForwardDouble where (==) = on (==) primal -instance Ord (ForwardDouble s) where +instance Ord ForwardDouble where compare = on compare primal -instance Num (ForwardDouble s) where +instance Num ForwardDouble where fromInteger 0 = zero fromInteger n = auto (fromInteger n) (+) = (<+>) -- binary (+) 1 1 @@ -120,13 +119,13 @@ abs = lift1 abs signum signum a = lift1 signum (const zero) a -instance Fractional (ForwardDouble s) where +instance Fractional ForwardDouble where fromRational 0 = zero fromRational r = auto (fromRational r) x / y = x * recip y recip = lift1_ recip (const . negate . join (*)) -instance Floating (ForwardDouble s) where +instance Floating ForwardDouble where pi = auto pi exp = lift1_ exp const log = lift1 log recip @@ -149,7 +148,7 @@ acosh = lift1 acosh $ \x -> recip (sqrt (join (*) x - 1)) atanh = lift1 atanh $ \x -> recip (1 - join (*) x) -instance Enum (ForwardDouble s) where +instance Enum ForwardDouble where succ = lift1 succ (const 1) pred = lift1 pred (const 1) toEnum = auto . toEnum @@ -159,10 +158,10 @@ enumFromThen a b = zipWith (fromBy a delta) [0..] $ enumFromThen (primal a) (primal b) where delta = b - a enumFromThenTo a b c = zipWith (fromBy a delta) [0..] $ enumFromThenTo (primal a) (primal b) (primal c) where delta = b - a -instance Real (ForwardDouble s) where +instance Real ForwardDouble where toRational = toRational . primal -instance RealFloat (ForwardDouble s) where +instance RealFloat ForwardDouble where floatRadix = floatRadix . primal floatDigits = floatDigits . primal floatRange = floatRange . primal @@ -178,7 +177,7 @@ significand x = unary significand (scaleFloat (- floatDigits x) 1) x atan2 = lift2 atan2 $ \vx vy -> let r = recip (join (*) vx + join (*) vy) in (vy * r, negate vx * r) -instance RealFrac (ForwardDouble s) where +instance RealFrac ForwardDouble where properFraction a = (w, a `withPrimal` pb) where pa = primal a (w, pb) = properFraction pa @@ -187,34 +186,34 @@ ceiling = ceiling . primal floor = floor . primal -instance Erf (ForwardDouble s) where +instance Erf ForwardDouble where erf = lift1 erf $ \x -> (2 / sqrt pi) * exp (negate x * x) erfc = lift1 erfc $ \x -> ((-2) / sqrt pi) * exp (negate x * x) normcdf = lift1 normcdf $ \x -> ((-1) / sqrt pi) * exp (x * x * fromRational (- recip 2) / sqrt 2) -instance InvErf (ForwardDouble s) where +instance InvErf ForwardDouble where inverf = lift1 inverfc $ \x -> recip $ (2 / sqrt pi) * exp (negate x * x) inverfc = lift1 inverfc $ \x -> recip $ negate (2 / sqrt pi) * exp (negate x * x) invnormcdf = lift1 invnormcdf $ \x -> recip $ ((-1) / sqrt pi) * exp (x * x * fromRational (- recip 2) / sqrt 2) -bind :: (Traversable f) => (f (ForwardDouble s) -> b) -> f Double -> f b +bind :: Traversable f => (f ForwardDouble -> b) -> f Double -> f b bind f as = snd $ mapAccumL outer (0 :: Int) as where outer !i _ = (i + 1, f $ snd $ mapAccumL (inner i) 0 as) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) -bind' :: (Traversable f) => (f (ForwardDouble s) -> b) -> f Double -> (b, f b) +bind' :: Traversable f => (f ForwardDouble -> b) -> f Double -> (b, f b) bind' f as = dropIx $ mapAccumL outer (0 :: Int, b0) as where outer (!i, _) _ = let b = f $ snd $ mapAccumL (inner i) (0 :: Int) as in ((i + 1, b), b) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) b0 = f (auto <$> as) dropIx ((_,b),bs) = (b,bs) -bindWith :: (Traversable f) => (Double -> b -> c) -> (f (ForwardDouble s) -> b) -> f Double -> f c +bindWith :: Traversable f => (Double -> b -> c) -> (f ForwardDouble -> b) -> f Double -> f c bindWith g f as = snd $ mapAccumL outer (0 :: Int) as where outer !i a = (i + 1, g a $ f $ snd $ mapAccumL (inner i) 0 as) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) -bindWith' :: (Traversable f) => (Double -> b -> c) -> (f (ForwardDouble s) -> b) -> f Double -> (b, f c) +bindWith' :: Traversable f => (Double -> b -> c) -> (f ForwardDouble -> b) -> f Double -> (b, f c) bindWith' g f as = dropIx $ mapAccumL outer (0 :: Int, b0) as where outer (!i, _) a = let b = f $ snd $ mapAccumL (inner i) (0 :: Int) as in ((i + 1, b), g a b) inner !i !j a = (j + 1, if i == j then bundle a 1 else auto a) diff --git a/src/Numeric/AD/Internal/Identity.hs b/src/Numeric/AD/Internal/Identity.hs --- a/src/Numeric/AD/Internal/Identity.hs +++ b/src/Numeric/AD/Internal/Identity.hs @@ -29,30 +29,29 @@ import Data.Typeable (Typeable) import Numeric.AD.Mode -newtype Id a s = Id { runId :: a } deriving +newtype Id a = Id { runId :: a } deriving (Eq, Ord, Show, Enum, Bounded, Num, Real, Fractional, Floating, RealFrac, RealFloat, Monoid, Data, Typeable, Erf, InvErf) -type instance Scalar (Id a s) = a - -probe :: a -> Id a s +probe :: a -> Id a probe = Id -unprobe :: Id a s -> a +unprobe :: Id a -> a unprobe = runId -pid :: Functor f => f a -> f (Id a s) +pid :: Functor f => f a -> f (Id a) pid = fmap probe -unpid :: Functor f => f (Id a s) -> f a +unpid :: Functor f => f (Id a) -> f a unpid = fmap unprobe -probed :: Functor f => f a -> f (Id a s) +probed :: Functor f => f a -> f (Id a) probed = pid -unprobed :: Functor f => f (Id a s) -> f a +unprobed :: Functor f => f (Id a) -> f a unprobed = unpid -instance Num a => Mode (Id a s) where +instance Num a => Mode (Id a) where + type Scalar (Id a) = a auto = Id Id a ^* b = Id (a * b) a *^ Id b = Id (a * b) diff --git a/src/Numeric/AD/Internal/Kahn.hs b/src/Numeric/AD/Internal/Kahn.hs --- a/src/Numeric/AD/Internal/Kahn.hs +++ b/src/Numeric/AD/Internal/Kahn.hs @@ -80,14 +80,10 @@ deriving (Show, Data, Typeable) -- | @Kahn@ is a 'Mode' using reverse-mode automatic differentiation that provides fast 'diffFU', 'diff2FU', 'grad', 'grad2' and a fast 'jacobian' when you have a significantly smaller number of outputs than inputs. -newtype Kahn a s = Kahn (Tape a (Kahn a s)) deriving (Show, Typeable) - -type instance Scalar (Kahn a s) = a - --- deriving instance (Data (Tape a (Kahn a)) => Data (Kahn a) +newtype Kahn a = Kahn (Tape a (Kahn a)) deriving (Show, Typeable) -instance MuRef (Kahn a s) where - type DeRef (Kahn a s) = Tape a +instance MuRef (Kahn a) where + type DeRef (Kahn a) = Tape a mapDeRef _ (Kahn Zero) = pure Zero mapDeRef _ (Kahn (Lift a)) = pure (Lift a) @@ -95,7 +91,9 @@ mapDeRef f (Kahn (Binary a dadb dadc b c)) = Binary a dadb dadc <$> f b <*> f c mapDeRef f (Kahn (Unary a dadb b)) = Unary a dadb <$> f b -instance Num a => Mode (Kahn a s) where +instance Num a => Mode (Kahn a) where + type Scalar (Kahn a) = a + isKnownZero (Kahn Zero) = True isKnownZero _ = False @@ -109,24 +107,24 @@ a ^* b = lift1 (* b) (\_ -> auto b) a a ^/ b = lift1 (/ b) (\_ -> auto (recip b)) a -(<+>) :: Num a => Kahn a s -> Kahn a s -> Kahn a s +(<+>) :: Num a => Kahn a -> Kahn a -> Kahn a (<+>) = binary (+) 1 1 -(<**>) :: Floating a => Kahn a s -> Kahn a s -> Kahn a s +(<**>) :: Floating a => Kahn a -> Kahn a -> Kahn a Kahn Zero <**> y = auto (0 ** primal y) _ <**> Kahn Zero = auto 1 x <**> Kahn (Lift y) = lift1 (**y) (\z -> y *^ z ** Id (y-1)) x x <**> y = lift2_ (**) (\z xi yi -> (yi * z / xi, z * xi)) x y -primal :: Num a => Kahn a s -> a +primal :: Num a => Kahn a -> a primal (Kahn Zero) = 0 primal (Kahn (Lift a)) = a primal (Kahn (Var a _)) = a primal (Kahn (Binary a _ _ _ _)) = a primal (Kahn (Unary a _ _)) = a -instance Num a => Jacobian (Kahn a s) where - type D (Kahn a s) = Id a s +instance Num a => Jacobian (Kahn a) where + type D (Kahn a) = Id a unary f _ (Kahn Zero) = Kahn (Lift (f 0)) unary f _ (Kahn (Lift a)) = Kahn (Lift (f a)) @@ -158,14 +156,14 @@ a = f pb pc (dadb, dadc) = df (Id a) (Id pb) (Id pc) -#define HEAD Kahn a s +#define HEAD Kahn a #include -derivative :: Num a => Kahn a s -> a +derivative :: Num a => Kahn a -> a derivative = sum . map snd . partials {-# INLINE derivative #-} -derivative' :: Num a => Kahn a s -> (a, a) +derivative' :: Num a => Kahn a -> (a, a) derivative' r = (primal r, derivative r) {-# INLINE derivative' #-} @@ -208,8 +206,8 @@ -- | This returns a list of contributions to the partials. -- The variable ids returned in the list are likely /not/ unique! -{-# SPECIALIZE partials :: Kahn Double s -> [(Int, Double)] #-} -partials :: forall s a . Num a => Kahn a s -> [(Int, a)] +{-# SPECIALIZE partials :: Kahn Double -> [(Int, Double)] #-} +partials :: forall a. Num a => Kahn a -> [(Int, a)] partials tape = [ let v = sensitivities ! ix in seq v (ident, v) | (ix, Var _ ident) <- xs ] where Reified.Graph xs start = unsafePerformIO $ reifyGraph tape g = array xsBounds [ (i, successors t) | (i, t) <- xs ] @@ -233,26 +231,26 @@ successors _ = [] -- | Return an 'Array' of 'partials' given bounds for the variable IDs. -partialArray :: Num a => (Int, Int) -> Kahn a s -> Array Int a +partialArray :: Num a => (Int, Int) -> Kahn a -> Array Int a partialArray vbounds tape = accumArray (+) 0 vbounds (partials tape) {-# INLINE partialArray #-} -- | Return an 'IntMap' of sparse partials -partialMap :: Num a => Kahn a s -> IntMap a +partialMap :: Num a => Kahn a -> IntMap a partialMap = fromListWith (+) . partials {-# INLINE partialMap #-} class Num a => Grad i o o' a | i -> a o o', o -> a i o', o' -> a i o where - pack :: i -> [Kahn a ()] -> Kahn a () + pack :: i -> [Kahn a] -> Kahn a unpack :: ([a] -> [a]) -> o unpack' :: ([a] -> (a, [a])) -> o' -instance Num a => Grad (Kahn a ()) [a] (a, [a]) a where +instance Num a => Grad (Kahn a) [a] (a, [a]) a where pack i _ = i unpack f = f [] unpack' f = f [] -instance Grad i o o' a => Grad (Kahn a () -> i) (a -> o) (a -> o') a where +instance Grad i o o' a => Grad (Kahn a -> i) (a -> o) (a -> o') a where pack f (a:as) = pack (f a) as pack _ [] = error "Grad.pack: logic error" unpack f a = unpack (f . (a:)) @@ -269,26 +267,26 @@ r = f vs (vs,bds) = bind as -var :: a -> Int -> Kahn a s +var :: a -> Int -> Kahn a var a v = Kahn (Var a v) -varId :: Kahn a s -> Int +varId :: Kahn a -> Int varId (Kahn (Var _ v)) = v varId _ = error "varId: not a Var" -bind :: Traversable f => f a -> (f (Kahn a s), (Int,Int)) +bind :: Traversable f => f a -> (f (Kahn a), (Int,Int)) bind xs = (r,(0,hi)) where (r,hi) = runState (mapM freshVar xs) 0 freshVar a = state $ \s -> let s' = s + 1 in s' `seq` (var a s, s') -unbind :: Functor f => f (Kahn a s) -> Array Int a -> f a +unbind :: Functor f => f (Kahn a) -> Array Int a -> f a unbind xs ys = fmap (\v -> ys ! varId v) xs -unbindWith :: (Functor f, Num a) => (a -> b -> c) -> f (Kahn a s) -> Array Int b -> f c +unbindWith :: (Functor f, Num a) => (a -> b -> c) -> f (Kahn a) -> Array Int b -> f c unbindWith f xs ys = fmap (\v -> f (primal v) (ys ! varId v)) xs -unbindMap :: (Functor f, Num a) => f (Kahn a s) -> IntMap a -> f a +unbindMap :: (Functor f, Num a) => f (Kahn a) -> IntMap a -> f a unbindMap xs ys = fmap (\v -> findWithDefault 0 (varId v) ys) xs -unbindMapWithDefault :: (Functor f, Num a) => b -> (a -> b -> c) -> f (Kahn a s) -> IntMap b -> f c +unbindMapWithDefault :: (Functor f, Num a) => b -> (a -> b -> c) -> f (Kahn a) -> IntMap b -> f c unbindMapWithDefault z f xs ys = fmap (\v -> f (primal v) $ findWithDefault z (varId v) ys) xs diff --git a/src/Numeric/AD/Internal/On.hs b/src/Numeric/AD/Internal/On.hs --- a/src/Numeric/AD/Internal/On.hs +++ b/src/Numeric/AD/Internal/On.hs @@ -41,9 +41,8 @@ , InvErf, RealFloat, Typeable ) -type instance Scalar (On t) = Scalar (Scalar t) - instance (Mode t, Mode (Scalar t)) => Mode (On t) where + type Scalar (On t) = Scalar (Scalar t) auto = On . auto . auto a *^ On b = On (auto a *^ b) On a ^* b = On (a ^* auto b) diff --git a/src/Numeric/AD/Internal/Or.hs b/src/Numeric/AD/Internal/Or.hs --- a/src/Numeric/AD/Internal/Or.hs +++ b/src/Numeric/AD/Internal/Or.hs @@ -35,25 +35,25 @@ #endif import Numeric.AD.Mode -runL :: Or a b F -> a +runL :: Or F a b -> a runL (L a) = a -runR :: Or a b T -> b +runR :: Or T a b -> b runR (R b) = b ------------------------------------------------------------------------------ -- On ------------------------------------------------------------------------------ -chosen :: (a -> r) -> (b -> r) -> Or a b s -> r +chosen :: (a -> r) -> (b -> r) -> Or s a b -> r chosen f _ (L a) = f a chosen _ g (R b) = g b -unary :: (a -> a) -> (b -> b) -> Or a b s -> Or a b s +unary :: (a -> a) -> (b -> b) -> Or s a b -> Or s a b unary f _ (L a) = L (f a) unary _ g (R a) = R (g a) -binary :: (a -> a -> a) -> (b -> b -> b) -> Or a b s -> Or a b s -> Or a b s +binary :: (a -> a -> a) -> (b -> b -> b) -> Or s a b -> Or s a b -> Or s a b binary f _ (L a) (L b) = L (f a b) binary _ g (R a) (R b) = R (g a b) binary _ _ _ _ = impossible @@ -62,7 +62,7 @@ data T class Chosen s where - choose :: a -> b -> Or a b s + choose :: a -> b -> Or s a b instance Chosen F where choose x _ = L x @@ -72,9 +72,9 @@ #ifndef HLINT -- | The choice between two AD modes is an AD mode in its own right -data Or a b s where - L :: a -> Or a b F - R :: b -> Or a b T +data Or s a b where + L :: a -> Or F a b + R :: b -> Or T a b #if __GLASGOW_HASKELL__ >= 707 deriving Typeable #endif @@ -83,17 +83,17 @@ impossible :: a impossible = error "Numeric.AD.Internal.Or: impossible case" -instance (Eq a, Eq b) => Eq (Or a b s) where +instance (Eq a, Eq b) => Eq (Or s a b) where L a == L b = a == b R a == R b = a == b _ == _ = impossible -instance (Ord a, Ord b) => Ord (Or a b s) where +instance (Ord a, Ord b) => Ord (Or s a b) where L a `compare` L b = compare a b R a `compare` R b = compare a b _ `compare` _ = impossible -instance (Enum a, Enum b, Chosen s) => Enum (Or a b s) where +instance (Enum a, Enum b, Chosen s) => Enum (Or s a b) where pred = unary pred pred succ = unary succ succ toEnum i = choose (toEnum i) (toEnum i) @@ -110,11 +110,11 @@ enumFromThenTo (R a) (R b) (R c) = R <$> enumFromThenTo a b c enumFromThenTo _ _ _ = impossible -instance (Bounded a, Bounded b, Chosen s) => Bounded (Or a b s) where +instance (Bounded a, Bounded b, Chosen s) => Bounded (Or s a b) where maxBound = choose maxBound maxBound minBound = choose minBound minBound -instance (Num a, Num b, Chosen s) => Num (Or a b s) where +instance (Num a, Num b, Chosen s) => Num (Or s a b) where (+) = binary (+) (+) (-) = binary (-) (-) (*) = binary (*) (*) @@ -123,15 +123,15 @@ signum = unary signum signum fromInteger = choose <$> fromInteger <*> fromInteger -instance (Real a, Real b, Chosen s) => Real (Or a b s) where +instance (Real a, Real b, Chosen s) => Real (Or s a b) where toRational = chosen toRational toRational -instance (Fractional a, Fractional b, Chosen s) => Fractional (Or a b s) where +instance (Fractional a, Fractional b, Chosen s) => Fractional (Or s a b) where (/) = binary (/) (/) recip = unary recip recip fromRational = choose <$> fromRational <*> fromRational -instance (RealFrac a, RealFrac b, Chosen s) => RealFrac (Or a b s) where +instance (RealFrac a, RealFrac b, Chosen s) => RealFrac (Or s a b) where properFraction (L a) = case properFraction a of (b, c) -> (b, L c) properFraction (R a) = case properFraction a of @@ -141,7 +141,7 @@ ceiling = chosen ceiling ceiling floor = chosen floor floor -instance (Floating a, Floating b, Chosen s) => Floating (Or a b s) where +instance (Floating a, Floating b, Chosen s) => Floating (Or s a b) where pi = choose pi pi exp = unary exp exp sqrt = unary sqrt sqrt @@ -161,18 +161,18 @@ atanh = unary atanh atanh acosh = unary acosh acosh -instance (Erf a, Erf b, Chosen s) => Erf (Or a b s) where +instance (Erf a, Erf b, Chosen s) => Erf (Or s a b) where erf = unary erf erf erfc = unary erfc erfc erfcx = unary erfcx erfcx normcdf = unary normcdf normcdf -instance (InvErf a, InvErf b, Chosen s) => InvErf (Or a b s) where +instance (InvErf a, InvErf b, Chosen s) => InvErf (Or s a b) where inverf = unary inverf inverf inverfc = unary inverfc inverfc invnormcdf = unary invnormcdf invnormcdf -instance (RealFloat a, RealFloat b, Chosen s) => RealFloat (Or a b s) where +instance (RealFloat a, RealFloat b, Chosen s) => RealFloat (Or s a b) where floatRadix = chosen floatRadix floatRadix floatDigits = chosen floatDigits floatDigits floatRange = chosen floatRange floatRange @@ -188,9 +188,9 @@ isIEEE = chosen isIEEE isIEEE atan2 = binary atan2 atan2 -type instance Scalar (Or a b s) = Scalar a -instance (Mode a, Mode b, Chosen s, Scalar a ~ Scalar b) => Mode (Or a b s) where +instance (Mode a, Mode b, Chosen s, Scalar a ~ Scalar b) => Mode (Or s a b) where + type Scalar (Or s a b) = Scalar a auto = choose <$> auto <*> auto isKnownConstant = chosen isKnownConstant isKnownConstant isKnownZero = chosen isKnownZero isKnownZero diff --git a/src/Numeric/AD/Internal/Reverse.hs b/src/Numeric/AD/Internal/Reverse.hs --- a/src/Numeric/AD/Internal/Reverse.hs +++ b/src/Numeric/AD/Internal/Reverse.hs @@ -114,27 +114,27 @@ -- | This is used to create a new entry on the chain given a unary function, its derivative with respect to its input, -- the variable ID of its input, and the value of its input. Used by 'unary' and 'binary' internally. -unarily :: forall s a. Reifies s Tape => (a -> a) -> a -> Int -> a -> Reverse a s +unarily :: forall s a. Reifies s Tape => (a -> a) -> a -> Int -> a -> Reverse s a unarily f di i b = Reverse (unsafePerformIO (modifyTape (Proxy :: Proxy s) (un i di))) $! f b {-# INLINE unarily #-} -- | This is used to create a new entry on the chain given a binary function, its derivatives with respect to its inputs, -- their variable IDs and values. Used by 'binary' internally. -binarily :: forall s a. Reifies s Tape => (a -> a -> a) -> a -> a -> Int -> a -> Int -> a -> Reverse a s +binarily :: forall s a. Reifies s Tape => (a -> a -> a) -> a -> a -> Int -> a -> Int -> a -> Reverse s a binarily f di dj i b j c = Reverse (unsafePerformIO (modifyTape (Proxy :: Proxy s) (bin i j di dj))) $! f b c {-# INLINE binarily #-} #ifndef HLINT -data Reverse a s where - Zero :: Reverse a s - Lift :: a -> Reverse a s - Reverse :: {-# UNPACK #-} !Int -> a -> Reverse a s +data Reverse s a where + Zero :: Reverse s a + Lift :: a -> Reverse s a + Reverse :: {-# UNPACK #-} !Int -> a -> Reverse s a deriving (Show, Typeable) #endif -type instance Scalar (Reverse a s) = a +instance (Reifies s Tape, Num a) => Mode (Reverse s a) where + type Scalar (Reverse s a) = a -instance (Num a, Reifies s Tape) => Mode (Reverse a s) where isKnownZero Zero = True isKnownZero _ = False @@ -147,22 +147,22 @@ a ^* b = lift1 (* b) (\_ -> auto b) a a ^/ b = lift1 (/ b) (\_ -> auto (recip b)) a -(<+>) :: (Reifies s Tape, Num a) => Reverse a s -> Reverse a s -> Reverse a s +(<+>) :: (Reifies s Tape, Num a) => Reverse s a -> Reverse s a -> Reverse s a (<+>) = binary (+) 1 1 -(<**>) :: (Reifies s Tape, Floating a) => Reverse a s -> Reverse a s -> Reverse a s +(<**>) :: (Reifies s Tape, Floating a) => Reverse s a -> Reverse s a -> Reverse s a Zero <**> y = auto (0 ** primal y) _ <**> Zero = auto 1 x <**> Lift y = lift1 (**y) (\z -> y *^ z ** Id (y - 1)) x x <**> y = lift2_ (**) (\z xi yi -> (yi * z / xi, z * log xi)) x y -primal :: Num a => Reverse a s -> a +primal :: Num a => Reverse s a -> a primal Zero = 0 primal (Lift a) = a primal (Reverse _ a) = a -instance (Reifies s Tape, Num a) => Jacobian (Reverse a s) where - type D (Reverse a s) = Id a s +instance (Reifies s Tape, Num a) => Jacobian (Reverse s a) where + type D (Reverse s a) = Id a unary f _ (Zero) = Lift (f 0) unary f _ (Lift a) = Lift (f a) @@ -197,16 +197,16 @@ #define BODY1(x) (Reifies s Tape,x) #define BODY2(x,y) (Reifies s Tape,x,y) -#define HEAD Reverse a s +#define HEAD Reverse s a #include "instances.h" -- | Helper that extracts the derivative of a chain when the chain was constructed with 1 variable. -derivativeOf :: (Reifies s Tape, Num a) => Proxy s -> Reverse a s -> a +derivativeOf :: (Reifies s Tape, Num a) => Proxy s -> Reverse s a -> a derivativeOf _ = sum . partials {-# INLINE derivativeOf #-} -- | Helper that extracts both the primal and derivative of a chain when the chain was constructed with 1 variable. -derivativeOf' :: (Reifies s Tape, Num a) => Proxy s -> Reverse a s -> (a, a) +derivativeOf' :: (Reifies s Tape, Num a) => Proxy s -> Reverse s a -> (a, a) derivativeOf' p r = (primal r, derivativeOf p r) {-# INLINE derivativeOf' #-} @@ -227,8 +227,8 @@ (backPropagate $! k - 1) xs ss -- | Extract the partials from the current chain for a given AD variable. -{-# SPECIALIZE partials :: Reifies s Tape => Reverse Double s -> [Double] #-} -partials :: forall s a. (Reifies s Tape, Num a) => Reverse a s -> [a] +{-# SPECIALIZE partials :: Reifies s Tape => Reverse s Double -> [Double] #-} +partials :: forall s a. (Reifies s Tape, Num a) => Reverse s a -> [a] partials Zero = [] partials (Lift _) = [] partials (Reverse k _) = map (sensitivities !) [0..vs] where @@ -242,12 +242,12 @@ return (v, as) -- | Return an 'Array' of 'partials' given bounds for the variable IDs. -partialArrayOf :: (Reifies s Tape, Num a) => Proxy s -> (Int, Int) -> Reverse a s -> Array Int a +partialArrayOf :: (Reifies s Tape, Num a) => Proxy s -> (Int, Int) -> Reverse s a -> Array Int a partialArrayOf _ vbounds = accumArray (+) 0 vbounds . zip [0..] . partials {-# INLINE partialArrayOf #-} -- | Return an 'IntMap' of sparse partials -partialMapOf :: (Reifies s Tape, Num a) => Proxy s -> Reverse a s -> IntMap a +partialMapOf :: (Reifies s Tape, Num a) => Proxy s -> Reverse s a -> IntMap a partialMapOf _ = fromDistinctAscList . zip [0..] . partials {-# INLINE partialMapOf #-} @@ -258,26 +258,26 @@ return (reify (Tape h) k) {-# NOINLINE reifyTape #-} -var :: a -> Int -> Reverse a s +var :: a -> Int -> Reverse s a var a v = Reverse v a -varId :: Reverse a s -> Int +varId :: Reverse s a -> Int varId (Reverse v _) = v varId _ = error "varId: not a Var" -bind :: Traversable f => f a -> (f (Reverse a s), (Int,Int)) +bind :: Traversable f => f a -> (f (Reverse s a), (Int,Int)) bind xs = (r,(0,hi)) where (r,hi) = runState (mapM freshVar xs) 0 freshVar a = state $ \s -> let s' = s + 1 in s' `seq` (var a s, s') -unbind :: Functor f => f (Reverse a s) -> Array Int a -> f a +unbind :: Functor f => f (Reverse s a) -> Array Int a -> f a unbind xs ys = fmap (\v -> ys ! varId v) xs -unbindWith :: (Functor f, Num a) => (a -> b -> c) -> f (Reverse a s) -> Array Int b -> f c +unbindWith :: (Functor f, Num a) => (a -> b -> c) -> f (Reverse s a) -> Array Int b -> f c unbindWith f xs ys = fmap (\v -> f (primal v) (ys ! varId v)) xs -unbindMap :: (Functor f, Num a) => f (Reverse a s) -> IntMap a -> f a +unbindMap :: (Functor f, Num a) => f (Reverse s a) -> IntMap a -> f a unbindMap xs ys = fmap (\v -> findWithDefault 0 (varId v) ys) xs -unbindMapWithDefault :: (Functor f, Num a) => b -> (a -> b -> c) -> f (Reverse a s) -> IntMap b -> f c +unbindMapWithDefault :: (Functor f, Num a) => b -> (a -> b -> c) -> f (Reverse s a) -> IntMap b -> f c unbindMapWithDefault z f xs ys = fmap (\v -> f (primal v) $ findWithDefault z (varId v) ys) xs diff --git a/src/Numeric/AD/Internal/Sparse.hs b/src/Numeric/AD/Internal/Sparse.hs --- a/src/Numeric/AD/Internal/Sparse.hs +++ b/src/Numeric/AD/Internal/Sparse.hs @@ -63,19 +63,16 @@ -- which it was found. This should be key for efficiently computing sparse hessians. -- there are only (n + k - 1) choose k distinct nth partial derivatives of a -- function with k inputs. -data Sparse a s - = Sparse !a (IntMap (Sparse a s)) +data Sparse a + = Sparse !a (IntMap (Sparse a)) | Zero deriving (Show, Data, Typeable) -type instance Scalar (Sparse a s) = a - --- | drop keys below a given value dropMap :: Int -> IntMap a -> IntMap a dropMap n = snd . IntMap.split (n - 1) {-# INLINE dropMap #-} -times :: Num a => Sparse a s -> Int -> Sparse a s -> Sparse a s +times :: Num a => Sparse a -> Int -> Sparse a -> Sparse a times Zero _ _ = Zero times _ _ Zero = Zero times (Sparse a as) n (Sparse b bs) = Sparse (a * b) $ @@ -84,12 +81,12 @@ (fmap (a *^) (dropMap n bs)) {-# INLINE times #-} -vars :: (Traversable f, Num a) => f a -> f (Sparse a s) +vars :: (Traversable f, Num a) => f a -> f (Sparse a) vars = snd . mapAccumL var 0 where var !n a = (n + 1, Sparse a $ singleton n $ auto 1) {-# INLINE vars #-} -apply :: (Traversable f, Num a) => (f (Sparse a s) -> b) -> f a -> b +apply :: (Traversable f, Num a) => (f (Sparse a) -> b) -> f a -> b apply f = f . vars {-# INLINE apply #-} @@ -97,17 +94,17 @@ skeleton = snd . mapAccumL (\ !n _ -> (n + 1, n)) 0 {-# INLINE skeleton #-} -d :: (Traversable f, Num a) => f b -> Sparse a s -> f a +d :: (Traversable f, Num a) => f b -> Sparse a -> f a d fs (Zero) = 0 <$ fs d fs (Sparse _ da) = snd $ mapAccumL (\ !n _ -> (n + 1, maybe 0 primal $ lookup n da)) 0 fs {-# INLINE d #-} -d' :: (Traversable f, Num a) => f a -> Sparse a s -> (a, f a) +d' :: (Traversable f, Num a) => f a -> Sparse a -> (a, f a) d' fs Zero = (0, 0 <$ fs) d' fs (Sparse a da) = (a, snd $ mapAccumL (\ !n _ -> (n + 1, maybe 0 primal $ lookup n da)) 0 fs) {-# INLINE d' #-} -ds :: (Traversable f, Num a) => f b -> Sparse a s -> Cofree f a +ds :: (Traversable f, Num a) => f b -> Sparse a -> Cofree f a ds fs Zero = r where r = 0 :< (r <$ fs) ds fs (as@(Sparse a _)) = a :< (go emptyIndex <$> fns) where fns = skeleton fs @@ -142,13 +139,13 @@ {-# INLINE vds #-} -} -partial :: Num a => [Int] -> Sparse a s -> a +partial :: Num a => [Int] -> Sparse a -> a partial [] (Sparse a _) = a partial (n:ns) (Sparse _ da) = partial ns $ findWithDefault (auto 0) n da partial _ Zero = 0 {-# INLINE partial #-} -spartial :: Num a => [Int] -> Sparse a s -> Maybe a +spartial :: Num a => [Int] -> Sparse a -> Maybe a spartial [] (Sparse a _) = Just a spartial (n:ns) (Sparse _ da) = do a' <- lookup n da @@ -156,21 +153,21 @@ spartial _ Zero = Nothing {-# INLINE spartial #-} -primal :: Num a => Sparse a s -> a +primal :: Num a => Sparse a -> a primal (Sparse a _) = a primal Zero = 0 -(<**>) :: Floating a => Sparse a s -> Sparse a s -> Sparse a s +(<**>) :: Floating a => Sparse a -> Sparse a -> Sparse a Zero <**> y = auto (0 ** primal y) _ <**> Zero = auto 1 x <**> y@(Sparse b bs) | IntMap.null bs = lift1 (**b) (\z -> b *^ z <**> Sparse (b-1) IntMap.empty) x | otherwise = lift2_ (**) (\z xi yi -> (yi * z / xi, z * log xi)) x y -instance Num a => Mode (Sparse a s) where +instance Num a => Mode (Sparse a) where + type Scalar (Sparse a) = a auto a = Sparse a IntMap.empty zero = Zero - Zero ^* _ = Zero Sparse a as ^* b = Sparse (a * b) $ fmap (^* b) as _ *^ Zero = Zero @@ -180,13 +177,13 @@ infixr 6 <+> -(<+>) :: Num a => Sparse a s -> Sparse a s -> Sparse a s +(<+>) :: Num a => Sparse a -> Sparse a -> Sparse a Zero <+> a = a a <+> Zero = a Sparse a as <+> Sparse b bs = Sparse (a + b) $ unionWith (<+>) as bs -instance Num a => Jacobian (Sparse a s) where - type D (Sparse a s) = Sparse a s +instance Num a => Jacobian (Sparse a) where + type D (Sparse a) = Sparse a unary f _ Zero = auto (f 0) unary f dadb (Sparse pb bs) = Sparse (f pb) $ mapWithKey (times dadb) bs @@ -224,20 +221,20 @@ (mapWithKey (times dadb) db) (mapWithKey (times dadc) dc) -#define HEAD Sparse a s +#define HEAD Sparse a #include "instances.h" class Num a => Grad i o o' a | i -> a o o', o -> a i o', o' -> a i o where - pack :: i -> [Sparse a ()] -> Sparse a () + pack :: i -> [Sparse a] -> Sparse a unpack :: ([a] -> [a]) -> o unpack' :: ([a] -> (a, [a])) -> o' -instance Num a => Grad (Sparse a ()) [a] (a, [a]) a where +instance Num a => Grad (Sparse a) [a] (a, [a]) a where pack i _ = i unpack f = f [] unpack' f = f [] -instance Grad i o o' a => Grad (Sparse a () -> i) (a -> o) (a -> o') a where +instance Grad i o o' a => Grad (Sparse a -> i) (a -> o) (a -> o') a where pack f (a:as) = pack (f a) as pack _ [] = error "Grad.pack: logic error" unpack f a = unpack (f . (a:)) @@ -254,14 +251,14 @@ {-# INLINE vgrad' #-} class Num a => Grads i o a | i -> a o, o -> a i where - packs :: i -> [Sparse a ()] -> Sparse a () + packs :: i -> [Sparse a] -> Sparse a unpacks :: ([a] -> Cofree [] a) -> o -instance Num a => Grads (Sparse a ()) (Cofree [] a) a where +instance Num a => Grads (Sparse a) (Cofree [] a) a where packs i _ = i unpacks f = f [] -instance Grads i o a => Grads (Sparse a () -> i) (a -> o) a where +instance Grads i o a => Grads (Sparse a -> i) (a -> o) a where packs f (a:as) = packs (f a) as packs _ [] = error "Grad.pack: logic error" unpacks f a = unpacks (f . (a:)) diff --git a/src/Numeric/AD/Internal/Tower.hs b/src/Numeric/AD/Internal/Tower.hs --- a/src/Numeric/AD/Internal/Tower.hs +++ b/src/Numeric/AD/Internal/Tower.hs @@ -5,7 +5,6 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE DeriveDataTypeable #-} {-# OPTIONS_GHC -fno-warn-name-shadowing #-} {-# OPTIONS_HADDOCK not-home #-} @@ -47,11 +46,9 @@ import Numeric.AD.Mode -- | @Tower@ is an AD 'Mode' that calculates a tangent tower by forward AD, and provides fast 'diffsUU', 'diffsUF' -newtype Tower a s = Tower { getTower :: [a] } deriving (Data, Typeable) - -type instance Scalar (Tower a s) = a +newtype Tower a = Tower { getTower :: [a] } deriving (Data, Typeable) -instance Show a => Show (Tower a s) where +instance Show a => Show (Tower a) where showsPrec n (Tower as) = showParen (n > 10) $ showString "Tower " . showList as -- Local combinators @@ -86,58 +83,58 @@ d' _ = (0, 0) {-# INLINE d' #-} -tangents :: Tower a s -> Tower a s +tangents :: Tower a -> Tower a tangents (Tower []) = Tower [] tangents (Tower (_:xs)) = Tower xs {-# INLINE tangents #-} -truncated :: Tower a s -> Bool +truncated :: Tower a -> Bool truncated (Tower []) = True truncated _ = False {-# INLINE truncated #-} -bundle :: a -> Tower a s -> Tower a s +bundle :: a -> Tower a -> Tower a bundle a (Tower as) = Tower (a:as) {-# INLINE bundle #-} -withD :: (a, a) -> Tower a s +withD :: (a, a) -> Tower a withD (a, da) = Tower [a,da] {-# INLINE withD #-} -apply :: Num a => (Tower a s -> b) -> a -> b +apply :: Num a => (Tower a -> b) -> a -> b apply f a = f (Tower [a,1]) {-# INLINE apply #-} -getADTower :: Tower a s -> [a] +getADTower :: Tower a -> [a] getADTower = getTower {-# INLINE getADTower #-} -tower :: [a] -> Tower a s +tower :: [a] -> Tower a tower = Tower -primal :: Num a => Tower a s -> a +primal :: Num a => Tower a -> a primal (Tower (x:_)) = x primal _ = 0 -instance Num a => Mode (Tower a s) where +instance Num a => Mode (Tower a) where + type Scalar (Tower a) = a auto a = Tower [a] zero = Tower [] - a *^ Tower bs = Tower (map (a*) bs) Tower as ^* b = Tower (map (*b) as) Tower as ^/ b = Tower (map (/b) as) infixr 6 <+> -(<+>) :: forall a s. Num a => Tower a s -> Tower a s -> Tower a s +(<+>) :: Num a => Tower a -> Tower a -> Tower a Tower [] <+> bs = bs as <+> Tower [] = as Tower (a:as) <+> Tower (b:bs) = Tower (c:cs) where c = a + b - Tower cs = Tower as <+> (Tower bs :: Tower a s) + Tower cs = Tower as <+> Tower bs -instance Num a => Jacobian (Tower a s) where - type D (Tower a s) = Tower a s +instance Num a => Jacobian (Tower a) where + type D (Tower a) = Tower a unary f dadb b = bundle (f (primal b)) (tangents b * dadb) lift1 f df b = bundle (f (primal b)) (tangents b * df b) lift1_ f df b = a where @@ -159,11 +156,11 @@ a = bundle a0 da (dadb, dadc) = df a b c -(<**>) :: Floating a => Tower a s -> Tower a s -> Tower a s +(<**>) :: Floating a => Tower a -> Tower a -> Tower a Tower [] <**> y = auto (0 ** primal y) _ <**> Tower [] = auto 1 x <**> Tower [y] = lift1 (**y) (\z -> y *^ z <**> Tower [y-1]) x x <**> y = lift2_ (**) (\z xi yi -> (yi * z / xi, z * log xi)) x y -#define HEAD Tower a s +#define HEAD Tower a #include diff --git a/src/Numeric/AD/Internal/Type.hs b/src/Numeric/AD/Internal/Type.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/AD/Internal/Type.hs @@ -0,0 +1,32 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +----------------------------------------------------------------------------- +---- | +---- Copyright : (c) Edward Kmett 2010-2014 +---- License : BSD3 +---- Maintainer : ekmett@gmail.com +---- Stability : experimental +---- Portability : GHC only +---- +------------------------------------------------------------------------------- +module Numeric.AD.Internal.Type + ( AD(..) + ) where + +import Data.Number.Erf +import Numeric.AD.Mode +import Data.Typeable + +newtype AD s a = AD { runAD :: a } + deriving (Eq,Ord,Show,Read,Bounded,Num,Real,Fractional,Floating,Enum,RealFrac,RealFloat,Erf,InvErf,Typeable) + +instance Mode a => Mode (AD s a) where + type Scalar (AD s a) = Scalar a + isKnownConstant = isKnownConstant . runAD + isKnownZero = isKnownZero . runAD + zero = AD zero + auto = AD . auto + AD a ^* b = AD (a ^* b) + a *^ AD b = AD (a *^ b) + AD a ^/ b = AD (a ^/ b) diff --git a/src/Numeric/AD/Mode.hs b/src/Numeric/AD/Mode.hs --- a/src/Numeric/AD/Mode.hs +++ b/src/Numeric/AD/Mode.hs @@ -22,16 +22,20 @@ ( -- * AD modes Mode(..) - , Scalar ) where -type family Scalar (t :: *) :: * +import Numeric.Natural +import Data.Complex +import Data.Int +import Data.Ratio +import Data.Word infixr 7 *^ infixl 7 ^* infixr 7 ^/ class (Num t, Num (Scalar t)) => Mode t where + type Scalar t -- | allowed to return False for items with a zero derivative, but we'll give more NaNs than strictly necessary isKnownConstant :: t -> Bool isKnownConstant _ = False @@ -59,3 +63,115 @@ -- @'zero' = 'lift' 0@ zero :: t zero = auto 0 + +instance Mode Double where + type Scalar Double = Double + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Float where + type Scalar Float = Float + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Int where + type Scalar Int = Int + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Integer where + type Scalar Integer = Integer + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Int8 where + type Scalar Int8 = Int8 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Int16 where + type Scalar Int16 = Int16 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Int32 where + type Scalar Int32 = Int32 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Int64 where + type Scalar Int64 = Int64 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Natural where + type Scalar Natural = Natural + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Word where + type Scalar Word = Word + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Word8 where + type Scalar Word8 = Word8 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Word16 where + type Scalar Word16 = Word16 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Word32 where + type Scalar Word32 = Word32 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Mode Word64 where + type Scalar Word64 = Word64 + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance RealFloat a => Mode (Complex a) where + type Scalar (Complex a) = Complex a + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) + +instance Integral a => Mode (Ratio a) where + type Scalar (Ratio a) = Ratio a + isKnownConstant _ = True + isKnownZero x = 0 == x + auto = id + (^/) = (/) diff --git a/src/Numeric/AD/Mode/Directed.hs b/src/Numeric/AD/Mode/Directed.hs deleted file mode 100644 --- a/src/Numeric/AD/Mode/Directed.hs +++ /dev/null @@ -1,93 +0,0 @@ -{-# LANGUAGE RankNTypes #-} ------------------------------------------------------------------------------ --- | --- Copyright : (c) Edward Kmett 2010-2014 --- License : BSD3 --- Maintainer : ekmett@gmail.com --- Stability : experimental --- Portability : GHC only --- --- Allows the choice of AD 'Mode' to be specified at the term level for --- benchmarking or more complicated usage patterns. ------------------------------------------------------------------------------ - -module Numeric.AD.Mode.Directed - ( - -- * Gradients - grad - , grad' - -- * Jacobians - , jacobian - , jacobian' - -- * Derivatives - , diff - , diff' - -- * Exposed Types - , Direction(..) - ) where - -import Prelude hiding (reverse) -import Numeric.AD.Mode -import Data.Traversable (Traversable) -import qualified Numeric.AD.Mode.Kahn as K -import qualified Numeric.AD.Mode.Forward as F -import qualified Numeric.AD.Mode.Tower as T -import qualified Numeric.AD.Mode.Reverse as R -import qualified Numeric.AD as M -import Data.Ix - -data Direction - = Forward - | Kahn - | Reverse - | Tower - | Mixed - deriving (Show, Eq, Ord, Read, Bounded, Enum, Ix) - -diff :: Num a => Direction -> (forall t. Mode t => t -> t) -> a -> a -diff Forward f a = F.diff f a -diff Kahn f a = K.diff f a -diff Reverse f a = R.diff f a -diff Tower f a = T.diff f a -diff Mixed f a = F.diff f a -{-# INLINE diff #-} - -diff' :: Num a => Direction -> (forall t. Mode t => t -> t) -> a -> (a, a) -diff' Forward f a = F.diff' f a -diff' Kahn f a = K.diff' f a -diff' Reverse f a = R.diff' f a -diff' Tower f a = T.diff' f a -diff' Mixed f a = F.diff' f a -{-# INLINE diff' #-} - -jacobian :: (Traversable f, Traversable g, Num a) => Direction -> (forall t. Mode t => f (t) -> g (t)) -> f a -> g (f a) -jacobian Forward f a = F.jacobian f a -jacobian Kahn f a = K.jacobian f a -jacobian Reverse f a = R.jacobian f a -jacobian Tower f a = F.jacobian f a -- error "jacobian Tower: unimplemented" -jacobian Mixed f a = M.jacobian f a -{-# INLINE jacobian #-} - -jacobian' :: (Traversable f, Traversable g, Num a) => Direction -> (forall t. Mode t => f (t) -> g (t)) -> f a -> g (a, f a) -jacobian' Forward f a = F.jacobian' f a -jacobian' Kahn f a = K.jacobian' f a -jacobian' Reverse f a = R.jacobian' f a -jacobian' Tower f a = F.jacobian' f a -- error "jacobian' Tower: unimplemented" -jacobian' Mixed f a = M.jacobian' f a -{-# INLINE jacobian' #-} - -grad :: (Traversable f, Num a) => Direction -> (forall t. Mode t => f (t) -> t) -> f a -> f a -grad Forward f a = F.grad f a -grad Kahn f a = K.grad f a -grad Reverse f a = R.grad f a -grad Tower f a = F.grad f a -- error "grad Tower: unimplemented" -grad Mixed f a = M.grad f a -{-# INLINE grad #-} - -grad' :: (Traversable f, Num a) => Direction -> (forall t. Mode t => f (t) -> t) -> f a -> (a, f a) -grad' Forward f a = F.grad' f a -grad' Kahn f a = K.grad' f a -grad' Reverse f a = R.grad' f a -grad' Tower f a = F.grad' f a -- error "grad' Tower: unimplemented" -grad' Mixed f a = M.grad' f a -{-# INLINE grad' #-} diff --git a/src/Numeric/AD/Mode/Forward.hs b/src/Numeric/AD/Mode/Forward.hs --- a/src/Numeric/AD/Mode/Forward.hs +++ b/src/Numeric/AD/Mode/Forward.hs @@ -13,7 +13,9 @@ ----------------------------------------------------------------------------- module Numeric.AD.Mode.Forward - ( Forward + ( AD + , Forward + , auto -- * Gradient , grad , grad' @@ -43,36 +45,38 @@ ) where import Data.Traversable (Traversable) -import Control.Applicative import Numeric.AD.Internal.Forward import Numeric.AD.Internal.On +import Numeric.AD.Internal.Type +import qualified Numeric.AD.Rank1.Forward as Rank1 +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) => (forall s. f (Forward a s) -> Forward a s) -> f (a, a) -> a -du f = tangent . f . fmap (uncurry bundle) +du :: (Functor f, Num a) => (forall s. f (AD s (Forward a)) -> AD s (Forward a)) -> f (a, a) -> a +du f = Rank1.du (runAD.f.fmap AD) {-# 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) => (forall s. f (Forward a s) -> Forward a s) -> f (a, a) -> (a, a) -du' f = unbundle . f . fmap (uncurry bundle) +du' :: (Functor f, Num a) => (forall s. f (AD s (Forward a)) -> AD s (Forward a)) -> f (a, a) -> (a, a) +du' f = Rank1.du' (runAD.f.fmap AD) {-# 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) => (forall s. f (Forward a s) -> g (Forward a s)) -> f (a, a) -> g a -duF f = fmap tangent . f . fmap (uncurry bundle) +duF :: (Functor f, Functor g, Num a) => (forall s. f (AD s (Forward a)) -> g (AD s (Forward a))) -> f (a, a) -> g a +duF f = Rank1.duF (fmap runAD.f.fmap AD) {-# 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) => (forall s. f (Forward a s) -> g (Forward a s)) -> f (a, a) -> g (a, a) -duF' f = fmap unbundle . f . fmap (uncurry bundle) +duF' :: (Functor f, Functor g, Num a) => (forall s. f (AD s (Forward a)) -> g (AD s (Forward a))) -> f (a, a) -> g (a, a) +duF' f = Rank1.duF' (fmap runAD.f.fmap AD) {-# 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 => (forall s. Forward a s -> Forward a s) -> a -> a -diff f a = tangent $ apply f a +diff :: Num a => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> a +diff f = Rank1.diff (runAD.f.AD) {-# INLINE diff #-} -- | The 'diff'' function calculates the result and first derivative of scalar-to-scalar function by 'Forward' mode 'AD' @@ -88,92 +92,79 @@ -- >>> diff' exp 0 -- (1.0,1.0) -diff' :: Num a => (forall s. Forward a s -> Forward a s) -> a -> (a, a) -diff' f a = unbundle $ apply f a +diff' :: Num a => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> (a, a) +diff' f = Rank1.diff' (runAD.f.AD) {-# 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) => (forall s. Forward a s -> f (Forward a s)) -> a -> f a -diffF f a = tangent <$> apply f a +diffF :: (Functor f, Num a) => (forall s. AD s (Forward a) -> f (AD s (Forward a))) -> a -> f a +diffF f = Rank1.diffF (fmap runAD.f.AD) {-# 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) => (forall s. Forward a s -> f (Forward a s)) -> a -> f (a, a) -diffF' f a = unbundle <$> apply f a +diffF' :: (Functor f, Num a) => (forall s. AD s (Forward a) -> f (AD s (Forward a))) -> a -> f (a, a) +diffF' f = Rank1.diffF' (fmap runAD.f.AD) {-# INLINE diffF' #-} -- | A fast, simple, transposed Jacobian computed with forward-mode AD. -jacobianT :: (Traversable f, Functor g, Num a) => (forall s. f (Forward a s) -> g (Forward a s)) -> f a -> f (g a) -jacobianT f = bind (fmap tangent . f) +jacobianT :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Forward a)) -> g (AD s (Forward a))) -> f a -> f (g a) +jacobianT f = Rank1.jacobianT (fmap runAD.f.fmap AD) {-# 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) -> (forall s. f (Forward a s) -> g (Forward a s)) -> f a -> f (g b) -jacobianWithT g f = bindWith g' f where - g' a ga = g a . tangent <$> ga +jacobianWithT :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Forward a)) -> g (AD s (Forward a))) -> f a -> f (g b) +jacobianWithT g f = Rank1.jacobianWithT g (fmap runAD.f.fmap AD) {-# INLINE jacobianWithT #-} -#ifdef HLINT -{-# ANN jacobianWithT "HLint: ignore Eta reduce" #-} -#endif -- | 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) => (forall s . f (Forward a s) -> g (Forward a s)) -> f a -> g (f a) -jacobian f as = transposeWith (const id) t p where - (p, t) = bind' (fmap tangent . f) as +jacobian :: (Traversable f, Traversable g, Num a) => (forall s . f (AD s (Forward a)) -> g (AD s (Forward a))) -> f a -> g (f a) +jacobian f = Rank1.jacobian (fmap runAD.f.fmap AD) {-# 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) -> (forall s. f (Forward a s) -> g (Forward a s)) -> f a -> g (f b) -jacobianWith g f as = transposeWith (const id) t p where - (p, t) = bindWith' g' f as - g' a ga = g a . tangent <$> ga +jacobianWith :: (Traversable f, Traversable g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Forward a)) -> g (AD s (Forward a))) -> f a -> g (f b) +jacobianWith g f = Rank1.jacobianWith g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith #-} -- | Compute the Jacobian using 'Forward' mode 'AD' along with the actual answer. -jacobian' :: (Traversable f, Traversable g, Num a) => (forall s. f (Forward a s) -> g (Forward a s)) -> f a -> g (a, f a) -jacobian' f as = transposeWith row t p where - (p, t) = bind' f as - row x as' = (primal x, tangent <$> as') +jacobian' :: (Traversable f, Traversable g, Num a) => (forall s. f (AD s (Forward a)) -> g (AD s (Forward a))) -> f a -> g (a, f a) +jacobian' f = Rank1.jacobian' (fmap runAD.f.fmap AD) {-# 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) -> (forall s. f (Forward a s) -> g (Forward a s)) -> f a -> g (a, f b) -jacobianWith' g f as = transposeWith row t p where - (p, t) = bindWith' g' f as - row x as' = (primal x, as') - g' a ga = g a . tangent <$> ga +jacobianWith' :: (Traversable f, Traversable g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Forward a)) -> g (AD s (Forward a))) -> f a -> g (a, f b) +jacobianWith' g f = Rank1.jacobianWith' g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith' #-} -- | Compute the gradient of a function using forward mode AD. -- -- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Wengert.grad' for @n@ inputs, in exchange for better space utilization. -grad :: (Traversable f, Num a) => (forall s. f (Forward a s) -> Forward a s) -> f a -> f a -grad f = bind (tangent . f) +grad :: (Traversable f, Num a) => (forall s. f (AD s (Forward a)) -> AD s (Forward a)) -> f a -> f a +grad f = Rank1.grad (runAD.f.fmap AD) {-# INLINE grad #-} -- | Compute the gradient and answer to a function using forward mode AD. -- -- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Wengert.grad'' for @n@ inputs, in exchange for better space utilization. -grad' :: (Traversable f, Num a) => (forall s. f (Forward a s) -> Forward a s) -> f a -> (a, f a) -grad' f as = (primal b, tangent <$> bs) where - (b, bs) = bind' f as +grad' :: (Traversable f, Num a) => (forall s. f (AD s (Forward a)) -> AD s (Forward a)) -> f a -> (a, f a) +grad' f = Rank1.grad' (runAD.f.fmap AD) {-# 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.Wengert.gradWith' for @n@ inputs, in exchange for better space utilization. -gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (Forward a s) -> Forward a s) -> f a -> f b -gradWith g f = bindWith g (tangent . f) +gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Forward a)) -> AD s (Forward a)) -> f a -> f b +gradWith g f = Rank1.gradWith g (runAD.f.fmap AD) {-# INLINE gradWith #-} -- | Compute the gradient of a function using forward mode AD and the answer, and combine the result with the input using a @@ -183,17 +174,17 @@ -- -- >>> gradWith' (,) sum [0..4] -- (10,[(0,1),(1,1),(2,1),(3,1),(4,1)]) -gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (Forward a s) -> Forward a s) -> f a -> (a, f b) -gradWith' g f as = (primal $ f (Lift <$> as), bindWith g (tangent . f) as) +gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Forward a)) -> AD s (Forward a)) -> f a -> (a, f b) +gradWith' g f = Rank1.gradWith' g (runAD.f.fmap AD) {-# INLINE gradWith' #-} -- | Compute the product of a vector with the Hessian using forward-on-forward-mode AD. -- -hessianProduct :: (Traversable f, Num a) => (forall s s'. f (On (Forward (Forward a s') s)) -> On (Forward (Forward a s') s)) -> f (a, a) -> f a -hessianProduct f = duF $ grad $ off . f . fmap On +hessianProduct :: (Traversable f, Num a) => (forall s. f (AD s (On (Forward (Forward a)))) -> AD s (On (Forward (Forward a)))) -> f (a, a) -> f a +hessianProduct f = Rank1.hessianProduct (runAD.f.fmap AD) {-# INLINE hessianProduct #-} -- | Compute the gradient and hessian product using forward-on-forward-mode AD. -hessianProduct' :: (Traversable f, Num a) => (forall s s'. f (On (Forward (Forward a s') s)) -> On (Forward (Forward a s') s)) -> f (a, a) -> f (a, a) -hessianProduct' f = duF' $ grad $ off . f . fmap On +hessianProduct' :: (Traversable f, Num a) => (forall s. f (AD s (On (Forward (Forward a)))) -> AD s (On (Forward (Forward a)))) -> f (a, a) -> f (a, a) +hessianProduct' f = Rank1.hessianProduct' (runAD.f.fmap AD) {-# INLINE hessianProduct' #-} diff --git a/src/Numeric/AD/Mode/Forward/Double.hs b/src/Numeric/AD/Mode/Forward/Double.hs --- a/src/Numeric/AD/Mode/Forward/Double.hs +++ b/src/Numeric/AD/Mode/Forward/Double.hs @@ -1,6 +1,7 @@ -{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE Rank2Types #-} module Numeric.AD.Mode.Forward.Double - ( ForwardDouble + ( AD + , ForwardDouble -- * Gradient , grad , grad' @@ -26,37 +27,37 @@ , duF' ) where -import Control.Applicative import Data.Traversable (Traversable) -import Numeric.AD.Mode -import Numeric.AD.Internal.Forward.Double +import Numeric.AD.Internal.Type (AD(AD), runAD) +import Numeric.AD.Internal.Forward.Double (ForwardDouble) +import qualified Numeric.AD.Rank1.Forward.Double as Rank1 -- | Compute the directional derivative of a function given a zipped up 'Functor' of the input values and their derivatives -du :: Functor f => (forall s. f (ForwardDouble s) -> ForwardDouble s) -> f (Double, Double) -> Double -du f = tangent . f . fmap (uncurry bundle) +du :: Functor f => (forall s. f (AD s ForwardDouble) -> AD s ForwardDouble) -> f (Double, Double) -> Double +du f = Rank1.du (runAD . f . fmap AD) {-# 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 => (forall s. f (ForwardDouble s) -> ForwardDouble s) -> f (Double, Double) -> (Double, Double) -du' f = unbundle . f . fmap (uncurry bundle) +du' :: Functor f => (forall s. f (AD s ForwardDouble) -> AD s ForwardDouble) -> f (Double, Double) -> (Double, Double) +du' f = Rank1.du' (runAD . f . fmap AD) {-# 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) => (forall s. f (ForwardDouble s) -> g (ForwardDouble s)) -> f (Double, Double) -> g Double -duF f = fmap tangent . f . fmap (uncurry bundle) +duF :: (Functor f, Functor g) => (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f (Double, Double) -> g Double +duF f = Rank1.duF (fmap runAD . f . fmap AD) {-# 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) => (forall s. f (ForwardDouble s) -> g (ForwardDouble s)) -> f (Double, Double) -> g (Double, Double) -duF' f = fmap unbundle . f . fmap (uncurry bundle) +duF' :: (Functor f, Functor g) => (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f (Double, Double) -> g (Double, Double) +duF' f = Rank1.duF' (fmap runAD . f . fmap AD) {-# 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 :: (forall s. ForwardDouble s -> ForwardDouble s) -> Double -> Double -diff f a = tangent $ apply f a +diff :: (forall s. AD s ForwardDouble -> AD s ForwardDouble) -> Double -> Double +diff f = Rank1.diff (runAD.f.AD) {-# INLINE diff #-} -- | The 'diff'' function calculates the result and first derivative of scalar-to-scalar function by 'Forward' mode 'AD' @@ -71,91 +72,79 @@ -- -- >>> diff' exp 0 -- (1.0,1.0) -diff' :: (forall s. ForwardDouble s -> ForwardDouble s) -> Double -> (Double, Double) -diff' f a = unbundle $ apply f a +diff' :: (forall s. AD s ForwardDouble -> AD s ForwardDouble) -> Double -> (Double, Double) +diff' f = Rank1.diff' (runAD.f.AD) {-# 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 => (forall s. ForwardDouble s -> f (ForwardDouble s)) -> Double -> f Double -diffF f a = tangent <$> apply f a +diffF :: Functor f => (forall s. AD s ForwardDouble -> f (AD s ForwardDouble)) -> Double -> f Double +diffF f = Rank1.diffF (fmap runAD.f.AD) {-# 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 => (forall s. ForwardDouble s -> f (ForwardDouble s)) -> Double -> f (Double, Double) -diffF' f a = unbundle <$> apply f a +diffF' :: Functor f => (forall s. AD s ForwardDouble -> f (AD s ForwardDouble)) -> Double -> f (Double, Double) +diffF' f = Rank1.diffF' (fmap runAD.f.AD) {-# INLINE diffF' #-} -- | A fast, simple, transposed Jacobian computed with forward-mode AD. -jacobianT :: (Traversable f, Functor g) => (forall s. f (ForwardDouble s) -> g (ForwardDouble s)) -> f Double -> f (g Double) -jacobianT f = bind (fmap tangent . f) +jacobianT :: (Traversable f, Functor g) => (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f Double -> f (g Double) +jacobianT f = Rank1.jacobianT (fmap runAD.f.fmap AD) {-# 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) -> (forall s. f (ForwardDouble s) -> g (ForwardDouble s)) -> f Double -> f (g b) -jacobianWithT g f = bindWith g' f where - g' a ga = g a . tangent <$> ga +jacobianWithT :: (Traversable f, Functor g) => (Double -> Double -> b) -> (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f Double -> f (g b) +jacobianWithT g f = Rank1.jacobianWithT g (fmap runAD.f.fmap AD) {-# INLINE jacobianWithT #-} -{-# ANN jacobianWithT "HLint: ignore Eta reduce" #-} -- | 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) => (forall s . f (ForwardDouble s) -> g (ForwardDouble s)) -> f Double -> g (f Double) -jacobian f as = transposeWith (const id) t p where - (p, t) = bind' (fmap tangent . f) as +jacobian :: (Traversable f, Traversable g) => (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f Double -> g (f Double) +jacobian f = Rank1.jacobian (fmap runAD.f.fmap AD) {-# 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) -> (forall s. f (ForwardDouble s) -> g (ForwardDouble s)) -> f Double -> g (f b) -jacobianWith g f as = transposeWith (const id) t p where - (p, t) = bindWith' g' f as - g' a ga = g a . tangent <$> ga +jacobianWith :: (Traversable f, Traversable g) => (Double -> Double -> b) -> (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f Double -> g (f b) +jacobianWith g f = Rank1.jacobianWith g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith #-} -- | Compute the Jacobian using 'Forward' mode 'AD' along with the actual answer. -jacobian' :: (Traversable f, Traversable g) => (forall s. f (ForwardDouble s) -> g (ForwardDouble s)) -> f Double -> g (Double, f Double) -jacobian' f as = transposeWith row t p where - (p, t) = bind' f as - row x as' = (primal x, tangent <$> as') +jacobian' :: (Traversable f, Traversable g) => (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f Double -> g (Double, f Double) +jacobian' f = Rank1.jacobian' (fmap runAD.f.fmap AD) {-# 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) -> (forall s. f (ForwardDouble s) -> g (ForwardDouble s)) -> f Double -> g (Double, f b) -jacobianWith' g f as = transposeWith row t p where - (p, t) = bindWith' g' f as - row x as' = (primal x, as') - g' a ga = g a . tangent <$> ga +jacobianWith' :: (Traversable f, Traversable g) => (Double -> Double -> b) -> (forall s. f (AD s ForwardDouble) -> g (AD s ForwardDouble)) -> f Double -> g (Double, f b) +jacobianWith' g f = Rank1.jacobianWith' g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith' #-} -- | Compute the gradient of a function using forward mode AD. -- -- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Wengert.grad' for @n@ inputs, in exchange for better space utilization. -grad :: (Traversable f) => (forall s. f (ForwardDouble s) -> ForwardDouble s) -> f Double -> f Double -grad f = bind (tangent . f) +grad :: Traversable f => (forall s. f (AD s ForwardDouble) -> AD s ForwardDouble) -> f Double -> f Double +grad f = Rank1.grad (runAD.f.fmap AD) {-# INLINE grad #-} -- | Compute the gradient and answer to a function using forward mode AD. -- -- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Wengert.grad'' for @n@ inputs, in exchange for better space utilization. -grad' :: (Traversable f) => (forall s. f (ForwardDouble s) -> ForwardDouble s) -> f Double -> (Double, f Double) -grad' f as = (primal b, tangent <$> bs) - where - (b, bs) = bind' f as +grad' :: Traversable f => (forall s. f (AD s ForwardDouble) -> AD s ForwardDouble) -> f Double -> (Double, f Double) +grad' f = Rank1.grad' (runAD.f.fmap AD) {-# 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.Wengert.gradWith' for @n@ inputs, in exchange for better space utilization. -gradWith :: (Traversable f) => (Double -> Double -> b) -> (forall s. f (ForwardDouble s) -> ForwardDouble s) -> f Double -> f b -gradWith g f = bindWith g (tangent . f) +gradWith :: Traversable f => (Double -> Double -> b) -> (forall s. f (AD s ForwardDouble) -> AD s ForwardDouble) -> f Double -> f b +gradWith g f = Rank1.gradWith g (runAD.f.fmap AD) {-# INLINE gradWith #-} -- | Compute the gradient of a function using forward mode AD and the answer, and combine the result with the input using a @@ -165,6 +154,6 @@ -- -- >>> 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) -> (forall s. f (ForwardDouble s) -> ForwardDouble s) -> f Double -> (Double, f b) -gradWith' g f as = (primal $ f (auto <$> as), bindWith g (tangent . f) as) +gradWith' :: Traversable f => (Double -> Double -> b) -> (forall s. f (AD s ForwardDouble) -> AD s ForwardDouble) -> f Double -> (Double, f b) +gradWith' g f = Rank1.gradWith' g (runAD.f.fmap AD) {-# INLINE gradWith' #-} diff --git a/src/Numeric/AD/Mode/Kahn.hs b/src/Numeric/AD/Mode/Kahn.hs --- a/src/Numeric/AD/Mode/Kahn.hs +++ b/src/Numeric/AD/Mode/Kahn.hs @@ -24,7 +24,7 @@ ----------------------------------------------------------------------------- module Numeric.AD.Mode.Kahn - ( Kahn + ( AD, Kahn, auto -- * Gradient , grad , grad' @@ -43,35 +43,30 @@ , diff' , diffF , diffF' - -- * Unsafe Variadic Gradient - -- $vgrad - , vgrad, vgrad' - , Grad ) where -import Control.Applicative ((<$>)) -import Data.Functor.Compose import Data.Traversable (Traversable) +import Numeric.AD.Internal.Kahn (Kahn) import Numeric.AD.Internal.On -import Numeric.AD.Internal.Kahn +import Numeric.AD.Internal.Type (AD(..)) +import Numeric.AD.Mode +import qualified Numeric.AD.Rank1.Kahn as Rank1 + -- | The 'grad' function calculates the gradient of a non-scalar-to-scalar function with kahn-mode AD in a single pass. -- -- >>> grad (\[x,y,z] -> x*y+z) [1,2,3] -- [2,1,1] -grad :: (Traversable f, Num a) => (forall s. f (Kahn a s) -> Kahn a s) -> f a -> f a -grad f as = unbind vs (partialArray bds $ f vs) where - (vs,bds) = bind as +grad :: (Traversable f, Num a) => (forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> f a +grad f = Rank1.grad (runAD.f.fmap AD) {-# INLINE grad #-} -- | The 'grad'' function calculates the result and gradient of a non-scalar-to-scalar function with kahn-mode AD in a single pass. -- -- >>> grad' (\[x,y,z] -> 4*x*exp y+cos z) [1,2,3] -- (28.566231899122155,[29.5562243957226,29.5562243957226,-0.1411200080598672]) -grad' :: (Traversable f, Num a) => (forall s. f (Kahn a s) -> Kahn a s) -> f a -> (a, f a) -grad' f as = (primal r, unbind vs $ partialArray bds r) where - (vs, bds) = bind as - r = f vs +grad' :: (Traversable f, Num a) => (forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> (a, f a) +grad' f = Rank1.grad' (runAD.f.fmap AD) {-# INLINE grad' #-} -- | @'grad' g f@ function calculates the gradient of a non-scalar-to-scalar function @f@ with kahn-mode AD in a single pass. @@ -83,19 +78,16 @@ -- @ -- -- -gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (Kahn a s) -> Kahn a s) -> f a -> f b -gradWith g f as = unbindWith g vs (partialArray bds $ f vs) where - (vs,bds) = bind as +gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> f b +gradWith g f = Rank1.gradWith g (runAD.f.fmap AD) {-# INLINE gradWith #-} -- | @'grad'' g f@ calculates the result and gradient of a non-scalar-to-scalar function @f@ with kahn-mode AD in a single pass -- the gradient is combined element-wise with the argument using the function @g@. -- -- @'grad'' == 'gradWith'' (\_ dx -> dx)@ -gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (Kahn a s) -> Kahn a s) -> f a -> (a, f b) -gradWith' g f as = (primal r, unbindWith g vs $ partialArray bds r) where - (vs, bds) = bind as - r = f vs +gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> (a, f b) +gradWith' g f = Rank1.gradWith' g (runAD.f.fmap AD) {-# INLINE gradWith' #-} -- | The 'jacobian' function calculates the jacobian of a non-scalar-to-non-scalar function with kahn AD lazily in @m@ passes for @m@ outputs. @@ -105,9 +97,8 @@ -- -- >>> jacobian (\[x,y] -> [exp y,cos x,x+y]) [1,2] -- [[0.0,7.38905609893065],[-0.8414709848078965,0.0],[1.0,1.0]] -jacobian :: (Traversable f, Functor g, Num a) => (forall s. f (Kahn a s) -> g (Kahn a s)) -> f a -> g (f a) -jacobian f as = unbind vs . partialArray bds <$> f vs where - (vs, bds) = bind as +jacobian :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (f a) +jacobian f = Rank1.jacobian (fmap runAD.f.fmap AD) {-# INLINE jacobian #-} -- | The 'jacobian'' function calculates both the result and the Jacobian of a nonscalar-to-nonscalar function, using @m@ invocations of kahn AD, @@ -116,10 +107,8 @@ -- -- ghci> jacobian' (\[x,y] -> [y,x,x*y]) [2,1] -- [(1,[0,1]),(2,[1,0]),(2,[1,2])] -jacobian' :: (Traversable f, Functor g, Num a) => (forall s. f (Kahn a s) -> g (Kahn a s)) -> f a -> g (a, f a) -jacobian' f as = row <$> f vs where - (vs, bds) = bind as - row a = (primal a, unbind vs (partialArray bds a)) +jacobian' :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (a, f a) +jacobian' f = Rank1.jacobian' (fmap runAD.f.fmap AD) {-# INLINE jacobian' #-} -- | 'jacobianWith g f' calculates the Jacobian of a non-scalar-to-non-scalar function @f@ with kahn AD lazily in @m@ passes for @m@ outputs. @@ -130,9 +119,8 @@ -- 'jacobian' = 'jacobianWith' (\_ dx -> dx) -- 'jacobianWith' 'const' = (\f x -> 'const' x '<$>' f x) -- @ -jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (Kahn a s) -> g (Kahn a s)) -> f a -> g (f b) -jacobianWith g f as = unbindWith g vs . partialArray bds <$> f vs where - (vs, bds) = bind as +jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (f b) +jacobianWith g f = Rank1.jacobianWith g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith #-} -- | 'jacobianWith' g f' calculates both the result and the Jacobian of a nonscalar-to-nonscalar function @f@, using @m@ invocations of kahn AD, @@ -141,10 +129,8 @@ -- Instead of returning the Jacobian matrix, the elements of the matrix are combined with the input using the @g@. -- -- @'jacobian'' == 'jacobianWith'' (\_ dx -> dx)@ -jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (Kahn a s) -> g (Kahn a s)) -> f a -> g (a, f b) -jacobianWith' g f as = row <$> f vs where - (vs, bds) = bind as - row a = (primal a, unbindWith g vs (partialArray bds a)) +jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (a, f b) +jacobianWith' g f = Rank1.jacobianWith' g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith' #-} -- | Compute the derivative of a function. @@ -154,8 +140,8 @@ -- -- >>> cos 0 -- 1.0 -diff :: Num a => (forall s. Kahn a s -> Kahn a s) -> a -> a -diff f a = derivative $ f (var a 0) +diff :: Num a => (forall s. AD s (Kahn a) -> AD s (Kahn a)) -> a -> a +diff f = Rank1.diff (runAD.f.AD) {-# INLINE diff #-} -- | The 'diff'' function calculates the value and derivative, as a @@ -164,16 +150,16 @@ -- -- >>> diff' sin 0 -- (0.0,1.0) -diff' :: Num a => (forall s. Kahn a s -> Kahn a s) -> a -> (a, a) -diff' f a = derivative' $ f (var a 0) +diff' :: Num a => (forall s. AD s (Kahn a) -> AD s (Kahn a)) -> a -> (a, a) +diff' f = Rank1.diff' (runAD.f.AD) {-# INLINE diff' #-} -- | Compute the derivatives of a function that returns a vector with regards to its single input. -- -- >>> diffF (\a -> [sin a, cos a]) 0 -- [1.0,0.0] -diffF :: (Functor f, Num a) => (forall s. Kahn a s -> f (Kahn a s)) -> a -> f a -diffF f a = derivative <$> f (var a 0) +diffF :: (Functor f, Num a) => (forall s. AD s (Kahn a) -> f (AD s (Kahn a))) -> a -> f a +diffF f = Rank1.diffF (fmap runAD.f.AD) {-# INLINE diffF #-} -- | Compute the derivatives of a function that returns a vector with regards to its single input @@ -181,19 +167,18 @@ -- -- >>> diffF' (\a -> [sin a, cos a]) 0 -- [(0.0,1.0),(1.0,0.0)] -diffF' :: (Functor f, Num a) => (forall s. Kahn a s -> f (Kahn a s)) -> a -> f (a, a) -diffF' f a = derivative' <$> f (var a 0) +diffF' :: (Functor f, Num a) => (forall s. AD s (Kahn a) -> f (AD s (Kahn a))) -> a -> f (a, a) +diffF' f = Rank1.diffF' (fmap runAD.f.AD) {-# INLINE diffF' #-} - -- | Compute the 'hessian' via the 'jacobian' of the gradient. gradient is computed in 'Kahn' mode and then the 'jacobian' is computed in 'Kahn' mode. -- -- However, since the @'grad' f :: f a -> f a@ is square this is not as fast as using the forward-mode 'jacobian' of a reverse mode gradient provided by 'Numeric.AD.hessian'. -- -- >>> hessian (\[x,y] -> x*y) [1,2] -- [[0,1],[1,0]] -hessian :: (Traversable f, Num a) => (forall s s'. f (On (Kahn (Kahn a s') s)) -> (On (Kahn (Kahn a s') s))) -> f a -> f (f a) -hessian f = jacobian (grad (off . f . fmap On)) +hessian :: (Traversable f, Num a) => (forall s. f (AD s (On (Kahn (Kahn a)))) -> AD s (On (Kahn (Kahn a)))) -> f a -> f (f a) +hessian f = Rank1.hessian (runAD.f.fmap AD) -- | Compute the order 3 Hessian tensor on a non-scalar-to-non-scalar function via the 'Kahn'-mode Jacobian of the 'Kahn'-mode Jacobian of the function. -- @@ -201,15 +186,5 @@ -- -- >>> hessianF (\[x,y] -> [x*y,x+y,exp x*cos y]) [1,2] -- [[[0.0,1.0],[1.0,0.0]],[[0.0,0.0],[0.0,0.0]],[[-1.1312043837568135,-2.4717266720048188],[-2.4717266720048188,1.1312043837568135]]] -hessianF :: (Traversable f, Functor g, Num a) => (forall s s'. f (On (Kahn (Kahn a s') s)) -> g (On (Kahn (Kahn a s') s))) -> f a -> g (f (f a)) -hessianF f = getCompose . jacobian (Compose . jacobian (fmap off . f . fmap On)) - - --- $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! +hessianF :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (On (Kahn (Kahn a)))) -> g (AD s (On (Kahn (Kahn a))))) -> f a -> g (f (f a)) +hessianF f = Rank1.hessianF (fmap runAD.f.fmap AD) diff --git a/src/Numeric/AD/Mode/Reverse.hs b/src/Numeric/AD/Mode/Reverse.hs --- a/src/Numeric/AD/Mode/Reverse.hs +++ b/src/Numeric/AD/Mode/Reverse.hs @@ -20,7 +20,7 @@ ----------------------------------------------------------------------------- module Numeric.AD.Mode.Reverse - ( Reverse + ( Reverse, auto -- * Gradient , grad , grad' @@ -50,13 +50,15 @@ import Data.Traversable (Traversable) import Numeric.AD.Internal.On import Numeric.AD.Internal.Reverse +import Numeric.AD.Mode + -- | The 'grad' function calculates the gradient of a non-scalar-to-scalar function with reverse-mode AD in a single pass. -- -- -- >>> grad (\[x,y,z] -> x*y+z) [1,2,3] -- [2,1,1] -grad :: (Traversable f, Num a) => (forall s. Reifies s Tape => f (Reverse a s) -> Reverse a s) -> f a -> f a +grad :: (Traversable f, Num a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> f a grad f as = reifyTape (snd bds) $ \p -> unbind vs $! partialArrayOf p bds $! f vs where (vs, bds) = bind as {-# INLINE grad #-} @@ -65,7 +67,7 @@ -- -- >>> grad' (\[x,y,z] -> x*y+z) [1,2,3] -- (5,[2,1,1]) -grad' :: (Traversable f, Num a) => (forall s. Reifies s Tape => f (Reverse a s) -> Reverse a s) -> f a -> (a, f a) +grad' :: (Traversable f, Num a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> (a, f a) grad' f as = reifyTape (snd bds) $ \p -> case f vs of r -> (primal r, unbind vs $! partialArrayOf p bds $! r) where (vs, bds) = bind as @@ -78,7 +80,7 @@ -- 'grad' == 'gradWith' (\_ dx -> dx) -- 'id' == 'gradWith' 'const' -- @ -gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse a s) -> Reverse a s) -> f a -> f b +gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> f b gradWith g f as = reifyTape (snd bds) $ \p -> unbindWith g vs $! partialArrayOf p bds $! f vs where (vs,bds) = bind as {-# INLINE gradWith #-} @@ -89,7 +91,7 @@ -- @ -- 'grad'' == 'gradWith'' (\_ dx -> dx) -- @ -gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse a s) -> Reverse a s) -> f a -> (a, f b) +gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> (a, f b) gradWith' g f as = reifyTape (snd bds) $ \p -> case f vs of r -> (primal r, unbindWith g vs $! partialArrayOf p bds $! r) where (vs, bds) = bind as @@ -99,7 +101,7 @@ -- -- >>> jacobian (\[x,y] -> [y,x,x*y]) [2,1] -- [[0,1],[1,0],[1,2]] -jacobian :: (Traversable f, Functor g, Num a) => (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (f a) +jacobian :: (Traversable f, Functor g, Num a) => (forall s. Reifies s Tape => f (Reverse s a) -> g (Reverse s a)) -> f a -> g (f a) jacobian f as = reifyTape (snd bds) $ \p -> unbind vs . partialArrayOf p bds <$> f vs where (vs, bds) = bind as {-# INLINE jacobian #-} @@ -110,7 +112,7 @@ -- -- >>> jacobian' (\[x,y] -> [y,x,x*y]) [2,1] -- [(1,[0,1]),(2,[1,0]),(2,[1,2])] -jacobian' :: (Traversable f, Functor g, Num a) => (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (a, f a) +jacobian' :: (Traversable f, Functor g, Num a) => (forall s. Reifies s Tape => f (Reverse s a) -> g (Reverse s a)) -> f a -> g (a, f a) jacobian' f as = reifyTape (snd bds) $ \p -> let row a = (primal a, unbind vs $! partialArrayOf p bds $! a) in row <$> f vs @@ -125,7 +127,7 @@ -- 'jacobian' == 'jacobianWith' (\_ dx -> dx) -- 'jacobianWith' 'const' == (\f x -> 'const' x '<$>' f x) -- @ -jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (f b) +jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse s a) -> g (Reverse s a)) -> f a -> g (f b) jacobianWith g f as = reifyTape (snd bds) $ \p -> unbindWith g vs . partialArrayOf p bds <$> f vs where (vs, bds) = bind as {-# INLINE jacobianWith #-} @@ -137,7 +139,7 @@ -- -- @'jacobian'' == 'jacobianWith'' (\_ dx -> dx)@ -- -jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse a s) -> g (Reverse a s)) -> f a -> g (a, f b) +jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. Reifies s Tape => f (Reverse s a) -> g (Reverse s a)) -> f a -> g (a, f b) jacobianWith' g f as = reifyTape (snd bds) $ \p -> let row a = (primal a, unbindWith g vs $! partialArrayOf p bds $! a) in row <$> f vs @@ -148,7 +150,7 @@ -- -- >>> diff sin 0 -- 1.0 -diff :: Num a => (forall s. Reifies s Tape => Reverse a s -> Reverse a s) -> a -> a +diff :: Num a => (forall s. Reifies s Tape => Reverse s a -> Reverse s a) -> a -> a diff f a = reifyTape 1 $ \p -> derivativeOf p $! f (var a 0) {-# INLINE diff #-} @@ -159,7 +161,7 @@ -- -- >>> diff' exp 0 -- (1.0,1.0) -diff' :: Num a => (forall s. Reifies s Tape => Reverse a s -> Reverse a s) -> a -> (a, a) +diff' :: Num a => (forall s. Reifies s Tape => Reverse s a -> Reverse s a) -> a -> (a, a) diff' f a = reifyTape 1 $ \p -> derivativeOf' p $! f (var a 0) {-# INLINE diff' #-} @@ -168,7 +170,7 @@ -- >>> diffF (\a -> [sin a, cos a]) 0 -- [1.0,0.0] -- -diffF :: (Functor f, Num a) => (forall s. Reifies s Tape => Reverse a s -> f (Reverse a s)) -> a -> f a +diffF :: (Functor f, Num a) => (forall s. Reifies s Tape => Reverse s a -> f (Reverse s a)) -> a -> f a diffF f a = reifyTape 1 $ \p -> derivativeOf p <$> f (var a 0) {-# INLINE diffF #-} @@ -176,7 +178,7 @@ -- -- >>> diffF' (\a -> [sin a, cos a]) 0 -- [(0.0,1.0),(1.0,0.0)] -diffF' :: (Functor f, Num a) => (forall s. Reifies s Tape => Reverse a s -> f (Reverse a s)) -> a -> f (a, a) +diffF' :: (Functor f, Num a) => (forall s. Reifies s Tape => Reverse s a -> f (Reverse s a)) -> a -> f (a, a) diffF' f a = reifyTape 1 $ \p -> derivativeOf' p <$> f (var a 0) {-# INLINE diffF' #-} @@ -186,7 +188,7 @@ -- -- >>> hessian (\[x,y] -> x*y) [1,2] -- [[0,1],[1,0]] -hessian :: (Traversable f, Num a) => (forall s s'. (Reifies s Tape, Reifies s' Tape) => f (On (Reverse (Reverse a s') s)) -> (On (Reverse (Reverse a s') s))) -> f a -> f (f a) +hessian :: (Traversable f, Num a) => (forall s s'. (Reifies s Tape, Reifies s' Tape) => f (On (Reverse s (Reverse s' a))) -> (On (Reverse s (Reverse s' a)))) -> f a -> f (f a) hessian f = jacobian (grad (off . f . fmap On)) {-# INLINE hessian #-} @@ -196,6 +198,6 @@ -- -- >>> hessianF (\[x,y] -> [x*y,x+y,exp x*cos y]) [1,2] -- [[[0.0,1.0],[1.0,0.0]],[[0.0,0.0],[0.0,0.0]],[[-1.1312043837568135,-2.4717266720048188],[-2.4717266720048188,1.1312043837568135]]] -hessianF :: (Traversable f, Functor g, Num a) => (forall s s'. (Reifies s Tape, Reifies s' Tape) => f (On (Reverse (Reverse a s') s)) -> g (On (Reverse (Reverse a s') s))) -> f a -> g (f (f a)) +hessianF :: (Traversable f, Functor g, Num a) => (forall s s'. (Reifies s Tape, Reifies s' Tape) => f (On (Reverse s (Reverse s' a))) -> g (On (Reverse s (Reverse s' a)))) -> f a -> g (f (f a)) hessianF f = getCompose . jacobian (Compose . jacobian (fmap off . f . fmap On)) {-# INLINE hessianF #-} diff --git a/src/Numeric/AD/Mode/Sparse.hs b/src/Numeric/AD/Mode/Sparse.hs --- a/src/Numeric/AD/Mode/Sparse.hs +++ b/src/Numeric/AD/Mode/Sparse.hs @@ -12,21 +12,13 @@ ----------------------------------------------------------------------------- module Numeric.AD.Mode.Sparse - ( Sparse + ( AD, Sparse, auto -- * Sparse Gradients , grad , grad' + , grads , gradWith , gradWith' - -- * Variadic Gradients - -- $vgrad - , Grad - , vgrad - -- * Higher-Order Gradients - , grads - -- * Variadic Higher-Order Gradients - , Grads - , vgrads -- * Sparse Jacobians (synonyms) , jacobian @@ -41,89 +33,68 @@ , hessianF , hessianF' - ) where -import Control.Comonad -import Data.Traversable -import Control.Comonad.Cofree -import Numeric.AD.Jet -import Numeric.AD.Internal.Sparse -import Numeric.AD.Internal.Combinators +import Control.Comonad.Cofree (Cofree) +import Data.Traversable (Traversable) +import Numeric.AD.Internal.Sparse (Sparse) +import qualified Numeric.AD.Rank1.Sparse as Rank1 +import Numeric.AD.Internal.Type +import Numeric.AD.Mode -second :: (a -> b) -> (c, a) -> (c, b) -second g (a,b) = (a, g b) -{-# INLINE second #-} -grad :: (Traversable f, Num a) => (forall s. f (Sparse a s) -> Sparse a s) -> f a -> f a -grad f as = d as $ apply f as +grad :: (Traversable f, Num a) => (forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> f a +grad f = Rank1.grad (runAD.f.fmap AD) {-# INLINE grad #-} -grad' :: (Traversable f, Num a) => (forall s. f (Sparse a s) -> Sparse a s) -> f a -> (a, f a) -grad' f as = d' as $ apply f as +grad' :: (Traversable f, Num a) => (forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> (a, f a) +grad' f = Rank1.grad' (runAD.f.fmap AD) {-# INLINE grad' #-} -gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (Sparse a s) -> Sparse a s) -> f a -> f b -gradWith g f as = zipWithT g as $ grad f as +gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> f b +gradWith g f = Rank1.gradWith g (runAD.f.fmap AD) {-# INLINE gradWith #-} -gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (Sparse a s) -> Sparse a s) -> f a -> (a, f b) -gradWith' g f as = second (zipWithT g as) $ grad' f as +gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> (a, f b) +gradWith' g f = Rank1.gradWith' g (runAD.f.fmap AD) {-# INLINE gradWith' #-} -jacobian :: (Traversable f, Functor g, Num a) => (forall s. f (Sparse a s) -> g (Sparse a s)) -> f a -> g (f a) -jacobian f as = d as <$> apply f as +jacobian :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Sparse a)) -> g (AD s (Sparse a))) -> f a -> g (f a) +jacobian f = Rank1.jacobian (fmap runAD.f.fmap AD) {-# INLINE jacobian #-} -jacobian' :: (Traversable f, Functor g, Num a) => (forall s. f (Sparse a s) -> g (Sparse a s)) -> f a -> g (a, f a) -jacobian' f as = d' as <$> apply f as +jacobian' :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Sparse a)) -> g (AD s (Sparse a))) -> f a -> g (a, f a) +jacobian' f = Rank1.jacobian' (fmap runAD.f.fmap AD) {-# INLINE jacobian' #-} -jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (Sparse a s) -> g (Sparse a s)) -> f a -> g (f b) -jacobianWith g f as = zipWithT g as <$> jacobian f as +jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Sparse a)) -> g (AD s (Sparse a))) -> f a -> g (f b) +jacobianWith g f = Rank1.jacobianWith g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith #-} -jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (Sparse a s) -> g (Sparse a s)) -> f a -> g (a, f b) -jacobianWith' g f as = second (zipWithT g as) <$> jacobian' f as +jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Sparse a)) -> g (AD s (Sparse a))) -> f a -> g (a, f b) +jacobianWith' g f = Rank1.jacobianWith' g (fmap runAD.f.fmap AD) {-# INLINE jacobianWith' #-} -grads :: (Traversable f, Num a) => (forall s. f (Sparse a s) -> Sparse a s) -> f a -> Cofree f a -grads f as = ds as $ apply f as +grads :: (Traversable f, Num a) => (forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> Cofree f a +grads f = Rank1.grads (runAD.f.fmap AD) {-# INLINE grads #-} -jacobians :: (Traversable f, Functor g, Num a) => (forall s. f (Sparse a s) -> g (Sparse a s)) -> f a -> g (Cofree f a) -jacobians f as = ds as <$> apply f as +jacobians :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Sparse a)) -> g (AD s (Sparse a))) -> f a -> g (Cofree f a) +jacobians f = Rank1.jacobians (fmap runAD.f.fmap AD) {-# INLINE jacobians #-} -d2 :: Functor f => Cofree f a -> f (f a) -d2 = headJet . tailJet . tailJet . jet -{-# INLINE d2 #-} - -d2' :: Functor f => Cofree f a -> (a, f (a, f a)) -d2' (a :< as) = (a, fmap (\(da :< das) -> (da, extract <$> das)) as) -{-# INLINE d2' #-} - -hessian :: (Traversable f, Num a) => (forall s. f (Sparse a s) -> Sparse a s) -> f a -> f (f a) -hessian f as = d2 $ grads f as +hessian :: (Traversable f, Num a) => (forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> f (f a) +hessian f = Rank1.hessian (runAD.f.fmap AD) {-# INLINE hessian #-} -hessian' :: (Traversable f, Num a) => (forall s. f (Sparse a s) -> Sparse a s) -> f a -> (a, f (a, f a)) -hessian' f as = d2' $ grads f as +hessian' :: (Traversable f, Num a) => (forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> (a, f (a, f a)) +hessian' f = Rank1.hessian' (runAD.f.fmap AD) {-# INLINE hessian' #-} -hessianF :: (Traversable f, Functor g, Num a) => (forall s. f (Sparse a s) -> g (Sparse a s)) -> f a -> g (f (f a)) -hessianF f as = d2 <$> jacobians f as +hessianF :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Sparse a)) -> g (AD s (Sparse a))) -> f a -> g (f (f a)) +hessianF f = Rank1.hessianF (fmap runAD.f.fmap AD) {-# INLINE hessianF #-} -hessianF' :: (Traversable f, Functor g, Num a) => (forall s. f (Sparse a s) -> g (Sparse a s)) -> f a -> g (a, f (a, f a)) -hessianF' f as = d2' <$> jacobians f as +hessianF' :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Sparse a)) -> g (AD s (Sparse a))) -> f a -> g (a, f (a, f a)) +hessianF' f = Rank1.hessianF' (fmap runAD.f.fmap AD) {-# 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! diff --git a/src/Numeric/AD/Mode/Tower.hs b/src/Numeric/AD/Mode/Tower.hs --- a/src/Numeric/AD/Mode/Tower.hs +++ b/src/Numeric/AD/Mode/Tower.hs @@ -1,5 +1,4 @@ {-# LANGUAGE Rank2Types #-} -{-# LANGUAGE BangPatterns #-} ----------------------------------------------------------------------------- -- | -- Copyright : (c) Edward Kmett 2010-2014 @@ -11,9 +10,10 @@ -- Higher order derivatives via a \"dual number tower\". -- ----------------------------------------------------------------------------- - module Numeric.AD.Mode.Tower - ( Tower + ( AD + , Tower + , auto -- * Taylor Series , taylor , taylor0 @@ -38,78 +38,79 @@ , dus0F -- answer and all zero padded directional derivatives of (a -> a) ) where -import Control.Applicative ((<$>)) -import Numeric.AD.Internal.Tower +import qualified Numeric.AD.Rank1.Tower as Rank1 +import Numeric.AD.Internal.Tower (Tower) +import Numeric.AD.Internal.Type (AD(..)) +import Numeric.AD.Mode -diffs :: Num a => (forall s. Tower a s -> Tower a s) -> a -> [a] -diffs f a = getADTower $ apply f a + +diffs :: Num a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> [a] +diffs f = Rank1.diffs (runAD.f.AD) {-# INLINE diffs #-} -diffs0 :: Num a => (forall s. Tower a s -> Tower a s) -> a -> [a] -diffs0 f a = zeroPad (diffs f a) +diffs0 :: Num a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> [a] +diffs0 f = Rank1.diffs0 (runAD.f.AD) {-# INLINE diffs0 #-} -diffsF :: (Functor f, Num a) => (forall s. Tower a s -> f (Tower a s)) -> a -> f [a] -diffsF f a = getADTower <$> apply f a +diffsF :: (Functor f, Num a) => (forall s. AD s (Tower a) -> f (AD s (Tower a))) -> a -> f [a] +diffsF f = Rank1.diffsF (fmap runAD.f.AD) {-# INLINE diffsF #-} -diffs0F :: (Functor f, Num a) => (forall s. Tower a s -> f (Tower a s)) -> a -> f [a] -diffs0F f a = (zeroPad . getADTower) <$> apply f a +diffs0F :: (Functor f, Num a) => (forall s. AD s (Tower a) -> f (AD s (Tower a))) -> a -> f [a] +diffs0F f = Rank1.diffs0F (fmap runAD.f.AD) {-# INLINE diffs0F #-} -taylor :: Fractional a => (forall s. Tower a s -> Tower a s) -> a -> a -> [a] -taylor f x dx = go 1 1 (diffs f x) where - go !n !acc (a:as) = a * acc : go (n + 1) (acc * dx / n) as - go _ _ [] = [] +taylor :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a -> [a] +taylor f = Rank1.taylor (runAD.f.AD) -taylor0 :: Fractional a => (forall s. Tower a s -> Tower a s) -> a -> a -> [a] -taylor0 f x dx = zeroPad (taylor f x dx) +taylor0 :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a -> [a] +taylor0 f = Rank1.taylor0 (runAD.f.AD) {-# INLINE taylor0 #-} -maclaurin :: Fractional a => (forall s. Tower a s -> Tower a s) -> a -> [a] -maclaurin f = taylor f 0 +maclaurin :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> [a] +maclaurin f = Rank1.maclaurin (runAD.f.AD) {-# INLINE maclaurin #-} -maclaurin0 :: Fractional a => (forall s. Tower a s -> Tower a s) -> a -> [a] -maclaurin0 f = taylor0 f 0 +maclaurin0 :: Fractional a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> [a] +maclaurin0 f = Rank1.maclaurin0 (runAD.f.AD) {-# INLINE maclaurin0 #-} -diff :: Num a => (forall s. Tower a s -> Tower a s) -> a -> a -diff f = d . diffs f +diff :: Num a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> a +diff f = Rank1.diff (runAD.f.AD) {-# INLINE diff #-} -diff' :: Num a => (forall s. Tower a s -> Tower a s) -> a -> (a, a) -diff' f = d' . diffs f +diff' :: Num a => (forall s. AD s (Tower a) -> AD s (Tower a)) -> a -> (a, a) +diff' f = Rank1.diff' (runAD.f.AD) {-# INLINE diff' #-} -du :: (Functor f, Num a) => (forall s. f (Tower a s) -> Tower a s) -> f (a, a) -> a -du f = d . getADTower . f . fmap withD +du :: (Functor f, Num a) => (forall s. f (AD s (Tower a)) -> AD s (Tower a)) -> f (a, a) -> a +du f = Rank1.du (runAD.f. fmap AD) {-# INLINE du #-} -du' :: (Functor f, Num a) => (forall s. f (Tower a s) -> Tower a s) -> f (a, a) -> (a, a) -du' f = d' . getADTower . f . fmap withD +du' :: (Functor f, Num a) => (forall s. f (AD s (Tower a)) -> AD s (Tower a)) -> f (a, a) -> (a, a) +du' f = Rank1.du' (runAD.f.fmap AD) {-# INLINE du' #-} -duF :: (Functor f, Functor g, Num a) => (forall s. f (Tower a s) -> g (Tower a s)) -> f (a, a) -> g a -duF f = fmap (d . getADTower) . f . fmap withD +duF :: (Functor f, Functor g, Num a) => (forall s. f (AD s (Tower a)) -> g (AD s (Tower a))) -> f (a, a) -> g a +duF f = Rank1.duF (fmap runAD.f.fmap AD) {-# INLINE duF #-} -duF' :: (Functor f, Functor g, Num a) => (forall s. f (Tower a s) -> g (Tower a s)) -> f (a, a) -> g (a, a) -duF' f = fmap (d' . getADTower) . f . fmap withD +duF' :: (Functor f, Functor g, Num a) => (forall s. f (AD s (Tower a)) -> g (AD s (Tower a))) -> f (a, a) -> g (a, a) +duF' f = Rank1.duF' (fmap runAD.f.fmap AD) {-# INLINE duF' #-} -dus :: (Functor f, Num a) => (forall s. f (Tower a s) -> Tower a s) -> f [a] -> [a] -dus f = getADTower . f . fmap tower +dus :: (Functor f, Num a) => (forall s. f (AD s (Tower a)) -> AD s (Tower a)) -> f [a] -> [a] +dus f = Rank1.dus (runAD.f.fmap AD) {-# INLINE dus #-} -dus0 :: (Functor f, Num a) => (forall s. f (Tower a s) -> Tower a s) -> f [a] -> [a] -dus0 f = zeroPad . getADTower . f . fmap tower +dus0 :: (Functor f, Num a) => (forall s. f (AD s (Tower a)) -> AD s (Tower a)) -> f [a] -> [a] +dus0 f = Rank1.dus0 (runAD.f.fmap AD) {-# INLINE dus0 #-} -dusF :: (Functor f, Functor g, Num a) => (forall s. f (Tower a s) -> g (Tower a s)) -> f [a] -> g [a] -dusF f = fmap getADTower . f . fmap tower +dusF :: (Functor f, Functor g, Num a) => (forall s. f (AD s (Tower a)) -> g (AD s (Tower a))) -> f [a] -> g [a] +dusF f = Rank1.dusF (fmap runAD.f.fmap AD) {-# INLINE dusF #-} -dus0F :: (Functor f, Functor g, Num a) => (forall s. f (Tower a s) -> g (Tower a s)) -> f [a] -> g [a] -dus0F f = fmap getADTower . f . fmap tower +dus0F :: (Functor f, Functor g, Num a) => (forall s. f (AD s (Tower a)) -> g (AD s (Tower a))) -> f [a] -> g [a] +dus0F f = Rank1.dus0F (fmap runAD.f.fmap AD) {-# INLINE dus0F #-} diff --git a/src/Numeric/AD/Newton.hs b/src/Numeric/AD/Newton.hs --- a/src/Numeric/AD/Newton.hs +++ b/src/Numeric/AD/Newton.hs @@ -5,7 +5,7 @@ {-# LANGUAGE TypeFamilies #-} ----------------------------------------------------------------------------- -- | --- Copyright : (c) Edward Kmett 2010 +-- Copyright : (c) Edward Kmett 2010-2014 -- License : BSD3 -- Maintainer : ekmett@gmail.com -- Stability : experimental @@ -27,19 +27,20 @@ , conjugateGradientAscent ) where -import Prelude hiding (all, mapM, sum) import Data.Foldable (all, sum) import Data.Reflection (Reifies) import Data.Traversable -import Numeric.AD.Mode -import Numeric.AD.Mode.Forward (diff, diff') -import Numeric.AD.Mode.Reverse as Reverse (gradWith') -import Numeric.AD.Mode.Kahn as Kahn (Kahn, grad) import Numeric.AD.Internal.Combinators import Numeric.AD.Internal.Forward (Forward) -import Numeric.AD.Internal.Or import Numeric.AD.Internal.On +import Numeric.AD.Internal.Or import Numeric.AD.Internal.Reverse (Reverse, Tape) +import Numeric.AD.Internal.Type (AD(..)) +import Numeric.AD.Mode +import Numeric.AD.Mode.Reverse as Reverse (gradWith') +import Numeric.AD.Rank1.Kahn as Kahn (Kahn, grad) +import qualified Numeric.AD.Rank1.Newton as Rank1 +import Prelude hiding (all, mapM, sum) -- $setup -- >>> import Data.Complex @@ -56,11 +57,8 @@ -- -- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1) -- 0.0 :+ 1.0 -findZero :: (Fractional a, Eq a) => (forall s. Forward a s -> Forward a s) -> a -> [a] -findZero f = go where - go x = x : if x == xn then [] else go xn where - (y,y') = diff' f x - xn = x - y/y' +findZero :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a] +findZero f = Rank1.findZero (runAD.f.AD) {-# INLINE findZero #-} -- | The 'inverse' function inverts a scalar function using @@ -72,8 +70,8 @@ -- -- >>> last $ take 10 $ inverse sqrt 1 (sqrt 10) -- 10.0 -inverse :: (Fractional a, Eq a) => (forall s. Forward a s -> Forward a s) -> a -> a -> [a] -inverse f x0 y = findZero (\x -> f x - auto y) x0 +inverse :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> a -> [a] +inverse f = Rank1.inverse (runAD.f.AD) {-# INLINE inverse #-} -- | The 'fixedPoint' function find a fixedpoint of a scalar @@ -85,8 +83,8 @@ -- -- >>> last $ take 10 $ fixedPoint cos 1 -- 0.7390851332151607 -fixedPoint :: (Fractional a, Eq a) => (forall s. Forward a s -> Forward a s) -> a -> [a] -fixedPoint f = findZero (\x -> f x - x) +fixedPoint :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a] +fixedPoint f = Rank1.fixedPoint (runAD.f.AD) {-# INLINE fixedPoint #-} -- | The 'extremum' function finds an extremum of a scalar @@ -96,8 +94,8 @@ -- -- >>> last $ take 10 $ extremum cos 1 -- 0.0 -extremum :: (Fractional a, Eq a) => (forall s s'. On (Forward (Forward a s') s) -> On (Forward (Forward a s') s)) -> a -> [a] -extremum f = findZero (diff (off . f . On)) +extremum :: (Fractional a, Eq a) => (forall s. AD s (On (Forward (Forward a))) -> AD s (On (Forward (Forward a)))) -> a -> [a] +extremum f = Rank1.extremum (runAD.f.AD) {-# INLINE extremum #-} -- | The 'gradientDescent' function performs a multivariate @@ -107,7 +105,7 @@ -- increasingly accurate results. (Modulo the usual caveats.) -- -- It uses reverse mode automatic differentiation to compute the gradient. -gradientDescent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse a s) -> Reverse a s) -> f a -> [f a] +gradientDescent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> [f a] gradientDescent f x0 = go x0 fx0 xgx0 0.1 (0 :: Int) where (fx0, xgx0) = Reverse.gradWith' (,) f x0 @@ -125,7 +123,7 @@ {-# INLINE gradientDescent #-} -- | Perform a gradient descent using reverse mode automatic differentiation to compute the gradient. -gradientAscent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse a s) -> Reverse a s) -> f a -> [f a] +gradientAscent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> [f a] gradientAscent f = gradientDescent (negate . f) {-# INLINE gradientAscent #-} @@ -139,21 +137,21 @@ -- True conjugateGradientDescent :: (Traversable f, Ord a, Fractional a) - => (forall s1 s2 s3 s4. Chosen s4 => f (Or (On (Forward (Forward a s1) s2)) (Kahn a s3) s4) -> Or (On (Forward (Forward a s1) s2)) (Kahn a s3) s4) + => (forall s. Chosen s => f (Or s (On (Forward (Forward a))) (Kahn a)) -> Or s (On (Forward (Forward a))) (Kahn a)) -> f a -> [f a] conjugateGradientDescent f = conjugateGradientAscent (negate . f) {-# INLINE conjugateGradientDescent #-} -lfu :: Functor f => (f (Or a b F) -> Or a b F) -> f a -> a +lfu :: Functor f => (f (Or F a b) -> Or F a b) -> f a -> a lfu f = runL . f . fmap L -rfu :: Functor f => (f (Or a b T) -> Or a b T) -> f b -> b +rfu :: Functor f => (f (Or T a b) -> Or T a b) -> f b -> b rfu f = runR . f . fmap R -- | Perform a conjugate gradient ascent using reverse mode automatic differentiation to compute the gradient. conjugateGradientAscent :: (Traversable f, Ord a, Fractional a) - => (forall s1 s2 s3 s4. Chosen s4 => f (Or (On (Forward (Forward a s1) s2)) (Kahn a s3) s4) -> Or (On (Forward (Forward a s1) s2)) (Kahn a s3) s4) + => (forall s. Chosen s => f (Or s (On (Forward (Forward a))) (Kahn a)) -> Or s (On (Forward (Forward a))) (Kahn a)) -> f a -> [f a] conjugateGradientAscent f x0 = takeWhile (all (\a -> a == a)) (go x0 d0 d0 delta0) where @@ -162,7 +160,7 @@ delta0 = dot d0 d0 go xi _ri di deltai = xi : go xi1 ri1 di1 deltai1 where - ai = last $ take 20 $ extremum (\a -> lfu f $ zipWithT (\x d -> auto x + a * auto d) xi di) 0 + ai = last $ take 20 $ Rank1.extremum (\a -> lfu f $ zipWithT (\x d -> auto x + a * auto d) xi di) 0 xi1 = zipWithT (\x d -> x + ai*d) xi di ri1 = Kahn.grad (rfu f) xi1 deltai1 = dot ri1 ri1 diff --git a/src/Numeric/AD/Rank1/Forward.hs b/src/Numeric/AD/Rank1/Forward.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/AD/Rank1/Forward.hs @@ -0,0 +1,201 @@ +{-# LANGUAGE CPP #-} +----------------------------------------------------------------------------- +-- | +-- Copyright : (c) Edward Kmett 2010-2014 +-- 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 Data.Traversable (Traversable) +import Control.Applicative +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 f = tangent . f . fmap (uncurry 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' f = unbundle . f . fmap (uncurry 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 f = fmap tangent . f . fmap (uncurry 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' f = fmap unbundle . f . fmap (uncurry 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 f a = tangent $ apply f 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' f a = unbundle $ apply f 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 f a = tangent <$> apply f 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' f a = unbundle <$> apply f 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 f = bind (fmap tangent . 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 g f = bindWith g' f where + g' a ga = g a . tangent <$> ga +{-# INLINE jacobianWithT #-} +#ifdef HLINT +{-# ANN jacobianWithT "HLint: ignore Eta reduce" #-} +#endif + +-- | 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 f as = transposeWith (const id) t p where + (p, t) = bind' (fmap tangent . f) 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 g f as = transposeWith (const id) t p where + (p, t) = bindWith' g' f as + g' a ga = g a . tangent <$> 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' f as = transposeWith row t p where + (p, t) = bind' f as + row x as' = (primal x, tangent <$> 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' g f as = transposeWith row t p where + (p, t) = bindWith' g' f as + row x as' = (primal x, as') + g' a ga = g a . tangent <$> ga +{-# INLINE jacobianWith' #-} + +-- | Compute the gradient of a function using forward mode AD. +-- +-- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Wengert.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 f = bind (tangent . 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.Wengert.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' f as = (primal b, tangent <$> bs) where + (b, bs) = bind' f 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.Wengert.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 g f = bindWith g (tangent . 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.Wengert.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' g f as = (primal $ f (Lift <$> as), bindWith g (tangent . f) 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 f = duF $ grad $ off . f . fmap 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' f = duF' $ grad $ off . f . fmap On +{-# INLINE hessianProduct' #-} diff --git a/src/Numeric/AD/Rank1/Forward/Double.hs b/src/Numeric/AD/Rank1/Forward/Double.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/AD/Rank1/Forward/Double.hs @@ -0,0 +1,169 @@ +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 Control.Applicative +import Data.Traversable (Traversable) +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 f = tangent . f . fmap (uncurry 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' f = unbundle . f . fmap (uncurry 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 f = fmap tangent . f . fmap (uncurry 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' f = fmap unbundle . f . fmap (uncurry 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 f a = tangent $ apply f 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' f a = unbundle $ apply f 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 f a = tangent <$> apply f 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' f a = unbundle <$> apply f 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 f = bind (fmap tangent . 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 g f = bindWith g' f where + g' a ga = g a . tangent <$> ga +{-# INLINE jacobianWithT #-} +{-# ANN jacobianWithT "HLint: ignore Eta reduce" #-} + +-- | 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 f as = transposeWith (const id) t p where + (p, t) = bind' (fmap tangent . f) 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 g f as = transposeWith (const id) t p where + (p, t) = bindWith' g' f as + g' a ga = g a . tangent <$> 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' f as = transposeWith row t p where + (p, t) = bind' f as + row x as' = (primal x, tangent <$> 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' g f as = transposeWith row t p where + (p, t) = bindWith' g' f as + row x as' = (primal x, as') + g' a ga = g a . tangent <$> ga +{-# INLINE jacobianWith' #-} + +-- | Compute the gradient of a function using forward mode AD. +-- +-- Note, this performs /O(n)/ worse than 'Numeric.AD.Mode.Wengert.grad' for @n@ inputs, in exchange for better space utilization. +grad :: Traversable f => (f ForwardDouble -> ForwardDouble) -> f Double -> f Double +grad f = bind (tangent . 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.Wengert.grad'' for @n@ inputs, in exchange for better space utilization. +grad' :: Traversable f => (f ForwardDouble -> ForwardDouble) -> f Double -> (Double, f Double) +grad' f as = (primal b, tangent <$> bs) + where + (b, bs) = bind' f 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.Wengert.gradWith' for @n@ inputs, in exchange for better space utilization. +gradWith :: Traversable f => (Double -> Double -> b) -> (f ForwardDouble -> ForwardDouble) -> f Double -> f b +gradWith g f = bindWith g (tangent . 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.Wengert.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' g f as = (primal $ f (auto <$> as), bindWith g (tangent . f) as) +{-# INLINE gradWith' #-} diff --git a/src/Numeric/AD/Rank1/Kahn.hs b/src/Numeric/AD/Rank1/Kahn.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/AD/Rank1/Kahn.hs @@ -0,0 +1,217 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UndecidableInstances #-} +----------------------------------------------------------------------------- +-- | +-- Copyright : (c) Edward Kmett 2010-2014 +-- License : BSD3 +-- Maintainer : ekmett@gmail.com +-- Stability : experimental +-- Portability : GHC only +-- +-- This module provides reverse-mode Automatic Differentiation using post-hoc linear time +-- topological sorting. +-- +-- For reverse mode AD we use 'System.Mem.StableName.StableName' to recover sharing information from +-- the tape to avoid combinatorial explosion, and thus run asymptotically faster +-- than it could without such sharing information, but the use of side-effects +-- contained herein is benign. +-- +----------------------------------------------------------------------------- + +module Numeric.AD.Rank1.Kahn + ( Kahn + , auto + -- * Gradient + , grad + , grad' + , gradWith + , gradWith' + -- * Jacobian + , jacobian + , jacobian' + , jacobianWith + , jacobianWith' + -- * Hessian + , hessian + , hessianF + -- * Derivatives + , diff + , diff' + , diffF + , diffF' + -- * Unsafe Variadic Gradient + -- $vgrad + , vgrad, vgrad' + , Grad + ) where + +import Control.Applicative ((<$>)) +import Data.Functor.Compose +import Data.Traversable (Traversable) +import Numeric.AD.Internal.On +import Numeric.AD.Internal.Kahn +import Numeric.AD.Mode + + +-- | The 'grad' function calculates the gradient of a non-scalar-to-scalar function with kahn-mode AD in a single pass. +-- +-- >>> grad (\[x,y,z] -> x*y+z) [1,2,3] +-- [2,1,1] +grad :: (Traversable f, Num a) => (f (Kahn a) -> Kahn a) -> f a -> f a +grad f as = unbind vs (partialArray bds $ f vs) where + (vs,bds) = bind as +{-# INLINE grad #-} + +-- | The 'grad'' function calculates the result and gradient of a non-scalar-to-scalar function with kahn-mode AD in a single pass. +-- +-- >>> grad' (\[x,y,z] -> 4*x*exp y+cos z) [1,2,3] +-- (28.566231899122155,[29.5562243957226,29.5562243957226,-0.1411200080598672]) +grad' :: (Traversable f, Num a) => (f (Kahn a) -> Kahn a) -> f a -> (a, f a) +grad' f as = (primal r, unbind vs $ partialArray bds r) where + (vs, bds) = bind as + r = f vs +{-# INLINE grad' #-} + +-- | @'grad' g f@ function calculates the gradient of a non-scalar-to-scalar function @f@ with kahn-mode AD in a single pass. +-- The gradient is combined element-wise with the argument using the function @g@. +-- +-- @ +-- 'grad' = 'gradWith' (\_ dx -> dx) +-- 'id' = 'gradWith' const +-- @ +-- +-- +gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (f (Kahn a) -> Kahn a) -> f a -> f b +gradWith g f as = unbindWith g vs (partialArray bds $ f vs) where + (vs,bds) = bind as +{-# INLINE gradWith #-} + +-- | @'grad'' g f@ calculates the result and gradient of a non-scalar-to-scalar function @f@ with kahn-mode AD in a single pass +-- the gradient is combined element-wise with the argument using the function @g@. +-- +-- @'grad'' == 'gradWith'' (\_ dx -> dx)@ +gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (f (Kahn a) -> Kahn a) -> f a -> (a, f b) +gradWith' g f as = (primal r, unbindWith g vs $ partialArray bds r) where + (vs, bds) = bind as + r = f vs +{-# INLINE gradWith' #-} + +-- | The 'jacobian' function calculates the jacobian of a non-scalar-to-non-scalar function with kahn AD lazily in @m@ passes for @m@ outputs. +-- +-- >>> jacobian (\[x,y] -> [y,x,x*y]) [2,1] +-- [[0,1],[1,0],[1,2]] +-- +-- >>> jacobian (\[x,y] -> [exp y,cos x,x+y]) [1,2] +-- [[0.0,7.38905609893065],[-0.8414709848078965,0.0],[1.0,1.0]] +jacobian :: (Traversable f, Functor g, Num a) => (f (Kahn a) -> g (Kahn a)) -> f a -> g (f a) +jacobian f as = unbind vs . partialArray bds <$> f vs where + (vs, bds) = bind as +{-# INLINE jacobian #-} + +-- | The 'jacobian'' function calculates both the result and the Jacobian of a nonscalar-to-nonscalar function, using @m@ invocations of kahn AD, +-- where @m@ is the output dimensionality. Applying @fmap snd@ to the result will recover the result of 'jacobian' +-- | An alias for 'gradF'' +-- +-- ghci> jacobian' (\[x,y] -> [y,x,x*y]) [2,1] +-- [(1,[0,1]),(2,[1,0]),(2,[1,2])] +jacobian' :: (Traversable f, Functor g, Num a) => (f (Kahn a) -> g (Kahn a)) -> f a -> g (a, f a) +jacobian' f as = row <$> f vs where + (vs, bds) = bind as + row a = (primal a, unbind vs (partialArray bds a)) +{-# INLINE jacobian' #-} + +-- | 'jacobianWith g f' calculates the Jacobian of a non-scalar-to-non-scalar function @f@ with kahn AD lazily in @m@ passes for @m@ outputs. +-- +-- Instead of returning the Jacobian matrix, the elements of the matrix are combined with the input using the @g@. +-- +-- @ +-- 'jacobian' = 'jacobianWith' (\_ dx -> dx) +-- 'jacobianWith' 'const' = (\f x -> 'const' x '<$>' f x) +-- @ +jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (f (Kahn a) -> g (Kahn a)) -> f a -> g (f b) +jacobianWith g f as = unbindWith g vs . partialArray bds <$> f vs where + (vs, bds) = bind as +{-# INLINE jacobianWith #-} + +-- | 'jacobianWith' g f' calculates both the result and the Jacobian of a nonscalar-to-nonscalar function @f@, using @m@ invocations of kahn AD, +-- where @m@ is the output dimensionality. Applying @fmap snd@ to the result will recover the result of 'jacobianWith' +-- +-- Instead of returning the Jacobian matrix, the elements of the matrix are combined with the input using the @g@. +-- +-- @'jacobian'' == 'jacobianWith'' (\_ dx -> dx)@ +jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (f (Kahn a) -> g (Kahn a)) -> f a -> g (a, f b) +jacobianWith' g f as = row <$> f vs where + (vs, bds) = bind as + row a = (primal a, unbindWith g vs (partialArray bds a)) +{-# INLINE jacobianWith' #-} + +-- | Compute the derivative of a function. +-- +-- >>> diff sin 0 +-- 1.0 +-- +-- >>> cos 0 +-- 1.0 +diff :: Num a => (Kahn a -> Kahn a) -> a -> a +diff f a = derivative $ f (var a 0) +{-# INLINE diff #-} + +-- | The 'diff'' function calculates the value and derivative, as a +-- pair, of a scalar-to-scalar function. +-- +-- +-- >>> diff' sin 0 +-- (0.0,1.0) +diff' :: Num a => (Kahn a -> Kahn a) -> a -> (a, a) +diff' f a = derivative' $ f (var a 0) +{-# INLINE diff' #-} + +-- | Compute the derivatives of a function that returns a vector with regards to its single input. +-- +-- >>> diffF (\a -> [sin a, cos a]) 0 +-- [1.0,0.0] +diffF :: (Functor f, Num a) => (Kahn a -> f (Kahn a)) -> a -> f a +diffF f a = derivative <$> f (var a 0) +{-# INLINE diffF #-} + +-- | Compute the derivatives of a function that returns a vector with regards to its single input +-- as well as the primal answer. +-- +-- >>> diffF' (\a -> [sin a, cos a]) 0 +-- [(0.0,1.0),(1.0,0.0)] +diffF' :: (Functor f, Num a) => (Kahn a -> f (Kahn a)) -> a -> f (a, a) +diffF' f a = derivative' <$> f (var a 0) +{-# INLINE diffF' #-} + + +-- | Compute the 'hessian' via the 'jacobian' of the gradient. gradient is computed in 'Kahn' mode and then the 'jacobian' is computed in 'Kahn' mode. +-- +-- However, since the @'grad' f :: f a -> f a@ is square this is not as fast as using the forward-mode 'jacobian' of a reverse mode gradient provided by 'Numeric.AD.hessian'. +-- +-- >>> hessian (\[x,y] -> x*y) [1,2] +-- [[0,1],[1,0]] +hessian :: (Traversable f, Num a) => (f (On (Kahn (Kahn a))) -> (On (Kahn (Kahn a)))) -> f a -> f (f a) +hessian f = jacobian (grad (off . f . fmap On)) + +-- | Compute the order 3 Hessian tensor on a non-scalar-to-non-scalar function via the 'Kahn'-mode Jacobian of the 'Kahn'-mode Jacobian of the function. +-- +-- Less efficient than 'Numeric.AD.Mode.Mixed.hessianF'. +-- +-- >>> hessianF (\[x,y] -> [x*y,x+y,exp x*cos y]) [1,2] +-- [[[0.0,1.0],[1.0,0.0]],[[0.0,0.0],[0.0,0.0]],[[-1.1312043837568135,-2.4717266720048188],[-2.4717266720048188,1.1312043837568135]]] +hessianF :: (Traversable f, Functor g, Num a) => (f (On (Kahn (Kahn a))) -> g (On (Kahn (Kahn a)))) -> f a -> g (f (f a)) +hessianF f = getCompose . jacobian (Compose . jacobian (fmap off . f . fmap On)) + +-- $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! diff --git a/src/Numeric/AD/Rank1/Newton.hs b/src/Numeric/AD/Rank1/Newton.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/AD/Rank1/Newton.hs @@ -0,0 +1,121 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +----------------------------------------------------------------------------- +-- | +-- Copyright : (c) Edward Kmett 2010-2014 +-- License : BSD3 +-- Maintainer : ekmett@gmail.com +-- Stability : experimental +-- Portability : GHC only +-- +----------------------------------------------------------------------------- + +module Numeric.AD.Rank1.Newton + ( + -- * Newton's Method (Forward) + findZero + , inverse + , fixedPoint + , extremum + -- * Gradient Ascent/Descent (Kahn) + , gradientDescent + , gradientAscent + ) where + +import Prelude hiding (all, mapM) +import Data.Foldable (all) +import Data.Traversable +import Numeric.AD.Mode +import Numeric.AD.Rank1.Forward (Forward, diff, diff') +import Numeric.AD.Rank1.Kahn as Kahn (Kahn, gradWith') +import Numeric.AD.Internal.On + +-- $setup +-- >>> import Data.Complex + +-- | The 'findZero' function finds a zero of a scalar function using +-- Newton's method; its output is a stream of increasingly accurate +-- results. (Modulo the usual caveats.) If the stream becomes constant +-- ("it converges"), no further elements are returned. +-- +-- Examples: +-- +-- >>> take 10 $ findZero (\x->x^2-4) 1 +-- [1.0,2.5,2.05,2.000609756097561,2.0000000929222947,2.000000000000002,2.0] +-- +-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1) +-- 0.0 :+ 1.0 +findZero :: (Fractional a, Eq a) => (Forward a -> Forward a) -> a -> [a] +findZero f = go where + go x = x : if x == xn then [] else go xn where + (y,y') = diff' f x + xn = x - y/y' +{-# INLINE findZero #-} + +-- | The 'inverse' function inverts a scalar function using +-- Newton's method; its output is a stream of increasingly accurate +-- results. (Modulo the usual caveats.) If the stream becomes +-- constant ("it converges"), no further elements are returned. +-- +-- Example: +-- +-- >>> last $ take 10 $ inverse sqrt 1 (sqrt 10) +-- 10.0 +inverse :: (Fractional a, Eq a) => (Forward a -> Forward a) -> a -> a -> [a] +inverse f x0 y = findZero (\x -> f x - auto y) x0 +{-# INLINE inverse #-} + +-- | The 'fixedPoint' function find a fixedpoint of a scalar +-- function using Newton's method; its output is a stream of +-- increasingly accurate results. (Modulo the usual caveats.) +-- +-- If the stream becomes constant ("it converges"), no further +-- elements are returned. +-- +-- >>> last $ take 10 $ fixedPoint cos 1 +-- 0.7390851332151607 +fixedPoint :: (Fractional a, Eq a) => (Forward a -> Forward a) -> a -> [a] +fixedPoint f = findZero (\x -> f x - x) +{-# INLINE fixedPoint #-} + +-- | The 'extremum' function finds an extremum of a scalar +-- function using Newton's method; produces a stream of increasingly +-- accurate results. (Modulo the usual caveats.) If the stream +-- becomes constant ("it converges"), no further elements are returned. +-- +-- >>> last $ take 10 $ extremum cos 1 +-- 0.0 +extremum :: (Fractional a, Eq a) => (On (Forward (Forward a)) -> On (Forward (Forward a))) -> a -> [a] +extremum f = findZero (diff (off . f . On)) +{-# INLINE extremum #-} + +-- | The 'gradientDescent' function performs a multivariate +-- optimization, based on the naive-gradient-descent in the file +-- @stalingrad\/examples\/flow-tests\/pre-saddle-1a.vlad@ from the +-- VLAD compiler Stalingrad sources. Its output is a stream of +-- increasingly accurate results. (Modulo the usual caveats.) +-- +-- It uses reverse mode automatic differentiation to compute the gradient. +gradientDescent :: (Traversable f, Fractional a, Ord a) => (f (Kahn a) -> Kahn a) -> f a -> [f a] +gradientDescent f x0 = go x0 fx0 xgx0 0.1 (0 :: Int) + where + (fx0, xgx0) = Kahn.gradWith' (,) f x0 + go x fx xgx !eta !i + | eta == 0 = [] -- step size is 0 + | fx1 > fx = go x fx xgx (eta/2) 0 -- we stepped too far + | zeroGrad xgx = [] -- gradient is 0 + | otherwise = x1 : if i == 10 + then go x1 fx1 xgx1 (eta*2) 0 + else go x1 fx1 xgx1 eta (i+1) + where + zeroGrad = all (\(_,g) -> g == 0) + x1 = fmap (\(xi,gxi) -> xi - eta * gxi) xgx + (fx1, xgx1) = Kahn.gradWith' (,) f x1 +{-# INLINE gradientDescent #-} + +-- | Perform a gradient descent using reverse mode automatic differentiation to compute the gradient. +gradientAscent :: (Traversable f, Fractional a, Ord a) => (f (Kahn a) -> Kahn a) -> f a -> [f a] +gradientAscent f = gradientDescent (negate . f) +{-# INLINE gradientAscent #-} diff --git a/src/Numeric/AD/Rank1/Sparse.hs b/src/Numeric/AD/Rank1/Sparse.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/AD/Rank1/Sparse.hs @@ -0,0 +1,130 @@ +----------------------------------------------------------------------------- +-- | +-- Copyright : (c) Edward Kmett 2010-2014 +-- 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 Data.Traversable +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 g (a,b) = (a, g b) +{-# INLINE second #-} + +grad :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> f a +grad f as = d as $ apply f as +{-# INLINE grad #-} + +grad' :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> (a, f a) +grad' f as = d' as $ apply f as +{-# INLINE grad' #-} + +gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (f (Sparse a) -> Sparse a) -> f a -> f b +gradWith g f as = zipWithT g as $ grad f as +{-# INLINE gradWith #-} + +gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (f (Sparse a) -> Sparse a) -> f a -> (a, f b) +gradWith' g f as = second (zipWithT g as) $ grad' f as +{-# INLINE gradWith' #-} + +jacobian :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (f a) +jacobian f as = d as <$> apply f as +{-# INLINE jacobian #-} + +jacobian' :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (a, f a) +jacobian' f as = d' as <$> apply f 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 g f as = zipWithT g as <$> jacobian f 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' g f as = second (zipWithT g as) <$> jacobian' f as +{-# INLINE jacobianWith' #-} + +grads :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> Cofree f a +grads f as = ds as $ apply f as +{-# INLINE grads #-} + +jacobians :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (Cofree f a) +jacobians f as = ds as <$> apply f as +{-# INLINE jacobians #-} + +d2 :: Functor f => Cofree f a -> f (f a) +d2 = headJet . tailJet . tailJet . jet +{-# INLINE d2 #-} + +d2' :: Functor f => Cofree f a -> (a, f (a, f a)) +d2' (a :< as) = (a, fmap (\(da :< das) -> (da, extract <$> das)) as) +{-# INLINE d2' #-} + +hessian :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> f (f a) +hessian f as = d2 $ grads f as +{-# INLINE hessian #-} + +hessian' :: (Traversable f, Num a) => (f (Sparse a) -> Sparse a) -> f a -> (a, f (a, f a)) +hessian' f as = d2' $ grads f as +{-# INLINE hessian' #-} + +hessianF :: (Traversable f, Functor g, Num a) => (f (Sparse a) -> g (Sparse a)) -> f a -> g (f (f a)) +hessianF f as = d2 <$> jacobians f 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' f as = d2' <$> jacobians f 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! diff --git a/src/Numeric/AD/Rank1/Tower.hs b/src/Numeric/AD/Rank1/Tower.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/AD/Rank1/Tower.hs @@ -0,0 +1,117 @@ +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE BangPatterns #-} +----------------------------------------------------------------------------- +-- | +-- Copyright : (c) Edward Kmett 2010-2014 +-- License : BSD3 +-- Maintainer : ekmett@gmail.com +-- Stability : experimental +-- Portability : GHC only +-- +-- Higher order derivatives via a \"dual number tower\". +-- +----------------------------------------------------------------------------- + +module Numeric.AD.Rank1.Tower + ( Tower + , auto + -- * Taylor Series + , taylor + , taylor0 + -- * Maclaurin Series + , maclaurin + , maclaurin0 + -- * Derivatives + , diff -- first derivative of (a -> a) + , diff' -- answer and first derivative of (a -> a) + , diffs -- answer and all derivatives of (a -> a) + , diffs0 -- zero padded derivatives of (a -> a) + , diffsF -- answer and all derivatives of (a -> f a) + , diffs0F -- zero padded derivatives of (a -> f a) + -- * Directional Derivatives + , du -- directional derivative of (a -> a) + , du' -- answer and directional derivative of (a -> a) + , dus -- answer and all directional derivatives of (a -> a) + , dus0 -- answer and all zero padded directional derivatives of (a -> a) + , duF -- directional derivative of (a -> f a) + , duF' -- answer and directional derivative of (a -> f a) + , dusF -- answer and all directional derivatives of (a -> f a) + , dus0F -- answer and all zero padded directional derivatives of (a -> a) + ) where + +import Control.Applicative ((<$>)) +import Numeric.AD.Internal.Tower +import Numeric.AD.Mode + +diffs :: Num a => (Tower a -> Tower a) -> a -> [a] +diffs f a = getADTower $ apply f a +{-# INLINE diffs #-} + +diffs0 :: Num a => (Tower a -> Tower a) -> a -> [a] +diffs0 f a = zeroPad (diffs f a) +{-# INLINE diffs0 #-} + +diffsF :: (Functor f, Num a) => (Tower a -> f (Tower a)) -> a -> f [a] +diffsF f a = getADTower <$> apply f a +{-# INLINE diffsF #-} + +diffs0F :: (Functor f, Num a) => (Tower a -> f (Tower a)) -> a -> f [a] +diffs0F f a = (zeroPad . getADTower) <$> apply f a +{-# INLINE diffs0F #-} + +taylor :: Fractional a => (Tower a -> Tower a) -> a -> a -> [a] +taylor f x dx = go 1 1 (diffs f x) where + go !n !acc (a:as) = a * acc : go (n + 1) (acc * dx / n) as + go _ _ [] = [] + +taylor0 :: Fractional a => (Tower a -> Tower a) -> a -> a -> [a] +taylor0 f x dx = zeroPad (taylor f x dx) +{-# INLINE taylor0 #-} + +maclaurin :: Fractional a => (Tower a -> Tower a) -> a -> [a] +maclaurin f = taylor f 0 +{-# INLINE maclaurin #-} + +maclaurin0 :: Fractional a => (Tower a -> Tower a) -> a -> [a] +maclaurin0 f = taylor0 f 0 +{-# INLINE maclaurin0 #-} + +diff :: Num a => (Tower a -> Tower a) -> a -> a +diff f = d . diffs f +{-# INLINE diff #-} + +diff' :: Num a => (Tower a -> Tower a) -> a -> (a, a) +diff' f = d' . diffs f +{-# INLINE diff' #-} + +du :: (Functor f, Num a) => (f (Tower a) -> Tower a) -> f (a, a) -> a +du f = d . getADTower . f . fmap withD +{-# INLINE du #-} + +du' :: (Functor f, Num a) => (f (Tower a) -> Tower a) -> f (a, a) -> (a, a) +du' f = d' . getADTower . f . fmap withD +{-# INLINE du' #-} + +duF :: (Functor f, Functor g, Num a) => (f (Tower a) -> g (Tower a)) -> f (a, a) -> g a +duF f = fmap (d . getADTower) . f . fmap withD +{-# INLINE duF #-} + +duF' :: (Functor f, Functor g, Num a) => (f (Tower a) -> g (Tower a)) -> f (a, a) -> g (a, a) +duF' f = fmap (d' . getADTower) . f . fmap withD +{-# INLINE duF' #-} + +dus :: (Functor f, Num a) => (f (Tower a) -> Tower a) -> f [a] -> [a] +dus f = getADTower . f . fmap tower +{-# INLINE dus #-} + +dus0 :: (Functor f, Num a) => (f (Tower a) -> Tower a) -> f [a] -> [a] +dus0 f = zeroPad . getADTower . f . fmap tower +{-# INLINE dus0 #-} + +dusF :: (Functor f, Functor g, Num a) => (f (Tower a) -> g (Tower a)) -> f [a] -> g [a] +dusF f = fmap getADTower . f . fmap tower +{-# INLINE dusF #-} + +dus0F :: (Functor f, Functor g, Num a) => (f (Tower a) -> g (Tower a)) -> f [a] -> g [a] +dus0F f = fmap getADTower . f . fmap tower +{-# INLINE dus0F #-}