{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-----------------------------------------------------------------------------
-- |
-- Module : Numeric.Optimization.Backprop.ToVector
-- Copyright : (c) Masahiro Sakai 2023
-- License : BSD-style
--
-- Maintainer : masahiro.sakai@gmail.com
-- Stability : provisional
-- Portability : non-portable
--
-- Conversion between a type and 'VS.Vector' 'Double'.
--
-----------------------------------------------------------------------------
module Numeric.Optimization.Backprop.ToVector
(
-- * ToVector class
ToVector (..)
, toVector
-- * Utilities for defining ToVector class
-- ** Generics
, GToVector (..)
-- ** @Foldable@/@Traversable@-based definition
, dimFoldable
, writeToMVectorFoldable
, updateFromVectorTraversable
-- ** @MonoFoldable@/@MonoTraversable@-based definition
, dimMonoFoldable
, writeToMVectorMonoFoldable
, updateFromVectorMonoTraversable
) where
import Control.Monad.Primitive
import Control.Monad.State
import qualified Data.MonoTraversable as MT
import Data.Traversable (mapAccumL)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import qualified Data.Vector.Unboxed as VU
import GHC.Generics
import qualified Data.Functor.Identity as Functor
import qualified Data.Functor.Compose as Functor
import qualified Data.Functor.Const as Functor
import qualified Data.Functor.Product as Functor
import qualified Data.Functor.Sum as Functor
import Data.IntMap (IntMap)
import Data.List.NonEmpty (NonEmpty)
import Data.Map (Map)
import Data.Monoid
import qualified Data.Semigroup as SG
import Data.Sequence (Seq)
import Data.Void
-- ------------------------------------------------------------------------
-- | Type that can be converted to @'VS.Vector' 'Double'@ and back.
--
-- Laws that should be satisfied:
--
-- * @'VS.length' . 'toVector' = dim@
--
-- * @updateFromVector a ('toVector' a) = a@
--
-- * @updateFromVector (updateFromVector a v1) v2 = updateFromVector a v2@
class ToVector a where
-- | Dimention of the resulting vector.
dim :: a -> Int
-- | Destination passing style version of 'toVector'.
writeToMVector :: PrimMonad m => a -> VSM.MVector (PrimState m) Double -> m ()
-- | Converting @'VS.Vector' 'Double'@ back to a value
updateFromVector :: a -> VS.Vector Double -> a
default dim :: (Generic a, GToVector (Rep a)) => a -> Int
dim x = gDim (from x)
default writeToMVector :: (Generic a, GToVector (Rep a), PrimMonad m) => a -> VSM.MVector (PrimState m) Double -> m ()
writeToMVector x vec = gWriteToMVector (from x) vec
default updateFromVector :: (Generic a, GToVector (Rep a)) => a -> VS.Vector Double -> a
updateFromVector x v = to (gUpdateFromVector (from x) v)
-- | Converting a value to @'VS.Vector' 'Double'@.
toVector :: ToVector a => a -> VS.Vector Double
toVector x = VS.create $ do
vec <- VSM.new (dim x)
writeToMVector x vec
return vec
-- ------------------------------------------------------------------------
-- | Implementation of 'dim' for the type of the form @f a@ for @'Foldable' f@.
dimFoldable :: (Foldable f, ToVector a) => f a -> Int
dimFoldable = getSum . foldMap (Sum . dim)
-- | Implementation of 'writeToMVector' for the type of the form @f a@ for @'Foldable' f@.
writeToMVectorFoldable :: (Foldable f, ToVector a, PrimMonad m) => f a -> VSM.MVector (PrimState m) Double -> m ()
writeToMVectorFoldable xs vec = foldM_ f vec xs
where
f vec' x =
case VSM.splitAt (dim x) vec' of
(vec1, vec2) -> do
writeToMVector x vec1
return vec2
-- | Implementation of 'updateFromVectorTraversable' for the type of the form @f a@ for @'Traversable' f@.
updateFromVectorTraversable :: (Traversable f, ToVector a) => f a -> VS.Vector Double -> f a
updateFromVectorTraversable xs v0 = flip evalState v0 $ do
forM xs $ \x -> do
v <- get
case VS.splitAt (dim x) v of
(v1, v2) -> do
put v2
return (updateFromVector x v1)
-- ------------------------------------------------------------------------
-- | Implementation of 'dim' for a 'MT.MonoFoldable' type
dimMonoFoldable :: (MT.MonoFoldable a, ToVector (MT.Element a)) => a -> Int
dimMonoFoldable = getSum . MT.ofoldMap (Sum . dim)
-- | Implementation of 'writeToMVector' for a 'MT.MonoFoldable' type
writeToMVectorMonoFoldable :: (MT.MonoFoldable a, ToVector (MT.Element a), PrimMonad m) => a -> VSM.MVector (PrimState m) Double -> m ()
writeToMVectorMonoFoldable xs vec = MT.ofoldM f vec xs >> return ()
where
f vec' x =
case VSM.splitAt (dim x) vec' of
(vec1, vec2) -> do
writeToMVector x vec1
return vec2
-- | Implementation of 'updateFromVector' for a 'MT.MonoTraversable' type
updateFromVectorMonoTraversable :: (MT.MonoTraversable a, ToVector (MT.Element a)) => a -> VS.Vector Double -> a
updateFromVectorMonoTraversable xs v0 = flip evalState v0 $ do
MT.oforM xs $ \x -> do
v <- get
case VS.splitAt (dim x) v of
(v1, v2) -> do
put v2
return (updateFromVector x v1)
-- ------------------------------------------------------------------------
-- | Class of generic representation types that can be converted to/from 'VS.Vector' 'Double'.
class GToVector f where
gDim :: f p -> Int
gWriteToMVector :: PrimMonad m => f p -> VSM.MVector (PrimState m) Double -> m ()
gUpdateFromVector :: f p -> VS.Vector Double -> f p
instance GToVector V1 where
gDim x = case x of { }
gWriteToMVector _x _vec = return ()
gUpdateFromVector x _v = case x of { }
instance GToVector U1 where
gDim _ = 0
gWriteToMVector _x _vec = return ()
gUpdateFromVector x _v = x
instance (GToVector f, GToVector g) => GToVector (f :+: g) where
gDim (L1 x) = gDim x
gDim (R1 x) = gDim x
gWriteToMVector (L1 x) vec = gWriteToMVector x vec
gWriteToMVector (R1 x) vec = gWriteToMVector x vec
gUpdateFromVector (L1 x) v = L1 (gUpdateFromVector x v)
gUpdateFromVector (R1 x) v = R1 (gUpdateFromVector x v)
instance (GToVector f, GToVector g) => GToVector (f :*: g) where
gDim (a :*: b) = gDim a + gDim b
gWriteToMVector (a :*: b) vec =
case VSM.splitAt (gDim a) vec of
(vec1, vec2) -> do
gWriteToMVector a vec1
gWriteToMVector b vec2
gUpdateFromVector (a :*: b) v =
case VS.splitAt (gDim a) v of
(vec1, vec2) -> (gUpdateFromVector a vec1 :*: gUpdateFromVector b vec2)
instance (ToVector c) => GToVector (K1 i c) where
gDim (K1 x) = dim x
gWriteToMVector (K1 x) vec = writeToMVector x vec
gUpdateFromVector (K1 x) v = K1 (updateFromVector x v)
instance (GToVector f) => GToVector (M1 i t f) where
gDim (M1 x) = gDim x
gWriteToMVector (M1 x) vec = gWriteToMVector x vec
gUpdateFromVector (M1 x) v = M1 (gUpdateFromVector x v)
-- ------------------------------------------------------------------------
instance ToVector Double where
dim _ = 1
writeToMVector x vec = VSM.write vec 0 x
updateFromVector _x v = v VS.! 0
instance (a ~ Double) => ToVector (VS.Vector a) where
dim x = VS.length x
#if MIN_VERSION_vector(0,12,2)
writeToMVector x vec = VS.imapM_ (VSM.write vec) x
#else
writeToMVector x vec = flip evalStateT 0 $ VS.mapM_ (\e -> do{ i <- get; VSM.write vec i e; put (i+1) }) x
#endif
updateFromVector _x v = v
instance (a ~ Double) => ToVector (VU.Vector a) where
dim x = VU.length x
#if MIN_VERSION_vector(0,12,2)
writeToMVector x vec = VU.imapM_ (VSM.write vec) x
#else
writeToMVector x vec = flip evalStateT 0 $ VU.mapM_ (\e -> do{ i <- get; VSM.write vec i e; put (i+1) }) x
#endif
updateFromVector _x v = VG.convert v
instance (ToVector a) => ToVector (V.Vector a) where
dim xs = V.sum (V.map dim xs)
writeToMVector xs vec = V.foldM_ f vec xs
where
f vec' x =
case VSM.splitAt (dim x) vec' of
(vec1, vec2) -> do
writeToMVector x vec1
return vec2
updateFromVector xs v = snd $ mapAccumL f v xs
where
f v' x =
case VS.splitAt (dim x) v' of
(v1, v2) -> (v2, updateFromVector x v1)
instance ToVector Void
instance ToVector ()
instance (ToVector a, ToVector b) => ToVector (a, b)
instance (ToVector a, ToVector b, ToVector c) => ToVector (a, b, c)
instance (ToVector a, ToVector b, ToVector c, ToVector d) => ToVector (a, b, c, d)
instance (ToVector a, ToVector b, ToVector c, ToVector d, ToVector e) => ToVector (a, b, c, d, e)
instance (ToVector a) => ToVector (Maybe a)
instance ToVector a => ToVector (SG.Min a)
instance ToVector a => ToVector (SG.Max a)
instance ToVector a => ToVector (SG.First a)
instance ToVector a => ToVector (SG.Last a)
instance ToVector a => ToVector (SG.WrappedMonoid a)
#if !MIN_VERSION_base(4,16,0)
instance ToVector a => ToVector (SG.Option a)
#endif
instance (ToVector a, ToVector b) => ToVector (SG.Arg a b)
instance ToVector a => ToVector (Dual a)
instance ToVector a => ToVector (Sum a)
instance ToVector a => ToVector (Product a)
instance ToVector a => ToVector (First a)
instance ToVector a => ToVector (Last a)
instance ToVector (f a) => ToVector (Alt f a)
instance ToVector (f a) => ToVector (Ap f a)
instance ToVector a => ToVector (Functor.Identity a)
instance ToVector (f (g a)) => ToVector (Functor.Compose f g a)
instance ToVector w => ToVector (Functor.Const w a)
instance (ToVector (f a), ToVector (g a)) => ToVector (Functor.Product f g a)
instance (ToVector (f a), ToVector (g a)) => ToVector (Functor.Sum f g a)
instance ToVector a => ToVector [a] where
dim = dimFoldable
writeToMVector = writeToMVectorFoldable
updateFromVector = updateFromVectorTraversable
instance ToVector a => ToVector (NonEmpty a) where
dim = dimFoldable
writeToMVector = writeToMVectorFoldable
updateFromVector = updateFromVectorTraversable
instance ToVector a => ToVector (Map k a) where
dim = dimFoldable
writeToMVector = writeToMVectorFoldable
updateFromVector = updateFromVectorTraversable
instance ToVector a => ToVector (IntMap a) where
dim = dimFoldable
writeToMVector = writeToMVectorFoldable
updateFromVector = updateFromVectorTraversable
instance ToVector a => ToVector (Seq a) where
dim = dimFoldable
writeToMVector = writeToMVectorFoldable
updateFromVector = updateFromVectorTraversable
-- ------------------------------------------------------------------------