Num.hs

{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}
{-# LANGUAGE PatternSynonyms  #-}
{-# LANGUAGE RankNTypes       #-}
{-# OPTIONS_HADDOCK not-home  #-}

-- |
-- Module      : Numeric.Backprop.Num
-- Copyright   : (c) Justin Le 2023
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Provides the exact same API as "Numeric.Backprop", except requiring
-- 'Num' instances for all types involved instead of 'Backprop' instances.
--
-- This was the original API of the library (for version 0.1).
--
-- 'Num' is strictly more powerful than 'Backprop', and is a stronger
-- constraint on types than is necessary for proper backpropagating.  In
-- particular, 'fromInteger' is a problem for many types, preventing useful
-- backpropagation for lists, variable-length vectors (like "Data.Vector")
-- and variable-size matrices from linear algebra libraries like /hmatrix/
-- and /accelerate/.
--
-- However, this module might be useful in situations where you are working
-- with external types with 'Num' instances, and you want to avoid writing
-- orphan instances for external types.
--
-- If you have external types that are not 'Num' instances, consider
-- instead "Numeric.Backprop.External".
--
-- If you need a 'Num' instance for tuples, you can use the orphan
-- instances in the <https://hackage.haskell.org/package/NumInstances
-- NumInstances> package (in particular, "Data.NumInstances.Tuple") if you
-- are writing an application and do not have to worry about orphan
-- instances.
--
-- See "Numeric.Backprop" for fuller documentation on using these
-- functions.
--
-- @since 0.2.0.0

module Numeric.Backprop.Num (
    -- * Types
    BVar, W
    -- * Running
  , backprop, E.evalBP, gradBP, backpropWith
    -- ** Multiple inputs
  , E.evalBP0
  , backprop2, E.evalBP2, gradBP2, backpropWith2
  , backpropN, E.evalBPN, gradBPN, backpropWithN
    -- * Manipulating 'BVar'
  , E.constVar, E.auto, E.coerceVar
  , (^^.), (.~~), (%~~), (^^?), (^^..), (^^?!)
  , viewVar, setVar, overVar
  , sequenceVar, collectVar
  , previewVar, toListOfVar
    -- ** With Isomorphisms
  , isoVar, isoVar2, isoVar3, isoVarN
    -- ** With 'Op's
  , liftOp
  , liftOp1, liftOp2, liftOp3
    -- * 'Op'
  , Op(..)
    -- ** Creation
  , op0, opConst, idOp
  , bpOp
    -- *** Giving gradients directly
  , op1, op2, op3
    -- *** From Isomorphisms
  , opCoerce, opTup, opIso, opIsoN, opLens
    -- *** No gradients
  , noGrad1, noGrad
    -- * Utility
  , Reifies
  ) where

import           Data.Functor.Identity
import           Data.Maybe
import           Data.Reflection
import           Data.Vinyl
import           Lens.Micro
import           Numeric.Backprop.Explicit (BVar, W)
import           Numeric.Backprop.Op
import qualified Numeric.Backprop.Explicit as E

-- | 'Numeric.Backprop.backpropN', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- The @'RPureConstrained' 'Num' as@ in the constraint says that every
-- value in the type-level list @as@ must have a 'Num' instance.  This
-- means you can use, say, @'[Double, Float, Int]@, but not @'[Double,
-- Bool, String]@.
--
-- If you stick to /concerete/, monomorphic usage of this (with specific
-- types, typed into source code, known at compile-time), then
-- @'AllPureConstrained' 'Num' as@ should be fulfilled automatically.
--
backpropN
    :: (RPureConstrained Num as, Num b)
    => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b)
    -> Rec Identity as
    -> (b, Rec Identity as)
backpropN = E.backpropN E.zfNums E.ofNum
{-# INLINE backpropN #-}

-- | 'Numeric.Backprop.backpropWithN', but with 'Num' constraints instead
-- of 'Backprop' constraints.
--
-- See 'backpropN' for information on the 'AllConstrained' constraint.
--
-- Note that argument order changed in v0.2.4.
--
-- @since 0.2.0.0
backpropWithN
    :: RPureConstrained Num as
    => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b)
    -> Rec Identity as
    -> (b, b -> Rec Identity as)
backpropWithN = E.backpropWithN E.zfNums
{-# INLINE backpropWithN #-}

-- | 'Numeric.Backprop.backprop', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- See module documentation for "Numeric.Backprop.Num" for information on
-- using this with tuples.
backprop
    :: (Num a, Num b)
    => (forall s. Reifies s W => BVar s a -> BVar s b)
    -> a
    -> (b, a)
backprop = E.backprop E.zfNum E.ofNum
{-# INLINE backprop #-}

-- | 'Numeric.Backprop.backpropWith', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- See module documentation for "Numeric.Backprop.Num" for information on
-- using this with tuples.
--
-- Note that argument order changed in v0.2.4.
--
-- @since 0.2.0.0
backpropWith
    :: Num a
    => (forall s. Reifies s W => BVar s a -> BVar s b)
    -> a
    -> (b, b -> a)
backpropWith = E.backpropWith E.zfNum
{-# INLINE backpropWith #-}

-- | 'Numeric.Backprop.gradBP', but with 'Num' constraints instead of
-- 'Backprop' constraints.
gradBP
    :: (Num a, Num b)
    => (forall s. Reifies s W => BVar s a -> BVar s b)
    -> a
    -> a
gradBP = E.gradBP E.zfNum E.ofNum
{-# INLINE gradBP #-}

-- | 'Numeric.Backprop.gradBPN', but with 'Num' constraints instead of
-- 'Backprop' constraints.
gradBPN
    :: (RPureConstrained Num as, Num b)
    => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b)
    -> Rec Identity as
    -> Rec Identity as
gradBPN = E.gradBPN E.zfNums E.ofNum
{-# INLINE gradBPN #-}

-- | 'Numeric.Backprop.backprop2', but with 'Num' constraints instead of
-- 'Backprop' constraints.
backprop2
    :: (Num a, Num b, Num c)
    => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c)
    -> a
    -> b
    -> (c, (a, b))
backprop2 = E.backprop2 E.zfNum E.zfNum E.ofNum
{-# INLINE backprop2 #-}

-- | 'Numeric.Backprop.backpropWith2', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- Note that argument order changed in v0.2.4.
--
-- @since 0.2.0.0
backpropWith2
    :: (Num a, Num b)
    => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c)
    -> a
    -> b
    -> (c, c -> (a, b)) -- ^ Takes function giving gradient of final result given the output of function
backpropWith2 = E.backpropWith2 E.zfNum E.zfNum
{-# INLINE backpropWith2 #-}

-- | 'Numeric.Backprop.gradBP2', but with 'Num' constraints instead of
-- 'Backprop' constraints.
gradBP2
    :: (Num a, Num b, Num c)
    => (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c)
    -> a
    -> b
    -> (a, b)
gradBP2 = E.gradBP2 E.zfNum E.zfNum E.ofNum
{-# INLINE gradBP2 #-}

-- | 'Numeric.Backprop.bpOp', but with 'Num' constraints instead of
-- 'Backprop' constraints.
bpOp
    :: RPureConstrained Num as
    => (forall s. Reifies s W => Rec (BVar s) as -> BVar s b)
    -> Op as b
bpOp = E.bpOp E.zfNums
{-# INLINE bpOp #-}

-- | 'Numeric.Backprop.^^.', but with 'Num' constraints instead of
-- 'Backprop' constraints.
(^^.)
    :: forall b a s. (Num a, Num b, Reifies s W)
    => BVar s b
    -> Lens' b a
    -> BVar s a
x ^^. l = viewVar l x
infixl 8 ^^.
{-# INLINE (^^.) #-}

-- | 'Numeric.Backprop.viewVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
viewVar
    :: forall b a s. (Num a, Num b, Reifies s W)
    => Lens' b a
    -> BVar s b
    -> BVar s a
viewVar = E.viewVar E.afNum E.zfNum
{-# INLINE viewVar #-}


-- | 'Numeric.Backprop..~~', but with 'Num' constraints instead of
-- 'Backprop' constraints.
(.~~)
    :: (Num a, Num b, Reifies s W)
    => Lens' b a
    -> BVar s a
    -> BVar s b
    -> BVar s b
l .~~ x = setVar l x
infixl 8 .~~
{-# INLINE (.~~) #-}

-- | 'Numeric.Backprop.setVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
setVar
    :: forall a b s. (Num a, Num b, Reifies s W)
    => Lens' b a
    -> BVar s a
    -> BVar s b
    -> BVar s b
setVar = E.setVar E.afNum E.afNum E.zfNum
{-# INLINE setVar #-}

-- | 'Numeric.Backprop.%~~', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- @since 0.2.4.0
--
(%~~)
    :: (Num a, Num b, Reifies s W)
    => Lens' b a
    -> (BVar s a -> BVar s a)
    -> BVar s b
    -> BVar s b
l %~~ f = overVar l f
infixr 4 %~~
{-# INLINE (%~~) #-}

-- | 'Numeric.Backprop.overVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- @since 0.2.4.0
overVar
    :: (Num a, Num b, Reifies s W)
    => Lens' b a
    -> (BVar s a -> BVar s a)
    -> BVar s b
    -> BVar s b
overVar = E.overVar E.afNum E.afNum E.zfNum E.zfNum
{-# INLINE overVar #-}

-- | 'Numeric.Backprop.^^?', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- Note that many automatically-generated prisms by the /lens/ package use
-- tuples, which cannot work this this by default (because tuples do not
-- have a 'Num' instance).
--
-- If you are writing an application or don't have to worry about orphan
-- instances, you can pull in the orphan instances from
-- <https://hackage.haskell.org/package/NumInstances NumInstances>.
-- Alternatively, you can chain those prisms with conversions to the
-- anonymous canonical strict tuple types in "Numeric.Backprop.Tuple",
-- which do have 'Num' instances.
--
-- @
-- myPrism                   :: 'Prism'' c (a, b)
-- myPrism . 'iso' 'tupT2' 't2Tup' :: 'Prism'' c ('T2' a b)
-- @
(^^?)
    :: forall b a s. (Num b, Num a, Reifies s W)
    => BVar s b
    -> Traversal' b a
    -> Maybe (BVar s a)
v ^^? t = previewVar t v
infixl 8 ^^?
{-# INLINE (^^?) #-}

-- | 'Numeric.Backprop.^^?!', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- Like 'Numeric.Backprop.^^?!', is *UNSAFE*.
--
-- @since 0.2.1.0
(^^?!)
    :: forall b a s. (Num b, Num a, Reifies s W)
    => BVar s b
    -> Traversal' b a
    -> BVar s a
v ^^?! t = fromMaybe (error e) (previewVar t v)
  where
    e = "Numeric.Backprop.Num.^^?!: Empty traversal"
infixl 8 ^^?!
{-# INLINE (^^?!) #-}

-- | 'Numeric.Backprop.previewVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- See documentation for '^^?' for more information and important notes.
previewVar
    :: forall b a s. (Num b, Num a, Reifies s W)
    => Traversal' b a
    -> BVar s b
    -> Maybe (BVar s a)
previewVar = E.previewVar E.afNum E.zfNum
{-# INLINE previewVar #-}

-- | 'Numeric.Backprop.^^..', but with 'Num' constraints instead of
-- 'Backprop' constraints.
(^^..)
    :: forall b a s. (Num b, Num a, Reifies s W)
    => BVar s b
    -> Traversal' b a
    -> [BVar s a]
v ^^.. t = toListOfVar t v
{-# INLINE (^^..) #-}

-- | 'Numeric.Backprop.toListOfVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
toListOfVar
    :: forall b a s. (Num b, Num a, Reifies s W)
    => Traversal' b a
    -> BVar s b
    -> [BVar s a]
toListOfVar = E.toListOfVar E.afNum E.zfNum
{-# INLINE toListOfVar #-}

-- | 'Numeric.Backprop.sequenceVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- Since v0.2.4, requires a 'Num' constraint on @t a@.
sequenceVar
    :: (Traversable t, Num a, Reifies s W)
    => BVar s (t a)
    -> t (BVar s a)
sequenceVar = E.sequenceVar E.afNum E.zfNum
{-# INLINE sequenceVar #-}

-- | 'Numeric.Backprop.collectVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
--
-- Prior to v0.2.3, required a 'Num' constraint on @t a@.
collectVar
    :: (Foldable t, Functor t, Num a, Reifies s W)
    => t (BVar s a)
    -> BVar s (t a)
collectVar = E.collectVar E.afNum E.zfNum
{-# INLINE collectVar #-}

-- | 'Numeric.Backprop.liftOp', but with 'Num' constraints instead of
-- 'Backprop' constraints.
liftOp
    :: (RPureConstrained Num as, Reifies s W)
    => Op as b
    -> Rec (BVar s) as
    -> BVar s b
liftOp = E.liftOp E.afNums
{-# INLINE liftOp #-}

-- | 'Numeric.Backprop.liftOp1', but with 'Num' constraints instead of
-- 'Backprop' constraints.
liftOp1
    :: (Num a, Reifies s W)
    => Op '[a] b
    -> BVar s a
    -> BVar s b
liftOp1 = E.liftOp1 E.afNum
{-# INLINE liftOp1 #-}

-- | 'Numeric.Backprop.liftOp2', but with 'Num' constraints instead of
-- 'Backprop' constraints.
liftOp2
    :: (Num a, Num b, Reifies s W)
    => Op '[a,b] c
    -> BVar s a
    -> BVar s b
    -> BVar s c
liftOp2 = E.liftOp2 E.afNum E.afNum
{-# INLINE liftOp2 #-}

-- | 'Numeric.Backprop.liftOp3', but with 'Num' constraints instead of
-- 'Backprop' constraints.
liftOp3
    :: (Num a, Num b, Num c, Reifies s W)
    => Op '[a,b,c] d
    -> BVar s a
    -> BVar s b
    -> BVar s c
    -> BVar s d
liftOp3 = E.liftOp3 E.afNum E.afNum E.afNum
{-# INLINE liftOp3 #-}

-- | 'Numeric.Backprop.isoVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
isoVar
    :: (Num a, Reifies s W)
    => (a -> b)
    -> (b -> a)
    -> BVar s a
    -> BVar s b
isoVar = E.isoVar E.afNum
{-# INLINE isoVar #-}

-- | 'Numeric.Backprop.isoVar', but with 'Num' constraints instead of
-- 'Backprop' constraints.
isoVar2
    :: (Num a, Num b, Reifies s W)
    => (a -> b -> c)
    -> (c -> (a, b))
    -> BVar s a
    -> BVar s b
    -> BVar s c
isoVar2 = E.isoVar2 E.afNum E.afNum
{-# INLINE isoVar2 #-}

-- | 'Numeric.Backprop.isoVar3', but with 'Num' constraints instead of
-- 'Backprop' constraints.
isoVar3
    :: (Num a, Num b, Num c, Reifies s W)
    => (a -> b -> c -> d)
    -> (d -> (a, b, c))
    -> BVar s a
    -> BVar s b
    -> BVar s c
    -> BVar s d
isoVar3 = E.isoVar3 E.afNum E.afNum E.afNum
{-# INLINE isoVar3 #-}

-- | 'Numeric.Backprop.isoVarN', but with 'Num' constraints instead of
-- 'Backprop' constraints.
isoVarN
    :: (RPureConstrained Num as, Reifies s W)
    => (Rec Identity as -> b)
    -> (b -> Rec Identity as)
    -> Rec (BVar s) as
    -> BVar s b
isoVarN = E.isoVarN E.afNums
{-# INLINE isoVarN #-}