{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Numeric.Optimization.AD
(
minimize
, minimizeReverse
, minimizeSparse
, Constraint (..)
, Method (..)
, isSupportedMethod
, Params (..)
, Result (..)
, Statistics (..)
, OptimizationException (..)
, Default (..)
, AD
, auto
, Reverse
, Reifies
, Tape
, Sparse
) where
import Control.Monad.Primitive
import Data.Default.Class
import Data.Foldable (foldlM, toList)
import Data.Functor.Contravariant
import Data.Reflection (Reifies)
import Data.Traversable
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import Numeric.AD (AD, auto)
import Numeric.AD.Internal.Reverse (Tape)
import Numeric.AD.Mode.Reverse (Reverse)
import qualified Numeric.AD.Mode.Reverse as Reverse
import Numeric.AD.Mode.Sparse (Sparse)
import qualified Numeric.AD.Mode.Sparse as Sparse
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.Optimization as Opt
import Numeric.Optimization hiding (minimize, IsProblem (..))
data ProblemReverse f
= ProblemReverse
(forall s. Reifies s Tape => f (Reverse s Double) -> Reverse s Double)
(Maybe (V.Vector (Double, Double)))
[Constraint]
Int
(f Int)
instance Traversable f => Opt.IsProblem (ProblemReverse f) where
func :: ProblemReverse f -> Vector Double -> Double
func (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x =
forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a -> (a, f a)
Reverse.grad' forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x)
bounds :: ProblemReverse f -> Maybe (Vector (Double, Double))
bounds (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
_f Maybe (Vector (Double, Double))
bounds [Constraint]
_constraints Int
_size f Int
_template) = Maybe (Vector (Double, Double))
bounds
constraints :: ProblemReverse f -> [Constraint]
constraints (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
_f Maybe (Vector (Double, Double))
_bounds [Constraint]
constraints Int
_size f Int
_template) = [Constraint]
constraints
instance Traversable f => Opt.HasGrad (ProblemReverse f) where
grad :: ProblemReverse f -> Vector Double -> Vector Double
grad (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
func Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
size f Int
template) =
forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a -> f a
Reverse.grad forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
func forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template
grad'M :: forall (m :: * -> *).
PrimMonad m =>
ProblemReverse f
-> Vector Double -> MVector (PrimState m) Double -> m Double
grad'M (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x MVector (PrimState m) Double
gvec = do
case forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a -> (a, f a)
Reverse.grad' forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x) of
(Double
y, f Double
g) -> do
forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f Double
g MVector (PrimState m) Double
gvec
forall (m :: * -> *) a. Monad m => a -> m a
return Double
y
instance Traversable f => Opt.Optionally (Opt.HasGrad (ProblemReverse f)) where
optionalDict :: Maybe (Dict (HasGrad (ProblemReverse f)))
optionalDict = forall (c :: Constraint). c => Maybe (Dict c)
hasOptionalDict
instance Opt.Optionally (Opt.HasHessian (ProblemReverse f)) where
optionalDict :: Maybe (Dict (HasHessian (ProblemReverse f)))
optionalDict = forall a. Maybe a
Nothing
minimizeReverse
:: forall f. Traversable f
=> Method
-> Params (f Double)
-> (forall s. Reifies s Tape => f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeReverse :: forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeReverse Method
method Params (f Double)
params forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (f (Double, Double))
bounds [Constraint]
constraints f Double
x0 = do
let size :: Int
template :: f Int
(Int
size, f Int
template) = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
i Double
_ -> Int
i seq :: forall a b. a -> b -> b
`seq` (Int
iforall a. Num a => a -> a -> a
+Int
1, Int
i)) Int
0 f Double
x0
bounds' :: Maybe (V.Vector (Double, Double))
bounds' :: Maybe (Vector (Double, Double))
bounds' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size) Maybe (f (Double, Double))
bounds
prob :: ProblemReverse f
prob = forall (f :: * -> *).
(forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double)
-> Maybe (Vector (Double, Double))
-> [Constraint]
-> Int
-> f Int
-> ProblemReverse f
ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (Vector (Double, Double))
bounds' [Constraint]
constraints Int
size f Int
template
params' :: Params (Vector Double)
params' = forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Params (f Double)
params
Result (Vector Double)
result <- forall prob.
(IsProblem prob, Optionally (HasGrad prob),
Optionally (HasHessian prob)) =>
Method
-> Params (Vector Double)
-> prob
-> Vector Double
-> IO (Result (Vector Double))
Opt.minimize Method
method Params (Vector Double)
params' ProblemReverse f
prob (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size f Double
x0)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Result (Vector Double)
result
data ProblemSparse f
= ProblemSparse
(forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))
(Maybe (V.Vector (Double, Double)))
[Constraint]
Int
(f Int)
instance Traversable f => Opt.IsProblem (ProblemSparse f) where
func :: ProblemSparse f -> Vector Double -> Double
func (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x =
forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a))
-> f a -> (a, f a)
Sparse.grad' forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x)
bounds :: ProblemSparse f -> Maybe (Vector (Double, Double))
bounds (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
_f Maybe (Vector (Double, Double))
bounds [Constraint]
_constraints Int
_size f Int
_template) = Maybe (Vector (Double, Double))
bounds
constraints :: ProblemSparse f -> [Constraint]
constraints (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
_f Maybe (Vector (Double, Double))
_bounds [Constraint]
constraints Int
_size f Int
_template) = [Constraint]
constraints
instance Traversable f => Opt.HasGrad (ProblemSparse f) where
grad :: ProblemSparse f -> Vector Double -> Vector Double
grad (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
size f Int
template) =
forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> f a
Sparse.grad forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template
grad'M :: forall (m :: * -> *).
PrimMonad m =>
ProblemSparse f
-> Vector Double -> MVector (PrimState m) Double -> m Double
grad'M (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x MVector (PrimState m) Double
gvec = do
case forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a))
-> f a -> (a, f a)
Sparse.grad' forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x) of
(Double
y, f Double
g) -> do
forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f Double
g MVector (PrimState m) Double
gvec
forall (m :: * -> *) a. Monad m => a -> m a
return Double
y
instance Traversable f => Opt.HasHessian (ProblemSparse f) where
hessian :: ProblemSparse f -> Vector Double -> Matrix Double
hessian (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
size f Int
template) =
forall {a} {t :: * -> *} {t :: * -> *}.
(Storable a, Foldable t, Foldable t) =>
Int -> t (t a) -> Matrix a
toMatrix Int
size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a))
-> f a -> f (f a)
Sparse.hessian forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template
where
toMatrix :: Int -> t (t a) -> Matrix a
toMatrix Int
n t (t a)
xss = (Int
n forall a. Storable a => Int -> Int -> [a] -> Matrix a
LA.>< Int
n) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t (t a)
xss
instance Traversable f => Opt.Optionally (Opt.HasGrad (ProblemSparse f)) where
optionalDict :: Maybe (Dict (HasGrad (ProblemSparse f)))
optionalDict = forall (c :: Constraint). c => Maybe (Dict c)
hasOptionalDict
instance Traversable f => Opt.Optionally (Opt.HasHessian (ProblemSparse f)) where
optionalDict :: Maybe (Dict (HasHessian (ProblemSparse f)))
optionalDict = forall (c :: Constraint). c => Maybe (Dict c)
hasOptionalDict
minimizeSparse
:: forall f. Traversable f
=> Method
-> Params (f Double)
-> (forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeSparse :: forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeSparse Method
method Params (f Double)
params forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (f (Double, Double))
bounds [Constraint]
constraints f Double
x0 = do
let size :: Int
template :: f Int
(Int
size, f Int
template) = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
i Double
_ -> Int
i seq :: forall a b. a -> b -> b
`seq` (Int
iforall a. Num a => a -> a -> a
+Int
1, Int
i)) Int
0 f Double
x0
bounds' :: Maybe (V.Vector (Double, Double))
bounds' :: Maybe (Vector (Double, Double))
bounds' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size) Maybe (f (Double, Double))
bounds
prob :: ProblemSparse f
prob = forall (f :: * -> *).
(forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))
-> Maybe (Vector (Double, Double))
-> [Constraint]
-> Int
-> f Int
-> ProblemSparse f
ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (Vector (Double, Double))
bounds' [Constraint]
constraints Int
size f Int
template
params' :: Params (Vector Double)
params' = forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Params (f Double)
params
Result (Vector Double)
result <- forall prob.
(IsProblem prob, Optionally (HasGrad prob),
Optionally (HasHessian prob)) =>
Method
-> Params (Vector Double)
-> prob
-> Vector Double
-> IO (Result (Vector Double))
Opt.minimize Method
method Params (Vector Double)
params' ProblemSparse f
prob (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size f Double
x0)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Result (Vector Double)
result
minimize
:: forall f. Traversable f
=> Method
-> Params (f Double)
-> (forall s. Reifies s Tape => f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimize :: forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimize = forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeReverse
fromVector :: (Functor f, VG.Vector v a) => f Int -> v a -> f a
fromVector :: forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template v a
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (v a
x forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!) f Int
template
toVector :: (Traversable f, VG.Vector v a) => Int -> f a -> v a
toVector :: forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size f a
x = forall (v :: * -> *) a.
Vector v a =>
(forall s. ST s (Mutable v s a)) -> v a
VG.create forall a b. (a -> b) -> a -> b
$ do
Mutable v s a
vec <- forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VGM.new Int
size
forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f a
x Mutable v s a
vec
forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s a
vec
writeToMVector :: (PrimMonad m, VGM.MVector mv a, Traversable f) => f a -> mv (PrimState m) a -> m ()
writeToMVector :: forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f a
x mv (PrimState m) a
vec = do
Int
_ <- forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM (\Int
i a
v -> forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write mv (PrimState m) a
vec Int
i a
v forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (Int
iforall a. Num a => a -> a -> a
+Int
1)) Int
0 f a
x
forall (m :: * -> *) a. Monad m => a -> m a
return ()