module Numeric.AD.Internal.Forward
( Forward(..)
, tangent
, bundle
, unbundle
, apply
, bind
, bind'
, bindWith
, bindWith'
, transposeWith
) where
import Language.Haskell.TH
import Data.Typeable
import Data.Traversable (Traversable, mapAccumL)
import Data.Foldable (Foldable, toList)
import Data.Data
import Control.Applicative
import Numeric.AD.Internal.Types
import Numeric.AD.Internal.Classes
import Numeric.AD.Internal.Identity
data Forward a
= Forward !a a
| Lift !a
| Zero
deriving (Show, Data, Typeable)
tangent :: Num a => AD Forward a -> a
tangent (AD (Forward _ da)) = da
tangent _ = 0
unbundle :: Num a => AD Forward a -> (a, a)
unbundle (AD (Forward a da)) = (a, da)
unbundle (AD Zero) = (0,0)
unbundle (AD (Lift a)) = (a, 0)
bundle :: a -> a -> AD Forward a
bundle a da = AD (Forward a da)
apply :: Num a => (AD Forward a -> b) -> a -> b
apply f a = f (bundle a 1)
instance Primal Forward where
primal (Forward a _) = a
primal (Lift a) = a
primal Zero = 0
instance Lifted Forward => Mode Forward where
auto = Lift
zero = Zero
isKnownZero Zero = True
isKnownZero _ = False
isKnownConstant Forward{} = False
isKnownConstant _ = True
Zero <+> a = a
a <+> Zero = a
Forward a da <+> Forward b db = Forward (a + b) (da + db)
Forward a da <+> Lift b = Forward (a + b) da
Lift a <+> Forward b db = Forward (a + b) db
Lift a <+> Lift b = Lift (a + b)
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 *! log1 xi)) x y
a *^ Forward b db = Forward (a * b) (a * db)
a *^ Lift b = Lift (a * b)
_ *^ Zero = Zero
Forward a da ^* b = Forward (a * b) (da * b)
Lift a ^* b = Lift (a * b)
Zero ^* _ = Zero
Forward a da ^/ b = Forward (a / b) (da / b)
Lift a ^/ b = Lift (a / b)
Zero ^/ _ = Zero
instance Lifted Forward => Jacobian Forward where
type D Forward = Id
unary f (Id dadb) (Forward b db) = Forward (f b) (dadb * db)
unary f _ (Lift b) = Lift (f b)
unary f _ Zero = Lift (f 0)
lift1 f _ Zero = Lift (f 0)
lift1 f _ (Lift b) = Lift (f b)
lift1 f df (Forward b db) = Forward (f b) (dadb * db)
where
Id dadb = df (Id b)
lift1_ f _ Zero = Lift (f 0)
lift1_ f _ (Lift b) = Lift (f b)
lift1_ f df (Forward b db) = Forward a da
where
a = f b
Id da = df (Id a) (Id b) ^* db
binary f _ _ Zero Zero = Lift (f 0 0)
binary f _ _ Zero (Lift c) = Lift (f 0 c)
binary f _ _ (Lift b) Zero = Lift (f b 0)
binary f _ _ (Lift b) (Lift c) = Lift (f b c)
binary f _ (Id dadc) Zero (Forward c dc) = Forward (f 0 c) $ dc * dadc
binary f _ (Id dadc) (Lift b) (Forward c dc) = Forward (f b c) $ dc * dadc
binary f (Id dadb) _ (Forward b db) Zero = Forward (f b 0) $ dadb * db
binary f (Id dadb) _ (Forward b db) (Lift c) = Forward (f b c) $ dadb * db
binary f (Id dadb) (Id dadc) (Forward b db) (Forward c dc) = Forward (f b c) $ dadb * db + dc * dadc
lift2 f _ Zero Zero = Lift (f 0 0)
lift2 f _ Zero (Lift c) = Lift (f 0 c)
lift2 f _ (Lift b) Zero = Lift (f b 0)
lift2 f _ (Lift b) (Lift c) = Lift (f b c)
lift2 f df Zero (Forward c dc) = Forward (f 0 c) $ dc * runId (snd (df (Id 0) (Id c)))
lift2 f df (Lift b) (Forward c dc) = Forward (f b c) $ dc * runId (snd (df (Id b) (Id c)))
lift2 f df (Forward b db) Zero = Forward (f b 0) $ runId (fst (df (Id b) (Id 0))) * db
lift2 f df (Forward b db) (Lift c) = Forward (f b c) $ runId (fst (df (Id b) (Id c))) * db
lift2 f df (Forward b db) (Forward c dc) = Forward a da
where
a = f b c
(Id dadb, Id dadc) = df (Id b) (Id c)
da = dadb * db + dc * dadc
lift2_ f _ Zero Zero = Lift (f 0 0)
lift2_ f _ Zero (Lift c) = Lift (f 0 c)
lift2_ f _ (Lift b) Zero = Lift (f b 0)
lift2_ f _ (Lift b) (Lift c) = Lift (f b c)
lift2_ f df Zero (Forward c dc) = Forward a $ dc * runId (snd (df (Id a) (Id 0) (Id c))) where a = f 0 c
lift2_ f df (Lift b) (Forward c dc) = Forward a $ dc * runId (snd (df (Id a) (Id b) (Id c))) where a = f b c
lift2_ f df (Forward b db) Zero = Forward a $ runId (fst (df (Id a) (Id b) (Id 0))) * db where a = f b 0
lift2_ f df (Forward b db) (Lift c) = Forward a $ runId (fst (df (Id a) (Id b) (Id c))) * db where a = f b c
lift2_ f df (Forward b db) (Forward c dc) = Forward a da
where
a = f b c
(Id dadb, Id dadc) = df (Id a) (Id b) (Id c)
da = dadb * db + dc * dadc
deriveLifted id $ conT ''Forward
bind :: (Traversable f, Num a) => (f (AD 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 (AD 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 (AD 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 (AD 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)
b0 = f (auto <$> as)
dropIx ((_,b),bs) = (b,bs)
transposeWith :: (Functor f, Foldable f, Traversable g) => (b -> f a -> c) -> f (g a) -> g b -> g c
transposeWith f as = snd . mapAccumL go xss0
where
go xss b = (tail <$> xss, f b (head <$> xss))
xss0 = toList <$> as